From 8c5fa0c0f9ec195636372bc8861f2f88402717f4 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 15 Aug 2026 20:05:43 +0200 Subject: [PATCH 001/126] feat: add show budget window usage --- .../key_management_endpoints.py | 62 ++++++ .../test_key_management_endpoints.py | 206 ++++++++++++++++++ 2 files changed, 268 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e190e8b19d..22edbf33f94 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -106,6 +106,7 @@ from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( get_ui_settings_cached, @@ -3516,6 +3517,61 @@ async def _build_model_max_budget_usage( return result +async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None: + """ + Attach current-window spend to each entry in key_info["budget_limits"], in place. + + Per-window spend is not persisted in the DB; it lives in the cross-pod spend + counters (spend:key:{hashed_token}:window:{budget_duration}) that + _virtual_key_multi_budget_check enforces against, so we read the same + counters via get_current_spend. Passing max_budget + window_start makes the + read re-check against the authoritative spend-log aggregate when the counter + is stale-low (e.g. after a Redis flush), same as the enforcement path. + """ + from litellm.proxy.proxy_server import get_current_spend + + budget_limits: Any = key_info.get("budget_limits") + if isinstance(budget_limits, str): + try: + budget_limits = json.loads(budget_limits) + except (TypeError, ValueError): + return + key_info["budget_limits"] = budget_limits + if not isinstance(budget_limits, list): + return + + for idx, window in enumerate(budget_limits): + w: dict | None = None + if isinstance(window, dict): + w = window + elif hasattr(window, "model_dump"): + try: + w = window.model_dump() + except Exception: # noqa: BLE001 + continue + budget_limits[idx] = w + if not w: + continue + duration: Any = w.get("budget_duration") + max_budget: Any = w.get("max_budget") + if not duration: + continue + try: + max_budget = float(max_budget) if max_budget is not None else None + except (TypeError, ValueError): + max_budget = None + counter_key: Final = f"spend:key:{api_key_hash}:window:{duration}" + spend: Final = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=max_budget, + window_entity_type="Key", + window_entity_id=api_key_hash, + window_start=get_budget_window_start(w), + ) + w["current_spend"] = round(spend, 4) + + @router.post( "/v2/key/info", tags=["key management"], @@ -3597,6 +3653,8 @@ async def info_key_fn_v2( model_max_budget=model_max_budget, user_api_key_cache=user_api_key_cache, ) + if k_token_hash: + await _attach_budget_limits_usage(key_info=k_dict, api_key_hash=k_token_hash) filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3633,6 +3691,9 @@ async def info_key_fn( - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - model_max_budget_usage: dict | None - Current-window spend per model, present only when the key has per-model budgets + - budget_limits: list | None - Concurrent budget windows. Each entry includes + current_spend: spend accumulated in the window so far (read from the same cross-pod + spend counter the budget enforcement uses) - models: list - Model_name's the key is allowed to call - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} @@ -3709,6 +3770,7 @@ async def info_key_fn( model_max_budget=model_max_budget, user_api_key_cache=user_api_key_cache, ) + await _attach_budget_limits_usage(key_info=key_info, api_key_hash=key_token_hash) # Attach object_permission if object_permission_id is set key_info = await attach_object_permission_to_dict(key_info, prisma_client) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bdf09a95e4b..59f7a2f1e8e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13789,6 +13789,212 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): mock_prisma_client.db.query_raw.assert_not_awaited() +@pytest.mark.asyncio +async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): + """ + /key/info should attach current_spend to each budget_limits window, read from + the same spend counter (spend:key:{token}:window:{duration}) that budget + enforcement uses. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + test_key_token = "hashed_token_window_test" + budget_limits = [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + } + ] + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=0.73) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) + mock_key_info.token = test_key_token + mock_key_info.object_permission_id = None + mock_key_info.user_id = "user-w" + mock_key_info.team_id = None + mock_key_info.litellm_budget_table = None + mock_key_info.model_dump.return_value = { + "token": test_key_token, + "budget_limits": [dict(w) for w in budget_limits], + "user_id": "user-w", + "team_id": None, + "object_permission_id": None, + "litellm_budget_table": None, + } + mock_key_info.dict.return_value = mock_key_info.model_dump.return_value + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_info + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test-window-key", + ) + + result = await info_key_fn( + key="sk-test-window-key", + user_api_key_dict=user_api_key_dict, + ) + + windows = result["info"]["budget_limits"] + assert len(windows) == 1 + assert windows[0]["current_spend"] == 0.73 + assert windows[0]["max_budget"] == 2.0 + assert windows[0]["budget_duration"] == "1h" + + mock_get_current_spend.assert_awaited_once() + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == f"spend:key:{test_key_token}:window:1h" + assert call_kwargs["max_budget"] == 2.0 + assert call_kwargs["window_entity_type"] == "Key" + assert call_kwargs["window_entity_id"] == test_key_token + assert call_kwargs["window_start"] is not None + + +@pytest.mark.asyncio +async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): + """Keys without budget_limits should not trigger window spend lookups.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + test_key_token = "hashed_token_no_windows" + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=0.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) + mock_key_info.token = test_key_token + mock_key_info.object_permission_id = None + mock_key_info.user_id = "user-nw" + mock_key_info.team_id = None + mock_key_info.litellm_budget_table = None + mock_key_info.model_dump.return_value = { + "token": test_key_token, + "budget_limits": None, + "user_id": "user-nw", + "team_id": None, + "object_permission_id": None, + "litellm_budget_table": None, + } + mock_key_info.dict.return_value = mock_key_info.model_dump.return_value + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_info + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test-no-window-key", + ) + + result = await info_key_fn( + key="sk-test-no-window-key", + user_api_key_dict=user_api_key_dict, + ) + + assert result["info"]["budget_limits"] is None + mock_get_current_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): + """/v2/key/info should attach current_spend to each budget_limits window.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import ( + info_key_fn_v2, + ) + + test_key_token = "hashed_token_v2_window_test" + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=1.25) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_key_token + mock_key.user_id = "user-v2-w" + mock_key.team_id = None + mock_key.model_dump.return_value = { + "token": test_key_token, + "budget_limits": [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + }, + { + "reset_at": "2026-08-16T00:00:00+00:00", + "max_budget": 20.0, + "budget_duration": "1d", + }, + ], + "user_id": "user-v2-w", + "team_id": None, + "litellm_budget_table": None, + } + mock_key.dict.return_value = mock_key.model_dump.return_value + + mock_prisma_client.get_data = AsyncMock(return_value=[mock_key]) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin-v2-w", + ) + + result = await info_key_fn_v2( + data=KeyRequest(keys=[test_key_token]), + user_api_key_dict=user_api_key_dict, + ) + + assert len(result["info"]) == 1 + windows = result["info"][0]["budget_limits"] + assert len(windows) == 2 + assert windows[0]["current_spend"] == 1.25 + assert windows[1]["current_spend"] == 1.25 + assert mock_get_current_spend.await_count == 2 + counter_keys = { + call.kwargs["counter_key"] for call in mock_get_current_spend.await_args_list + } + assert counter_keys == { + f"spend:key:{test_key_token}:window:1h", + f"spend:key:{test_key_token}:window:1d", + } + + @pytest.mark.asyncio async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): """Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match.""" From 34bc22fffd3459954afc42ede850f622d1344520 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 15 Aug 2026 21:10:12 +0200 Subject: [PATCH 002/126] fix: address pr comments --- .../key_management_endpoints.py | 25 +++-- .../test_key_management_endpoints.py | 103 ++++++++++++++++++ 2 files changed, 119 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 22edbf33f94..5e4642d66a9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,6 +3517,20 @@ async def _build_model_max_budget_usage( return result +def _budget_window_to_dict(window: object) -> dict | None: + """Coerce a budget_limits entry to a dict; None when the entry is unusable.""" + if isinstance(window, dict): + return window + model_dump = getattr(window, "model_dump", None) + if callable(model_dump): + try: + dumped: Any = model_dump() + except Exception: # noqa: BLE001 + return None + return dumped if isinstance(dumped, dict) else None + return None + + async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None: """ Attach current-window spend to each entry in key_info["budget_limits"], in place. @@ -3541,17 +3555,10 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None return for idx, window in enumerate(budget_limits): - w: dict | None = None - if isinstance(window, dict): - w = window - elif hasattr(window, "model_dump"): - try: - w = window.model_dump() - except Exception: # noqa: BLE001 - continue - budget_limits[idx] = w + w: Final = _budget_window_to_dict(window) if not w: continue + budget_limits[idx] = w duration: Any = w.get("budget_duration") max_budget: Any = w.get("max_budget") if not duration: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 59f7a2f1e8e..5df7c1cb28d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13995,6 +13995,109 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): } +@pytest.mark.asyncio +async def test_attach_budget_limits_usage_json_string_input(monkeypatch): + """budget_limits stored as a JSON string should be parsed and annotated.""" + import json as json_module + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _attach_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.5) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + key_info = { + "budget_limits": json_module.dumps( + [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] + ) + } + await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + + assert isinstance(key_info["budget_limits"], list) + assert key_info["budget_limits"][0]["current_spend"] == 0.5 + + +@pytest.mark.asyncio +async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): + """Invalid JSON strings, non-list values, and malformed windows are skipped.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _attach_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + # invalid JSON string + key_info = {"budget_limits": "{not json"} + await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + assert key_info["budget_limits"] == "{not json" + + # non-list value + key_info = {"budget_limits": {"budget_duration": "1h"}} + await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + + # windows that are falsy, missing budget_duration, or not dict-like + key_info = { + "budget_limits": [ + {}, + {"max_budget": 2.0}, + {"budget_duration": "1h", "max_budget": "not-a-number"}, + 42, + ] + } + await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + + # only the well-formed window (with unparseable max_budget coerced to None) + # triggers a spend lookup + mock_get_current_spend.assert_awaited_once() + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == "spend:key:hash-1:window:1h" + assert call_kwargs["max_budget"] is None + assert key_info["budget_limits"][2]["current_spend"] == 0.0 + assert key_info["budget_limits"][3] == 42 + + +@pytest.mark.asyncio +async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch): + """Window objects with model_dump() are converted to dicts in place.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _attach_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=1.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + good_window = MagicMock() + good_window.model_dump.return_value = { + "budget_duration": "7d", + "max_budget": 10.0, + "reset_at": None, + } + bad_window = MagicMock() + bad_window.model_dump.side_effect = ValueError("boom") + + key_info = {"budget_limits": [good_window, bad_window]} + await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-2") + + # good window converted to dict and annotated; failing window left as-is + assert isinstance(key_info["budget_limits"][0], dict) + assert key_info["budget_limits"][0]["current_spend"] == 1.0 + assert key_info["budget_limits"][1] is bad_window + mock_get_current_spend.assert_awaited_once() + + @pytest.mark.asyncio async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): """Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match.""" From db67bafd59a0e131d70cfc4a5b1a3e494ca63373 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 15 Aug 2026 21:26:55 +0200 Subject: [PATCH 003/126] fix: drop Final from loop locals to stay within the basedpyright budget --- .../proxy/management_endpoints/key_management_endpoints.py | 6 +++--- ui/litellm-dashboard/src/lib/http/schema.d.ts | 3 +++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5e4642d66a9..51b9ebf91d8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3555,7 +3555,7 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None return for idx, window in enumerate(budget_limits): - w: Final = _budget_window_to_dict(window) + w = _budget_window_to_dict(window) if not w: continue budget_limits[idx] = w @@ -3567,8 +3567,8 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None max_budget = float(max_budget) if max_budget is not None else None except (TypeError, ValueError): max_budget = None - counter_key: Final = f"spend:key:{api_key_hash}:window:{duration}" - spend: Final = await get_current_spend( + counter_key = f"spend:key:{api_key_hash}:window:{duration}" + spend = await get_current_spend( counter_key=counter_key, fallback_spend=0.0, max_budget=max_budget, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3b4b96bdd4e..f4703202b18 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7020,6 +7020,9 @@ export interface paths { * - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - model_max_budget_usage: dict | None - Current-window spend per model, present only when * the key has per-model budgets + * - budget_limits: list | None - Concurrent budget windows. Each entry includes + * current_spend: spend accumulated in the window so far (read from the same cross-pod + * spend counter the budget enforcement uses) * - models: list - Model_name's the key is allowed to call * - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits * - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} From 31130036c0013b8c0ad34e6cffbd58b24ee8b78a Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 17 Aug 2026 16:21:31 +0200 Subject: [PATCH 004/126] refactor: build budget window usage without in-place mutation Replace _attach_budget_limits_usage, which rewrote the caller's key_info dict, with _budget_limits_with_usage returning a new list. Callers assign the result once. Keeps the response shape and spend-counter read path identical while following the repo's no-mutation rule. --- .../key_management_endpoints.py | 119 +++++++++++------- .../test_key_management_endpoints.py | 83 ++++++------ 2 files changed, 121 insertions(+), 81 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 51b9ebf91d8..f6f739f3eae 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,23 +3517,44 @@ async def _build_model_max_budget_usage( return result -def _budget_window_to_dict(window: object) -> dict | None: +def _budget_window_to_dict(window: object) -> Mapping[str, object] | None: """Coerce a budget_limits entry to a dict; None when the entry is unusable.""" if isinstance(window, dict): return window - model_dump = getattr(window, "model_dump", None) - if callable(model_dump): + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return None + try: + dumped: Final = model_dump() + except Exception: # noqa: BLE001 # model_dump implementations can raise arbitrary errors + return None + return dumped if isinstance(dumped, dict) else None + + +def _coerce_budget_limits(budget_limits: object) -> Sequence[object] | None: + """Coerce budget_limits to a sequence of windows, parsing JSON strings; None when unusable.""" + if isinstance(budget_limits, str): try: - dumped: Any = model_dump() - except Exception: # noqa: BLE001 + parsed: Final = json.loads(budget_limits) + except (TypeError, ValueError): + return None + return parsed if isinstance(parsed, list) else None + return budget_limits if isinstance(budget_limits, list) else None + + +def _parse_window_max_budget(value: object) -> float | None: + """Coerce a window's max_budget to float; None when absent or unparseable.""" + if isinstance(value, (int, float, str)): + try: + return float(value) + except (TypeError, ValueError): return None - return dumped if isinstance(dumped, dict) else None return None -async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None: +async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: str) -> Mapping[str, object]: """ - Attach current-window spend to each entry in key_info["budget_limits"], in place. + Return a copy of a budget window with current-window spend attached. Per-window spend is not persisted in the DB; it lives in the cross-pod spend counters (spend:key:{hashed_token}:window:{budget_duration}) that @@ -3544,39 +3565,43 @@ async def _attach_budget_limits_usage(key_info: dict, api_key_hash: str) -> None """ from litellm.proxy.proxy_server import get_current_spend - budget_limits: Any = key_info.get("budget_limits") - if isinstance(budget_limits, str): - try: - budget_limits = json.loads(budget_limits) - except (TypeError, ValueError): - return - key_info["budget_limits"] = budget_limits - if not isinstance(budget_limits, list): - return + duration: Final = window.get("budget_duration") + if not duration: + return dict(window) # mutable-ok: per-window response copy, built once per window + spend: Final = await get_current_spend( + counter_key=f"spend:key:{api_key_hash}:window:{duration}", + fallback_spend=0.0, + max_budget=_parse_window_max_budget(window.get("max_budget")), + window_entity_type="Key", + window_entity_id=api_key_hash, + window_start=get_budget_window_start(window), + ) + return {**window, "current_spend": round(spend, 4)} # mutable-ok: per-window response copy, built once per window - for idx, window in enumerate(budget_limits): - w = _budget_window_to_dict(window) - if not w: - continue - budget_limits[idx] = w - duration: Any = w.get("budget_duration") - max_budget: Any = w.get("max_budget") - if not duration: - continue - try: - max_budget = float(max_budget) if max_budget is not None else None - except (TypeError, ValueError): - max_budget = None - counter_key = f"spend:key:{api_key_hash}:window:{duration}" - spend = await get_current_spend( - counter_key=counter_key, - fallback_spend=0.0, - max_budget=max_budget, - window_entity_type="Key", - window_entity_id=api_key_hash, - window_start=get_budget_window_start(w), - ) - w["current_spend"] = round(spend, 4) + +async def _budget_limits_entry_with_usage(window: object, api_key_hash: str) -> object: + """Return the window as an enriched dict when dict-coercible; the original entry otherwise.""" + coerced: Final = _budget_window_to_dict(window) + if not coerced: + return window + return await _budget_window_with_usage(window=coerced, api_key_hash=api_key_hash) + + +async def _budget_limits_with_usage(budget_limits: object, api_key_hash: str) -> Sequence[object] | None: + """ + Return budget_limits as window dicts with current-window spend attached. + + None when budget_limits is not a usable (possibly JSON-encoded) list; the + caller keeps the original value then. Entries that are not dict-coercible + are preserved as-is. + """ + windows: Final = _coerce_budget_limits(budget_limits) + if windows is None: + return None + return [ # mutable-ok: entries are awaited, so they cannot be built inside a frozen wrapper + await _budget_limits_entry_with_usage(window=window, api_key_hash=api_key_hash) + for window in windows + ] @router.post( @@ -3661,7 +3686,12 @@ async def info_key_fn_v2( user_api_key_cache=user_api_key_cache, ) if k_token_hash: - await _attach_budget_limits_usage(key_info=k_dict, api_key_hash=k_token_hash) + budget_limits_usage = await _budget_limits_with_usage( + budget_limits=k_dict.get("budget_limits"), + api_key_hash=k_token_hash, + ) + if budget_limits_usage is not None: + k_dict["budget_limits"] = budget_limits_usage filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3777,7 +3807,12 @@ async def info_key_fn( model_max_budget=model_max_budget, user_api_key_cache=user_api_key_cache, ) - await _attach_budget_limits_usage(key_info=key_info, api_key_hash=key_token_hash) + budget_limits_usage: Final = await _budget_limits_with_usage( + budget_limits=key_info.get("budget_limits"), + api_key_hash=key_token_hash, + ) + if budget_limits_usage is not None: + key_info["budget_limits"] = budget_limits_usage # Attach object_permission if object_permission_id is set key_info = await attach_object_permission_to_dict(key_info, prisma_client) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 5df7c1cb28d..1f3b9042ed7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13996,13 +13996,13 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): @pytest.mark.asyncio -async def test_attach_budget_limits_usage_json_string_input(monkeypatch): +async def test_budget_limits_with_usage_json_string_input(monkeypatch): """budget_limits stored as a JSON string should be parsed and annotated.""" import json as json_module from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _attach_budget_limits_usage, + _budget_limits_with_usage, ) mock_get_current_spend = AsyncMock(return_value=0.5) @@ -14010,24 +14010,29 @@ async def test_attach_budget_limits_usage_json_string_input(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - key_info = { - "budget_limits": json_module.dumps( - [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] - ) - } - await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + raw = json_module.dumps( + [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] + ) + result = await _budget_limits_with_usage(budget_limits=raw, api_key_hash="hash-1") - assert isinstance(key_info["budget_limits"], list) - assert key_info["budget_limits"][0]["current_spend"] == 0.5 + assert result == [ + { + "budget_duration": "1h", + "max_budget": 2.0, + "reset_at": None, + "current_spend": 0.5, + } + ] + mock_get_current_spend.assert_awaited_once() @pytest.mark.asyncio -async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): +async def test_budget_limits_with_usage_skips_unusable_inputs(monkeypatch): """Invalid JSON strings, non-list values, and malformed windows are skipped.""" from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _attach_budget_limits_usage, + _budget_limits_with_usage, ) mock_get_current_spend = AsyncMock(return_value=0.0) @@ -14035,25 +14040,18 @@ async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - # invalid JSON string - key_info = {"budget_limits": "{not json"} - await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") - assert key_info["budget_limits"] == "{not json" - - # non-list value - key_info = {"budget_limits": {"budget_duration": "1h"}} - await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + # invalid JSON string and non-list values return None: callers keep the original + assert await _budget_limits_with_usage(budget_limits="{not json", api_key_hash="hash-1") is None + assert await _budget_limits_with_usage(budget_limits={"budget_duration": "1h"}, api_key_hash="hash-1") is None # windows that are falsy, missing budget_duration, or not dict-like - key_info = { - "budget_limits": [ - {}, - {"max_budget": 2.0}, - {"budget_duration": "1h", "max_budget": "not-a-number"}, - 42, - ] - } - await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-1") + windows = [ + {}, + {"max_budget": 2.0}, + {"budget_duration": "1h", "max_budget": "not-a-number"}, + 42, + ] + result = await _budget_limits_with_usage(budget_limits=windows, api_key_hash="hash-1") # only the well-formed window (with unparseable max_budget coerced to None) # triggers a spend lookup @@ -14061,17 +14059,22 @@ async def test_attach_budget_limits_usage_skips_unusable_inputs(monkeypatch): call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-1:window:1h" assert call_kwargs["max_budget"] is None - assert key_info["budget_limits"][2]["current_spend"] == 0.0 - assert key_info["budget_limits"][3] == 42 + assert result is not None + assert result[0] == {} + assert result[1] == {"max_budget": 2.0} + assert result[2] == {"budget_duration": "1h", "max_budget": "not-a-number", "current_spend": 0.0} + assert result[3] == 42 + # input is not mutated + assert windows[2] == {"budget_duration": "1h", "max_budget": "not-a-number"} @pytest.mark.asyncio -async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch): - """Window objects with model_dump() are converted to dicts in place.""" +async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): + """Window objects with model_dump() are converted to dicts; failing windows pass through.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _attach_budget_limits_usage, + _budget_limits_with_usage, ) mock_get_current_spend = AsyncMock(return_value=1.0) @@ -14088,13 +14091,15 @@ async def test_attach_budget_limits_usage_pydantic_windows(monkeypatch): bad_window = MagicMock() bad_window.model_dump.side_effect = ValueError("boom") - key_info = {"budget_limits": [good_window, bad_window]} - await _attach_budget_limits_usage(key_info=key_info, api_key_hash="hash-2") + result = await _budget_limits_with_usage( + budget_limits=[good_window, bad_window], api_key_hash="hash-2" + ) # good window converted to dict and annotated; failing window left as-is - assert isinstance(key_info["budget_limits"][0], dict) - assert key_info["budget_limits"][0]["current_spend"] == 1.0 - assert key_info["budget_limits"][1] is bad_window + assert result is not None + assert isinstance(result[0], dict) + assert result[0]["current_spend"] == 1.0 + assert result[1] is bad_window mock_get_current_spend.assert_awaited_once() From 9b30e73d07ebccdf4e6ff24981110e238c080d4d Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 17 Aug 2026 20:48:12 +0200 Subject: [PATCH 005/126] fix: cap /v2/key/info batch size to bound spend-log query fan-out --- .../key_management_endpoints.py | 16 +++++++++ .../test_key_management_endpoints.py | 33 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f6f739f3eae..46d7f617352 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3604,6 +3604,11 @@ async def _budget_limits_with_usage(budget_limits: object, api_key_hash: str) -> ] +# Caps per-request fan-out: each key with budget windows costs one spend-counter +# read (worst case a SpendLogs aggregation) per window. +MAX_KEY_INFO_KEYS_PER_REQUEST: Final = 100 + + @router.post( "/v2/key/info", tags=["key management"], @@ -3643,6 +3648,17 @@ async def info_key_fn_v2( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail={"message": "Malformed request. No keys passed in."}, ) + requested_key_count: Final = len(data.keys or []) + len(data.key_aliases or []) + if requested_key_count > MAX_KEY_INFO_KEYS_PER_REQUEST: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "message": ( + f"Too many keys requested: {requested_key_count}. " + f"At most {MAX_KEY_INFO_KEYS_PER_REQUEST} keys and key_aliases combined per request." + ) + }, + ) # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query: Final = list(data.keys) if data.keys else [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 1f3b9042ed7..9194ffb603b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13995,6 +13995,39 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): } +@pytest.mark.asyncio +async def test_info_key_fn_v2_rejects_oversized_batch(monkeypatch): + """/v2/key/info must reject over-cap batches before doing any DB work.""" + from unittest.mock import AsyncMock + + from litellm.proxy._types import KeyRequest, ProxyException + from litellm.proxy.management_endpoints.key_management_endpoints import ( + MAX_KEY_INFO_KEYS_PER_REQUEST, + info_key_fn_v2, + ) + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin-batch-cap", + ) + + with pytest.raises(ProxyException) as exc_info: + await info_key_fn_v2( + data=KeyRequest( + keys=[f"hash-{i}" for i in range(MAX_KEY_INFO_KEYS_PER_REQUEST)], + key_aliases=["alias-over-cap"], + ), + user_api_key_dict=user_api_key_dict, + ) + + assert exc_info.value.code == "422" + mock_prisma_client.get_data.assert_not_awaited() + + @pytest.mark.asyncio async def test_budget_limits_with_usage_json_string_input(monkeypatch): """budget_limits stored as a JSON string should be parsed and annotated.""" From ffd23ab1319d0caf9c870be5849006cbc931f3d1 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 17 Aug 2026 21:00:24 +0200 Subject: [PATCH 006/126] style: apply ruff format to budget limits comprehension --- litellm/proxy/management_endpoints/key_management_endpoints.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 46d7f617352..170b431c1e2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3599,8 +3599,7 @@ async def _budget_limits_with_usage(budget_limits: object, api_key_hash: str) -> if windows is None: return None return [ # mutable-ok: entries are awaited, so they cannot be built inside a frozen wrapper - await _budget_limits_entry_with_usage(window=window, api_key_hash=api_key_hash) - for window in windows + await _budget_limits_entry_with_usage(window=window, api_key_hash=api_key_hash) for window in windows ] From 4e3e7d8cf6c0afe6171f784a1bba1b0080831edf Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 17 Aug 2026 21:12:36 +0200 Subject: [PATCH 007/126] test: cover soft budget row creation and windows without max_budget --- .../test_key_management_endpoints.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 9194ffb603b..c05380dd054 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -540,6 +540,57 @@ async def test_key_generation_with_object_permission(monkeypatch): assert key_insert_calls[0]["data"].get("object_permission_id") == "objperm123" +@pytest.mark.asyncio +async def test_generate_key_with_soft_budget_creates_budget_row(monkeypatch): + """soft_budget on /key/generate must create a budget table row and link its budget_id to the key.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.jsonify_object = lambda data: data + mock_prisma_client.db = MagicMock() + mock_budget_create = AsyncMock(return_value=MagicMock(budget_id="budget-soft-123")) + mock_prisma_client.db.litellm_budgettable = MagicMock() + mock_prisma_client.db.litellm_budgettable.create = mock_budget_create + + async def _insert_data_side_effect(*args, **kwargs): + if kwargs.get("table_name") == "user": + return MagicMock(models=[], spend=0) + return MagicMock( + token="hashed_token_soft", + litellm_budget_table=None, + object_permission=None, + ) + + mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + await generate_key_fn( + data=GenerateKeyRequest(soft_budget=5.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-1", + ), + ) + + mock_budget_create.assert_awaited_once() + created_budget = mock_budget_create.call_args.kwargs["data"] + assert created_budget["soft_budget"] == 5.0 + assert created_budget["created_by"] == "admin-1" + + key_insert_calls = [ + call.kwargs + for call in mock_prisma_client.insert_data.call_args_list + if call.kwargs.get("table_name") == "key" + ] + assert len(key_insert_calls) == 1 + assert key_insert_calls[0]["data"].get("budget_id") == "budget-soft-123" + + @pytest.mark.asyncio async def test_generate_key_debug_log_never_contains_raw_token(monkeypatch, caplog): """Regression for LIT-4356: /key/generate must never emit the raw virtual key @@ -14101,6 +14152,30 @@ async def test_budget_limits_with_usage_skips_unusable_inputs(monkeypatch): assert windows[2] == {"budget_duration": "1h", "max_budget": "not-a-number"} +@pytest.mark.asyncio +async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): + """A window with only budget_duration still gets current_spend, read without a budget ceiling.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _budget_limits_with_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.75) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + result = await _budget_limits_with_usage( + budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" + ) + + assert result == [{"budget_duration": "2d", "current_spend": 0.75}] + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == "spend:key:hash-no-max:window:2d" + assert call_kwargs["max_budget"] is None + + @pytest.mark.asyncio async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): """Window objects with model_dump() are converted to dicts; failing windows pass through.""" From f4406b4060cddb60731f804513455bea321e7407 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 24 Aug 2026 08:52:18 +0200 Subject: [PATCH 008/126] fix(proxy): clear LIT002 in /v2/key/info batch-cap guard The two `or []` fallbacks only feed len(), so tuples do the job without a mutable literal. The detail dict mirrors the sibling HTTPException payload above and is never mutated, so it carries a mutable-ok reason. --- .../proxy/management_endpoints/key_management_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2b37d0bc8fc..85f6d25f6dc 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3675,11 +3675,11 @@ async def info_key_fn_v2( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail={"message": "Malformed request. No keys passed in."}, ) - requested_key_count: Final = len(data.keys or []) + len(data.key_aliases or []) + requested_key_count: Final = len(data.keys or ()) + len(data.key_aliases or ()) if requested_key_count > MAX_KEY_INFO_KEYS_PER_REQUEST: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ + detail={ # mutable-ok: one-shot HTTPException payload matching the sibling detail dict above; never mutated after construction "message": ( f"Too many keys requested: {requested_key_count}. " f"At most {MAX_KEY_INFO_KEYS_PER_REQUEST} keys and key_aliases combined per request." From 41ab7e57e8ac4f222fa6191d1940d7e636debffd Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Mon, 24 Aug 2026 09:09:30 +0200 Subject: [PATCH 009/126] ci(test-unit): raise job timeout to 60m for the three 55m shards caching-local, proxy-extras and enterprise-package gave pytest 20m but capped the job at 55m; with a 35m setup ceiling plus 5m of runner overhead the job deadline could preempt pytest itself. --- .github/workflows/test-unit.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a7c67f2b35d..2dfca3d308f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -211,7 +211,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: proxy-extras artifact-name: proxy-extras @@ -219,7 +219,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: enterprise-package artifact-name: enterprise-package @@ -227,7 +227,7 @@ jobs: workers: 4 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: responses-caching-types artifact-name: responses-caching-types From 02bb3f9310e6633caefd8f739d2340ebd0d56cff Mon Sep 17 00:00:00 2001 From: Timik232 <100406268+Timik232@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:56:54 +0300 Subject: [PATCH 010/126] fix(streaming): keep response id stable across streamed chunks Providers that stream via GenericStreamingChunk (e.g. GigaChat) do not propagate an upstream response id, so every chunk of one streamed response got a freshly generated id. Pin CustomStreamWrapper.response_id from the first chunk it creates, mirroring the existing 'created' pinning (#11437). Clients that merge deltas by chunk id (e.g. goose) split one reply into one message per chunk. Fixes #38098 --- .../litellm_core_utils/streaming_handler.py | 2 + .../test_streaming_handler.py | 79 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f6340426c1b..fb693943b2b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -816,6 +816,8 @@ class CustomStreamWrapper: model_response: Final = ModelResponseStream(**args) if self.response_id is not None: model_response.id = self.response_id + elif model_response.id: + self.response_id = model_response.id if self.system_fingerprint is not None: model_response.system_fingerprint = self.system_fingerprint diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index b5e33a4e421..a76d427495c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4460,3 +4460,82 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp finally: trace_id_var.set("") session_id_var.set("") + + +class TestStableStreamingResponseId: + """ + All chunks of one streamed response must share the same top-level id + (OpenAI streaming contract). Providers streaming via GenericStreamingChunk + (e.g. GigaChat) do not propagate an upstream response id, so + CustomStreamWrapper must pin the id from the first chunk it creates, + mirroring the existing `created` pinning (issue #11437). + + Clients such as goose merge streamed deltas into one assistant message by + chunk id; per-chunk ids split a single reply into many messages. + """ + + def test_generic_chunks_share_one_id(self): + def _generic_chunks(): + return iter( + [ + { + "text": "Hello", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, + }, + { + "text": " world", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, + }, + { + "text": "", + "tool_use": None, + "is_finished": True, + "finish_reason": "stop", + "usage": { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + }, + "index": 0, + }, + ] + ) + + wrapper = CustomStreamWrapper( + completion_stream=_generic_chunks(), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + ids = [chunk.id for chunk in wrapper if chunk.id] + assert ids, "no chunks emitted" + assert len(set(ids)) == 1, f"chunk ids differ across one stream: {ids}" + + def test_creator_pins_id_from_first_chunk(self): + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + first = wrapper.model_response_creator() + assert wrapper.response_id == first.id + assert wrapper.model_response_creator().id == first.id + + def test_provider_supplied_id_still_wins(self): + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + wrapper.response_id = "chatcmpl-from-provider" + assert wrapper.model_response_creator().id == "chatcmpl-from-provider" From 4f6fd85ab1e36af0cb3f44b045c3767d0ebe89b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:10:05 -0700 Subject: [PATCH 011/126] feat(proxy): default to the v2 migration resolver, keep v1 as an opt-out The v2 resolver skips the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contend for one database during a rolling deploy. The standalone migration Job already defaulted to v2; this aligns the proxy-server path. v1 stays reachable two ways: --use_legacy_migration_resolver on the CLI, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, where prisma_migration.py calls run_server with a fixed argv and the env var is the only route in. --use_v2_migration_resolver still parses, so existing commands do not die on an unknown option. Because v2 fails fast where v1 retried every failed deploy, a database that is not accepting connections yet, or another instance holding the migration advisory lock, would now kill a boot that used to ride it out. Those two failures are retried, with Prisma's stderr logged each round, and still raise once the attempts are spent. Moves the resolver tests from litellm-proxy-extras/tests, which no CI job runs, into tests/litellm-proxy-extras, and repoints the dedicated Postgres CircleCI job at the legacy path so v1 keeps real-DB and proxy-boot coverage. --- .circleci/config.yml | 15 +- CLAUDE.md | 2 +- .../litellm_proxy_extras/utils.py | 81 ++++-- litellm-proxy-extras/tests/__init__.py | 0 litellm/proxy/proxy_cli.py | 22 +- .../test_setup_database_fail_fast.py | 238 +++++++++++++++++- .../test_basic_python_version.py | 14 +- tests/test_litellm/proxy/test_proxy_cli.py | 83 +++++- 8 files changed, 409 insertions(+), 46 deletions(-) delete mode 100644 litellm-proxy-extras/tests/__init__.py rename {litellm-proxy-extras/tests => tests/litellm-proxy-extras}/test_setup_database_fail_fast.py (50%) diff --git a/.circleci/config.yml b/.circleci/config.yml index 4615a6a5a7e..0fbbd5ec7f7 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1483,7 +1483,7 @@ jobs: - run: name: Run tests command: | - uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver" + uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver" installing_litellm_on_python_3_13: docker: @@ -1507,9 +1507,9 @@ jobs: - run: name: Run tests command: | - uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver" + uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver" - installing_litellm_on_python_v2_migration_resolver: + installing_litellm_on_python_legacy_migration_resolver: docker: - *python312_image - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 @@ -1536,10 +1536,10 @@ jobs: url: tcp://localhost:5432 timeout: "60" - run: - name: Run v2 migration resolver proxy smoke test + name: Run legacy migration resolver proxy smoke test command: | uv run --no-sync python -m pytest -vv \ - tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver + tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver helm_chart_testing: machine: @@ -2918,7 +2918,8 @@ jobs: command: | if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \ (grep -q "Database setup failed after multiple retries" docker_output.log || \ - grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then + grep -q "ERROR: Application startup failed. Exiting." docker_output.log || \ + grep -q "Database migration cannot proceed" docker_output.log); then echo "Expected error found. Test passed." else echo "Expected error not found. Test failed." @@ -3050,7 +3051,7 @@ workflows: filters: *main_branches - installing_litellm_on_python_3_13: filters: *main_branches - - installing_litellm_on_python_v2_migration_resolver: + - installing_litellm_on_python_legacy_migration_resolver: filters: *main_branches - helm_chart_testing: requires: diff --git a/CLAUDE.md b/CLAUDE.md index 930825aeb89..819f65059b7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -35,7 +35,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b27221c9beb..22da9b834ab 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -7,7 +7,8 @@ import subprocess import tempfile import time from pathlib import Path -from typing import Optional +from types import MappingProxyType +from typing import Final, Optional from litellm_proxy_extras._logging import logger from litellm_proxy_extras.replica_identity import ( @@ -50,6 +51,35 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile( re.IGNORECASE, ) +_MIGRATE_DEPLOY_ATTEMPTS: Final = 4 + +_TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType( + { + "deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)", + "P1001": "an unreachable database server", + "P1002": "a database server that timed out", + } +) + + +def _transient_deploy_failure(stderr: str) -> str | None: + """Describe why a failed `prisma migrate deploy` is worth retrying, or None. + + These are environment failures, not migration failures: the database is not + up yet, or another instance holds the migration lock. v1 retried every + failed deploy and absorbed them; failing fast on them instead would turn a + database that is ten seconds late into a dead proxy. + """ + return next( + ( + reason + for marker, reason in _TRANSIENT_DEPLOY_FAILURES.items() + if marker in stderr + ), + None, + ) + + PARTITIONED_SPEND_LOGS_PUSH_ERROR = ( "LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), " "so its primary key must include the partition key (\"startTime\"). `prisma db push` " @@ -648,7 +678,7 @@ class ProxyExtrasDBManager: @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ - v2 migration resolver (opt-in via --use_v2_migration_resolver). + v2 migration resolver (what the proxy CLI selects by default). Runs `prisma migrate deploy` and handles standard recovery paths (P3005 baseline, P3009/P3018 idempotent errors). Critically, it does @@ -692,7 +722,7 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - for attempt in range(4): + for attempt in range(_MIGRATE_DEPLOY_ATTEMPTS): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -807,16 +837,36 @@ class ProxyExtrasDBManager: f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" - ) from e + transient = _transient_deploy_failure(stderr) + if transient is None: + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from e + + if attempt == _MIGRATE_DEPLOY_ATTEMPTS - 1: + raise RuntimeError( + f"Database migration failed after " + f"{_MIGRATE_DEPLOY_ATTEMPTS} attempts on {transient}. " + "Check database connectivity and load." + f"\n\nPrisma error:\n{stderr}" + ) from e + + logger.info( + "prisma migrate deploy attempt %s failed on %s, retrying. " + "Prisma error:\n%s", + attempt + 1, + transient, + stderr, + ) + time.sleep(random.randrange(5, 15)) + continue raise RuntimeError( - "Database migration failed after 4 attempts (retry loop " - "exhausted by timeouts or repeated idempotent-recovery " - "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + f"Database migration failed after {_MIGRATE_DEPLOY_ATTEMPTS} " + "attempts (retry loop exhausted by timeouts or repeated " + "idempotent-recovery continues). Check database connectivity, " + "load, and _prisma_migrations ledger state." ) finally: os.chdir(original_dir) @@ -864,10 +914,11 @@ class ProxyExtrasDBManager: Args: use_migrate: Whether to use prisma migrate instead of db push - use_v2_resolver: Opt into the v2 migration resolver (safer during + use_v2_resolver: Run the v2 migration resolver (safer during rolling deploys; does not run the diff-and-force recovery - that causes schema thrashing). Defaults to False for - backwards compatibility. + that causes schema thrashing). Defaults to False here so + direct callers keep the old behavior; the proxy CLI passes + True, so the proxy's runtime default is v2. Returns: bool: True if setup was successful, False otherwise @@ -885,7 +936,7 @@ class ProxyExtrasDBManager: @staticmethod def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool: if use_v2_resolver: - logger.info("Using v2 migration resolver (--use_v2_migration_resolver)") + logger.info("Using v2 migration resolver") return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate) schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" diff --git a/litellm-proxy-extras/tests/__init__.py b/litellm-proxy-extras/tests/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8ac63ba25c9..c1d86344a1e 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -913,13 +913,14 @@ class ProxyInitializationHelpers: envvar="ENFORCE_PRISMA_MIGRATION_CHECK", ) @click.option( - "--use_v2_migration_resolver", - is_flag=True, - default=False, + "--use_v2_migration_resolver/--use_legacy_migration_resolver", + default=True, help=( - "Opt into the v2 migration resolver. Avoids the diff-and-force recovery " - "path that can cause schema thrashing during rolling deploys where two " - "LiteLLM versions contend for the same DB. Default is the v1 resolver." + "Which database migration resolver to run at startup. The default v2 " + "resolver avoids the diff-and-force recovery path that can cause schema " + "thrashing during rolling deploys where two LiteLLM versions contend for " + "the same DB. Pass --use_legacy_migration_resolver, or set " + "USE_V2_MIGRATION_RESOLVER=false, to fall back to v1." ), envvar="USE_V2_MIGRATION_RESOLVER", ) @@ -1310,10 +1311,11 @@ def run_server( else: if not use_v2_migration_resolver: print( - "\033[1;33mLiteLLM Proxy: Using default (v1) migration resolver. " - "If your deployment has seen schema thrashing during rolling " - "deploys, try --use_v2_migration_resolver (safer: avoids the " - "diff-and-force recovery that caused the thrash).\033[0m" + "\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration resolver. " + "The default v2 resolver is safer: it avoids the diff-and-force " + "recovery that caused schema thrashing during rolling deploys. " + "Remove --use_legacy_migration_resolver / " + "USE_V2_MIGRATION_RESOLVER=false to switch back to it.\033[0m" ) try: setup_ok: Final = PrismaManager.setup_database( diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py similarity index 50% rename from litellm-proxy-extras/tests/test_setup_database_fail_fast.py rename to tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 8d66bf872de..f637e139e86 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -1,15 +1,26 @@ -"""Regression tests for ProxyExtrasDBManager v2 migration resolver. +"""Regression tests for ProxyExtrasDBManager's v2 migration resolver. -The v2 resolver is opt-in via `--use_v2_migration_resolver` / the -`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1 -(default) behavior is unchanged from pre-fix. +v2 is what the proxy CLI selects by default; v1 stays reachable via +`--use_legacy_migration_resolver` or `USE_V2_MIGRATION_RESOLVER=false`. At the +library level the resolver is picked with the `use_v2_resolver` kwarg, which +still defaults to False so `migrations/run.py` and any direct caller keep their +own explicit choice. """ +import os import subprocess +import sys from unittest.mock import patch import pytest +sys.path.insert( + 0, + os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras") + ), +) + from litellm_proxy_extras.utils import ( ProxyExtrasDBManager, _max_migration_timestamp, @@ -240,3 +251,222 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" + + +_DEADLOCK_STDERR = ( + "Error: ERROR: deadlock detected\n" + "DETAIL: Process 277 waits for ExclusiveLock on advisory lock " + "[17556,0,72707369,1]; blocked by process 278.\n" + "Process 278 waits for ShareLock on virtual transaction 3/1041; " + "blocked by process 277." +) + + +class _DeployApplied: + stdout = "All migrations have been successfully applied." + stderr = "" + returncode = 0 + + +def _deploy_only(deploy_side_effect): + """subprocess.run stand-in that only intercepts `prisma migrate deploy`. + + Everything else the resolver shells out to, the Prisma toolchain check + above all, succeeds untouched, so a mock meant for the deploy call cannot + be silently consumed by an earlier subprocess call. + """ + deploys = {"n": 0} + + def _run(*args, **kwargs): + cmd = args[0] if args else kwargs.get("args", []) + if list(cmd)[-2:] == ["migrate", "deploy"]: + deploys["n"] += 1 + return deploy_side_effect(deploys["n"], cmd) + return _DeployApplied() + + return _run, deploys + + +def _prepare_v2_resolver(monkeypatch, tmp_path): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr("time.sleep", lambda *_a, **_k: None) + + +def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path): + """A deadlock on Prisma's migration advisory lock is transient and retried. + + Several proxy replicas booting against one database race `migrate deploy`, + and Postgres aborts one side. v1 retried any failed deploy, so it rode this + out; v2 classifies unrecognised stderr as unrecoverable and raises, which + with v2 as the default would take a replica's whole boot down. + """ + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output="" + ) + return _DeployApplied() + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert ok is True + assert deploys["n"] == 2, "the deadlocked deploy must be retried, not raised" + + +def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp_path): + """The deadlock retry stays bounded: a deadlock that never clears still + raises rather than looping forever or reporting a successful migration.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output="" + ) + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match="after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert deploys["n"] == 4 + + +@pytest.mark.parametrize( + "stderr", + [ + "Error: P1001: Can't reach database server at `db`:`5432`", + "Error: P1002: The database server was reached but timed out.", + ], +) +def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path, stderr): + """A database that is not accepting connections yet is retried, not fatal. + + A proxy and its database starting together race routinely, and v1 rode that + out by retrying every failed deploy. v2 treats unrecognised stderr as + unrecoverable, so without this the default flip would turn a database that + is a few seconds late into a dead proxy instead of a slow boot. + """ + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=stderr, output="" + ) + return _DeployApplied() + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert ok is True + assert deploys["n"] == 2, "an unreachable database must be retried, not raised" + + +def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_path): + """Retrying connectivity errors must not turn a genuinely unreachable + database into a silent success: after the attempts are spent it still + raises, so the proxy exits instead of serving without its database.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr="Error: P1001: Can't reach database server at `db`:`5432`", + output="", + ) + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match="after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert deploys["n"] == 4 + + +def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, caplog): + """Retrying must not swallow why the database was unreachable. + + Prisma's stderr is captured, so if the retry path neither logs it nor puts + it in the final error, an operator (and CI's bad-DATABASE_URL job, which + greps the boot log for the P1001 line) sees four silent retries and no + cause. + """ + _prepare_v2_resolver(monkeypatch, tmp_path) + stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`" + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=stderr, output="" + ) + + run, _ = _deploy_only(_side_effect) + with caplog.at_level("INFO", logger="litellm_proxy_extras"): + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError) as exc_info: + ProxyExtrasDBManager.setup_database( + use_migrate=True, use_v2_resolver=True + ) + + assert "P1001" in str(exc_info.value) + assert "P1001" in caplog.text + + +def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path): + """The transient classification must stay narrow: a genuinely broken + migration still fails fast on the first attempt rather than being retried + into the same error four times.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "The `20260101000000_genuinely_broken` migration failed to apply.\n" + 'Reason: syntax error at or near "BRKN" LINE 42' + ) + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=stderr, output="" + ) + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert deploys["n"] == 1 + + +def test_v1_still_runs_the_diff_and_force_recovery(monkeypatch, tmp_path): + """v1 remains the pre-existing diff-and-force resolver, unchanged by the + default flip: it still calls _resolve_all_migrations after a deploy that + applied something. Operators opting back in must get exactly the old path. + """ + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + class FakeResult: + stdout = "Applied migration.\n" + stderr = "" + + resolve_called = {"n": 0} + + def fake_resolve(*args, **kwargs): + resolve_called["n"] += 1 + + monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=False) + assert ok is True + assert resolve_called["n"] == 1 diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index fb06ed6b69d..506c58d26b4 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -305,14 +305,16 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None): def test_litellm_proxy_server_config_no_general_settings(): - """Exercises the default (v1) migration resolver.""" + """Exercises the default (v2) migration resolver.""" _run_proxy_server_smoke_test() -def test_litellm_proxy_server_config_no_general_settings_v2_resolver(): - """Exercises the opt-in v2 migration resolver. +def test_litellm_proxy_server_config_no_general_settings_legacy_resolver(): + """Exercises the legacy (v1) migration resolver against a real database. - Runs in a separate CI job against a local Postgres to avoid collisions - with the v1 variant when they share a database. + v2 is the default, so the no-arg test above already covers it. This one is + the only place the v1 opt-out gets real-DB migration plus proxy-boot + coverage, and it runs in a separate CI job against its own Postgres to + avoid collisions with the default variant. """ - _run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"]) + _run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"]) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6ea6f208bb5..ac1defe8d79 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1787,7 +1787,7 @@ class TestRunServerDbSetup: # use_prisma_db_push should be False (default), so use_migrate should be True run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) mock_setup_database.assert_called_with( - use_migrate=True, use_v2_resolver=False + use_migrate=True, use_v2_resolver=True ) # Reset mocks @@ -1802,7 +1802,7 @@ class TestRunServerDbSetup: standalone_mode=False, ) mock_setup_database.assert_called_with( - use_migrate=False, use_v2_resolver=False + use_migrate=False, use_v2_resolver=True ) @patch("subprocess.run") @@ -1869,7 +1869,7 @@ class TestRunServerDbSetup: ) assert exc_info.value.code == 1 mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=False + use_migrate=True, use_v2_resolver=True ) @patch("subprocess.run") @@ -1981,6 +1981,83 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=True ) + @pytest.mark.parametrize( + "argv, env_value, expected_v2", + [ + ([], None, True), + (["--use_v2_migration_resolver"], None, True), + (["--use_legacy_migration_resolver"], None, False), + ([], "false", False), + ([], "true", True), + (["--use_v2_migration_resolver"], "false", True), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_migration_resolver_default_and_opt_out( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + argv, + env_value, + expected_v2, + ): + """The proxy defaults to the v2 resolver, and v1 stays reachable. + + Both opt-out routes matter: --use_legacy_migration_resolver for a CLI + boot, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, + where litellm/proxy/prisma_migration.py calls run_server with a fixed + argv and an env var is the only way in. The deprecated + --use_v2_migration_resolver must still parse so existing commands do + not die on an unknown option, and an explicit flag still beats the env. + """ + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + mock_setup_database.return_value = True + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k + not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + } + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + if env_value is not None: + clean_env["USE_V2_MIGRATION_RESOLVER"] = env_value + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + ): + run_server.main( + ["--local", "--skip_server_startup", *argv], standalone_mode=False + ) + + mock_setup_database.assert_called_once_with( + use_migrate=True, use_v2_resolver=expected_v2 + ) + # --- Module-level helpers for worker startup hook tests --- From 7b36bfb96720651c756c911ad5bcdb7a819e2006 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:31:44 -0700 Subject: [PATCH 012/126] fix(proxy-extras): retry transient db push failures, drop a vacuous test `prisma db push` under v2 raised on the first failure while v1 retried it four times, so making v2 the default silently cost --use_prisma_db_push its retries. It now uses the same transient classification as migrate deploy. The classifier moves onto ProxyExtrasDBManager next to _is_permission_error and _is_idempotent_error, which do the same kind of stderr matching. Replaces a test that claimed to pin the transient classification but fed it a P3009 stderr, which an earlier branch catches, so it passed even when the classifier was mutated to treat everything as transient. The replacement uses an unclassified error and fails on that mutant. Drops a v1 test that duplicated test_v1_default_still_calls_resolve_all_migrations. --- .../litellm_proxy_extras/utils.py | 98 +++++++++------ .../test_setup_database_fail_fast.py | 113 ++++++++---------- tests/test_litellm/proxy/test_proxy_cli.py | 12 +- 3 files changed, 110 insertions(+), 113 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 22da9b834ab..a31016e0f17 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -51,9 +51,9 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile( re.IGNORECASE, ) -_MIGRATE_DEPLOY_ATTEMPTS: Final = 4 +_PRISMA_ATTEMPTS: Final = 4 -_TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType( +_TRANSIENT_PRISMA_FAILURES: Final = MappingProxyType( { "deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)", "P1001": "an unreachable database server", @@ -62,24 +62,6 @@ _TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType( ) -def _transient_deploy_failure(stderr: str) -> str | None: - """Describe why a failed `prisma migrate deploy` is worth retrying, or None. - - These are environment failures, not migration failures: the database is not - up yet, or another instance holds the migration lock. v1 retried every - failed deploy and absorbed them; failing fast on them instead would turn a - database that is ten seconds late into a dead proxy. - """ - return next( - ( - reason - for marker, reason in _TRANSIENT_DEPLOY_FAILURES.items() - if marker in stderr - ), - None, - ) - - PARTITIONED_SPEND_LOGS_PUSH_ERROR = ( "LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), " "so its primary key must include the partition key (\"startTime\"). `prisma db push` " @@ -304,6 +286,23 @@ class ProxyExtrasDBManager: env=prisma_env, ) + @staticmethod + def _transient_prisma_failure(stderr: str) -> str | None: + """Why a failed prisma command is worth retrying, or None. + + v1 retried every failure, so it absorbed a database that was not up yet + or another instance holding the migration lock. v2 fails fast, which is + right for a broken migration and wrong for these. + """ + return next( + ( + reason + for marker, reason in _TRANSIENT_PRISMA_FAILURES.items() + if marker in stderr + ), + None, + ) + @staticmethod def _is_permission_error(error_message: str) -> bool: """ @@ -699,20 +698,43 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - subprocess.run( - [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=prisma_command_timeout(), - check=True, - env=_get_prisma_env(), + for attempt in range(_PRISMA_ATTEMPTS): + try: + subprocess.run( + [_get_prisma_command(), "db", "push", "--accept-data-loss"], + timeout=prisma_command_timeout(), + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + return True + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as e: + stderr = e.stderr or "" + transient = ProxyExtrasDBManager._transient_prisma_failure( + stderr + ) + # Re-raise as RuntimeError so proxy_cli.py's + # `except RuntimeError` catches it and exits cleanly. + if transient is None or attempt == _PRISMA_ATTEMPTS - 1: + raise RuntimeError( + f"prisma db push failed.\n\nDetail: {e}" + f"\n\nPrisma error:\n{stderr}" + ) from e + logger.info( + "prisma db push attempt %s failed on %s, retrying. " + "Prisma error:\n%s", + attempt + 1, + transient, + stderr, + ) + time.sleep(random.randrange(5, 15)) + raise RuntimeError( + f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." ) - return True - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as e: - # Re-raise as RuntimeError so proxy_cli.py's - # `except RuntimeError` catches it and exits cleanly. - raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e finally: os.chdir(original_dir) @@ -722,7 +744,7 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - for attempt in range(_MIGRATE_DEPLOY_ATTEMPTS): + for attempt in range(_PRISMA_ATTEMPTS): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -837,17 +859,17 @@ class ProxyExtrasDBManager: f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - transient = _transient_deploy_failure(stderr) + transient = ProxyExtrasDBManager._transient_prisma_failure(stderr) if transient is None: raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - if attempt == _MIGRATE_DEPLOY_ATTEMPTS - 1: + if attempt == _PRISMA_ATTEMPTS - 1: raise RuntimeError( f"Database migration failed after " - f"{_MIGRATE_DEPLOY_ATTEMPTS} attempts on {transient}. " + f"{_PRISMA_ATTEMPTS} attempts on {transient}. " "Check database connectivity and load." f"\n\nPrisma error:\n{stderr}" ) from e @@ -863,7 +885,7 @@ class ProxyExtrasDBManager: continue raise RuntimeError( - f"Database migration failed after {_MIGRATE_DEPLOY_ATTEMPTS} " + f"Database migration failed after {_PRISMA_ATTEMPTS} " "attempts (retry loop exhausted by timeouts or repeated " "idempotent-recovery continues). Check database connectivity, " "load, and _prisma_migrations ledger state." diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index f637e139e86..deb78dcdb30 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -1,10 +1,7 @@ """Regression tests for ProxyExtrasDBManager's v2 migration resolver. -v2 is what the proxy CLI selects by default; v1 stays reachable via -`--use_legacy_migration_resolver` or `USE_V2_MIGRATION_RESOLVER=false`. At the -library level the resolver is picked with the `use_v2_resolver` kwarg, which -still defaults to False so `migrations/run.py` and any direct caller keep their -own explicit choice. +v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver` +kwarg, which still defaults to False for direct callers. """ import os @@ -271,9 +268,7 @@ class _DeployApplied: def _deploy_only(deploy_side_effect): """subprocess.run stand-in that only intercepts `prisma migrate deploy`. - Everything else the resolver shells out to, the Prisma toolchain check - above all, succeeds untouched, so a mock meant for the deploy call cannot - be silently consumed by an earlier subprocess call. + Scoped by argv so the Prisma toolchain check cannot consume the mock first. """ deploys = {"n": 0} @@ -298,13 +293,8 @@ def _prepare_v2_resolver(monkeypatch, tmp_path): def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path): - """A deadlock on Prisma's migration advisory lock is transient and retried. - - Several proxy replicas booting against one database race `migrate deploy`, - and Postgres aborts one side. v1 retried any failed deploy, so it rode this - out; v2 classifies unrecognised stderr as unrecoverable and raises, which - with v2 as the default would take a replica's whole boot down. - """ + """v2: replicas racing `migrate deploy` deadlock on Prisma's advisory + lock, which is transient and must be retried rather than kill the boot.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -323,8 +313,8 @@ def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path): def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp_path): - """The deadlock retry stays bounded: a deadlock that never clears still - raises rather than looping forever or reporting a successful migration.""" + """v2: the deadlock retry is bounded, so a deadlock that never clears + still raises instead of looping or reporting success.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -348,13 +338,7 @@ def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp ], ) def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path, stderr): - """A database that is not accepting connections yet is retried, not fatal. - - A proxy and its database starting together race routinely, and v1 rode that - out by retrying every failed deploy. v2 treats unrecognised stderr as - unrecoverable, so without this the default flip would turn a database that - is a few seconds late into a dead proxy instead of a slow boot. - """ + """v2: a database not accepting connections yet is retried, not fatal.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -373,9 +357,8 @@ def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_path): - """Retrying connectivity errors must not turn a genuinely unreachable - database into a silent success: after the attempts are spent it still - raises, so the proxy exits instead of serving without its database.""" + """v2: a genuinely unreachable database still raises once the attempts + are spent, rather than passing as a successful migration.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -395,13 +378,8 @@ def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, caplog): - """Retrying must not swallow why the database was unreachable. - - Prisma's stderr is captured, so if the retry path neither logs it nor puts - it in the final error, an operator (and CI's bad-DATABASE_URL job, which - greps the boot log for the P1001 line) sees four silent retries and no - cause. - """ + """v2: retrying must not swallow Prisma's stderr, which is captured and is + the only place the cause appears for an operator or a boot-log grep.""" _prepare_v2_resolver(monkeypatch, tmp_path) stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`" @@ -422,21 +400,47 @@ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, cap assert "P1001" in caplog.text -def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path): - """The transient classification must stay narrow: a genuinely broken - migration still fails fast on the first attempt rather than being retried - into the same error four times.""" +def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path): + """v2: `prisma db push` retries a transient failure like v1 did, so the + default flip does not cost --use_prisma_db_push its retries.""" _prepare_v2_resolver(monkeypatch, tmp_path) - stderr = ( - "Error: P3009\n" - "The `20260101000000_genuinely_broken` migration failed to apply.\n" - 'Reason: syntax error at or near "BRKN" LINE 42' + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + if pushes["n"] == 1: + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr="Error: P1001: Can't reach database server at `db`:`5432`", + output="", + ) + return _DeployApplied() + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False ) + with patch("subprocess.run", side_effect=_run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert ok is True + assert pushes["n"] == 2 + + +def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): + """v2: an unrecognised deploy failure still raises on the first attempt.""" + _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): raise subprocess.CalledProcessError( - returncode=1, cmd=cmd, stderr=stderr, output="" + returncode=1, + cmd=cmd, + stderr="Error: relation \"LiteLLM_SpendLogs\" does not exist", + output="", ) run, deploys = _deploy_only(_side_effect) @@ -447,26 +451,3 @@ def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path) assert deploys["n"] == 1 -def test_v1_still_runs_the_diff_and_force_recovery(monkeypatch, tmp_path): - """v1 remains the pre-existing diff-and-force resolver, unchanged by the - default flip: it still calls _resolve_all_migrations after a deploy that - applied something. Operators opting back in must get exactly the old path. - """ - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") - - class FakeResult: - stdout = "Applied migration.\n" - stderr = "" - - resolve_called = {"n": 0} - - def fake_resolve(*args, **kwargs): - resolve_called["n"] += 1 - - monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) - monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) - - ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=False) - assert ok is True - assert resolve_called["n"] == 1 diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index ac1defe8d79..dac36965a6c 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2008,15 +2008,9 @@ class TestRunServerDbSetup: env_value, expected_v2, ): - """The proxy defaults to the v2 resolver, and v1 stays reachable. - - Both opt-out routes matter: --use_legacy_migration_resolver for a CLI - boot, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, - where litellm/proxy/prisma_migration.py calls run_server with a fixed - argv and an env var is the only way in. The deprecated - --use_v2_migration_resolver must still parse so existing commands do - not die on an unknown option, and an explicit flag still beats the env. - """ + """The proxy defaults to v2, and both v1 opt-out routes work: the + flag for a CLI boot, USE_V2_MIGRATION_RESOLVER=false for deploys that + cannot pass one. An explicit flag beats the env var.""" from litellm.proxy.proxy_cli import run_server mock_subprocess_run.return_value = MagicMock(returncode=0) From fbd1339993252ca3c134494ef431a22ce591c83f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:38:27 -0700 Subject: [PATCH 013/126] fix(tests): satisfy the tests-tree ruff config and correct a stale comment Moving the resolver tests under tests/ brings them under ruff-tests.toml, which the package-internal directory they came from was never linted by, so a pre-existing pytest.raises pattern now needs to be a raw string (RUF043). Also corrects the comment on proxy_cli's RuntimeError handler: both resolvers raise on permission failures, not just v2. --- litellm/proxy/proxy_cli.py | 8 ++++---- .../litellm-proxy-extras/test_setup_database_fail_fast.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c1d86344a1e..86f6853a625 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1323,10 +1323,10 @@ def run_server( use_v2_resolver=use_v2_migration_resolver, ) except RuntimeError as e: - # Raised on unrecoverable migration errors: the v2 - # resolver's non-idempotent failures and permission - # issues, and any `prisma db push` against a - # partitioned LiteLLM_SpendLogs. + # Raised on unrecoverable migration errors: permission + # failures from either resolver, the v2 resolver's + # non-idempotent failures, and any `prisma db push` + # against a partitioned LiteLLM_SpendLogs. print( f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m", file=sys.stderr, diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index deb78dcdb30..0b59c5f9430 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -219,7 +219,7 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error( ) with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises( - RuntimeError, match="Failed to mark migration .* as applied" + RuntimeError, match=r"Failed to mark migration .* as applied" ): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) From c3fc86869d45e2b5057dce53ec7bec0617f84149 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:52:15 -0700 Subject: [PATCH 014/126] test: keep the resolver tests inside the test-quality ceilings The moved fail-fast test carried a sys.path.insert that the uv workspace makes unnecessary, and one pre-existing case asserted nothing beyond "did not raise", so it could not tell a swallowed error from a skipped query. Give it a liveness gate on the connect count instead. Fold the resolver default/opt-out matrix into the existing db-push flag test rather than standing up another patched test, so the flag pair, the env var, and their precedence are covered without new mock scaffolding. --- .../test_setup_database_fail_fast.py | 16 +-- tests/test_litellm/proxy/test_proxy_cli.py | 105 ++++++------------ 2 files changed, 38 insertions(+), 83 deletions(-) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 0b59c5f9430..964355492a1 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -4,20 +4,11 @@ v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver` kwarg, which still defaults to False for direct callers. """ -import os import subprocess -import sys from unittest.mock import patch import pytest -sys.path.insert( - 0, - os.path.abspath( - os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras") - ), -) - from litellm_proxy_extras.utils import ( ProxyExtrasDBManager, _max_migration_timestamp, @@ -175,13 +166,16 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): # Simulate an InsufficientPrivilege (subclass of DatabaseError). raise psycopg.errors.InsufficientPrivilege("permission denied") + connects = {"n": 0} + def _fake_connect(*a, **kw): + connects["n"] += 1 return _FakeConn() monkeypatch.setattr("psycopg.connect", _fake_connect) - # Must not raise. - ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) + assert ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) is None + assert connects["n"] == 1, "the failing query must actually have been reached" def test_v2_resolve_specific_migration_failure_raises_runtime_error( diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index dac36965a6c..4f205ada609 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1805,6 +1805,39 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=True ) + # Test 3+: the resolver default and both routes back to v1. The flag + # covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that + # cannot pass one, where prisma_migration.py fixes the argv. An + # explicit flag beats the env var. + for argv, env_value, expected_v2 in ( + ([], None, True), + (["--use_v2_migration_resolver"], None, True), + (["--use_legacy_migration_resolver"], None, False), + ([], "false", False), + ([], "true", True), + (["--use_v2_migration_resolver"], "false", True), + (["--use_legacy_migration_resolver"], "true", False), + ): + mock_setup_database.reset_mock() + mock_should_update_schema.reset_mock() + mock_should_update_schema.return_value = True + + resolver_env = ( + {"USE_V2_MIGRATION_RESOLVER": env_value} + if env_value is not None + else {} + ) + os.environ.pop("USE_V2_MIGRATION_RESOLVER", None) + with patch.dict(os.environ, resolver_env): + run_server.main( + ["--local", "--skip_server_startup", *argv], + standalone_mode=False, + ) + assert mock_setup_database.call_args.kwargs == { + "use_migrate": True, + "use_v2_resolver": expected_v2, + }, f"argv={argv} env={env_value}" + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") @@ -1981,78 +2014,6 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=True ) - @pytest.mark.parametrize( - "argv, env_value, expected_v2", - [ - ([], None, True), - (["--use_v2_migration_resolver"], None, True), - (["--use_legacy_migration_resolver"], None, False), - ([], "false", False), - ([], "true", True), - (["--use_v2_migration_resolver"], "false", True), - ], - ) - @patch("subprocess.run") - @patch("atexit.register") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") - @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") - def test_migration_resolver_default_and_opt_out( - self, - mock_should_update_schema, - mock_check_schema_diff, - mock_setup_database, - mock_atexit_register, - mock_subprocess_run, - argv, - env_value, - expected_v2, - ): - """The proxy defaults to v2, and both v1 opt-out routes work: the - flag for a CLI boot, USE_V2_MIGRATION_RESOLVER=false for deploys that - cannot pass one. An explicit flag beats the env var.""" - from litellm.proxy.proxy_cli import run_server - - mock_subprocess_run.return_value = MagicMock(returncode=0) - mock_should_update_schema.return_value = True - mock_setup_database.return_value = True - - mock_proxy_module = MagicMock( - app=MagicMock(), - ProxyConfig=MagicMock(), - KeyManagementSettings=MagicMock(), - save_worker_config=MagicMock(), - ) - - clean_env = { - k: v - for k, v in os.environ.items() - if k - not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") - } - clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" - if env_value is not None: - clean_env["USE_V2_MIGRATION_RESOLVER"] = env_value - - with ( - patch.dict(os.environ, clean_env, clear=True), - patch.dict( - "sys.modules", - { - "proxy_server": mock_proxy_module, - "litellm.proxy.proxy_server": mock_proxy_module, - }, - ), - ): - run_server.main( - ["--local", "--skip_server_startup", *argv], standalone_mode=False - ) - - mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=expected_v2 - ) - - # --- Module-level helpers for worker startup hook tests --- _dummy_hook_called = False From 6b3e30e6601ac4fb285e97525061ccdd93aaac16 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:56:15 -0700 Subject: [PATCH 015/126] fix(proxy-extras): bound the db push retries in the reachable branch The retry loop already raises on the final attempt, so the raise that followed the loop could never run. Drop it and cover the exhaustion path with a test that pins the attempt count and keeps the prisma error in the message, which is the only thing that tells an operator why the boot stopped. --- .../litellm_proxy_extras/utils.py | 3 -- .../test_setup_database_fail_fast.py | 35 +++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index a31016e0f17..f751c4087eb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -732,9 +732,6 @@ class ProxyExtrasDBManager: stderr, ) time.sleep(random.randrange(5, 15)) - raise RuntimeError( - f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." - ) finally: os.chdir(original_dir) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 964355492a1..ed82037590a 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -10,6 +10,7 @@ from unittest.mock import patch import pytest from litellm_proxy_extras.utils import ( + _PRISMA_ATTEMPTS, ProxyExtrasDBManager, _max_migration_timestamp, _migration_timestamp, @@ -425,6 +426,40 @@ def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path): assert pushes["n"] == 2 +def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error( + monkeypatch, tmp_path +): + """v2: a database that never comes back stops after _PRISMA_ATTEMPTS and + surfaces the prisma error, rather than retrying the boot forever.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr="Error: P1001: Can't reach database server at `db`:`5432`", + output="", + ) + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + with patch("subprocess.run", side_effect=_run): + with pytest.raises(RuntimeError) as exc: + ProxyExtrasDBManager.setup_database( + use_migrate=False, use_v2_resolver=True + ) + + assert pushes["n"] == _PRISMA_ATTEMPTS + assert "P1001" in str(exc.value) + + def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): """v2: an unrecognised deploy failure still raises on the first attempt.""" _prepare_v2_resolver(monkeypatch, tmp_path) From 2495673d01ee26c73d54adb6b78e30ff87cf0bc2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 02:14:57 -0700 Subject: [PATCH 016/126] fix(proxy-extras): stop a db push timeout crashing the migration job subprocess.run leaves stderr as bytes on TimeoutExpired even under text=True, unlike CalledProcessError. Classifying both in one handler meant a real `prisma db push` timeout died on a TypeError, which proxy_cli.py's `except RuntimeError` does not catch, so the migrations Job container ended on an unhandled traceback instead of a clean exit. Give the timeout its own handler and retry it, matching what the migrate deploy loop beside it already does. That puts a fallthrough back into the loop, so the trailing raise removed in the previous commit is reachable again and comes back with it. Also drop a comment restating why the resolver cases exist and widen the db push test's docstring, which had stopped describing what it covers. --- .../litellm_proxy_extras/utils.py | 14 ++- .../test_setup_database_fail_fast.py | 86 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 7 +- 3 files changed, 98 insertions(+), 9 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f751c4087eb..22e30dfa897 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -709,10 +709,13 @@ class ProxyExtrasDBManager: env=_get_prisma_env(), ) return True - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as e: + except subprocess.TimeoutExpired: + logger.info( + "prisma db push attempt %s timed out, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + except subprocess.CalledProcessError as e: stderr = e.stderr or "" transient = ProxyExtrasDBManager._transient_prisma_failure( stderr @@ -732,6 +735,9 @@ class ProxyExtrasDBManager: stderr, ) time.sleep(random.randrange(5, 15)) + raise RuntimeError( + f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." + ) finally: os.chdir(original_dir) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index ed82037590a..3996b91c2a6 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -460,6 +460,92 @@ def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error( assert "P1001" in str(exc.value) +def _db_push_only(push_side_effect): + """subprocess.run stand-in that only intercepts `prisma db push`.""" + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + return push_side_effect(pushes["n"], cmd) + + return _run, pushes + + +def _timed_out_for_real(): + """Capture what subprocess.run really puts on a TimeoutExpired. + + Under text=True it still leaves stderr as bytes, unlike CalledProcessError, + so hardcoding a str here would test a shape production never sees. Derived + at import, before any test patches subprocess.run. + """ + try: + subprocess.run( + ["sh", "-c", "echo 'Error: P1001 unreachable' >&2; sleep 5"], + timeout=0.2, + check=True, + capture_output=True, + text=True, + ) + except subprocess.TimeoutExpired as e: + return e + raise AssertionError("the helper command was supposed to time out") + + +_TIMEOUT_TEMPLATE = _timed_out_for_real() + + +def _real_timeout_expired(cmd): + return subprocess.TimeoutExpired( + cmd=cmd, + timeout=_TIMEOUT_TEMPLATE.timeout, + output=_TIMEOUT_TEMPLATE.stdout, + stderr=_TIMEOUT_TEMPLATE.stderr, + ) + + +def test_v2_db_push_retries_a_timeout(monkeypatch, tmp_path): + """v2: a `prisma db push` that times out is retried, not turned into a + TypeError by classifying its bytes stderr as if it were text.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise _real_timeout_expired(cmd) + return _DeployApplied() + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + run, pushes = _db_push_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert ok is True + assert pushes["n"] == 2 + + +def test_v2_db_push_timeouts_are_bounded(monkeypatch, tmp_path): + """v2: a `prisma db push` that never stops timing out gives up as a + RuntimeError, which is the only exception proxy_cli.py exits cleanly on.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise _real_timeout_expired(cmd) + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + run, pushes = _db_push_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match=r"prisma db push failed after \d+"): + ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert pushes["n"] == _PRISMA_ATTEMPTS + + def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): """v2: an unrecognised deploy failure still raises on the first attempt.""" _prepare_v2_resolver(monkeypatch, tmp_path) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4f205ada609..5e2dd358d75 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1737,7 +1737,8 @@ class TestRunServerDbSetup: mock_atexit_register, mock_subprocess_run, ): - """Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter""" + """Which resolver and which migration mode run_server hands setup_database, + across the db push flag, the v2/legacy flag pair and USE_V2_MIGRATION_RESOLVER.""" from litellm.proxy.proxy_cli import run_server # Mock subprocess.run to simulate prisma being available @@ -1805,10 +1806,6 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=True ) - # Test 3+: the resolver default and both routes back to v1. The flag - # covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that - # cannot pass one, where prisma_migration.py fixes the argv. An - # explicit flag beats the env var. for argv, env_value, expected_v2 in ( ([], None, True), (["--use_v2_migration_resolver"], None, True), From 5dfe32c889ca778b682bcab6e7772a4ebb14d47f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 02:16:20 -0700 Subject: [PATCH 017/126] docs(tests): name the entrypoint that actually reaches the db push branch The proxy CLI's --use_prisma_db_push never gets here; PrismaManager keeps its own db push loop and only delegates when use_migrate is true. The caller this covers is the migrations Job with USE_PRISMA_DB_PUSH=true. --- .../litellm-proxy-extras/test_setup_database_fail_fast.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 3996b91c2a6..ef447315a8c 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -396,8 +396,11 @@ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, cap def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path): - """v2: `prisma db push` retries a transient failure like v1 did, so the - default flip does not cost --use_prisma_db_push its retries.""" + """v2: `prisma db push` retries a transient failure like v1 did. + + Reached from the migrations Job (USE_PRISMA_DB_PUSH=true), not from the + proxy CLI, whose --use_prisma_db_push has its own loop in prisma_client. + """ _prepare_v2_resolver(monkeypatch, tmp_path) pushes = {"n": 0} From e22744c4393858357d34b2abfffb0249528880ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:06:29 -0700 Subject: [PATCH 018/126] fix(embeddings): omit encoding_format when the client omits it on OpenAI-compatible calls When no encoding_format is set on the call, the model config, or LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT, leave the field out of the upstream request instead of defaulting to float, and bypass the OpenAI SDK's own base64 default so nothing re-adds it on the wire. Downstreams that reject encoding_format, such as a second LiteLLM proxy fronting Bedrock Titan embeddings, now work when the client omits the field. Fixes #38661 --- litellm/llms/hosted_vllm/embedding/README.md | 5 +- litellm/llms/openai/openai.py | 72 ++++++++++++-------- litellm/main.py | 19 +++--- tests/test_litellm/test_main.py | 61 +++++++++++++++++ 4 files changed, 115 insertions(+), 42 deletions(-) diff --git a/litellm/llms/hosted_vllm/embedding/README.md b/litellm/llms/hosted_vllm/embedding/README.md index 2c58e16fc23..32b7ea5c560 100644 --- a/litellm/llms/hosted_vllm/embedding/README.md +++ b/litellm/llms/hosted_vllm/embedding/README.md @@ -9,8 +9,7 @@ For OpenAI-compatible embedding calls (including `openai/...` with a custom `api 1. Explicit value on the embedding call (`encoding_format=...`). 2. Model config (`litellm_params.encoding_format` on the proxy `model_list` entry). 3. Environment variable `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT` (e.g. in `.env` or container env). -4. Default **`float`**. -That avoids forwarding `encoding_format=None` to the provider/SDK where some servers behave poorly. +If none of those is set, or the winning value is the literal string `none`, the field is omitted from the upstream request entirely (LiteLLM also bypasses the OpenAI SDK's own base64 default), so OpenAI-compatible servers that reject `encoding_format` keep working. -To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). \ No newline at end of file +To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6e66c998acf..7b71e532b61 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -12,9 +12,14 @@ if TYPE_CHECKING: import openai from openai import AsyncOpenAI, OpenAI +from openai._base_client import make_request_options +from openai._constants import RAW_RESPONSE_HEADER +from openai._legacy_response import LegacyAPIResponse +from openai._types import RequestOptions +from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import overload import litellm @@ -322,6 +327,24 @@ class OpenAIChatCompletionResponseIterator(BaseModelResponseIterator): raise e +_EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None) +_EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _embedding_request_without_sdk_defaults( + data: Mapping[str, object], timeout: float | httpx.Timeout +) -> tuple[dict[str, object], RequestOptions]: + body: Final = {k: v for k, v in data.items() if k not in ("extra_headers", "extra_query", "extra_body")} + extra_headers: Final = _EXTRA_HEADERS_ADAPTER.validate_python(data.get("extra_headers")) + options: Final = make_request_options( + extra_headers={**(extra_headers or {}), RAW_RESPONSE_HEADER: "true"}, + extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")), + extra_body=data.get("extra_body"), + timeout=timeout, + ) + return body, options + + class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): def __init__(self) -> None: super().__init__() @@ -1148,19 +1171,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - ): - """ - Helper to: - - call embeddings.create.with_raw_response when litellm.return_response_headers is True - - call embeddings.create by default - """ - try: - raw_response = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + ) -> tuple[dict[str, str], CreateEmbeddingResponse]: + if "encoding_format" not in data: + body, options = _embedding_request_without_sdk_defaults(data, timeout) + bypass_response: Final = await openai_aclient.post( + "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse + ) + assert isinstance(bypass_response, LegacyAPIResponse) + return dict(bypass_response.headers), bypass_response.parse(to=CreateEmbeddingResponse) + raw_response: Final = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) + return dict(raw_response.headers), raw_response.parse() @track_llm_api_timing() def make_sync_openai_embedding_request( @@ -1169,20 +1189,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - ): - """ - Helper to: - - call embeddings.create.with_raw_response when litellm.return_response_headers is True - - call embeddings.create by default - """ - try: - raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) - - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + ) -> tuple[dict[str, str], CreateEmbeddingResponse]: + if "encoding_format" not in data: + body, options = _embedding_request_without_sdk_defaults(data, timeout) + bypass_response: Final = openai_client.post( + "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse + ) + assert isinstance(bypass_response, LegacyAPIResponse) + return dict(bypass_response.headers), bypass_response.parse(to=CreateEmbeddingResponse) + raw_response: Final = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) + return dict(raw_response.headers), raw_response.parse() async def aembedding( self, diff --git a/litellm/main.py b/litellm/main.py index cafa1e4718f..f40b0e5c227 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6289,18 +6289,15 @@ def embedding( if headers is not None and headers != {}: optional_params["extra_headers"] = headers - if encoding_format is not None: - optional_params["encoding_format"] = encoding_format + requested_encoding_format: Final = ( + encoding_format + or optional_params.get("encoding_format") + or get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") + ) + if requested_encoding_format is None or requested_encoding_format.strip().lower() == "none": + optional_params.pop("encoding_format", None) else: - env_fmt: Final = get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") - if env_fmt is not None and env_fmt.strip().lower() == "none": - optional_params.pop("encoding_format", None) - else: - _default_fmt: Final = optional_params.get("encoding_format") or env_fmt or "float" - if _default_fmt.strip().lower() == "none": - optional_params.pop("encoding_format", None) - else: - optional_params["encoding_format"] = _default_fmt + optional_params["encoding_format"] = requested_encoding_format api_version = None diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..4c591be0f61 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,3 +3181,64 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None + + +def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +def test_embedding_openai_env_var_sets_default_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "float" + + +@pytest.mark.asyncio +async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] From c254605e924455062e72305fdfaa2bd9c9feb089 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:10:21 -0700 Subject: [PATCH 019/126] test(embeddings): move encoding_format default coverage to wire-level assertions Consolidate the new regression tests into test_openai_embedding_encoding_format_default.py, replacing mocks that pinned the old float default with respx captures of the request body, and update the stale local_testing default-float test to assert omission --- tests/local_testing/test_embedding.py | 64 +++--- tests/test_litellm/test_main.py | 60 ------ ...penai_embedding_encoding_format_default.py | 188 ++++++++---------- 3 files changed, 109 insertions(+), 203 deletions(-) diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index aed2849f056..ee2ac14f498 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -3,6 +3,8 @@ import os import re import traceback +import httpx + import openai import pytest from dotenv import load_dotenv @@ -1255,56 +1257,42 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input): assert sent_data["input"] == expected_payload_input -def test_encoding_format_defaults_to_float_for_openai_sdk(monkeypatch): +def test_encoding_format_omitted_by_default_for_openai_sdk(monkeypatch): """ - When encoding_format is not provided, LiteLLM sends `float` for OpenAI-path embeddings. + When encoding_format is not provided, LiteLLM leaves it out of the upstream request. Optional global override: `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT`. """ monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - # Create a mock client instance - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance + captured_bodies = [] - # Mock the embeddings.with_raw_response.create method - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-ada-002", "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} - - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response + }, ) - # Call the embedding function without encoding_format - response = embedding( - model="text-embedding-ada-002", - input="Hello world", - ) + client = openai.OpenAI( + api_key="sk-test", http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) - # Get the call arguments to verify what was sent to OpenAI SDK - call_args = mock_client_instance.embeddings.with_raw_response.create.call_args - assert ( - call_args is not None - ), "OpenAI SDK embeddings.create should have been called" + response = embedding( + model="text-embedding-ada-002", + input="Hello world", + api_key="sk-test", + client=client, + ) - call_kwargs = call_args[1] # Get kwargs - - assert "encoding_format" in call_kwargs - assert ( - call_kwargs["encoding_format"] == "float" - ), "encoding_format should default to float when not provided by user" - - print("✅ PASS: encoding_format='float' is correctly passed to OpenAI SDK") + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert "encoding_format" not in captured_bodies[0], ( + "encoding_format should be omitted from the upstream request when not provided by user" + ) def test_encoding_format_explicit_value_preserved(): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4c591be0f61..84225c3feb1 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3182,63 +3182,3 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None - -def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: - return respx_mock.post("https://api.openai.com/v1/embeddings").mock( - return_value=httpx.Response( - 200, - json={ - "object": "list", - "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], - "model": "text-embedding-3-small", - "usage": {"prompt_tokens": 2, "total_tokens": 2}, - }, - ) - ) - - -def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert "encoding_format" not in request_body - assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] - - -def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - litellm.embedding( - model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" - ) - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert request_body["encoding_format"] == "base64" - - -def test_embedding_openai_env_var_sets_default_encoding_format( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert request_body["encoding_format"] == "float" - - -@pytest.mark.asyncio -async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert "encoding_format" not in request_body - assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/test_litellm/test_openai_embedding_encoding_format_default.py index 94e4e3c81e5..9842bf30585 100644 --- a/tests/test_litellm/test_openai_embedding_encoding_format_default.py +++ b/tests/test_litellm/test_openai_embedding_encoding_format_default.py @@ -1,124 +1,102 @@ -from unittest.mock import MagicMock, patch +import json +from typing import Final +import httpx import pytest +import respx -from litellm import embedding +import litellm -@pytest.mark.parametrize( - "set_env, env_value, expected", - [ - (False, None, "float"), - (True, "base64", "base64"), - ], -) -def test_openai_embedding_encoding_format_default( - monkeypatch, set_env, env_value, expected -): - monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - if set_env: - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) - - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } +def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) ) - mock_response.headers = {} - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - embedding( - model="text-embedding-ada-002", - input="Hello world", - ) +@pytest.fixture(autouse=True) +def clear_default_encoding_format_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert call_kwargs["encoding_format"] == expected + +def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +def test_embedding_openai_explicit_encoding_format_wins_over_env_var( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +@pytest.mark.parametrize("env_value", ["float", "base64"]) +def test_embedding_openai_env_var_sets_default_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, env_value: str +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == env_value @pytest.mark.parametrize("env_none", ["none", "NONE", " none "]) -def test_openai_embedding_encoding_format_env_none_omits_param( - monkeypatch, env_none -): - """LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT=none omits encoding_format (provider default).""" +def test_embedding_openai_env_none_omits_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, env_none: str +) -> None: monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_none) + mock_route: Final = _mock_openai_embedding_route(respx_mock) - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - - embedding( - model="text-embedding-ada-002", - input="Hello world", - ) - - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert "encoding_format" not in call_kwargs + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body -def test_openai_embedding_encoding_format_explicit_overrides_env(monkeypatch): - """Request `encoding_format` wins over LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT.""" - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") +@pytest.mark.asyncio +async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + mock_route: Final = _mock_openai_embedding_route(respx_mock) - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} + response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - - embedding( - model="text-embedding-ada-002", - input="Hello world", - encoding_format="base64", - ) - - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert call_kwargs["encoding_format"] == "base64" + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] From b65592e623805749827d03b1fcacf03173a5d16d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:18:19 -0700 Subject: [PATCH 020/126] test: drop stray trailing blank line in test_main.py --- tests/test_litellm/test_main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 84225c3feb1..8cf878d05d9 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,4 +3181,3 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None - From 14484d67fde20fd0ca31409c5143422655e0e1cd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 19:00:43 +0000 Subject: [PATCH 021/126] refactor(types): replace Any with real types across 54 more backend files Second pass over the highest-Any-density modules that the first pass left untouched: guardrail hooks, the gemini and anthropic transformation layers, the proxy spend-tracking and pass-through endpoints, and the caching clients. Untyped `response.json()` bodies and `dict[str, Any]` request payloads are described once at their boundary with a TypedDict or Protocol, so the fields read downstream resolve to real types instead of Any. No cast, no type: ignore, no noqa, and no new Any annotations. --- .../proxy/common_utils/check_batch_cost.py | 26 ++- .../proxy/hooks/managed_files.py | 23 +- litellm/caching/redis_cache.py | 57 +++-- litellm/caching/valkey_semantic_cache.py | 40 +++- .../transformation.py | 7 +- litellm/cost_calculator.py | 12 +- .../google_genai/adapters/transformation.py | 16 +- .../arize/arize_phoenix_prompt_manager.py | 66 ++++-- litellm/integrations/custom_logger.py | 75 ++++--- .../gitlab/gitlab_prompt_manager.py | 59 +++-- litellm/integrations/posthog.py | 62 +++-- .../llm_cost_calc/tool_call_cost_tracking.py | 44 ++-- .../prompt_templates/common_utils.py | 20 +- litellm/litellm_core_utils/token_counter.py | 43 ++-- .../chat/guardrail_translation/handler.py | 24 +- litellm/llms/anthropic/chat/handler.py | 34 +-- litellm/llms/anthropic/common_utils.py | 12 +- .../adapters/transformation.py | 66 ++++-- .../context_management/editors/compact.py | 65 ++++-- .../responses_adapters/transformation.py | 42 ++-- .../base_managed_resource.py | 59 +++-- litellm/llms/gemini/common_utils.py | 94 ++++---- litellm/llms/gemini/files/transformation.py | 38 +++- .../llms/gemini/realtime/transformation.py | 61 +++-- .../litellm_proxy/skills/code_execution.py | 130 +++++++++-- .../audio_transcription/audio_utils.py | 30 ++- litellm/llms/oci/common_utils.py | 41 ++-- .../llms/openai/containers/transformation.py | 32 +-- .../runwayml/text_to_speech/transformation.py | 31 ++- litellm/llms/vertex_ai/common_utils.py | 33 +-- .../mcp_server/mcp_server_manager.py | 14 +- .../mcp_server/oauth2_flow_backfill.py | 50 ++++- litellm/proxy/auth/auth_utils.py | 16 +- litellm/proxy/common_utils/callback_utils.py | 24 +- .../proxy/common_utils/custom_openapi_spec.py | 211 ++++++++++-------- .../proxy/common_utils/user_api_key_cache.py | 55 ++--- .../guardrail_hooks/bedrock_guardrails.py | 26 +-- .../cisco_ai_defense/cisco_ai_defense_mcp.py | 53 ++--- .../mcp_jwt_signer/mcp_jwt_signer.py | 33 ++- .../guardrail_hooks/noma/noma_v2.py | 21 +- .../guardrails/guardrail_hooks/presidio.py | 32 ++- .../guardrail_hooks/xecguard/xecguard.py | 4 +- .../proxy/hooks/mcp_semantic_filter/hook.py | 52 +++-- litellm/proxy/litellm_pre_call_utils.py | 84 +++---- .../key_management_endpoints.py | 6 +- .../policy_endpoints/endpoints.py | 12 +- .../llm_passthrough_endpoints.py | 63 +++--- .../pass_through_endpoints.py | 18 +- litellm/proxy/rag_endpoints/endpoints.py | 56 +++-- .../response_polling/background_streaming.py | 48 +++- .../spend_tracking/budget_reservation.py | 47 ++-- .../spend_tracking/spend_tracking_utils.py | 63 ++++-- .../proxy_setting_endpoints.py | 180 +++++++++------ litellm/responses/streaming_iterator.py | 36 +-- litellm/types/guardrail_base_init.py | 24 ++ 55 files changed, 1630 insertions(+), 940 deletions(-) create mode 100644 litellm/types/guardrail_base_init.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index aee3295d1da..0a75768f709 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -4,7 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple +from typing import TYPE_CHECKING, Final, List, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -86,7 +86,7 @@ class CheckBatchCost: return self.batch_processed_support_confirmed = True - async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: + async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]: """ Look up user email and key alias by user_id for enriching the S3 callback metadata. Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). @@ -96,8 +96,10 @@ class CheckBatchCost: if not user_id: return {} try: - user_row = await self.prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} + user_row: prisma_models.LiteLLM_UserTable | None = ( + await self.prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) ) if user_row is None: return {} @@ -114,8 +116,10 @@ class CheckBatchCost: if not api_key: return None try: - key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": api_key} + key_row: prisma_models.LiteLLM_VerificationToken | None = ( + await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) ) return getattr(key_row, "key_alias", None) if key_row is not None else None except Exception as e: @@ -127,8 +131,10 @@ class CheckBatchCost: if not team_id: return None try: - team_row = await self.prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row: prisma_models.LiteLLM_TeamTable | None = ( + await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) return getattr(team_row, "team_alias", None) if team_row is not None else None except Exception as e: @@ -137,7 +143,7 @@ class CheckBatchCost: async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str - ) -> Dict[str, Any]: + ) -> dict[str, object]: """ Rebuild the spend-tracking metadata for the key, team, and tags that created the batch so the batch-cost spend log is attributed the same way a non-batch request @@ -151,7 +157,7 @@ class CheckBatchCost: team_id = getattr(job, "team_id", None) request_tags = getattr(job, "request_tags", None) - metadata: Dict[str, Any] = { + metadata: dict[str, object] = { "user_api_key_user_id": job.created_by, "user_api_key": api_key, "user_api_key_team_id": team_id, diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 39f8de0b0cc..00528c9ade9 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -181,6 +181,10 @@ class _ManagedObjectTableActions(Protocol): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... +class _SchedulerWithJobLookup(Protocol): + def get_job(self, job_id: str) -> object: ... + + class _CursorPageArgs(TypedDict, total=False): cursor: Mapping[str, str] skip: int @@ -815,7 +819,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]: + def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]: """ Gets file ids from responses API input. @@ -840,7 +844,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check for direct input_file type if item.get("type") == "input_file": file_id = item.get("file_id") - if file_id: + if isinstance(file_id, str) and file_id: file_ids.append(file_id) # Check for input_file in content array @@ -849,7 +853,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for content_item in content: if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") - if file_id: + if isinstance(file_id, str) and file_id: file_ids.append(file_id) return file_ids @@ -1189,7 +1193,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Handle both output_file_id and error_file_id for file_attr in ["output_file_id", "error_file_id"]: - file_id_value = getattr(response, file_attr, None) + file_id_value: str | None = getattr(response, file_attr, None) if file_id_value and model_id: decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value) if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id: @@ -1458,7 +1462,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): import litellm.proxy.proxy_server as proxy_server_module # Check if the scheduler has the batch cost checking job registered - scheduler = getattr(proxy_server_module, "scheduler", None) + scheduler: Final[_SchedulerWithJobLookup | None] = getattr(proxy_server_module, "scheduler", None) if scheduler is None: return False @@ -1504,7 +1508,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) MAX_MATCHES_TO_RETURN = 10 - batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + batches = await _managed_object_table(self.prisma_client).find_many( where={ "file_purpose": "batch", "batch_processed": False, @@ -1514,11 +1518,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): order={"created_at": "desc"}, ) - referencing_batches = [] + referencing_batches: Final[list[dict[str, object]]] = [] for batch in batches: try: # Parse the batch file_object to check for file references - batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object + decoded_file_object = _decode_json_blob(batch.file_object) + batch_data: Mapping[str, object] = ( + decoded_file_object if isinstance(decoded_file_object, Mapping) else {} + ) # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f1c80eaacbe..2b04a075114 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,7 +18,7 @@ import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -58,6 +58,26 @@ else: Span = Any +class _AsyncRedisCommands(Protocol): + """Async redis commands this cache issues. + + redis-py's type stubs omit these methods on RedisCluster, so the union returned by + init_async_client() is untyped at every call site without this protocol. + """ + + def ping(self) -> Awaitable[bool]: ... + + def delete(self, *names: str) -> Awaitable[int]: ... + + def ttl(self, name: str) -> Awaitable[int]: ... + + def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... + + def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... + + def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + + def _get_call_stack_info(num_frames: int = 2) -> str: """ Get the function names from the previous 1-2 functions in the call stack. @@ -429,6 +449,9 @@ class RedisCache(BaseCache): self.redis_async_client = redis_async_client return redis_async_client + def _async_commands(self) -> _AsyncRedisCommands: + return self.init_async_client() + def check_and_fix_namespace(self, key: str) -> str: """ Make sure each key starts with the given namespace @@ -1055,19 +1078,17 @@ class RedisCache(BaseCache): await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) self.redis_batch_writing_buffer = [] - def _get_cache_logic(self, cached_response: Any): + def _get_cache_logic(self, cached_response: bytes | str | None): """ Common 'get_cache_logic' across sync + async redis client implementations """ if cached_response is None: - return cached_response - # cached_response is in `b{} convert it to ModelResponse - cached_response = cached_response.decode("utf-8") # Convert bytes to string + return None + decoded: Final = cached_response.decode("utf-8") if isinstance(cached_response, bytes) else cached_response try: - cached_response = json.loads(cached_response) # Convert string to dictionary + return json.loads(decoded) except Exception: - cached_response = ast.literal_eval(cached_response) - return cached_response + return ast.literal_eval(decoded) def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): try: @@ -1314,8 +1335,7 @@ class RedisCache(BaseCache): raise e async def ping(self) -> bool: - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ping` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() start_time: Final = time.time() print_verbose("Pinging Async Redis Cache") try: @@ -1349,8 +1369,7 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def delete_cache_keys(self, keys): - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() keys = [self.check_and_fix_namespace(key=key) for key in keys] # keys is a list, unpack it so it gets passed as individual elements to delete await _redis_client.delete(*keys) @@ -1415,8 +1434,7 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_delete_cache(self, key: str): - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) # keys is str return await _redis_client.delete(key) @@ -1523,8 +1541,7 @@ class RedisCache(BaseCache): Redis ref: https://redis.io/docs/latest/commands/ttl/ """ try: - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) ttl: Final = await _redis_client.ttl(key) if ttl <= -1: # -1 means the key does not exist, -2 key does not exist @@ -1554,7 +1571,7 @@ class RedisCache(BaseCache): Returns: int: The length of the list after the push operation """ - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) start_time: Final = time.time() try: @@ -1621,7 +1638,7 @@ class RedisCache(BaseCache): if len(rpush_list) == 0: return [] - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() start_time: Final = time.time() try: @@ -1678,7 +1695,7 @@ class RedisCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> Any | list[Any]: - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) start_time: Final = time.time() print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") @@ -1810,7 +1827,7 @@ class RedisCache(BaseCache): if len(lpop_list) == 0: return [] - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() start_time: Final = time.time() try: diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index c66f6873383..b63b2e0dc10 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -17,8 +17,9 @@ RedisSemanticCache since those are backend agnostic. import asyncio import hashlib import os +from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any, Final +from typing import Any, Final, Protocol from redis import Redis from redis.asyncio import Redis as AsyncRedis @@ -40,6 +41,19 @@ class _ValkeyCacheHit: distance: float +class _SearchDocumentLike(Protocol): + """A valkey-search result document, whose fields are addressed by configurable name.""" + + def __getattr__(self, name: str, /) -> str | bytes | int | float: ... + + +class _SearchResultLike(Protocol): + """The one field this backend reads off an ``FT.SEARCH`` reply.""" + + @property + def docs(self) -> Sequence[_SearchDocumentLike]: ... + + class ValkeySemanticCache(RedisSemanticCache): """Valkey-backed semantic cache for LLM responses.""" @@ -64,7 +78,7 @@ class ValkeySemanticCache(RedisSemanticCache): async_client: AsyncRedis | None = None, embedding_max_input_tokens: int | None = None, embedding_timeout: float | None = None, - **kwargs: Any, + **kwargs: object, ): if similarity_threshold is None: raise ValueError("similarity_threshold must be provided, passed None") @@ -192,7 +206,9 @@ class ValkeySemanticCache(RedisSemanticCache): def _doc_key(self, key: str) -> str: return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}" - def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict: + def _doc_mapping( + self, key: str, prompt: str, value_str: str, embedding: list[float] + ) -> dict[str | bytes, str | bytes]: return { self.CACHE_KEY_FIELD_NAME: self._scope_tag(key), self.PROMPT_FIELD_NAME: prompt, @@ -209,8 +225,8 @@ class ValkeySemanticCache(RedisSemanticCache): return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2) @classmethod - def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: - docs: Final = getattr(search_result, "docs", []) + def _first_hit(cls, search_result: _SearchResultLike) -> _ValkeyCacheHit | None: + docs: Final[Sequence[_SearchDocumentLike]] = getattr(search_result, "docs", ()) if not docs: return None doc: Final = docs[0] @@ -219,7 +235,7 @@ class ValkeySemanticCache(RedisSemanticCache): distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)), ) - def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any: + def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> object: if hit is None: kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None @@ -231,7 +247,7 @@ class ValkeySemanticCache(RedisSemanticCache): return None return self._get_cache_logic(cached_response=hit.response) - def set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + def set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -250,7 +266,7 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in Valkey semantic-cache set_cache: {e}") - def get_cache(self, key: str, **kwargs: Any) -> Any: + def get_cache(self, key: str, **kwargs: Any) -> object: print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -270,7 +286,7 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose(f"Error in Valkey semantic-cache get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -289,7 +305,7 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}") - async def async_get_cache(self, key: str, **kwargs: Any) -> Any: + async def async_get_cache(self, key: str, **kwargs: Any) -> object: print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -309,11 +325,11 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None: try: await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list]) except Exception as e: print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}") - async def _index_info(self) -> dict: + async def _index_info(self) -> Mapping[str, object]: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 17815976b4a..a3196e25581 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -200,7 +200,8 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch LiteLLMCompletionResponsesConfig, ) - is_custom: Final = item.get("type") == "custom_tool_call" + item_type: Final[object] = item.get("type") + is_custom: Final = item_type == "custom_tool_call" arguments: Final = (item.get("input") if is_custom else item.get("arguments")) or "" name: Final = item.get("name") or ("custom_tool" if is_custom else "") function_chunk: Final = ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments) @@ -210,7 +211,7 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch function=function_chunk, index=index, ) - raw_provider_fields: Final = item.get("provider_specific_fields") + raw_provider_fields: Final[object] = item.get("provider_specific_fields") if isinstance(raw_provider_fields, dict): provider_specific_fields = raw_provider_fields elif raw_provider_fields and hasattr(raw_provider_fields, "__dict__"): @@ -495,7 +496,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _merge_responses_api_request_into_request_data( self, - request_data: dict[str, Any], + request_data: dict[str, object], responses_api_request: "ResponsesAPIOptionalRequestParams", instructions: str | None, ) -> None: diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 37a79e2f6d4..da89ad8919f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2253,6 +2253,10 @@ def batch_cost_calculator( return total_prompt_cost, total_completion_cost +def _attribute_value(obj: object, name: str) -> object: + return getattr(obj, name) + + def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]: field_names: Final = list(type(prompt_tokens_details).model_fields) if getattr(prompt_tokens_details, "cache_write_tokens", None) is None: @@ -2278,7 +2282,7 @@ class BaseTokenUsageProcessor: for usage in usage_objects: # Handle direct attributes by checking what exists in the model for attr in dir(usage): - if not attr.startswith("_") and not callable(getattr(usage, attr)): + if not attr.startswith("_") and not callable(_attribute_value(usage, attr)): current_val = getattr(combined, attr, 0) new_val = getattr(usage, attr, 0) if ( @@ -2298,7 +2302,7 @@ class BaseTokenUsageProcessor: if ( hasattr(usage.prompt_tokens_details, attr) and not attr.startswith("_") - and not callable(getattr(usage.prompt_tokens_details, attr)) + and not callable(_attribute_value(usage.prompt_tokens_details, attr)) ): current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 @@ -2317,7 +2321,9 @@ class BaseTokenUsageProcessor: # Check what keys exist in the model's completion_tokens_details # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings for attr in type(usage.completion_tokens_details).model_fields: - if not attr.startswith("_") and not callable(getattr(usage.completion_tokens_details, attr)): + if not attr.startswith("_") and not callable( + _attribute_value(usage.completion_tokens_details, attr) + ): current_val = getattr(combined.completion_tokens_details, attr, 0) or 0 new_val = getattr(usage.completion_tokens_details, attr, 0) or 0 if isinstance(new_val, (int, float)): diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 7c86ceafd7f..2864663df93 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -23,7 +23,11 @@ from litellm.types.llms.openai import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( AdapterCompletionStreamWrapper, + ChatCompletionDeltaCustomToolCall, + ChatCompletionMessageCustomToolCall, Choices, + Delta, + Message, ModelResponse, ModelResponseStream, StreamingChoices, @@ -635,7 +639,7 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, - message: Any, + message: Message, ) -> list[_GenAIPart]: """Transform OpenAI message to Google GenAI parts format""" parts: Final[list[_GenAIPart]] = [] @@ -647,7 +651,11 @@ class GoogleGenAIAdapter: # Add tool calls if present if hasattr(message, "tool_calls") and message.tool_calls: for tool_call in message.tool_calls: - if hasattr(tool_call, "function") and tool_call.function: + if ( + hasattr(tool_call, "function") + and not isinstance(tool_call, ChatCompletionMessageCustomToolCall) + and tool_call.function + ): try: args = ( _decode_tool_call_arguments(tool_call.function.arguments) @@ -668,7 +676,7 @@ class GoogleGenAIAdapter: return parts if parts else [{"text": ""}] def _transform_openai_delta_to_google_genai_parts_with_accumulation( - self, delta: Any, wrapper: GoogleGenAIStreamWrapper + self, delta: Delta, wrapper: GoogleGenAIStreamWrapper ) -> list[_GenAIPart]: """Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls.""" @@ -685,7 +693,7 @@ class GoogleGenAIAdapter: tool_calls: Final = delta.tool_calls or [] for tool_call in tool_calls: - if not hasattr(tool_call, "function"): + if not hasattr(tool_call, "function") or isinstance(tool_call, ChatCompletionDeltaCustomToolCall): continue # 3. Use `index` as the primary key for accumulation diff --git a/litellm/integrations/arize/arize_phoenix_prompt_manager.py b/litellm/integrations/arize/arize_phoenix_prompt_manager.py index 71f4902bbe5..0c616e845f8 100644 --- a/litellm/integrations/arize/arize_phoenix_prompt_manager.py +++ b/litellm/integrations/arize/arize_phoenix_prompt_manager.py @@ -3,10 +3,12 @@ Arize Phoenix prompt manager that integrates with LiteLLM's prompt management sy Fetches prompt versions from Arize Phoenix and provides workspace-based access control. """ +from collections.abc import Mapping, Sequence from typing import Any, Final from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import ReadOnly, TypedDict from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import ( @@ -20,6 +22,31 @@ from litellm.types.utils import StandardCallbackDynamicParams from .arize_phoenix_client import ArizePhoenixClient +class ArizePhoenixContentPart(TypedDict, total=False): + type: ReadOnly[str] + text: ReadOnly[str] + + +class ArizePhoenixTemplateMessage(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[Sequence[ArizePhoenixContentPart]] + + +class ArizePhoenixTemplateBody(TypedDict, total=False): + messages: ReadOnly[Sequence[ArizePhoenixTemplateMessage]] + + +class ArizePhoenixPromptMetadata(TypedDict): + model_name: ReadOnly[str | None] + model_provider: ReadOnly[str | None] + description: ReadOnly[str] + template_type: ReadOnly[str | None] + template_format: ReadOnly[str] + invocation_parameters: ReadOnly[Mapping[str, Mapping[str, object]]] + temperature: ReadOnly[float | None] + max_tokens: ReadOnly[int | None] + + class ArizePhoenixPromptTemplate: """ Represents a prompt template loaded from Arize Phoenix. @@ -28,10 +55,10 @@ class ArizePhoenixPromptTemplate: def __init__( self, template_id: str, - messages: list[dict[str, Any]], - metadata: dict[str, Any], + messages: Sequence[ArizePhoenixTemplateMessage], + metadata: ArizePhoenixPromptMetadata, model: str | None = None, - ): + ) -> None: self.template_id = template_id self.messages = messages self.metadata = metadata @@ -43,7 +70,7 @@ class ArizePhoenixPromptTemplate: self.description = metadata.get("description", "") self.template_format = metadata.get("template_format", "MUSTACHE") - def __repr__(self): + def __repr__(self) -> str: return f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')" @@ -109,7 +136,7 @@ class ArizePhoenixTemplateManager: def _parse_prompt_data(self, data: dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate: """Parse Arize Phoenix prompt data and extract messages and metadata.""" - template_data: Final = data.get("template", {}) + template_data: Final[ArizePhoenixTemplateBody] = data.get("template", {}) messages: Final = template_data.get("messages", []) # Extract invocation parameters @@ -129,7 +156,7 @@ class ArizePhoenixTemplateManager: break # Build metadata dictionary - metadata: Final = { + metadata: Final[ArizePhoenixPromptMetadata] = { "model_name": data.get("model_name"), "model_provider": data.get("model_provider"), "description": data.get("description", ""), @@ -146,7 +173,9 @@ class ArizePhoenixTemplateManager: metadata=metadata, ) - def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> list[AllMessageValues]: + def render_template( + self, template_id: str, variables: Mapping[str, object] | None = None + ) -> list[AllMessageValues]: """Render a template with the given variables and return formatted messages.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") @@ -243,8 +272,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, - ) -> tuple[list[AllMessageValues], dict[str, Any]]: + prompt_variables: Mapping[str, object] | None = None, + ) -> tuple[list[AllMessageValues], dict[str, object]]: """ Get a prompt template and render it with variables. @@ -263,7 +292,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): rendered_messages: Final = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) # Extract metadata - metadata: Final = { + metadata: Final[dict[str, object]] = { "model": template.model, "temperature": template.temperature, "max_tokens": template.max_tokens, @@ -271,7 +300,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): # Add additional invocation parameters invocation_params: Final = template.invocation_parameters - provider_params = {} + provider_params: Mapping[str, object] = {} if "openai" in invocation_params: provider_params = invocation_params["openai"] @@ -289,12 +318,12 @@ class ArizePhoenixPromptManager(CustomPromptManagement): self, user_id: str | None, messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: dict[str, object] | str | None = None, + litellm_params: dict[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: dict[str, object] | None = None, **kwargs, - ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: + ) -> tuple[list[AllMessageValues], dict[str, object] | None]: """ Pre-call hook that processes the prompt template before making the LLM call. """ @@ -335,9 +364,9 @@ class ArizePhoenixPromptManager(CustomPromptManagement): except Exception as e: # Log error but don't fail the call - import litellm + from litellm._logging import verbose_proxy_logger - litellm._logging.verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e) + verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e) return messages, litellm_params def get_available_prompts(self) -> list[str]: @@ -393,7 +422,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement): rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Extract model from metadata (if specified) - template_model: Final = prompt_metadata.get("model") + raw_template_model: Final = prompt_metadata.get("model") + template_model: Final = raw_template_model if isinstance(raw_template_model, str) else None # Extract optional parameters from metadata optional_params: Final = {} diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 41caf732db0..83ef46efe40 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -2,7 +2,7 @@ # On success, logs events to Promptlayer import re import traceback -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Mapping, Sequence from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional from pydantic import BaseModel @@ -31,6 +31,9 @@ if TYPE_CHECKING: from litellm.caching.caching import DualCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import ( MCPPostCallResponseObject, @@ -39,7 +42,7 @@ if TYPE_CHECKING: ) from litellm.types.router import PreRoutingHookResponse - Span = _Span | Any + Span = _Span else: Span = Any LiteLLMLoggingObj = Any @@ -123,11 +126,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return [] callbacks: Final = AllCallbacks() - callback_info: Final = getattr(callbacks, lookup_name, None) + callback_info: Final[object] = getattr(callbacks, lookup_name, None) if callback_info is None: return [] - params: Final = getattr(callback_info, "litellm_callback_params", None) + params: Final[list[str] | None] = getattr(callback_info, "litellm_callback_params", None) if not params: return [] @@ -268,7 +271,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> list[dict]: return healthy_deployments - async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict | None: """ Allow modifying the request just before it's sent to the deployment. @@ -344,9 +349,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_post_call_streaming_deployment_hook( self, request_data: dict, - response_chunk: Any, + response_chunk: object, call_type: CallTypes | None, - ) -> Any | None: + ) -> object | None: """ Allow modifying streaming chunks just before they're returned to the user. @@ -378,7 +383,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ def translate_completion_output_params_streaming( - self, completion_stream: Any + self, completion_stream: object ) -> AdapterCompletionStreamWrapper | None: """ Translates the streaming chunk, from the OpenAI format to the custom format. @@ -418,9 +423,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self, data: dict, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: object, request_headers: dict[str, str] | None = None, - litellm_call_info: dict[str, Any] | None = None, + litellm_call_info: dict[str, object] | None = None, ) -> dict[str, str] | None: """ Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers. @@ -471,11 +476,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> Any: pass - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result - def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result @@ -581,7 +586,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -642,8 +647,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None", anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, @@ -711,8 +716,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None", anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, @@ -728,7 +733,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_post_agentic_loop_response_hook( self, - response: Any, + response: object, plan: AgenticLoopPlan, kwargs: dict, ) -> Any: @@ -767,7 +772,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_should_run_chat_completion_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -785,12 +790,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, + response: object, optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, kwargs: dict, - ) -> Any: + ) -> object: """ Hook to execute chat completion agentic loop based on context from should_run hook. """ @@ -800,7 +805,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, + response: object, optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, @@ -851,7 +856,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac - Converting to string and then truncating the logged content catches this 2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user """ - field_value: Final = standard_logging_object.get(field_name) + field_value: Final[object] = standard_logging_object.get(field_name) if field_value: str_value: Final = str(field_value) if len(str_value) > max_length: @@ -1005,8 +1010,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -1037,8 +1042,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -1056,10 +1061,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _redact_base64( self, - value: Any, + value: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> Any: + ) -> object: """Recursively redact inline base64 from any nested structure with a max recursion depth limit.""" if depth > max_depth: verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth) @@ -1079,7 +1084,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return value - def _should_keep_content(self, content: Any) -> bool: + def _should_keep_content(self, content: object) -> bool: """Return True if this content item should be retained.""" if not isinstance(content, dict): return True @@ -1090,16 +1095,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _process_messages( self, - messages: list[Any], + messages: Sequence[object], max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> list[dict[str, Any]]: - filtered_messages: Final[list[dict[str, Any]]] = [] + ) -> list[dict[str, object]]: + filtered_messages: Final[list[dict[str, object]]] = [] for msg in messages: if not isinstance(msg, dict): continue - contents: Any = msg.get("content") + contents: object = msg.get("content") if isinstance(contents, list): - cleaned: list[Any] = [] + cleaned: list[object] = [] for c in contents: if self._should_keep_content(content=c): cleaned.append(self._redact_base64(value=c, max_depth=max_depth)) diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index c41d9dd240f..d4602176650 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -2,10 +2,12 @@ GitLab prompt manager with configurable prompts folder. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, TypeVar from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import ReadOnly, TypedDict from litellm.integrations.custom_prompt_management import CustomPromptManagement @@ -24,6 +26,19 @@ from litellm.types.utils import StandardCallbackDynamicParams GITLAB_PREFIX: Final = "gitlab::" +_ResponseT = TypeVar("_ResponseT") + + +class GitLabCachedPrompt(TypedDict): + id: ReadOnly[str] + path: ReadOnly[str] + content: ReadOnly[str] + metadata: ReadOnly[Mapping[str, object]] + model: ReadOnly[str | None] + temperature: ReadOnly[float | None] + max_tokens: ReadOnly[int | None] + optional_params: ReadOnly[Mapping[str, object]] + def encode_prompt_id(raw_id: str) -> str: """Convert GitLab path IDs like 'invoice/extract' → 'gitlab::invoice::extract'""" @@ -206,7 +221,7 @@ class GitLabTemplateManager: result[key] = value.strip("\"'") return result - def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str: + def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str: if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") template: Final = self.prompts[template_id] @@ -313,7 +328,7 @@ class GitLabPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, *, ref: str | None = None, ) -> tuple[str, dict[str, Any]]: @@ -338,13 +353,13 @@ class GitLabPromptManager(CustomPromptManagement): self, user_id: str | None, messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: dict[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, prompt_version: str | None = None, **kwargs, - ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: + ) -> tuple[list[AllMessageValues], dict[str, object] | None]: if not prompt_id: return messages, litellm_params try: @@ -377,9 +392,9 @@ class GitLabPromptManager(CustomPromptManagement): return final_messages, litellm_params except Exception as e: - import litellm + from litellm._logging import verbose_proxy_logger - litellm._logging.verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e) + verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e) return messages, litellm_params def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]: @@ -435,14 +450,14 @@ class GitLabPromptManager(CustomPromptManagement): def post_call_hook( self, user_id: str | None, - response: Any, + response: _ResponseT, input_messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: Mapping[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, **kwargs, - ) -> Any: + ) -> _ResponseT: return response def get_available_prompts(self) -> list[str]: @@ -498,7 +513,7 @@ class GitLabPromptManager(CustomPromptManagement): messages: Final = self._parse_prompt_to_messages(rendered_prompt) template_model: Final = prompt_metadata.get("model") - optional_params: Final[dict[str, Any]] = {} + optional_params: Final[dict[str, object]] = {} for param in [ "temperature", "max_tokens", @@ -658,14 +673,14 @@ class GitLabPromptCache: self.template_manager: GitLabTemplateManager = self.prompt_manager.prompt_manager # In-memory stores - self._by_file: dict[str, dict[str, Any]] = {} - self._by_id: dict[str, dict[str, Any]] = {} + self._by_file: dict[str, GitLabCachedPrompt] = {} + self._by_id: dict[str, GitLabCachedPrompt] = {} # ------------------------- # Public API # ------------------------- - def load_all(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]: + def load_all(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]: """ Scan GitLab for all .prompt files under prompts_path, load and parse each, and return the mapping of repo file path -> JSON-like dict. @@ -695,7 +710,7 @@ class GitLabPromptCache: return self._by_id - def reload(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]: + def reload(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]: """Clear the cache and re-load from GitLab.""" self._by_file.clear() self._by_id.clear() @@ -709,11 +724,11 @@ class GitLabPromptCache: """Return the template IDs (relative to prompts_path, without extension) currently cached.""" return list(self._by_id.keys()) - def get_by_file(self, file_path: str) -> dict[str, Any] | None: + def get_by_file(self, file_path: str) -> GitLabCachedPrompt | None: """Get a cached prompt JSON by repo file path.""" return self._by_file.get(file_path) - def get_by_id(self, prompt_id: str) -> dict[str, Any] | None: + def get_by_id(self, prompt_id: str) -> GitLabCachedPrompt | None: """Get a cached prompt JSON by prompt ID (relative to prompts_path).""" if prompt_id in self._by_id: return self._by_id[prompt_id] @@ -728,7 +743,7 @@ class GitLabPromptCache: # Internals # ------------------------- - def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> dict[str, Any]: + def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> GitLabCachedPrompt: """ Normalize a GitLabPromptTemplate into a JSON-like dict that is easy to serialize. """ diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index db9610a5a3c..4f7dff952e6 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -12,7 +12,10 @@ For batching specific details see CustomBatchLogger class import asyncio import atexit import os -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -34,6 +37,21 @@ from litellm.types.integrations.posthog import ( from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload +class PostHogBatchPayload(TypedDict): + api_key: ReadOnly[str] + batch: ReadOnly[Sequence[PostHogEventPayload]] + + +class PostHogLiteLLMParams(TypedDict, total=False): + metadata: ReadOnly[Mapping[str, object]] + + +class PostHogLogKwargs(TypedDict, total=False): + standard_logging_object: ReadOnly[StandardLoggingPayload] + standard_callback_dynamic_params: ReadOnly[StandardCallbackDynamicParams] + litellm_params: ReadOnly[PostHogLiteLLMParams] + + class PostHogLogger(CustomBatchLogger): def __init__(self, **kwargs): """ @@ -137,7 +155,7 @@ class PostHogLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() - def create_posthog_event_payload(self, kwargs: dict[str, Any]) -> PostHogEventPayload: + def create_posthog_event_payload(self, kwargs: PostHogLogKwargs) -> PostHogEventPayload: """ Helper function to create a PostHog event payload for logging @@ -171,11 +189,11 @@ class PostHogLogger(CustomBatchLogger): def _create_posthog_properties( self, standard_logging_object: StandardLoggingPayload, - kwargs: dict[str, Any], + kwargs: PostHogLogKwargs, event_name: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Create PostHog properties following LLM Analytics spec""" - properties: Final = {} + properties: Final[dict[str, object]] = {} # Core model information properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "") @@ -211,16 +229,19 @@ class PostHogLogger(CustomBatchLogger): properties["$ai_error"] = error_str # Add trace properties - self._add_trace_properties(properties, kwargs) + self._add_trace_properties(properties, standard_logging_object, kwargs) # Add custom metadata fields self._add_custom_metadata_properties(properties, kwargs) return properties - def _add_trace_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]): - standard_logging_object: Final = self._safe_get(kwargs, "standard_logging_object", {}) - + def _add_trace_properties( + self, + properties: dict[str, object], + standard_logging_object: StandardLoggingPayload, + kwargs: PostHogLogKwargs, + ) -> None: trace_id: Final = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid()) properties["$ai_trace_id"] = trace_id @@ -232,7 +253,7 @@ class PostHogLogger(CustomBatchLogger): if parent_id: properties["$ai_parent_id"] = parent_id - def _add_custom_metadata_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]): + def _add_custom_metadata_properties(self, properties: dict[str, object], kwargs: PostHogLogKwargs) -> None: """Add custom metadata fields to PostHog properties""" metadata: Final = self._extract_metadata(kwargs) if not isinstance(metadata, dict): @@ -277,7 +298,7 @@ class PostHogLogger(CustomBatchLogger): if key not in litellm_internal_fields: properties[key] = value - def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: dict[str, Any]) -> str: + def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: PostHogLogKwargs) -> str: metadata: Final = self._extract_metadata(kwargs) user_id: Final = self._safe_get(metadata, "user_id") if user_id: @@ -291,7 +312,7 @@ class PostHogLogger(CustomBatchLogger): return self._safe_uuid() - def _get_credentials_for_request(self, kwargs: dict[str, Any]) -> tuple[str | None, str | None]: + def _get_credentials_for_request(self, kwargs: PostHogLogKwargs) -> tuple[str | None, str | None]: """ Get PostHog credentials for this request. @@ -334,7 +355,7 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted") # Group events by credentials for batch sending - batches_by_credentials: Final[dict[tuple[str, str], list]] = {} + batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {} for item in self.log_queue: key = (item["api_key"], item["api_url"]) if key not in batches_by_credentials: @@ -380,18 +401,19 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.error("PostHog: Failed to initialize async components: %s", e) raise - def _extract_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any]: - litellm_params: Final = kwargs.get("litellm_params", {}) or {} - return litellm_params.get("metadata", {}) or {} + def _extract_metadata(self, kwargs: PostHogLogKwargs) -> Mapping[str, object]: + litellm_params: Final[PostHogLiteLLMParams] = kwargs.get("litellm_params", {}) or {} + metadata: Final[Mapping[str, object]] = litellm_params.get("metadata", {}) or {} + return metadata def _safe_uuid(self) -> str: return str(uuid.uuid4()) - def _create_posthog_payload(self, events: list, api_key: str) -> dict[str, Any]: + def _create_posthog_payload(self, events: Sequence[PostHogEventPayload], api_key: str) -> PostHogBatchPayload: return {"api_key": api_key, "batch": events} - def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any: - if obj is None or not hasattr(obj, "get"): + def _safe_get(self, obj: Mapping[str, object] | None, key: str, default: object = None) -> object: + if not isinstance(obj, Mapping): return default return obj.get(key, default) @@ -412,7 +434,7 @@ class PostHogLogger(CustomBatchLogger): try: # Group events by credentials (same logic as async_send_batch) - batches_by_credentials: Final[dict[tuple[str, str], list]] = {} + batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {} for item in self.log_queue: key = (item["api_key"], item["api_url"]) if key not in batches_by_credentials: diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 9a2c4e244fb..864bbac70c3 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -3,7 +3,7 @@ Helper utilities for tracking the cost of built-in tools. """ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS @@ -14,6 +14,7 @@ from litellm.types.llms.openai import ( WebSearchOptions, ) from litellm.types.utils import ( + ChatCompletionAnnotation, Message, ModelInfo, ModelResponse, @@ -47,7 +48,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def get_cost_for_built_in_tools( model: str, - response_object: Any, + response_object: object, usage: Usage | None = None, custom_llm_provider: str | None = None, standard_built_in_tools_params: StandardBuiltInToolsParams | None = None, @@ -199,8 +200,7 @@ class StandardBuiltInToolCostTracking: model_info: Final = StandardBuiltInToolCostTracking._safe_get_model_info( model=model, custom_llm_provider=custom_llm_provider ) - file_search_raw: Final[Any] = standard_built_in_tools_params.get("file_search", {}) - file_search_usage: Final[FileSearchTool | None] = FileSearchTool(**file_search_raw) if file_search_raw else None + file_search_usage: Final[FileSearchTool | None] = standard_built_in_tools_params.get("file_search") or None # Convert model_info to dict and extract usage parameters model_info_dict: Final = dict(model_info) if model_info is not None else None @@ -243,7 +243,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def _extract_file_search_params( - file_search_usage: Any, + file_search_usage: object, ) -> tuple[float | None, float | None]: """Extract and convert file search parameters safely.""" storage_gb = None @@ -333,7 +333,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def _extract_token_counts( - computer_use_usage: Any, + computer_use_usage: object, ) -> tuple[int | None, int | None]: """Extract and convert token counts safely.""" input_tokens = None @@ -349,9 +349,9 @@ class StandardBuiltInToolCostTracking: return input_tokens, output_tokens @staticmethod - def _safe_convert_to_int(value: Any) -> int | None: + def _safe_convert_to_int(value: object) -> int | None: """Safely convert a value to int.""" - if value is not None: + if isinstance(value, (int, float, str)): try: return int(value) except (TypeError, ValueError): @@ -379,7 +379,7 @@ class StandardBuiltInToolCostTracking: return usage.model_copy(update={"server_tool_use": server_tool_use}) @staticmethod - def response_object_includes_web_search_call(response_object: Any, usage: Usage | None = None) -> bool: + def response_object_includes_web_search_call(response_object: object, usage: Usage | None = None) -> bool: """ Check if the response object includes a web search call. @@ -448,7 +448,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def response_object_includes_file_search_call( - response_object: Any, + response_object: object, ) -> bool: """ Check if the response object includes a file search call. @@ -479,11 +479,11 @@ class StandardBuiltInToolCostTracking: message: Message | None = getattr(choice, "message", None) if message is None: continue - if annotations := getattr(message, "annotations", None): - if len(annotations) > 0: - for annotation in annotations: - if annotation.get("type", None) == annotation_type: - return True + annotations: list[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) + if annotations: + for annotation in annotations: + if annotation.get("type", None) == annotation_type: + return True return False @staticmethod @@ -524,10 +524,8 @@ class StandardBuiltInToolCostTracking: if model_info is None: return 0.0 - search_context_raw: Final[Any] = model_info.get("search_context_cost_per_query", {}) - search_context_pricing: Final[SearchContextCostPerQuery] = ( - SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery() - ) + search_context_raw: Final = model_info.get("search_context_cost_per_query") + search_context_pricing: Final[SearchContextCostPerQuery] = search_context_raw or SearchContextCostPerQuery() if web_search_options.get("search_context_size", None) == "low": return search_context_pricing.get("search_context_size_low", 0.0) elif web_search_options.get("search_context_size", None) == "medium": @@ -547,10 +545,8 @@ class StandardBuiltInToolCostTracking: """ if model_info is None: return 0.0 - search_context_raw: Final[Any] = model_info.get("search_context_cost_per_query", {}) or {} - search_context_pricing: Final[SearchContextCostPerQuery] = ( - SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery() - ) + search_context_raw: Final = model_info.get("search_context_cost_per_query") + search_context_pricing: Final[SearchContextCostPerQuery] = search_context_raw or SearchContextCostPerQuery() return search_context_pricing.get("search_context_size_medium", 0.0) @staticmethod @@ -716,7 +712,7 @@ class StandardBuiltInToolCostTracking: response_object: ModelResponse, ) -> bool: for _choice in response_object.choices: - message = getattr(_choice, "message", None) + message: Message | None = getattr(_choice, "message", None) if ( message is not None and hasattr(message, "annotations") diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 748347fe938..fa1b57c894f 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -519,10 +519,10 @@ def update_messages_with_model_file_ids( def update_responses_input_with_model_file_ids( - input: Any, + input: object, model_id: str | None = None, model_file_id_mapping: dict[str, dict[str, str]] | None = None, -) -> str | list[dict[str, Any]]: +) -> object: """ Updates responses API input with provider-specific file IDs. File IDs are always inside the content array, not as direct input_file items. @@ -603,8 +603,8 @@ def update_responses_input_with_model_file_ids( def _decode_vector_store_ids_in_tools( - tools: list[dict[str, Any]] | None, -) -> list[dict[str, Any]] | None: + tools: list[dict[str, object]] | None, +) -> list[dict[str, object]] | None: """ Decodes unified (LiteLLM-managed) vector_store_ids in file_search tools to provider-native IDs. Non-unified IDs are passed through unchanged. @@ -656,10 +656,10 @@ def _decode_vector_store_ids_in_tools( def update_responses_tools_with_model_file_ids( - tools: list[dict[str, Any]] | None, + tools: list[dict[str, object]] | None, model_id: str | None = None, model_file_id_mapping: dict[str, dict[str, str]] | None = None, -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """ Updates responses API tools with provider-specific file IDs. @@ -852,7 +852,7 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData: # --------------------------------------------------------------------------- -def _estimate_json_bytes(obj: Any) -> int: +def _estimate_json_bytes(obj: object) -> int: """Estimate the JSON-serialised byte size of ``obj`` without materialising JSON. Walks iteratively (no recursion stack risk). @@ -1747,7 +1747,7 @@ def hoist_images_from_tool_messages( ] -def _attempt_json_repair(s: str) -> Any | None: +def _attempt_json_repair(s: str) -> object | None: """ Attempt to repair truncated JSON produced by LLM tool calls. @@ -1863,7 +1863,7 @@ def parse_tool_call_arguments( raise ValueError(error_message) from original_error -def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: +def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: """ Split a string that contains one or more concatenated JSON objects into a list of parsed dicts. @@ -1899,7 +1899,7 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: return [] decoder: Final = json.JSONDecoder() - results: Final[list[dict[str, Any]]] = [] + results: Final[list[dict[str, object]]] = [] idx = 0 length: Final = len(raw) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 858b078d626..f51e2122fca 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -4,8 +4,9 @@ import base64 import io import struct from collections.abc import Callable, Mapping -from typing import Any, Final, Literal, cast +from typing import Final, Literal, cast +import httpx import tiktoken import litellm @@ -164,6 +165,10 @@ def calculate_tiles_needed( return total_tiles +def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: + return struct.unpack(fmt, buffer) + + def get_image_type(image_data: bytes) -> str | None: """take an image (really only the first ~100 bytes max are needed) and return 'png' 'gif' 'jpeg' 'webp' 'heic' or None. method added to @@ -203,9 +208,9 @@ def get_image_dimensions( if data.startswith(("http://", "https://")): try: client: Final = _get_httpx_client() - response: Final = safe_get(client, data) + response: Final[httpx.Response] = safe_get(client, data) max_bytes: Final = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) - content_length: Final = response.headers.get("Content-Length") + content_length: Final[str | None] = response.headers.get("Content-Length") if content_length is not None and int(content_length) > max_bytes: pass # skip download; img_data stays None else: @@ -222,10 +227,10 @@ def get_image_dimensions( img_type: Final = get_image_type(img_data) if img_type == "png": - w, h = struct.unpack(">LL", img_data[16:24]) + w, h = _unpack_ints(">LL", img_data[16:24]) return w, h elif img_type == "gif": - w, h = struct.unpack("H", fhandle.read(2))[0] - 2 + size = _unpack_ints(">H", fhandle.read(2))[0] - 2 fhandle.seek(1, 1) - h, w = struct.unpack(">HH", fhandle.read(4)) + h, w = _unpack_ints(">HH", fhandle.read(4)) return w, h elif img_type == "webp": # For WebP, the dimensions are stored at different offsets depending on the format # Check for VP8X (extended format) if img_data[12:16] == b"VP8X": - w = struct.unpack("> 14) & 0x3FFF) + 1 return w, h @@ -413,8 +418,8 @@ def token_counter( def _count_function_call_tokens( key: str, - value: Any, - message: Mapping[str, Any], + value: object, + message: Mapping[str, object], count_function: TokenCounterFunction, ) -> int: """ @@ -580,7 +585,7 @@ def _fix_model_name(model: str) -> str: def _count_image_tokens( - image_url: Any, + image_url: object, use_default_image_token_count: bool, ) -> int: """ @@ -620,7 +625,7 @@ def _count_image_tokens( raise ValueError(f"Invalid image_url type: {type(image_url).__name__}. Expected str or dict with 'url' field.") -def _validate_anthropic_content(content: Mapping[str, Any]) -> type: +def _validate_anthropic_content(content: Mapping[str, object]) -> type: """ Validate and determine which Anthropic TypedDict applies. @@ -635,7 +640,7 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type: "tool_result": AnthropicMessagesToolResultParam, } - expected_cls: Final = mapping.get(content_type) + expected_cls: Final = mapping.get(content_type) if isinstance(content_type, str) else None if expected_cls is None: raise ValueError(f"Unknown Anthropic content type: '{content_type}'") @@ -647,7 +652,7 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type: def _count_anthropic_content( - content: Mapping[str, Any], + content: Mapping[str, object], count_function: TokenCounterFunction, use_default_image_token_count: bool, default_token_count: int | None, @@ -662,7 +667,7 @@ def _count_anthropic_content( avoiding hardcoded field names. """ typeddict_cls: Final = _validate_anthropic_content(content) - type_hints: Final = getattr(typeddict_cls, "__annotations__", {}) + type_hints: Final[Mapping[str, object]] = getattr(typeddict_cls, "__annotations__", {}) tokens = 0 # Fields to skip (metadata/identifiers that don't contribute to prompt tokens) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 721a6653597..850cc74bab6 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -18,7 +18,7 @@ from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, cast -from typing_extensions import assert_never +from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -111,6 +111,16 @@ class ExtractedInput: EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) +class _AnthropicSSEDelta(TypedDict, total=False): + type: ReadOnly[str] + text: ReadOnly[str] + stop_reason: ReadOnly[str | None] + + +class _AnthropicSSEEvent(TypedDict, total=False): + delta: ReadOnly[_AnthropicSSEDelta] + + class AnthropicMessagesHandler(BaseTranslation): """Process Anthropic messages with guardrails. @@ -747,7 +757,7 @@ class AnthropicMessagesHandler(BaseTranslation): if scan_only_tool_results: return EMPTY_EXTRACTED_INPUT - text_str: Final = content_item.get("text", None) + text_str: Final[str | None] = content_item.get("text", None) return ExtractedInput( scanned=( () if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),) @@ -1156,8 +1166,8 @@ class AnthropicMessagesHandler(BaseTranslation): # Only process content_block_delta events if event_type == "content_block_delta" and data_line: try: - data = json.loads(data_line) - delta = data.get("delta", {}) + data: _AnthropicSSEEvent = json.loads(data_line) + delta: _AnthropicSSEDelta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") except json.JSONDecodeError: @@ -1219,9 +1229,9 @@ class AnthropicMessagesHandler(BaseTranslation): # Check for message_delta event with stop_reason if event_type == "message_delta" and data_line: try: - data = json.loads(data_line) - delta = data.get("delta", {}) - stop_reason = delta.get("stop_reason") + data: _AnthropicSSEEvent = json.loads(data_line) + delta: _AnthropicSSEDelta = data.get("delta", {}) + stop_reason: str | None = delta.get("stop_reason") if stop_reason is not None: return True except json.JSONDecodeError: diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index cd47cdd57d6..c82be07a5c5 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -66,6 +66,10 @@ if TYPE_CHECKING: from litellm.llms.base_llm.chat.transformation import BaseConfig +def _loads_stream_chunk(payload: str) -> dict[str, object]: + return json.loads(payload) + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -78,7 +82,7 @@ async def make_call( json_mode: bool, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, -) -> tuple[Any, httpx.Headers]: +) -> tuple["ModelResponseIterator", httpx.Headers]: if client is None: client = litellm.module_level_aclient @@ -93,7 +97,7 @@ async def make_call( ) except httpx.HTTPStatusError as e: error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise AnthropicError( @@ -138,7 +142,7 @@ def make_sync_call( json_mode: bool, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, -) -> tuple[Any, httpx.Headers]: +) -> tuple["ModelResponseIterator", httpx.Headers]: if client is None: client = litellm.module_level_client # re-use a module level client @@ -153,7 +157,7 @@ def make_sync_call( ) except httpx.HTTPStatusError as e: error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise AnthropicError( @@ -292,7 +296,7 @@ class AnthropicChatCompletion(BaseLLM): status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) if error_response and hasattr(error_response, "text"): @@ -593,7 +597,7 @@ class AnthropicChatCompletion(BaseLLM): status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) if error_response and hasattr(error_response, "text"): @@ -664,10 +668,10 @@ class ModelResponseIterator: # Accumulate web_search_tool_result blocks for multi-turn reconstruction # See: https://github.com/BerriAI/litellm/issues/17737 - self.web_search_results: list[dict[str, Any]] = [] + self.web_search_results: list[dict[str, object]] = [] # Accumulate compaction blocks for multi-turn reconstruction - self.compaction_blocks: list[dict[str, Any]] = [] + self.compaction_blocks: list[dict[str, object]] = [] # Accumulate streamed thinking text so final usage can split reasoning # tokens from regular output tokens. @@ -727,7 +731,7 @@ class ModelResponseIterator: str, ChatCompletionToolCallChunk | None, list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], - dict[str, Any], + dict[str, object], str | None, ]: """ @@ -735,7 +739,7 @@ class ModelResponseIterator: """ text = "" tool_use: ChatCompletionToolCallChunk | None = None - provider_specific_fields: Final = {} + provider_specific_fields: Final[dict[str, object]] = {} reasoning_content: str | None = None content_block: Final = ContentBlockDelta(**chunk) thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = [] @@ -809,8 +813,8 @@ class ModelResponseIterator: def _handle_redacted_thinking_content( self, content_block_start: ContentBlockStart, - provider_specific_fields: dict[str, Any], - ) -> tuple[list[ChatCompletionRedactedThinkingBlock], dict[str, Any]]: + provider_specific_fields: dict[str, object], + ) -> tuple[list[ChatCompletionRedactedThinkingBlock], dict[str, object]]: """ Handle the redacted thinking content """ @@ -878,7 +882,7 @@ class ModelResponseIterator: tool_use: ChatCompletionToolCallChunk | None = None finish_reason = "" usage: Usage | None = None - provider_specific_fields: dict[str, Any] = {} + provider_specific_fields: dict[str, object] = {} reasoning_content: str | None = None thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None = None @@ -1212,7 +1216,7 @@ class ModelResponseIterator: # Try to parse as valid JSON first try: - data_json: Final = json.loads(data_str) + data_json: Final = _loads_stream_chunk(data_str) return self.chunk_parser(chunk=data_json) except json.JSONDecodeError: # Switch to accumulation mode and start accumulating @@ -1330,7 +1334,7 @@ class ModelResponseIterator: str_line = str_line[index:] if str_line.startswith("data:"): - data_json: Final = json.loads(str_line[5:]) + data_json: Final = _loads_stream_chunk(str_line[5:]) return self.chunk_parser(chunk=data_json) else: return ModelResponseStream(id=self.response_id) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index c73376ba498..5058bb460d6 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -865,13 +865,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): f"Failed to fetch models from Anthropic. Status code: {response.status_code}, Response: {response.text}" ) - models: Final = response.json()["data"] + models: Final[Sequence[Mapping[str, str]]] = response.json()["data"] - litellm_model_names: Final = [] - for model in models: - stripped_model_name = model["id"] - litellm_model_name = "anthropic/" + stripped_model_name - litellm_model_names.append(litellm_model_name) + litellm_model_names: Final = ["anthropic/" + model["id"] for model in models] return litellm_model_names def get_token_counter(self) -> BaseTokenCounter | None: @@ -1064,7 +1060,7 @@ def strip_empty_text_blocks_from_anthropic_messages( return out -def _is_empty_text_block(block: Any) -> bool: +def _is_empty_text_block(block: object) -> bool: if not isinstance(block, dict) or block.get("type") != "text": return False text: Final = block.get("text") @@ -1084,7 +1080,7 @@ def normalize_anthropic_tool_use_id(raw_id: str) -> str: return sanitized or "tool_use_id" -def _sanitize_tool_use_id_content_block(block: Any) -> Any: +def _sanitize_tool_use_id_content_block(block: object) -> object: if not isinstance(block, dict): return block block_type: Final = block.get("type") diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index d7b527824ea..90f87e38842 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1,7 +1,7 @@ import copy import hashlib import json -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, cast import litellm @@ -18,6 +18,24 @@ TOOL_NAME_PREFIX_LENGTH: Final = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LE PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: Final = frozenset({"litellm_proxy"}) +def _optional_attr(source: object, name: str) -> object: + return getattr(source, name, None) + + +def _as_string_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +def _thought_signature(provider_specific_fields: object) -> str | None: + fields: Final = _as_string_mapping(provider_specific_fields) + if fields is None: + return None + signature: Final = fields.get("thought_signature") + return signature if isinstance(signature, str) else None + + def truncate_tool_name(name: str) -> str: """ Truncate tool names that exceed OpenAI's 64-character limit. @@ -40,7 +58,7 @@ def truncate_tool_name(name: str) -> str: def create_tool_name_mapping( - tools: list[dict[str, Any]], + tools: Sequence[Mapping[str, object]], ) -> dict[str, str]: """ Create a mapping of truncated tool names to original names. @@ -54,6 +72,8 @@ def create_tool_name_mapping( mapping: Final[dict[str, str]] = {} for tool in tools: original_name = tool.get("name", "") + if not isinstance(original_name, str): + continue truncated_name = truncate_tool_name(original_name) if truncated_name != original_name: mapping[truncated_name] = original_name @@ -263,44 +283,44 @@ class LiteLLMAnthropicMessagesAdapter: ### FOR [BETA] `/v1/messages` endpoint support - def _extract_signature_from_tool_call(self, tool_call: Any) -> str | None: + def _extract_signature_from_tool_call(self, tool_call: object) -> str | None: """ Extract signature from a tool call's provider_specific_fields. Only checks provider_specific_fields, not thinking blocks. """ - signature = None + fields: Final = _optional_attr(tool_call, "provider_specific_fields") + if fields: + return _thought_signature(fields) - if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: - if "thought_signature" in tool_call.provider_specific_fields: - signature = tool_call.provider_specific_fields["thought_signature"] - elif hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields: - if "thought_signature" in tool_call.function.provider_specific_fields: - signature = tool_call.function.provider_specific_fields["thought_signature"] + function_fields: Final = _optional_attr(_optional_attr(tool_call, "function"), "provider_specific_fields") + if function_fields: + return _thought_signature(function_fields) - return signature + return None - def _extract_signature_from_tool_use_content(self, content: dict[str, Any]) -> str | None: + def _extract_signature_from_tool_use_content(self, content: Mapping[str, object]) -> str | None: """ Extract signature from a tool_use content block's provider_specific_fields. """ - provider_specific_fields: Final = content.get("provider_specific_fields", {}) + provider_specific_fields: Final = _as_string_mapping(content.get("provider_specific_fields", {})) if provider_specific_fields: - return provider_specific_fields.get("signature") + signature: Final = provider_specific_fields.get("signature") + return signature if isinstance(signature, str) else None return None def _add_cache_control_if_applicable( self, - source: Any, - target: Any, + source: object, + target: object, model: str | None, ) -> None: """ Extract cache_control from source and add to target if it should be preserved. - This method accepts Any type to support both regular dicts and TypedDict objects. - TypedDict objects (like ChatCompletionTextObject, ChatCompletionImageObject, etc.) - are dicts at runtime but have specific types at type-check time. Using Any allows - this method to work with both while maintaining runtime correctness. + This method accepts an unconstrained type to support both regular dicts and + TypedDict objects. TypedDict objects (like ChatCompletionTextObject, + ChatCompletionImageObject, etc.) are dicts at runtime but have specific types at + type-check time, so the widest parameter type works with both. Args: source: Dict or TypedDict containing potential cache_control field @@ -801,7 +821,7 @@ class LiteLLMAnthropicMessagesAdapter: return new_tools, tool_name_mapping - def translate_anthropic_output_format_to_openai(self, output_format: Any) -> dict[str, object] | None: + def translate_anthropic_output_format_to_openai(self, output_format: object) -> dict[str, object] | None: """ Translate Anthropic's output_format to OpenAI's response_format. @@ -1326,7 +1346,7 @@ class LiteLLMAnthropicMessagesAdapter: @classmethod def _first_positive_prompt_tokens_detail_value(cls, usage: Usage, field_names: tuple[str, ...]) -> int: - prompt_tokens_details: Final = getattr(usage, "prompt_tokens_details", None) + prompt_tokens_details: Final = _optional_attr(usage, "prompt_tokens_details") if prompt_tokens_details is None: return 0 @@ -1334,7 +1354,7 @@ class LiteLLMAnthropicMessagesAdapter: if isinstance(prompt_tokens_details, dict): value = cls._positive_int(prompt_tokens_details.get(field_name)) else: - value = cls._positive_int(getattr(prompt_tokens_details, field_name, None)) + value = cls._positive_int(_optional_attr(prompt_tokens_details, field_name)) if value > 0: return value return 0 diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index c8cbbba8784..4551ff5213f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -14,7 +14,18 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: import re from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast +from typing import ( + TYPE_CHECKING, + Final, + Literal, + NotRequired, + Optional, + Protocol, + TypedDict, + Union, + cast, + runtime_checkable, +) from typing_extensions import ReadOnly @@ -159,11 +170,11 @@ async def _check_summary_model_access( return True key_models: Final = list(getattr(user_api_key_auth, "models", None) or []) - team_id: Final = getattr(user_api_key_auth, "team_id", None) - team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) + team_model_aliases: Final[dict[str, str] | None] = getattr(user_api_key_auth, "team_model_aliases", None) team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or []) - user_id: Final = getattr(user_api_key_auth, "user_id", None) - project_id: Final = getattr(user_api_key_auth, "project_id", None) + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) + project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None) checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = ( ("key", key_models), @@ -371,8 +382,10 @@ async def _check_summary_model_budget( ) return False - end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) - end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None) + end_user_model_max_budget: Final[dict[str, object] | None] = getattr( + user_api_key_auth, "end_user_model_max_budget", None + ) + end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: try: await model_max_budget_limiter.is_end_user_within_model_budget( @@ -490,7 +503,7 @@ def _find_latest_compaction_index( def _slice_around_compaction_block( - messages: list[dict[str, Any]], + messages: list[dict[str, object]], ) -> tuple[list[dict[str, object]], dict[str, object] | None]: """Apply Anthropic's "drop everything before the compaction block" rule. @@ -505,7 +518,8 @@ def _slice_around_compaction_block( return messages, None original_msg: Final = messages[msg_idx] - original_content: Final = original_msg["content"] + raw_content: Final = original_msg.get("content") + original_content: Final[list[object]] = raw_content if isinstance(raw_content, list) else [] compaction_block: Final = cast(dict[str, object], original_content[blk_idx]) # Per Anthropic's contract everything before the compaction block is @@ -760,7 +774,7 @@ def _extract_summary_text(raw: str | None) -> str | None: def _system_to_openai_message( - system: str | list[dict[str, Any]] | None, + system: str | list[dict[str, object]] | None, ) -> dict[str, object] | None: """Translate Anthropic-shaped ``system`` to an OpenAI system message. @@ -772,8 +786,10 @@ def _system_to_openai_message( if isinstance(system, str): return {"role": "system", "content": system} if system else None if isinstance(system, list): - parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"] - joined: Final = "\n\n".join(part for part in parts if part) + parts: Final[list[object]] = [ + block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text" + ] + joined: Final = "\n\n".join(part for part in parts if isinstance(part, str) and part) return {"role": "system", "content": joined} if joined else None return None @@ -873,7 +889,7 @@ async def _call_summary_model( summary_model: str, summary_messages: list[dict[str, object]], metadata: Mapping[str, object], - llm_router: Any, + llm_router: Optional["Router"], allowed_model_region: str | None = None, max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS, ) -> Union["ModelResponse", "CustomStreamWrapper"]: @@ -927,11 +943,17 @@ async def _call_summary_model( return await litellm.acompletion(**call_kwargs) -def _extract_response_text(response: Any) -> str | None: +@runtime_checkable +class _ResponseWithChoices(Protocol): + choices: Sequence[object] + + +def _extract_response_text(response: object) -> str | None: + if not isinstance(response, _ResponseWithChoices) or not response.choices: + return None try: - choice: Final = response.choices[0] - message: Final = choice.message - content: Final = getattr(message, "content", None) + message: Final[object] = getattr(response.choices[0], "message", None) + content: Final[object] = getattr(message, "content", None) if isinstance(content, str): return content # Some providers return a list of content parts. @@ -946,13 +968,12 @@ def _extract_response_text(response: Any) -> str | None: def _extract_usage(response: object) -> tuple[int, int]: - usage: Final = getattr(response, "usage", None) + usage: Final[object] = getattr(response, "usage", None) if usage is None: return 0, 0 - return ( - int(getattr(usage, "prompt_tokens", 0) or 0), - int(getattr(usage, "completion_tokens", 0) or 0), - ) + prompt_tokens: Final[int | None] = getattr(usage, "prompt_tokens", 0) + completion_tokens: Final[int | None] = getattr(usage, "completion_tokens", 0) + return int(prompt_tokens or 0), int(completion_tokens or 0) def apply_client_compaction_block_history( diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index ace7fc25dc9..0eb0e38a46e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -179,14 +179,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ) @staticmethod - def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, Any]]) -> str: + def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: """Group a run of consecutive thinking blocks together; keep every other block alone.""" index, block = indexed_block return "thinking" if block.get("type") == "thinking" else f"block:{index}" @classmethod def _assistant_group_to_input_item( - cls, group: tuple[Mapping[str, Any], ...] + cls, group: tuple[Mapping[str, object], ...] ) -> dict[str, Any] | None: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") @@ -206,7 +206,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_messages_to_responses_input( self, messages: list[AllAnthropicPassThroughMessageValues], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Convert Anthropic messages list to Responses API `input` items. @@ -220,7 +220,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: assistant thinking -> reasoning assistant tool_use -> function_call """ - input_items: Final[list[dict[str, Any]]] = [] + input_items: Final[list[dict[str, object]]] = [] for m in messages: if m["role"] == "system": @@ -248,7 +248,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: } ) elif isinstance(content, list): - user_parts: list[dict[str, Any]] = [] + user_parts: list[Mapping[str, object]] = [] tool_image_parts: list[dict[str, Any]] = [] # mutable-ok: json content parts for block in content: if not isinstance(block, dict): @@ -379,9 +379,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_tools_to_responses_api( self, tools: list[AllAnthropicToolsValues], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """Convert Anthropic tool definitions to Responses API function tools.""" - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in tools: tool_dict = cast(dict[str, Any], tool) tool_type = tool_dict.get("type", "") @@ -392,7 +392,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: continue # Responses turns strict mode on when `strict` is omitted, silently rewriting # `required` to every property. Anthropic tools are non-strict unless asked. - func_tool: dict[str, Any] = { + func_tool: dict[str, object] = { "type": "function", "name": tool_name, "strict": bool(tool_dict.get("strict")), @@ -407,7 +407,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_tool_choice_to_responses_api( tool_choice: AnthropicMessagesToolChoice, - ) -> str | dict[str, Any]: + ) -> str | dict[str, object]: """Convert Anthropic tool_choice to Responses API tool_choice.""" tc_type: Final = tool_choice.get("type") if tc_type == "any": @@ -420,8 +420,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_context_management_to_responses_api( - context_management: dict[str, Any], - ) -> list[dict[str, Any]] | None: + context_management: dict[str, object], + ) -> list[dict[str, object]] | None: """ Convert Anthropic context_management dict to OpenAI Responses API array format. @@ -435,13 +435,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if not isinstance(edits, list): return None - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for edit in edits: if not isinstance(edit, dict): continue edit_type = edit.get("type", "") if edit_type == "compact_20260112": - entry: dict[str, Any] = {"type": "compaction"} + entry: dict[str, object] = {"type": "compaction"} trigger = edit.get("trigger") if isinstance(trigger, dict) and trigger.get("value") is not None: entry["compact_threshold"] = int(trigger["value"]) @@ -451,9 +451,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_thinking_to_reasoning( - thinking: dict[str, Any], - output_config: dict[str, Any] | None = None, - ) -> dict[str, Any] | None: + thinking: dict[str, object], + output_config: dict[str, object] | None = None, + ) -> dict[str, object] | None: """ Convert Anthropic thinking param to Responses API reasoning param. @@ -473,12 +473,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(output_config, dict) and output_config.get("effort"): effort = output_config["effort"] elif thinking_type == "enabled": - effort = reasoning_effort_from_thinking_budget(thinking.get("budget_tokens", 0)) + raw_budget: Final = thinking.get("budget_tokens", 0) + budget_tokens: Final = int(raw_budget) if isinstance(raw_budget, (int, float)) else 0 + effort = reasoning_effort_from_thinking_budget(budget_tokens) else: return None auto_summary: Final = is_reasoning_auto_summary_enabled() - result: Final[dict[str, Any]] = {"effort": effort} + result: Final[dict[str, object]] = {"effort": effort} summary: Final = thinking.get("summary") if summary: result["summary"] = summary @@ -570,7 +572,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: # output_format / output_config.format -> text format # output_format: {"type": "json_schema", "schema": {...}} # output_config: {"format": {"type": "json_schema", "schema": {...}}} - output_format: Any = anthropic_request.get("output_format") + output_format: object = anthropic_request.get("output_format") output_config = anthropic_request.get("output_config") if not isinstance(output_format, dict) and isinstance(output_config, dict): output_format = output_config.get("format") @@ -620,7 +622,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ResponseReasoningItem, ) - content: Final[list[dict[str, Any]]] = [] + content: Final[list[dict[str, object]]] = [] stop_reason: AnthropicFinishReason = "end_turn" for item in response.output: diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2a59eddf88a..4edc7fff1dc 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -5,7 +5,8 @@ import base64 import json from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar, cast +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Generic, Protocol, TypeVar, cast, runtime_checkable from litellm import verbose_logger from litellm.llms.base_llm.managed_resources.isolation import ( @@ -37,6 +38,30 @@ else: ResourceObjectType = TypeVar("ResourceObjectType") +@runtime_checkable +class _HasIdentifier(Protocol): + id: str + + +class _ManagedResourceRecord(Protocol[ResourceObjectType]): + unified_resource_id: str + resource_object: ResourceObjectType + + def model_dump(self) -> dict[str, object]: ... + + +class _ManagedResourceTable(Protocol[ResourceObjectType]): + async def create(self, *, data: Mapping[str, object]) -> object: ... + + async def find_first(self, *, where: Mapping[str, object]) -> _ManagedResourceRecord[ResourceObjectType] | None: ... + + async def find_many( + self, *, where: Mapping[str, object], take: int, order: Mapping[str, str] + ) -> list[_ManagedResourceRecord[ResourceObjectType]]: ... + + async def delete(self, *, where: Mapping[str, object]) -> object: ... + + class BaseManagedResource(ABC, Generic[ResourceObjectType]): """ Base class for managing resources with target_model_names support. @@ -63,6 +88,9 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client + def _resource_table(self) -> _ManagedResourceTable[ResourceObjectType]: + return getattr(self.prisma_client.db, self.table_name) + # ============================================================================ # ABSTRACT METHODS # ============================================================================ @@ -136,7 +164,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): litellm_parent_otel_span: Span | None, model_mappings: dict[str, str], user_api_key_dict: UserAPIKeyAuth, - additional_db_fields: dict[str, Any] | None = None, + additional_db_fields: Mapping[str, object] | None = None, ) -> None: """ Store unified resource ID with model mappings in cache and database. @@ -152,7 +180,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): verbose_logger.info("Storing LiteLLM Managed %s with id=%s in cache", self.resource_type, unified_resource_id) # Prepare cache data - cache_data: Final = { + cache_data: Final[dict[str, object]] = { "unified_resource_id": unified_resource_id, "resource_object": resource_object, "model_mappings": model_mappings, @@ -175,7 +203,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) # Prepare database data - db_data: Final = { + db_data: Final[dict[str, object]] = { "unified_resource_id": unified_resource_id, "model_mappings": json.dumps(model_mappings), "flat_model_resource_ids": list(model_mappings.values()), @@ -204,7 +232,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): db_data.update(additional_db_fields) # Store in database - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() result: Final = await table.create(data=db_data) verbose_logger.debug( @@ -239,7 +267,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): return result # Check database - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() db_object: Final = await table.find_first(where={"unified_resource_id": unified_resource_id}) if db_object: @@ -263,7 +291,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): The deleted resource object or None if not found """ # Get old value from database - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() initial_value: Final = await table.find_first(where={"unified_resource_id": unified_resource_id}) if initial_value is None: @@ -514,7 +542,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): user_api_key_dict: UserAPIKeyAuth, limit: int | None = None, after: str | None = None, - additional_filters: dict[str, Any] | None = None, + additional_filters: Mapping[str, object] | None = None, ) -> dict[str, Any]: """ List resources created by a user. @@ -532,7 +560,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): if owner_filter is None: return build_list_page([]) - where_clause: Final[dict[str, Any]] = {**owner_filter} + where_clause: Final[dict[str, object]] = {**owner_filter} if after: where_clause["id"] = {"gt": after} @@ -543,14 +571,14 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Fetch resources fetch_limit: Final = limit or 20 - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() resources: Final = await table.find_many( where=where_clause, take=fetch_limit, order={"created_at": "desc"}, ) - resource_objects: Final[list[Any]] = [] + resource_objects: Final[list[object]] = [] for resource in resources: try: # Stop once we have enough @@ -558,12 +586,13 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): break # Parse resource object - resource_data = resource.resource_object - if isinstance(resource_data, str): - resource_data = json.loads(resource_data) + stored_resource = resource.resource_object + resource_data: object = ( + json.loads(stored_resource) if isinstance(stored_resource, str) else stored_resource + ) # Set unified ID - if hasattr(resource_data, "id"): + if isinstance(resource_data, _HasIdentifier): resource_data.id = resource.unified_resource_id elif isinstance(resource_data, dict): resource_data["id"] = resource.unified_resource_id diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index bd2b124605c..78e6e6aaf82 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -2,7 +2,7 @@ import base64 import datetime import json import math -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -128,24 +128,35 @@ def is_gemini_image_model(model: str) -> bool: return "gemini" in base_model +def _parse_image_config_string(raw_image_config: str, model: str) -> object: + try: + return json.loads(raw_image_config) + except json.JSONDecodeError as exc: + raise litellm.UnsupportedParamsError( + model=model, + message="`imageConfig` must be valid JSON when provided as a string.", + ) from exc + + def map_openai_image_params_to_gemini( - params: dict[str, Any], + params: Mapping[str, object], model: str, supported_params: Sequence[str], - optional_params: dict[str, Any] | None = None, + optional_params: Mapping[str, object] | None = None, parse_image_config_string: bool = False, -) -> dict[str, Any]: - optional_params = optional_params or {} +) -> dict[str, object]: + already_mapped: Final[Mapping[str, object]] = optional_params or {} filtered_params: Final = {key: value for key, value in params.items() if key in supported_params} - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} - if "n" in filtered_params and "n" not in optional_params: + if "n" in filtered_params and "n" not in already_mapped: mapped_params["sampleCount"] = filtered_params["n"] - if "size" in filtered_params and "size" not in optional_params: + size_param: Final = filtered_params.get("size") + if isinstance(size_param, str) and "size" not in already_mapped: image_config: Final = map_openai_size_to_gemini_image_config( - filtered_params["size"], + size_param, model, ) if image_config is not None: @@ -156,33 +167,30 @@ def map_openai_image_params_to_gemini( if "imageSize" in image_config: mapped_params["imageSize"] = image_config["imageSize"] - image_config_param = filtered_params.get("imageConfig") - if isinstance(image_config_param, str) and parse_image_config_string: - try: - image_config_param = json.loads(image_config_param) - except json.JSONDecodeError as exc: - raise litellm.UnsupportedParamsError( - model=model, - message="`imageConfig` must be valid JSON when provided as a string.", - ) from exc + raw_image_config: Final = filtered_params.get("imageConfig") + image_config_param: Final[object] = ( + _parse_image_config_string(raw_image_config, model) + if isinstance(raw_image_config, str) and parse_image_config_string + else raw_image_config + ) if isinstance(image_config_param, dict): mapped_params["imageConfig"] = image_config_param for key, value in filtered_params.items(): - if key not in ("n", "size", "imageConfig", "tools", "web_search_options") and key not in optional_params: + if key not in ("n", "size", "imageConfig", "tools", "web_search_options") and key not in already_mapped: mapped_params[key] = value return mapped_params -def _dedupe_gemini_search_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: +def _dedupe_gemini_search_tools(tools: list[dict[str, object]]) -> list[dict[str, object]]: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) search_tool_keys: Final = VertexGeminiConfig._search_tool_keys() seen_search_keys: Final[set[str]] = set() - deduped_tools: Final[list[dict[str, Any]]] = [] + deduped_tools: Final[list[dict[str, object]]] = [] for tool in tools: if not isinstance(tool, dict): @@ -203,7 +211,7 @@ def _dedupe_gemini_search_tools(tools: list[dict[str, Any]]) -> list[dict[str, A return deduped_tools -def _has_gemini_search_tool(tools: list[Any]) -> bool: +def _has_gemini_search_tool(tools: list[object]) -> bool: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -213,9 +221,9 @@ def _has_gemini_search_tool(tools: list[Any]) -> bool: def map_gemini_image_tools_params( - non_default_params: dict[str, Any], - mapped_params: dict[str, Any], -) -> dict[str, Any]: + non_default_params: Mapping[str, object], + mapped_params: Mapping[str, object], +) -> dict[str, object]: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -239,21 +247,24 @@ def map_gemini_image_tools_params( gemini_config._drop_search_tools_mixed_with_functions(result) - if isinstance(result.get("tools"), list): - result["tools"] = _dedupe_gemini_search_tools(result["tools"]) + resolved_tools: Final = result.get("tools") + if isinstance(resolved_tools, list): + result["tools"] = _dedupe_gemini_search_tools(resolved_tools) return result def get_gemini_image_web_search_requests( - response_data: dict[str, Any], + response_data: Mapping[str, object], ) -> int | None: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) - grounding_metadata: Final[list[dict[str, Any]]] = [] - for candidate in response_data.get("candidates", []): + raw_candidates: Final = response_data.get("candidates") + candidates: Final[list[object]] = raw_candidates if isinstance(raw_candidates, list) else [] + grounding_metadata: Final[list[dict[str, object]]] = [] + for candidate in candidates: if not isinstance(candidate, dict): continue candidate_grounding = candidate.get("groundingMetadata") @@ -267,13 +278,14 @@ def get_gemini_image_web_search_requests( def get_gemini_image_generation_config( model: str, - optional_params: dict[str, Any], -) -> dict[str, Any]: - generation_config: Final[dict[str, Any]] = {"response_modalities": ["IMAGE", "TEXT"]} + optional_params: Mapping[str, object], +) -> dict[str, object]: + generation_config: Final[dict[str, object]] = {"response_modalities": ["IMAGE", "TEXT"]} - image_config: Final[dict[str, Any]] = {} - if isinstance(optional_params.get("imageConfig"), dict): - image_config.update(optional_params["imageConfig"]) + raw_image_config: Final = optional_params.get("imageConfig") + image_config: Final[dict[str, object]] = {} + if isinstance(raw_image_config, dict): + image_config.update(raw_image_config) if not supports_gemini_image_size(model): image_config.pop("imageSize", None) @@ -398,7 +410,7 @@ class GeminiModelInfo(BaseLLMModelInfo): f"Failed to fetch models from Gemini. Status code: {response.status_code}, Response: {response.json()}" ) - models: Final = response.json()["models"] + models: Final[list[dict[str, str]]] = response.json()["models"] litellm_model_names: Final = self.process_model_name(models) return litellm_model_names @@ -473,12 +485,12 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, object]] | None = None, + system: object | None = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index dee83407cb5..2c62e04c5a3 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -5,11 +5,13 @@ For vertex ai, check out the vertex_ai/files/handler.py file. """ import time -from typing import Any, Final, Literal +from collections.abc import Mapping +from typing import Final, Literal, TypedDict from urllib.parse import urlparse import httpx from openai.types.file_deleted import FileDeleted +from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data @@ -18,7 +20,6 @@ from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) -from litellm.types.llms.gemini import GeminiCreateFilesResponseObject from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -31,6 +32,25 @@ from litellm.types.utils import LlmProviders from ..common_utils import GeminiModelInfo +class _GeminiFileMetadata(TypedDict, total=False): + name: ReadOnly[str] + uri: ReadOnly[Required[str]] + displayName: ReadOnly[Required[str]] + mimeType: ReadOnly[str] + sizeBytes: ReadOnly[Required[str]] + createTime: ReadOnly[Required[str]] + updateTime: ReadOnly[str] + expirationTime: ReadOnly[str] + sha256Hash: ReadOnly[str] + state: ReadOnly[str] + source: ReadOnly[str] + error: ReadOnly[Mapping[str, object]] + + +class _GeminiCreateFileResponse(TypedDict): + file: ReadOnly[_GeminiFileMetadata] + + class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def __init__(self): pass @@ -41,14 +61,14 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def validate_environment( self, - headers: dict[Any, Any], + headers: dict[str, str], model: str, messages: list[AllMessageValues], - optional_params: dict[Any, Any], - litellm_params: dict[Any, Any], + optional_params: dict[str, object], + litellm_params: dict[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict[Any, Any]: + ) -> dict[str, str]: """ Validate environment and add Gemini API key to headers. Google AI Studio uses x-goog-api-key header for authentication. @@ -164,9 +184,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): Transform Gemini's file upload response into OpenAI-style FileObject """ try: - response_json: Final = raw_response.json() + response_json: Final[_GeminiCreateFileResponse] = raw_response.json() - response_object: Final = GeminiCreateFilesResponseObject(**response_json.get("file", {})) + response_object: Final = response_json["file"] # Extract file information from Gemini response @@ -262,7 +282,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ try: verbose_logger.debug("Retrieve file response: %s", raw_response.text) - response_json: Final = raw_response.json() + response_json: Final[_GeminiFileMetadata] = raw_response.json() verbose_logger.debug("Response JSON: %s", response_json) # Map Gemini state to OpenAI status gemini_state: Final = response_json.get("state", "STATE_UNSPECIFIED") diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 51801e91356..0b1dabbef33 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -7,6 +7,8 @@ from collections import OrderedDict from collections.abc import Mapping from typing import Any, Final, cast +from typing_extensions import ReadOnly, Required, TypedDict + import litellm from litellm import verbose_logger from litellm._uuid import uuid @@ -95,6 +97,23 @@ def _gemini_live_speech_config(voice: object) -> Mapping[str, object] | None: return VertexGeminiConfig()._map_audio_params({"voice": voice}) +class _GeminiLiveSetupEnvelope(TypedDict, total=False): + setup: ReadOnly[BidiGenerateContentSetup] + + +class _OpenAIRealtimeClientEvent(TypedDict, total=False): + type: ReadOnly[str] + audio: ReadOnly[Required[str]] + session: ReadOnly[dict[str, object]] + item: ReadOnly[dict[str, object]] + + +def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup: + envelope: Final[_GeminiLiveSetupEnvelope] = json.loads(session_configuration_request) + empty_setup: Final[BidiGenerateContentSetup] = {} + return envelope.get("setup", empty_setup) + + class GeminiRealtimeConfig(BaseRealtimeConfig): _TOOL_CALL_ID_TO_NAME_MAX = 256 # LRU cap for call_id→name mapping @@ -116,7 +135,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return True @staticmethod - def _usage_detail_alias(details: Any, defaults: dict[str, int]) -> dict[str, Any]: + def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: if not isinstance(details, dict): return dict(defaults) return { @@ -125,7 +144,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): } @staticmethod - def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, Any]) -> dict[str, Any]: + def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, Any]) -> dict[str, object]: usage_dict.setdefault( "input_token_details", GeminiRealtimeConfig._usage_detail_alias( @@ -208,8 +227,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if not session_configuration_request: return False try: - setup: Final = json.loads(session_configuration_request).get("setup", {}) - automatic_detection: Final = setup.get("realtimeInputConfig", {}).get("automaticActivityDetection", {}) + setup: Final = _parse_setup(session_configuration_request) + automatic_detection: Final[object] = setup.get("realtimeInputConfig", {}).get( + "automaticActivityDetection", {} + ) return isinstance(automatic_detection, dict) and automatic_detection.get("disabled") is True except (json.JSONDecodeError, TypeError, AttributeError): return False @@ -384,7 +405,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return bool(entry.get("gemini_native_audio") or entry.get("gemini_audio_only_live")) @staticmethod - def _coerce_response_modalities(model: str, modalities: list[Any]) -> list[str]: + def _coerce_response_modalities(model: str, modalities: list[object]) -> list[str]: """Map unsupported TEXT responseModalities to AUDIO for audio-only Live models.""" normalized: Final = [ modality.upper() if isinstance(modality, str) else str(modality).upper() for modality in modalities @@ -409,7 +430,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): def _handle_session_update( self, - json_message: dict, + json_message: _OpenAIRealtimeClientEvent, model: str, session_configuration_request: str | None, ) -> list[str]: @@ -423,7 +444,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): with a 1007, tearing the session down). To carry tools/instructions, send them on the first session.update before any conversation content. """ - session_payload = json_message.get("session") or {} + empty_session: Final[dict[str, object]] = {} + session_payload = json_message.get("session") or empty_session # Normalize GA-remapped fields (``output_modalities``, # nested ``audio.input.transcription``, # ``audio.input.turn_detection``) back to their flat beta keys so @@ -464,14 +486,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Gemini Realtime: Ignoring session.update (setup already sent)") return [] - def _handle_conversation_item(self, json_message: dict) -> list[str]: + def _handle_conversation_item(self, json_message: _OpenAIRealtimeClientEvent) -> list[str]: """ Handle conversation.item.create for user text or function call output. Converts OpenAI format to Gemini's clientContent (for user text) or toolResponse (for function outputs). """ - item: Final = json_message.get("item", {}) + empty_item: Final[dict[str, object]] = {} + item: Final = json_message.get("item", empty_item) item_type: Final = item.get("type") if item_type == "function_call_output": @@ -502,7 +525,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): call_id, ) - function_response: Final[dict[str, Any]] = {"response": output_dict} + function_response: Final[dict[str, object]] = {"response": output_dict} if self._include_function_response_id() and call_id: function_response["id"] = call_id if function_name: @@ -537,7 +560,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) -> list[str]: realtime_input_dict: BidiGenerateContentRealtimeInput = {} try: - json_message: Final = json.loads(message) + json_message: Final[_OpenAIRealtimeClientEvent] = json.loads(message) except json.JSONDecodeError: if isinstance(message, bytes): message_str = message.decode("utf-8", errors="replace") @@ -587,9 +610,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_configuration_request: str | None = None, ) -> OpenAIRealtimeStreamSessionEvents: if session_configuration_request: - session_configuration_request_dict: BidiGenerateContentSetup = json.loads( - session_configuration_request - ).get("setup", {}) + session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request) else: session_configuration_request_dict = {} @@ -640,7 +661,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_configuration_request_dict: BidiGenerateContentSetup = {} if session_configuration_request is not None: try: - session_configuration_request_dict = json.loads(session_configuration_request).get("setup", {}) + session_configuration_request_dict = _parse_setup(session_configuration_request) except json.JSONDecodeError: session_configuration_request_dict = {} generation_config: Final = session_configuration_request_dict.get("generationConfig", {}) @@ -908,9 +929,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return events @staticmethod - def get_nested_value(obj: dict, path: str) -> Any: + def get_nested_value(obj: dict, path: str) -> object | None: keys: Final = path.split(".") - current = obj + current: object = obj for key in keys: if isinstance(current, dict) and key in current: current = current[key] @@ -988,9 +1009,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): current_response_id = f"resp_{uuid.uuid4()}" if session_configuration_request: - session_configuration_request_dict: BidiGenerateContentSetup = json.loads( - session_configuration_request - ).get("setup", {}) + session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request) else: session_configuration_request_dict = {} @@ -1286,7 +1305,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_setup: BidiGenerateContentSetup = {} if session_configuration_request is not None: try: - session_setup = json.loads(session_configuration_request).get("setup", {}) + session_setup = _parse_setup(session_configuration_request) except (json.JSONDecodeError, TypeError): session_setup = {} tool_call_generation_config = session_setup.get("generationConfig", {}) or {} diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index d435994ce20..0b156379d0d 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -13,12 +13,111 @@ Generated files are returned directly in the response - no separate storage need import base64 import json +from collections.abc import Mapping, Sequence from enum import Enum -from typing import Any, Final +from typing import Any, Final, Protocol, TypedDict + +from typing_extensions import ReadOnly from litellm._logging import verbose_logger +class _ToolParameterSchema(TypedDict, total=False): + type: ReadOnly[str] + description: ReadOnly[str] + + +class _ToolArgumentSchema(TypedDict, total=False): + type: ReadOnly[str] + properties: ReadOnly[Mapping[str, _ToolParameterSchema]] + required: ReadOnly[Sequence[str]] + + +class _OpenAIToolFunction(TypedDict, total=False): + name: ReadOnly[str] + description: ReadOnly[str] + parameters: ReadOnly[_ToolArgumentSchema] + + +class _OpenAIToolSpec(TypedDict, total=False): + type: ReadOnly[str] + function: ReadOnly[_OpenAIToolFunction] + + +class _AnthropicToolSpec(TypedDict, total=False): + name: ReadOnly[str] + description: ReadOnly[str] + input_schema: ReadOnly[_ToolArgumentSchema] + + +class _CodeExecutionArguments(TypedDict, total=False): + code: ReadOnly[str] + + +class _GeneratedFile(TypedDict, total=False): + name: ReadOnly[str] + mime_type: ReadOnly[str] + content_base64: ReadOnly[str] + size: ReadOnly[int] + + +class _SandboxGeneratedFile(TypedDict): + name: ReadOnly[str] + mime_type: ReadOnly[str] + content_base64: ReadOnly[str] + + +class _SandboxExecutionResult(TypedDict): + success: ReadOnly[bool] + output: ReadOnly[str] + error: ReadOnly[str] + files: ReadOnly[Sequence[_SandboxGeneratedFile]] + + +class _ExecutionResult(TypedDict, total=False): + iteration: ReadOnly[int] + success: ReadOnly[bool] + output: ReadOnly[str] + error: ReadOnly[str] + files: ReadOnly[Sequence[str]] + + +class _ToolCallFunction(Protocol): + name: str + arguments: str + + +class _ToolCall(Protocol): + id: str + function: _ToolCallFunction + + +class _AssistantMessage(Protocol): + content: str | None + tool_calls: Sequence[_ToolCall] | None + + +class _ResponseChoice(Protocol): + message: _AssistantMessage + finish_reason: str | None + + +class _CompletionResponse(Protocol): + choices: Sequence[_ResponseChoice] + + +class _CodeExecutionOutcome(TypedDict, total=False): + response: ReadOnly[_CompletionResponse | None] + files: ReadOnly[Sequence[_GeneratedFile]] + execution_results: ReadOnly[Sequence[_ExecutionResult]] + messages: ReadOnly[Sequence[dict[str, object]]] + max_iterations_reached: ReadOnly[bool] + + +def _parse_code_execution_arguments(serialized_arguments: str) -> _CodeExecutionArguments: + return json.loads(serialized_arguments) + + class LiteLLMInternalTools(str, Enum): """ Enum for internal LiteLLM tools that are injected into requests. @@ -30,7 +129,7 @@ class LiteLLMInternalTools(str, Enum): CODE_EXECUTION = "litellm_code_execution" -def get_litellm_code_execution_tool() -> dict[str, Any]: +def get_litellm_code_execution_tool() -> _OpenAIToolSpec: """ Returns the litellm_code_execution tool definition in OpenAI format. @@ -51,7 +150,7 @@ def get_litellm_code_execution_tool() -> dict[str, Any]: } -def get_litellm_code_execution_tool_anthropic() -> dict[str, Any]: +def get_litellm_code_execution_tool_anthropic() -> _AnthropicToolSpec: """ Returns the litellm_code_execution tool definition in Anthropic/messages API format. @@ -98,12 +197,12 @@ class CodeExecutionHandler: async def execute_with_code_execution( self, model: str, - messages: list[dict], - tools: list[dict], + messages: list[dict[str, object]], + tools: list[_OpenAIToolSpec], skill_files: dict[str, bytes], skill_id: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> _CodeExecutionOutcome: """ Execute an LLM call with automatic code execution handling. @@ -134,8 +233,8 @@ class CodeExecutionHandler: ) current_messages: Final = list(messages) - generated_files: Final[list[dict[str, Any]]] = [] # Files returned directly - execution_results: Final[list[dict]] = [] + generated_files: Final[list[_GeneratedFile]] = [] # Files returned directly + execution_results: Final[list[_ExecutionResult]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) response: Any = None # Initialize to avoid possibly unbound error @@ -151,11 +250,12 @@ class CodeExecutionHandler: **kwargs, ) - assistant_message = response.choices[0].message - stop_reason = response.choices[0].finish_reason + choice: _ResponseChoice = response.choices[0] + assistant_message = choice.message + stop_reason = choice.finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -190,12 +290,12 @@ class CodeExecutionHandler: if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: # Execute code in sandbox try: - args = json.loads(tool_call.function.arguments) + args = _parse_code_execution_arguments(tool_call.function.arguments) code = args.get("code", "") verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code)) - exec_result = executor.execute( + exec_result: _SandboxExecutionResult = executor.execute( code=code, skill_files=skill_files, ) @@ -278,7 +378,7 @@ class CodeExecutionHandler: } -def has_code_execution_tool(tools: list[dict] | None) -> bool: +def has_code_execution_tool(tools: list[_OpenAIToolSpec] | None) -> bool: """Check if litellm_code_execution tool is in the tools list.""" if not tools: return False @@ -289,7 +389,7 @@ def has_code_execution_tool(tools: list[dict] | None) -> bool: return False -def add_code_execution_tool(tools: list[dict] | None) -> list[dict]: +def add_code_execution_tool(tools: list[_OpenAIToolSpec] | None) -> list[_OpenAIToolSpec]: """Add litellm_code_execution tool if not already present.""" tools = tools or [] if not has_code_execution_tool(tools): diff --git a/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py b/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py index 008a5a5780f..046b4e29a0a 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py +++ b/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py @@ -16,7 +16,7 @@ import io import os import tempfile from dataclasses import dataclass -from typing import Any, Final, cast +from typing import Final, Protocol, cast from litellm.llms.nvidia_riva.audio_transcription.transformation import ( RIVA_TARGET_NUM_CHANNELS, @@ -24,10 +24,30 @@ from litellm.llms.nvidia_riva.audio_transcription.transformation import ( ) from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException -# Keep this as Any: the module intentionally avoids importing numpy at module -# import time (optional dependency), and project-wide mypy config evaluates this -# file in contexts where conditional type aliases can degrade to "FloatArray?". -FloatArray = Any + +class FloatArray(Protocol): + """Structural view of the ``numpy.ndarray`` surface this module relies on.""" + + @property + def ndim(self) -> int: ... + + @property + def shape(self) -> tuple[int, ...]: ... + + @property + def size(self) -> int: ... + + def mean(self, axis: int) -> "FloatArray": ... + + def ravel(self) -> "FloatArray": ... + + def astype(self, dtype: object) -> "FloatArray": ... + + def tobytes(self) -> bytes: ... + + def __getitem__(self, key: object) -> "FloatArray": ... + + def __mul__(self, other: float) -> "FloatArray": ... _INSTALL_HINT = "Install Riva STT extras to enable automatic audio resampling: `pip install 'litellm[stt-nvidia-riva]'`" diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index 5c3962bc05d..3f703564b5a 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -5,10 +5,11 @@ import os import re from dataclasses import dataclass from email.utils import formatdate -from typing import Any, Final, Protocol +from typing import Final, Protocol from urllib.parse import urlparse import httpx +from pydantic import JsonValue from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -64,7 +65,7 @@ class OCISignerProtocol(Protocol): See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/signing.html """ - def do_request_sign(self, request: Any, *, enforce_content_headers: bool = False) -> None: + def do_request_sign(self, request: "OCIRequestWrapper", *, enforce_content_headers: bool = False) -> None: pass @@ -113,7 +114,7 @@ def build_signature_string(method: str, path: str, headers: dict, signed_headers return "\n".join(lines) -def load_private_key_from_str(key_str: str) -> Any: +def load_private_key_from_str(key_str: str) -> "rsa.RSAPrivateKey": _require_cryptography() key: Final = serialization.load_pem_private_key( key_str.encode("utf-8"), @@ -124,7 +125,7 @@ def load_private_key_from_str(key_str: str) -> Any: return key -def load_private_key_from_file(file_path: str) -> Any: +def load_private_key_from_file(file_path: str) -> "rsa.RSAPrivateKey": """Loads a private key from a file path.""" try: with open(file_path, "r", encoding="utf-8") as f: @@ -421,16 +422,17 @@ OCI_JSON_TO_PYTHON_TYPES: Final[dict[str, str]] = { } -def resolve_oci_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: +def resolve_oci_schema_refs(schema: JsonValue) -> JsonValue: """Inline all ``$ref``/``$defs`` references — OCI does not support JSON Schema ``$ref``.""" - defs: Final = schema.get("$defs", {}) - resolving_stack: Final[set] = set() + raw_defs: Final = schema.get("$defs") if isinstance(schema, dict) else None + defs: Final[dict[str, JsonValue]] = raw_defs if isinstance(raw_defs, dict) else {} + resolving_stack: Final[set[str]] = set() - def _resolve(obj: Any) -> Any: + def _resolve(obj: JsonValue) -> JsonValue: if isinstance(obj, dict): - if "$ref" in obj: - ref: Final = obj["$ref"] - if ref.startswith("#/$defs/"): + ref: Final = obj.get("$ref") + if ref is not None: + if isinstance(ref, str) and ref.startswith("#/$defs/"): key: Final = ref.split("/")[-1] if key in resolving_stack: return {"type": "object"} # break cycles @@ -451,7 +453,7 @@ def resolve_oci_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: return resolved -def resolve_oci_schema_anyof(obj: Any) -> Any: +def resolve_oci_schema_anyof(obj: JsonValue) -> JsonValue: """Resolve Pydantic v2 ``Optional[T]`` → ``anyOf`` patterns. Pydantic v2 emits ``{"anyOf": [{"type": "T"}, {"type": "null"}]}`` for @@ -459,10 +461,13 @@ def resolve_oci_schema_anyof(obj: Any) -> Any: first non-null branch and merge top-level metadata into it. """ if isinstance(obj, dict): - if "anyOf" in obj and "type" not in obj: - non_null: Final = [t for t in obj["anyOf"] if not (isinstance(t, dict) and t.get("type") == "null")] + raw_any_of: Final = obj.get("anyOf") + if raw_any_of is not None and "type" not in obj: + branches: Final = raw_any_of if isinstance(raw_any_of, list) else [] + non_null: Final = [t for t in branches if not (isinstance(t, dict) and t.get("type") == "null")] if non_null: - resolved: Final = {**obj, **non_null[0]} + first: Final = non_null[0] + resolved: Final[dict[str, JsonValue]] = {**obj, **first} if isinstance(first, dict) else {**obj} resolved.pop("anyOf", None) return resolve_oci_schema_anyof(resolved) return {k: resolve_oci_schema_anyof(v) for k, v in obj.items()} @@ -471,7 +476,7 @@ def resolve_oci_schema_anyof(obj: Any) -> Any: return obj -def sanitize_oci_schema(schema: Any) -> Any: +def sanitize_oci_schema(schema: JsonValue) -> JsonValue: """Recursively remove OCI-incompatible fields from a JSON schema. Strips ``title`` keys, removes ``None``-valued ``default`` entries, @@ -483,7 +488,7 @@ def sanitize_oci_schema(schema: Any) -> Any: if not isinstance(schema, dict): return schema - sanitized: Final[dict[str, Any]] = {} + sanitized: Final[dict[str, JsonValue]] = {} for key, value in schema.items(): if key == "title": continue @@ -513,7 +518,7 @@ def sanitize_oci_schema(schema: Any) -> Any: return sanitized -def enrich_cohere_param_description(description: str, param_schema: dict[str, Any]) -> str: +def enrich_cohere_param_description(description: str, param_schema: dict[str, JsonValue]) -> str: """Embed schema constraints into a Cohere parameter description. ``CohereParameterDefinition`` only has ``type``, ``description``, and diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 6fc50458aa3..d1e5e12d1ef 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -111,10 +111,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container creation response.""" - response_data: Final = raw_response.json() - - # Transform the response data - container_obj: Final = ContainerObject(**response_data) + container_obj: Final = ContainerObject.model_validate(raw_response.json()) # Add cost for container creation (OpenAI containers are code interpreter sessions) # https://platform.openai.com/docs/pricing @@ -171,10 +168,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerListResponse: """Transform the OpenAI container list response.""" - response_data: Final = raw_response.json() - - # Transform the response data - container_list: Final = ContainerListResponse(**response_data) + container_list: Final = ContainerListResponse.model_validate(raw_response.json()) return container_list @@ -191,7 +185,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -201,9 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container retrieve response.""" - response_data: Final = raw_response.json() - # Transform the response data - container_obj: Final = ContainerObject(**response_data) + container_obj: Final = ContainerObject.model_validate(raw_response.json()) return container_obj @@ -224,7 +216,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -234,10 +226,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteContainerResult: """Transform the OpenAI container delete response.""" - response_data: Final = raw_response.json() - - # Transform the response data - delete_result: Final = DeleteContainerResult(**response_data) + delete_result: Final = DeleteContainerResult.model_validate(raw_response.json()) return delete_result @@ -262,7 +251,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters - params: Final[dict[str, Any]] = {} + params: Final[dict[str, str]] = {} if after is not None: params["after"] = after if limit is not None: @@ -282,10 +271,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerFileListResponse: """Transform the OpenAI container file list response.""" - response_data: Final = raw_response.json() - - # Transform the response data - file_list: Final = ContainerFileListResponse(**response_data) + file_list: Final = ContainerFileListResponse.model_validate(raw_response.json()) return file_list @@ -308,7 +294,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content") # No query parameters needed - params: Final[dict[str, Any]] = {} + params: Final[dict[str, str]] = {} return url, params diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 1da8f0c66f0..19e6d8ff494 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -6,10 +6,11 @@ Maps OpenAI TTS spec to RunwayML Text-to-Speech API import asyncio import time -from collections.abc import Coroutine -from typing import TYPE_CHECKING, Any, Final, Union +from collections.abc import Coroutine, Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict, Union import httpx +from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger @@ -31,6 +32,14 @@ else: HttpxBinaryResponseContent = Any +class _RunwayTtsTaskResponse(TypedDict, total=False): + id: ReadOnly[str] + status: ReadOnly[str] + output: ReadOnly[Sequence[object]] + failure: ReadOnly[str] + failureCode: ReadOnly[str] + + class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Configuration for RunwayML Text-to-Speech @@ -64,7 +73,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): litellm_params_dict: dict, logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout, - extra_headers: dict[str, Any] | None, + extra_headers: dict[str, object] | None, base_llm_http_handler: Any, aspeech: bool, api_base: str | None, @@ -72,7 +81,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): **kwargs: Any, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Dispatch method to handle RunwayML TTS requests @@ -242,7 +251,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): raise TimeoutError(f"RunwayML TTS task polling timed out after {timeout_secs} seconds") @staticmethod - def _check_task_status(response_data: dict[str, Any]) -> str: + def _check_task_status(response_data: _RunwayTtsTaskResponse) -> str: """ Check RunwayML task status from response. @@ -314,7 +323,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayTtsTaskResponse = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -362,7 +371,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayTtsTaskResponse = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -453,7 +462,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -483,7 +492,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() verbose_logger.debug("RunwayML TTS polling complete, downloading audio") @@ -522,7 +531,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -552,7 +561,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() verbose_logger.debug("RunwayML TTS polling complete (async), downloading audio") diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 1de2337d8eb..48649cf3105 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -31,7 +31,7 @@ class VertexAIError(BaseLLMException): super().__init__(message=message, status_code=status_code, headers=headers) -def redact_vertex_ai_metadata_from_logged_object(obj: Any) -> None: +def redact_vertex_ai_metadata_from_logged_object(obj: object) -> None: if isinstance(obj, dict): for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: if field in obj: @@ -651,7 +651,7 @@ def _build_json_schema(parameters: dict) -> dict: return parameters -def _filter_anyof_fields(schema_dict: dict[str, Any]) -> dict[str, Any]: +def _filter_anyof_fields(schema_dict: dict[str, object]) -> dict[str, object]: """ When anyof is present, only keep the anyof field and its contents - otherwise VertexAI will throw an error - https://github.com/BerriAI/litellm/issues/11164 Filter out other fields in the same dict. @@ -704,7 +704,7 @@ def process_items(schema, depth=0): process_items(item, depth + 1) -def set_schema_property_ordering(schema: dict[str, Any], depth: int = 0) -> dict[str, Any]: +def set_schema_property_ordering(schema: dict[str, object], depth: int = 0) -> dict[str, object]: """ vertex ai and generativeai apis order output of fields alphabetically, unless you specify the order. python dicts retain order, so we just use that. Note that this field only applies to structured outputs, and not tools. @@ -731,7 +731,7 @@ def set_schema_property_ordering(schema: dict[str, Any], depth: int = 0) -> dict return schema -def filter_schema_fields(schema_dict: dict[str, Any], valid_fields: set[str], processed=None) -> dict[str, Any]: +def filter_schema_fields(schema_dict: dict[str, object], valid_fields: set[str], processed=None) -> dict[str, object]: """ Recursively filter a schema dictionary to keep only valid fields. """ @@ -905,7 +905,7 @@ def _convert_schema_types(schema, depth=0): "maxProperties", } - any_of: Final[list[dict[str, Any]]] = [] + any_of: Final[list[dict[str, object]]] = [] for t in type_val: if not isinstance(t, str): continue @@ -916,7 +916,7 @@ def _convert_schema_types(schema, depth=0): # For object/array types, include type-specific fields if t in ("object", "array"): - item_schema = {"type": t} + item_schema: dict[str, object] = {"type": t} # Move type-specific fields into this anyOf item for field in type_specific_fields: if field in schema: @@ -1110,11 +1110,11 @@ class VertexAITokenCounter(BaseTokenCounter): self, model_to_use: str, messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, object]] | None = None, + system: object | None = None, ) -> TokenCountResponse | None: import copy @@ -1131,25 +1131,26 @@ class VertexAITokenCounter(BaseTokenCounter): partner_models_handler: Final = VertexAIPartnerModels() # Extract vertex-specific params from litellm_params - vertex_project = count_tokens_params_request.get("vertex_project") or count_tokens_params_request.get( + partner_litellm_params: Final[dict[str, object]] = count_tokens_params_request + vertex_project = partner_litellm_params.get("vertex_project") or partner_litellm_params.get( "vertex_ai_project" ) - vertex_location = count_tokens_params_request.get("vertex_location") or count_tokens_params_request.get( + vertex_location = partner_litellm_params.get("vertex_location") or partner_litellm_params.get( "vertex_ai_location" ) # Count tokens not available on global location: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens - vertex_location = count_tokens_params_request.get("vertex_count_tokens_location") or vertex_location + vertex_location = partner_litellm_params.get("vertex_count_tokens_location") or vertex_location - vertex_credentials: Final = count_tokens_params_request.get( - "vertex_credentials" - ) or count_tokens_params_request.get("vertex_ai_credentials") + vertex_credentials: Final = partner_litellm_params.get("vertex_credentials") or partner_litellm_params.get( + "vertex_ai_credentials" + ) result = await partner_models_handler.count_tokens( model=model_to_use, messages=messages or [], - litellm_params=count_tokens_params_request, + litellm_params=partner_litellm_params, vertex_project=vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 308813039ca..24ae9b0a311 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast @@ -1210,7 +1210,7 @@ def _deserialize_json_dict(data: str | _StringMap | None) -> dict[str, str] | No return data -def _deserialize_json_list(data: Any) -> list[dict[str, Any]] | None: +def _deserialize_json_list(data: object) -> list[dict[str, Any]] | None: """Deserialize a JSON array stored in the DB (``env_vars`` and friends). Returns ``None`` for empty / null / unparseable input. Accepts strings @@ -1223,7 +1223,7 @@ def _deserialize_json_list(data: Any) -> list[dict[str, Any]] | None: return None if isinstance(data, str): try: - parsed: Final = json.loads(data) + parsed: Final[object] = json.loads(data) except (json.JSONDecodeError, TypeError): return None data = parsed @@ -1918,7 +1918,7 @@ class MCPServerManager: async def load_servers_from_config( self, - mcp_servers_config: dict[str, Any], + mcp_servers_config: dict[str, MCPServerConfig], mcp_aliases: dict[str, str] | None = None, ): """ @@ -3070,7 +3070,7 @@ class MCPServerManager: return {} cache_key: Final = "toolset_perms:" + ",".join(sorted(toolset_ids)) - cached: Final = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final[dict[str, list[str]] | None] = await user_api_key_cache.async_get_cache(key=cache_key) if cached is not None: return cached @@ -5154,7 +5154,7 @@ class MCPServerManager: # Wrapped so the bridge runs inside the task: the caller only holds the task and # gathers it later, so there is no other point that still sees a block here. - async def _run_during_call_hook() -> Mapping[str, Any] | None: + async def _run_during_call_hook() -> Mapping[str, object] | None: try: return await proxy_logging_obj.during_call_hook( user_api_key_dict=user_api_key_auth, @@ -5655,7 +5655,7 @@ class MCPServerManager: async def _gather_openapi_tool_tasks( self, - tasks: list[Any], + tasks: Sequence[Awaitable[object]], proxy_logging_obj: ProxyLogging | None, ) -> CallToolResult: """Await OpenAPI tool tasks and return the tool call result.""" diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py index c09106273e1..f2e9049c19c 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py @@ -34,9 +34,11 @@ validation error. Runs before the first registry load on every boot and is idemp a healed fleet has no null rows and the backfill exits after one query. """ -import json from collections import Counter -from typing import Any, Final, Literal +from collections.abc import Mapping +from typing import Final, Literal, Protocol + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials @@ -53,14 +55,46 @@ BackfillRule = Literal[ ] _BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill" +_CREDENTIALS_JSON: Final = TypeAdapter(dict[str, object]) -def _decrypted_credentials(raw_credentials: Any) -> MCPCredentials | None: +class _MCPServerRow(Protocol): + """The MCP server row fields this backfill reads, narrowing the untyped DB record once here.""" + + server_id: str + authorization_url: str | None + registration_url: str | None + token_url: str | None + credentials: object + + +class _MCPUserCredentialRow(Protocol): + """The per-user credential row fields this backfill reads.""" + + server_id: str + credential_b64: str + + +class _MCPServerTable(Protocol): + """The ``LiteLLM_MCPServerTable`` queries this backfill issues.""" + + async def find_many(self, *, where: Mapping[str, object]) -> list[_MCPServerRow]: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _MCPUserCredentialsTable(Protocol): + """The ``LiteLLM_MCPUserCredentials`` query this backfill issues.""" + + async def find_many(self, *, where: Mapping[str, object]) -> list[_MCPUserCredentialRow]: ... + + +def _decrypted_credentials(raw_credentials: object) -> MCPCredentials | None: if raw_credentials is None: return None if isinstance(raw_credentials, str): try: - parsed = json.loads(raw_credentials) + parsed: object = _CREDENTIALS_JSON.validate_json(raw_credentials) except (ValueError, TypeError): return None else: @@ -92,14 +126,16 @@ def classify_null_flow_row( async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]: """Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable ones, warn on the ambiguous ones, and return counts per rule.""" - null_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpservertable.find_many( + server_table: Final[_MCPServerTable] = prisma_client.db.litellm_mcpservertable + null_rows: Final = await server_table.find_many( where={"auth_type": "oauth2", "oauth2_flow": None}, ) if not null_rows: return {} server_ids: Final = [row.server_id for row in null_rows] - token_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpusercredentials.find_many( + user_credentials_table: Final[_MCPUserCredentialsTable] = prisma_client.db.litellm_mcpusercredentials + token_rows: Final = await user_credentials_table.find_many( where={"server_id": {"in": server_ids}}, ) server_ids_with_oauth_tokens: Final[set[str]] = { @@ -141,7 +177,7 @@ async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[Backfi stamped_flows: Final = {flow for _, (flow, _) in classified if flow is not None} for stamped_flow in stamped_flows: server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow] - await prisma_client.db.litellm_mcpservertable.update_many( + await server_table.update_many( where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None}, data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR}, ) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9b1a6ba5aa7..59aca5d8cfd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -956,7 +956,7 @@ def get_key_model_rpm_limit( # 2. Check model_max_budget if user_api_key_dict.model_max_budget: - model_rpm_limit: Final[dict[str, Any]] = {} + model_rpm_limit: Final[dict[str, int]] = {} for model, budget in user_api_key_dict.model_max_budget.items(): if isinstance(budget, dict) and budget.get("rpm_limit") is not None: model_rpm_limit[model] = budget["rpm_limit"] @@ -999,7 +999,7 @@ def get_key_model_tpm_limit( # 2. Check model_max_budget (iterate per-model like RPM does) if user_api_key_dict.model_max_budget: - model_tpm_limit: Final[dict[str, Any]] = {} + model_tpm_limit: Final[dict[str, int]] = {} for model, budget in user_api_key_dict.model_max_budget.items(): if isinstance(budget, dict) and budget.get("tpm_limit") is not None: model_tpm_limit[model] = budget["tpm_limit"] @@ -1062,7 +1062,7 @@ def _validated_output_token_estimates_per_model(raw: object) -> Mapping[str, int def _estimated_output_tokens_from_metadata( - metadata: Mapping[str, Any] | None, + metadata: Mapping[str, object] | None, model_name: str | None, ) -> int | None: """Resolve the per-model, then global, estimate out of one metadata blob. @@ -1628,7 +1628,7 @@ def _dedupe_model_candidates(candidates: list[str]) -> list[str]: return deduped -def _get_case_insensitive_mapping_value(mapping: Mapping[str, Any] | None, key: str) -> Any: +def _get_case_insensitive_mapping_value(mapping: Mapping[str, object] | None, key: str) -> object: if not mapping: return None if key in mapping: @@ -1732,8 +1732,8 @@ def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | Non def _extract_model_candidates_from_request( request_data: dict, route: str, - request_headers: Mapping[str, Any] | None = None, - request_query_params: Mapping[str, Any] | None = None, + request_headers: Mapping[str, object] | None = None, + request_query_params: Mapping[str, object] | None = None, llm_router: Router | None = None, ) -> list[str]: candidates: Final[list[str]] = [] @@ -1825,8 +1825,8 @@ def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool def get_model_from_request( request_data: dict, route: str, - request_headers: Mapping[str, Any] | None = None, - request_query_params: Mapping[str, Any] | None = None, + request_headers: Mapping[str, object] | None = None, + request_query_params: Mapping[str, object] | None = None, llm_router: Router | None = None, request: Request | None = None, ) -> str | list[str] | None: diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 9379a8577a3..1cc27f4784d 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -2,7 +2,7 @@ import copy import os from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, TypeVar from typing_extensions import assert_never @@ -31,6 +31,8 @@ from litellm.types.utils import ( StandardLoggingPayload, ) +_CallbackMetadataT: Final = TypeVar("_CallbackMetadataT") + _CALLBACK_VAR_MASKER: Final = SensitiveDataMasker() # Compound names that are credential-bearing but don't contain any of the # default sensitive segments (so SensitiveDataMasker won't flag them). @@ -525,7 +527,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( def sanitize_openai_provider_metadata( - metadata: dict[str, Any] | None, + metadata: dict[str, object] | None, ) -> dict[str, str] | None: """ Keep only provider-safe OpenAI metadata entries (string keys -> string values). @@ -533,8 +535,8 @@ def sanitize_openai_provider_metadata( Strips LiteLLM proxy-internal tracking fields that must not be forwarded to OpenAI batch/file APIs. """ - if not metadata: - return metadata + if metadata is None: + return None sanitized: Final[dict[str, str]] = {} for key, value in metadata.items(): if key in LITELLM_PROXY_INTERNAL_METADATA_KEYS: @@ -547,7 +549,7 @@ def sanitize_openai_provider_metadata( key, type(value).__name__, ) - return sanitized or None + return None if metadata and not sanitized else sanitized def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_name: str | None): @@ -644,13 +646,13 @@ def process_callback(_callback: str, callback_type: str, environment_variables: return {"name": _callback, "variables": env_vars_dict, "type": callback_type} -def normalize_callback_names(callbacks: Iterable[Any]) -> list[Any]: +def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object]: if callbacks is None: return [] return [c.lower() if isinstance(c, str) else c for c in callbacks] -def strip_callback_config(metadata: dict[str, Any] | None) -> dict[str, Any] | None: +def strip_callback_config(metadata: dict[str, object] | None) -> dict[str, object] | None: """Return key/team metadata without the slots that carry callback credentials.""" if not isinstance(metadata, dict): return metadata @@ -674,7 +676,9 @@ def decrypt_callback_vars(metadata: Any) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) -def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any: +def _transform_callback_vars( + metadata: _CallbackMetadataT, transform: Callable[[str, object], object] +) -> _CallbackMetadataT: if not isinstance(metadata, dict): return metadata out: Final = copy.deepcopy(metadata) @@ -704,7 +708,7 @@ def is_sensitive_callback_key( return _CALLBACK_VAR_MASKER.is_sensitive_key(key) -def _encrypt_if_plaintext(key: str, value: Any) -> Any: +def _encrypt_if_plaintext(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not is_sensitive_callback_key(key): @@ -725,7 +729,7 @@ def _encrypt_if_plaintext(key: str, value: Any) -> Any: return value -def _decrypt_or_passthrough(key: str, value: Any) -> Any: +def _decrypt_or_passthrough(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not value.startswith(_CALLBACK_VAR_ENCRYPTED_PREFIX): diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index bc7b80801fe..2a20e7b07ce 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -1,8 +1,12 @@ from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Final, TypeAlias, Union from litellm._logging import verbose_proxy_logger +JsonValue: TypeAlias = Union["JsonObject", "JsonArray", str, int, float, bool, None] +JsonObject: TypeAlias = dict[str, JsonValue] +JsonArray: TypeAlias = list[JsonValue] + class CustomOpenAPISpec: """ @@ -27,7 +31,20 @@ class CustomOpenAPISpec: RESPONSES_API_PATHS = ["/v1/responses", "/responses"] @staticmethod - def get_pydantic_schema(model_class) -> Mapping[str, object] | None: + def _as_object(node: JsonValue) -> JsonObject: + return node if isinstance(node, dict) else {} + + @staticmethod + def _as_array(node: JsonValue) -> JsonArray: + return node if isinstance(node, list) else [] + + @staticmethod + def _components_schemas(openapi_schema: JsonObject) -> JsonObject: + components: Final = CustomOpenAPISpec._as_object(openapi_schema.setdefault("components", {})) + return CustomOpenAPISpec._as_object(components.setdefault("schemas", {})) + + @staticmethod + def get_pydantic_schema(model_class) -> JsonObject | None: """ Get JSON schema from a Pydantic model, handling both v1 and v2 APIs. @@ -54,9 +71,7 @@ class CustomOpenAPISpec: return None @staticmethod - def add_schema_to_components( - openapi_schema: dict[str, Any], schema_name: str, schema_def: Mapping[str, object] - ) -> None: + def add_schema_to_components(openapi_schema: JsonObject, schema_name: str, schema_def: JsonObject) -> None: """ Add a schema definition to the OpenAPI components/schemas section. @@ -66,16 +81,25 @@ class CustomOpenAPISpec: schema_def: The schema definition """ # Ensure components/schemas structure exists - if "components" not in openapi_schema: - openapi_schema["components"] = {} - if "schemas" not in openapi_schema["components"]: - openapi_schema["components"]["schemas"] = {} + _ = CustomOpenAPISpec._components_schemas(openapi_schema) # Add the schema CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def}) @staticmethod - def add_request_body_to_paths(openapi_schema: dict[str, Any], paths: Sequence[str], schema_ref: str) -> None: + def _expanded_request_field(field_name: str, field_def: JsonValue) -> JsonValue: + expanded: Final = CustomOpenAPISpec._rewrite_defs_refs( + CustomOpenAPISpec._expand_field_definition(CustomOpenAPISpec._as_object(field_def)) + ) + if field_name != "messages": + return expanded + return { + **CustomOpenAPISpec._as_object(expanded), + "example": [{"role": "user", "content": "Hello, how are you?"}], + } + + @staticmethod + def add_request_body_to_paths(openapi_schema: JsonObject, paths: Sequence[str], schema_ref: str) -> None: """ Add request body with expanded form fields for better Swagger UI display. This keeps the request body but expands it to show individual fields in the UI. @@ -86,54 +110,58 @@ class CustomOpenAPISpec: schema_ref: Reference to the schema component (e.g., "#/components/schemas/ModelName") """ for path in paths: - if path in openapi_schema.get("paths", {}) and "post" in openapi_schema["paths"][path]: - # Get the actual schema to extract ALL field definitions - schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref - actual_schema = openapi_schema.get("components", {}).get("schemas", {}).get(schema_name, {}) - schema_properties = actual_schema.get("properties", {}) - required_fields = actual_schema.get("required", []) + path_item = CustomOpenAPISpec._as_object( + CustomOpenAPISpec._as_object(openapi_schema.get("paths")).get(path) + ) + if "post" not in path_item: + continue - # Extract $defs and add them to components/schemas - # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI - if "$defs" in actual_schema: - CustomOpenAPISpec._move_defs_to_components(openapi_schema, actual_schema["$defs"]) + post_operation = CustomOpenAPISpec._as_object(path_item["post"]) - # Create an expanded inline schema instead of just a $ref - # This makes Swagger UI show all individual fields in the request body editor - expanded_schema = { - "type": "object", - "required": required_fields, - "properties": {}, - } + # Get the actual schema to extract ALL field definitions + schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref + components = CustomOpenAPISpec._as_object(openapi_schema.get("components")) + actual_schema = CustomOpenAPISpec._as_object( + CustomOpenAPISpec._as_object(components.get("schemas")).get(schema_name) + ) + schema_properties = CustomOpenAPISpec._as_object(actual_schema.get("properties")) + required_fields = actual_schema.get("required", []) - # Add all properties with their full definitions - for field_name, field_def in schema_properties.items(): - expanded_field = CustomOpenAPISpec._expand_field_definition(field_def) + # Extract $defs and add them to components/schemas + # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI + if "$defs" in actual_schema: + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, CustomOpenAPISpec._as_object(actual_schema["$defs"]) + ) - # Rewrite $defs references to use components/schemas instead - expanded_field = CustomOpenAPISpec._rewrite_defs_refs(expanded_field) + # Create an expanded inline schema instead of just a $ref + # This makes Swagger UI show all individual fields in the request body editor + expanded_schema: JsonObject = { + "type": "object", + "required": required_fields, + "properties": { + field_name: CustomOpenAPISpec._expanded_request_field(field_name, field_def) + for field_name, field_def in schema_properties.items() + }, + } - # Add a simple example for the messages field - if field_name == "messages": - expanded_field["example"] = [{"role": "user", "content": "Hello, how are you?"}] + # Set the request body with the expanded schema + post_operation["requestBody"] = { + "required": True, + "content": {"application/json": {"schema": expanded_schema}}, + } - expanded_schema["properties"][field_name] = expanded_field - - # Set the request body with the expanded schema - openapi_schema["paths"][path]["post"]["requestBody"] = { - "required": True, - "content": {"application/json": {"schema": expanded_schema}}, - } - - # Keep any existing parameters (like path parameters) but remove conflicting query params - if "parameters" in openapi_schema["paths"][path]["post"]: - existing_params = openapi_schema["paths"][path]["post"]["parameters"] - # Only keep path parameters, remove query params that conflict with request body - filtered_params = [param for param in existing_params if param.get("in") == "path"] - openapi_schema["paths"][path]["post"]["parameters"] = filtered_params + # Keep any existing parameters (like path parameters) but remove conflicting query params + if "parameters" in post_operation: + # Only keep path parameters, remove query params that conflict with request body + post_operation["parameters"] = [ + param + for param in CustomOpenAPISpec._as_array(post_operation["parameters"]) + if CustomOpenAPISpec._as_object(param).get("in") == "path" + ] @staticmethod - def _move_defs_to_components(openapi_schema: dict[str, Any], defs: Mapping[str, Mapping[str, Any]]) -> None: + def _move_defs_to_components(openapi_schema: JsonObject, defs: Mapping[str, JsonValue]) -> None: """ Move $defs from Pydantic v2 schema to OpenAPI components/schemas. This makes the definitions resolvable in Swagger/OpenAPI viewers. @@ -146,23 +174,31 @@ class CustomOpenAPISpec: return # Ensure components/schemas exists - if "components" not in openapi_schema: - openapi_schema["components"] = {} - if "schemas" not in openapi_schema["components"]: - openapi_schema["components"]["schemas"] = {} + schemas: Final = CustomOpenAPISpec._components_schemas(openapi_schema) # Add each definition to components/schemas for def_name, def_schema in defs.items(): # Recursively rewrite any nested $defs references within this definition - rewritten_def = CustomOpenAPISpec._rewrite_defs_refs(def_schema) - openapi_schema["components"]["schemas"][def_name] = rewritten_def + schemas[def_name] = CustomOpenAPISpec._rewrite_defs_refs(def_schema) # If this definition also has $defs, process them recursively - if "$defs" in def_schema: - CustomOpenAPISpec._move_defs_to_components(openapi_schema, def_schema["$defs"]) + def_object = CustomOpenAPISpec._as_object(def_schema) + if "$defs" in def_object: + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, CustomOpenAPISpec._as_object(def_object["$defs"]) + ) @staticmethod - def _rewrite_defs_refs(schema: Any) -> Any: + def _rewritten_defs_entry(key: str, value: JsonValue) -> JsonValue: + if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): + # Rewrite the reference to use components/schemas + def_name: Final = value.replace("#/$defs/", "") + return f"#/components/schemas/{def_name}" + # Recursively process nested structures + return CustomOpenAPISpec._rewrite_defs_refs(value) + + @staticmethod + def _rewrite_defs_refs(schema: JsonValue) -> JsonValue: """ Recursively rewrite $ref values from #/$defs/... to #/components/schemas/... This converts Pydantic v2 references to OpenAPI-compatible references. @@ -174,26 +210,17 @@ class CustomOpenAPISpec: Schema with rewritten references """ if isinstance(schema, dict): - result: Final = {} - for key, value in schema.items(): - if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): - # Rewrite the reference to use components/schemas - def_name = value.replace("#/$defs/", "") - result[key] = f"#/components/schemas/{def_name}" - elif key == "$defs": - # Remove $defs from the schema since they're moved to components - continue - else: - # Recursively process nested structures - result[key] = CustomOpenAPISpec._rewrite_defs_refs(value) - return result - elif isinstance(schema, list): + return { + key: CustomOpenAPISpec._rewritten_defs_entry(key, value) + for key, value in schema.items() + if key != "$defs" + } + if isinstance(schema, list): return [CustomOpenAPISpec._rewrite_defs_refs(item) for item in schema] - else: - return schema + return schema @staticmethod - def _extract_field_schema(field_def: dict[str, Any]) -> dict[str, Any]: + def _extract_field_schema(field_def: JsonObject) -> JsonValue: """ Extract a simple schema from a Pydantic field definition for parameter display. @@ -209,10 +236,10 @@ class CustomOpenAPISpec: # Handle anyOf (Optional fields in Pydantic v2) if "anyOf" in field_def: - any_of: Final = field_def["anyOf"] + any_of: Final = CustomOpenAPISpec._as_array(field_def["anyOf"]) # Find the non-null type for option in any_of: - if option.get("type") != "null": + if CustomOpenAPISpec._as_object(option).get("type") != "null": return option # Fallback to string if all else fails return {"type": "string"} @@ -221,7 +248,7 @@ class CustomOpenAPISpec: return {"type": "string"} @staticmethod - def _expand_field_definition(field_def: dict[str, object]) -> dict[str, object]: + def _expand_field_definition(field_def: JsonObject) -> JsonObject: """ Expand a Pydantic field definition for inline use in OpenAPI schema. This creates a full field definition that Swagger UI can render as individual form fields. @@ -237,12 +264,12 @@ class CustomOpenAPISpec: @staticmethod def add_request_schema( - openapi_schema: dict[str, object], + openapi_schema: JsonObject, model_class: type, schema_name: str, paths: Sequence[str], operation_name: str, - ) -> dict[str, object]: + ) -> JsonObject: """ Generic method to add a request schema to OpenAPI specification. @@ -282,8 +309,8 @@ class CustomOpenAPISpec: @staticmethod def add_chat_completion_request_schema( - openapi_schema: dict[str, object], - ) -> dict[str, object]: + openapi_schema: JsonObject, + ) -> JsonObject: """ Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -309,7 +336,7 @@ class CustomOpenAPISpec: return openapi_schema @staticmethod - def add_embedding_request_schema(openapi_schema: dict[str, object]) -> dict[str, object]: + def add_embedding_request_schema(openapi_schema: JsonObject) -> JsonObject: """ Add EmbeddingRequest schema to embedding endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -336,8 +363,8 @@ class CustomOpenAPISpec: @staticmethod def add_responses_api_request_schema( - openapi_schema: dict[str, object], - ) -> dict[str, object]: + openapi_schema: JsonObject, + ) -> JsonObject: """ Add ResponsesAPIRequestParams schema to responses API endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -364,8 +391,8 @@ class CustomOpenAPISpec: @staticmethod def add_llm_api_request_schema_body( - openapi_schema: dict[str, object], - ) -> dict[str, object]: + openapi_schema: JsonObject, + ) -> JsonObject: """ Add LLM API request schema bodies to OpenAPI specification for documentation. @@ -376,12 +403,10 @@ class CustomOpenAPISpec: OpenAPI schema with added request body schemas """ # Add chat completion request schema - openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema) + with_chat_completions: Final = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema) # Add embedding request schema - openapi_schema = CustomOpenAPISpec.add_embedding_request_schema(openapi_schema) + with_embeddings: Final = CustomOpenAPISpec.add_embedding_request_schema(with_chat_completions) # Add responses API request schema - openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema(openapi_schema) - - return openapi_schema + return CustomOpenAPISpec.add_responses_api_request_schema(with_embeddings) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index b8df0105b7b..5820296a3cc 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Final, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload from pydantic import BaseModel @@ -9,6 +9,9 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +if TYPE_CHECKING: + from opentelemetry.trace import Span + T = TypeVar("T", bound=BaseModel) @@ -40,31 +43,32 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... def get_cache( self, - key, - parent_otel_span=None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, - ) -> Any | BaseModel | None: + **kwargs: object, + ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) cached: Final = super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) @@ -85,31 +89,32 @@ class UserApiKeyCache(DualCache): @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... async def async_get_cache( self, - key, - parent_otel_span=None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, - ) -> Any | BaseModel | None: + **kwargs: object, + ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) cached: Final = await super().async_get_cache( @@ -129,14 +134,14 @@ class UserApiKeyCache(DualCache): return None return decoded - def set_cache(self, key, value, local_only: bool = False, **kwargs): + def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs) -> None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index c70a2ee8a74..4d15fe96b64 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -232,7 +232,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` # routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail. - self.checks: dict[str, Any] | None = self._normalize_checks(checks) + self.checks: dict[str, object] | None = self._normalize_checks(checks) # Per-check block thresholds; a score >= threshold blocks. None => the # check is detect-only (logged, never blocks). self.content_filter_threshold = content_filter_threshold @@ -289,7 +289,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ] @staticmethod - def _normalize_checks(checks: BedrockChecksConfigModel | Mapping[str, object] | None) -> dict[str, Any] | None: + def _normalize_checks(checks: BedrockChecksConfigModel | Mapping[str, object] | None) -> dict[str, object] | None: """Normalize the configured `checks` into a plain dict for the API body. Accepts a pydantic ``BedrockChecksConfigModel`` or a raw dict; drops None / @@ -340,7 +340,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _create_bedrock_output_content_request( self, - response: Any | ModelResponse, + response: object, messages: list[AllMessageValues] | None = None, ) -> BedrockRequest: """ @@ -365,7 +365,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return bedrock_request def _build_response_content_items( - self, response: Any | ModelResponse, has_grounding: bool + self, response: object, has_grounding: bool ) -> list[BedrockContentItem]: """Build content item(s) from the model response. When the request supplied grounding, the response is qualified ``guard_content`` so Bedrock can score it. @@ -390,7 +390,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, source: Literal["INPUT", "OUTPUT"], messages: list[AllMessageValues] | None = None, - response: Any | ModelResponse | None = None, + response: object | None = None, ) -> BedrockRequest: """ Convert the litellm messages/response to the bedrock request format. @@ -911,7 +911,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _apply_guardrail_content_with_chunking( self, content: Sequence[BedrockContentItem], - base_request_data: Mapping[str, Any], + base_request_data: Mapping[str, object], credentials: "Credentials", aws_region_name: str, api_key: str | None, @@ -1049,7 +1049,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _post_apply_guardrail_content_with_retry( self, content: Sequence[BedrockContentItem], - base_request_data: Mapping[str, Any], + base_request_data: Mapping[str, object], credentials: "Credentials", aws_region_name: str, api_key: str | None, @@ -1099,7 +1099,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _post_apply_guardrail_content( self, content: Sequence[BedrockContentItem], - base_request_data: Mapping[str, Any], + base_request_data: Mapping[str, object], credentials: "Credentials", aws_region_name: str, api_key: str | None, @@ -1827,7 +1827,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() credentials, aws_region_name = self._load_credentials() - body: Final[dict[str, Any]] = {"messages": checks_messages, "checks": self.checks} + body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} api_key: Final[str | None] = request_data.get("api_key") if request_data else None prepared_request: Final = self._prepare_request( @@ -2309,7 +2309,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): guardrail_name=self.guardrail_name, ) - detail: Final[dict[str, Any]] = { + detail: Final[dict[str, object]] = { "error": "Violated guardrail policy", "bedrock_guardrail_response": bedrock_guardrail_output_text, } @@ -2853,7 +2853,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return updated_messages def _mask_content_list( - self, content_list: list[Any], masked_texts: list[str], masking_index: int + self, content_list: Sequence[object], masked_texts: list[str], masking_index: int ) -> tuple[list[Any], int]: """ Apply masking to a list of content items. @@ -2866,7 +2866,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Returns: Updated content list with masked items """ - new_content: Final[list[dict | str]] = [] + new_content: Final[list[dict[str, object] | str]] = [] for item in content_list: if isinstance(item, dict) and "text" in item: new_item = item.copy() @@ -2885,7 +2885,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _apply_masking_to_response( self, - response: ModelResponse | Any, + response: object, bedrock_guardrail_response: BedrockGuardrailResponse, ) -> None: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py index 8398ec9f141..5a6be1089b6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py @@ -5,8 +5,9 @@ The public guardrail class imports this private mixin from while preserving the existing public import path. """ +from collections.abc import Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException @@ -23,7 +24,7 @@ if TYPE_CHECKING: from .cisco_ai_defense import _ScanContext -def _serialize_mcp_content_item(item: object) -> dict[str, Any]: +def _serialize_mcp_content_item(item: object) -> dict[str, object]: """Serialize an MCP content item to a JSON-friendly dict. Handles raw dicts, MCP SDK Pydantic models, and simple ``.text`` objects. @@ -57,7 +58,7 @@ class _CiscoAIDefenseMcpMixin: def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: ... - async def _post_inspection(self, url: str, payload: dict[str, Any], surface: str) -> dict[str, Any]: ... + async def _post_inspection(self, url: str, payload: dict[str, object], surface: str) -> dict[str, object]: ... def _handle_api_error( self, @@ -67,16 +68,16 @@ class _CiscoAIDefenseMcpMixin: start_time: datetime | None = ..., surface: str = ..., direction: str = ..., - ) -> dict[str, Any]: ... + ) -> dict[str, object]: ... def _finalize_inspection( self, - inspect_response: dict[str, Any], + inspect_response: dict[str, object], request_data: dict, context: "_ScanContext", start_time: datetime, response_obj: object = ..., - ) -> dict[str, Any]: ... + ) -> dict[str, object]: ... # ------------------------------------------------------------------ # MCP post-tool hook (dispatcher contract) @@ -95,7 +96,7 @@ class _CiscoAIDefenseMcpMixin: if self.inspection_type != "mcp": return None - request_data: Final[dict[str, Any]] = {} + request_data: Final[dict[str, object]] = {} for key in ( "name", "litellm_call_id", @@ -188,9 +189,9 @@ class _CiscoAIDefenseMcpMixin: original_hidden: Final = getattr(original_response_obj, "hidden_params", None) if isinstance(original_hidden, HiddenParams): - hidden_params: Any = original_hidden + hidden_params: HiddenParams = original_hidden else: - response_cost: Final = getattr(original_hidden, "response_cost", None) + response_cost: Final[float | None] = getattr(original_hidden, "response_cost", None) hidden_params = HiddenParams(response_cost=response_cost) if response_cost is not None else HiddenParams() return MCPPostCallResponseObject( @@ -200,11 +201,11 @@ class _CiscoAIDefenseMcpMixin: @staticmethod def _replace_mcp_tool_response(response_obj: object, replacement_obj: object) -> bool: - replacement: Final = getattr(replacement_obj, "mcp_tool_call_response", None) + replacement: Final[list[object] | None] = getattr(replacement_obj, "mcp_tool_call_response", None) if replacement is None: return False - inner: Final = getattr(response_obj, "mcp_tool_call_response", None) + inner: Final[object | None] = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: if _CiscoAIDefenseMcpMixin._replace_mcp_tool_response(inner, replacement_obj): return True @@ -276,7 +277,7 @@ class _CiscoAIDefenseMcpMixin: self, data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: + ) -> dict[str, object]: del user_api_key_dict # carried via logging metadata, not the wire payload url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_mcp_request_payload(data=data) @@ -312,7 +313,7 @@ class _CiscoAIDefenseMcpMixin: response: object, user_api_key_dict: UserAPIKeyAuth | None = None, redact_response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: del user_api_key_dict # carried via logging metadata, not the wire payload url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_mcp_response_payload( @@ -349,7 +350,7 @@ class _CiscoAIDefenseMcpMixin: def _build_mcp_request_payload( self, data: dict, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """Build the JSON-RPC ``tools/call`` envelope sent to ``/inspect/mcp``. The Cisco AI Defense MCP inspect endpoint expects the JSON-RPC @@ -390,7 +391,7 @@ class _CiscoAIDefenseMcpMixin: self, request_data: dict, response: object, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """Build the MCP response-inspection body sent to ``/inspect/mcp``.""" request_payload: Final = self._build_mcp_request_payload(data=request_data) if request_payload is None: @@ -415,7 +416,7 @@ class _CiscoAIDefenseMcpMixin: return payload @staticmethod - def _hydrate_mcp_tool_context(request_data: dict[str, Any]) -> None: + def _hydrate_mcp_tool_context(request_data: dict[str, object]) -> None: metadata = request_data.get("mcp_tool_call_metadata") if metadata is None: nested: Final = request_data.get("metadata") or request_data.get("litellm_metadata") @@ -440,7 +441,7 @@ class _CiscoAIDefenseMcpMixin: request_data.setdefault("server_name", server_name) @staticmethod - def _normalize_mcp_response(response: object) -> dict[str, Any] | None: + def _normalize_mcp_response(response: object) -> dict[str, object] | None: """Normalize an MCP tool response into a JSON-RPC envelope. Handles JSON-RPC dicts, raw content lists, MCP SDK models, and @@ -502,10 +503,10 @@ class _CiscoAIDefenseMcpMixin: @staticmethod def _build_mcp_result( - content: list[Any], + content: Sequence[object], source: object = None, - ) -> dict[str, Any]: - result: Final[dict[str, Any]] = {"content": [_serialize_mcp_content_item(item) for item in content]} + ) -> dict[str, object]: + result: Final[dict[str, object]] = {"content": [_serialize_mcp_content_item(item) for item in content]} for key in ("structuredContent", "isError"): value = source.get(key) if isinstance(source, dict) else getattr(source, key, None) if value is not None and (key != "isError" or isinstance(value, bool)): @@ -522,7 +523,7 @@ class _CiscoAIDefenseMcpMixin: if response_obj is None: return False - inner: Final = getattr(response_obj, "mcp_tool_call_response", None) + inner: Final[object | None] = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: return _CiscoAIDefenseMcpMixin._set_mcp_tool_response_text(inner, text) @@ -559,7 +560,7 @@ class _CiscoAIDefenseMcpMixin: pass elif isinstance(response_obj, dict): result: Final = response_obj.get("result") - target: Final[dict[Any, Any]] = result if isinstance(result, dict) else response_obj + target: Final[dict[object, object]] = result if isinstance(result, dict) else response_obj if "structuredContent" in target: target["structuredContent"] = replacement replaced = True @@ -567,11 +568,11 @@ class _CiscoAIDefenseMcpMixin: return replaced @staticmethod - def _coerce_to_content_list(response_obj: object) -> list[Any] | None: + def _coerce_to_content_list(response_obj: object) -> list[object] | None: """Find the MCP content list inside supported response shapes.""" if response_obj is None: return None - inner: Final = getattr(response_obj, "mcp_tool_call_response", None) + inner: Final[object | None] = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: return _CiscoAIDefenseMcpMixin._coerce_to_content_list(inner) content: Final = getattr(response_obj, "content", None) @@ -594,8 +595,8 @@ class _CiscoAIDefenseMcpMixin: @staticmethod def _extract_sanitized_mcp_arguments( - inspect_response: dict[str, Any], - ) -> dict[str, Any] | None: + inspect_response: dict[str, object], + ) -> dict[str, object] | None: """Pull sanitized MCP tool-call arguments off the verdict. Cisco can return them at the top level (``params.arguments``) or diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index e2d7c06f7c5..a269ad31a6b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -80,7 +80,7 @@ import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey -from typing_extensions import NotRequired, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -89,6 +89,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -107,6 +108,19 @@ class _JWTDecodeKwargs(TypedDict): issuer: NotRequired[str] +class _DebugHeaderClaims(TypedDict, total=False): + sub: ReadOnly[object] + iss: ReadOnly[object] + exp: ReadOnly[object] + scope: ReadOnly[str] + + +class _SignedClaimSummary(TypedDict): + sub: ReadOnly[object] + act: ReadOnly[Mapping[str, object]] + exp: ReadOnly[object] + + # Module-level singleton for the JWKS discovery endpoint to access. _mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None @@ -265,7 +279,8 @@ class MCPJWTSigner(CustomGuardrail): **kwargs: Any, ) -> None: kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - super().__init__(**kwargs) + base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs + super().__init__(**base_kwargs) # --- Signing key setup --- key_material: Final = os.environ.get(self.SIGNING_KEY_ENV) @@ -677,7 +692,7 @@ class MCPJWTSigner(CustomGuardrail): data: dict, jwt_claims: Mapping[str, object] | None = None, call_type: CallTypesLiteral | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Build JWT claims for the outbound MCP access token. @@ -752,7 +767,7 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ @staticmethod - def _build_debug_header(claims: dict[str, Any], kid: str) -> str: + def _build_debug_header(claims: _DebugHeaderClaims, kid: str) -> str: """ Build the x-litellm-mcp-debug header value. @@ -873,16 +888,18 @@ class MCPJWTSigner(CustomGuardrail): # FR-9: Debug header # ------------------------------------------------------------------ if self.debug_headers: - new_headers["x-litellm-mcp-debug"] = self._build_debug_header(claims, self._kid) + debug_claims: Final[_DebugHeaderClaims] = claims + new_headers["x-litellm-mcp-debug"] = self._build_debug_header(debug_claims, self._kid) hook_data["extra_headers"] = new_headers + logged_claims: Final[_SignedClaimSummary] = claims verbose_proxy_logger.debug( "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d verified=%s channel=%s call_type=%s", - claims.get("sub"), - claims.get("act", {}).get("sub"), + logged_claims.get("sub"), + logged_claims.get("act", {}).get("sub"), hook_data.get("mcp_tool_name"), - claims["exp"], + logged_claims["exp"], jwt_claims is not None, bool(self.channel_token_audience), call_type, diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index e9cd6addef8..babb3f8aee2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -23,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage +from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus @@ -80,7 +81,8 @@ class NomaV2Guardrail(CustomGuardrail): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - super().__init__(**kwargs) + base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs + super().__init__(**base_kwargs) @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: @@ -111,7 +113,7 @@ class NomaV2Guardrail(CustomGuardrail): return parsed.hostname == _DEFAULT_API_BASE_HOSTNAME @staticmethod - def _get_non_empty_str(value: Any) -> str | None: + def _get_non_empty_str(value: object) -> str | None: if not isinstance(value, str): return None stripped: Final = value.strip() @@ -153,7 +155,7 @@ class NomaV2Guardrail(CustomGuardrail): else model_call_details ) - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "inputs": inputs, "request_data": payload_request_data, "input_type": input_type, @@ -165,7 +167,7 @@ class NomaV2Guardrail(CustomGuardrail): @staticmethod def _sanitize_payload_for_transport(payload: dict) -> dict: - def _default(obj: Any) -> Any: + def _default(obj: object) -> object: if hasattr(obj, "model_dump"): try: return obj.model_dump() @@ -178,7 +180,7 @@ class NomaV2Guardrail(CustomGuardrail): except (ValueError, TypeError): json_str = safe_dumps(payload) - safe_payload: Final = safe_json_loads(json_str, default={}) + safe_payload: Final[object] = safe_json_loads(json_str, default={}) if safe_payload == {} and payload: verbose_proxy_logger.warning( "Noma v2 guardrail: payload serialization failed, falling back to empty payload" @@ -215,7 +217,7 @@ class NomaV2Guardrail(CustomGuardrail): response.text, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[dict[str, object]] = response.json() verbose_proxy_logger.debug( "Noma v2 AIDR response parsed: %s", json.dumps(response_json, default=str), @@ -227,7 +229,7 @@ class NomaV2Guardrail(CustomGuardrail): request_data: dict, start_time: datetime, guardrail_status: GuardrailStatus, - guardrail_json_response: Any, + guardrail_json_response: str | dict[str, object], ) -> None: end_time: Final = datetime.now() duration: Final = (end_time - start_time).total_seconds() @@ -270,7 +272,7 @@ class NomaV2Guardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: start_time: Final = datetime.now() guardrail_status: GuardrailStatus = "success" - guardrail_json_response: Any = {} + guardrail_json_response: str | dict[str, object] = {} dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data) if not isinstance(dynamic_params, dict): dynamic_params = {} @@ -320,8 +322,9 @@ class NomaV2Guardrail(CustomGuardrail): except NomaBlockedMessage as e: guardrail_status = "guardrail_intervened" + blocked_detail: Final[dict[str, object]] = {"error": "blocked"} guardrail_json_response = ( - response_json if isinstance(response_json, dict) else getattr(e, "detail", {"error": "blocked"}) + response_json if isinstance(response_json, dict) else getattr(e, "detail", blocked_detail) ) raise except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index bcee45355e3..a942dd70611 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -11,10 +11,10 @@ import asyncio import json import threading -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable from contextlib import asynccontextmanager from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, cast import aiohttp from typing_extensions import NotRequired, ReadOnly @@ -63,6 +63,14 @@ class _PresidioAnonymizeResponse(TypedDict): items: ReadOnly[NotRequired[list[_PresidioAnonymizeItem]]] +class _JsonResponse(Protocol): + def json(self) -> Awaitable[object]: ... + + +async def _json_body(response: _JsonResponse) -> object: + return await response.json() + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -345,7 +353,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): f"expected application/json Content-Type but received '{content_type}'; body: '{error_body[:200]}'" ) - analyze_results: Final = await response.json() + analyze_results: Final = await _json_body(response) verbose_proxy_logger.debug("analyze_results: %s", analyze_results) # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) @@ -758,7 +766,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception as e: raise e - def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: from concurrent.futures import ThreadPoolExecutor def run_in_new_loop(): @@ -786,7 +794,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # No running event loop, we can safely run in this thread return run_in_new_loop() - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: """ Masks the input and output before logging to langfuse, datadog, etc. """ @@ -853,9 +861,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): and not isinstance(result.choices[0], StreamingChoices) ): await self._process_response_for_pii(response=result, request_data=kwargs, mode="mask") - elif self._is_anthropic_message_response(result): + elif isinstance(result, dict) and self._is_anthropic_message_response(result): await self._process_anthropic_response_for_pii( - response=cast(dict, result), # cast-ok: _is_anthropic_message_response narrows via isinstance + response=result, request_data=kwargs, mode="mask", ) @@ -1082,7 +1090,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _stream_apply_output_masking( self, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream | bytes, None]: """Apply Presidio masking to streaming output (apply_to_output=True path).""" @@ -1186,7 +1194,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return "\n".join(result_lines).encode("utf-8") - def _unmask_responses_api_completed_chunk(self, chunk: Any, pii_tokens: dict[str, str]) -> None: + def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None: """ Unmask PII tokens in-place for a ``response.completed`` Responses API event. @@ -1195,7 +1203,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): blocks; text blocks expose a ``.text`` string attribute. We walk the tree and replace every PII token with its original value. """ - response_obj: Final = getattr(chunk, "response", None) + response_obj: Final[object] = getattr(chunk, "response", None) if response_obj is None: return @@ -1211,7 +1219,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _stream_pii_unmasking( self, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream | bytes, None]: """Apply PII unmasking to streaming output (output_parse_pii=True path).""" @@ -1287,7 +1295,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream | bytes, None]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index ddb40dc3ca0..831df43692b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -310,7 +310,7 @@ class XecGuardGuardrail(CustomGuardrail): scan_type: str, suppress_errors: bool = False, ) -> dict | None: - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "model": self.xecguard_model, "scan_type": scan_type, "messages": messages, @@ -385,7 +385,7 @@ class XecGuardGuardrail(CustomGuardrail): def _build_full_history( self, request_data: dict, - inputs: Any, + inputs: GenericGuardrailAPIInputs, input_type: str, ) -> list[dict]: """Assemble the full message list that will be sent to XecGuard. diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 3ce406eef73..8b82842353c 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,10 +5,11 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Optional +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -30,6 +31,13 @@ if TYPE_CHECKING: from litellm.router import Router +class SemanticToolFilterConfig(TypedDict, total=False): + enabled: ReadOnly[bool] + embedding_model: ReadOnly[str] + top_k: ReadOnly[int] + similarity_threshold: ReadOnly[float] + + def _truncate_csv_at_tool_name_boundary(tool_names_csv: str, max_length: int) -> str: """Cap a CSV of tool names to max_length, dropping any name that does not fit whole.""" if len(tool_names_csv) <= max_length: @@ -68,7 +76,7 @@ class SemanticToolFilterHook(CustomLogger): semantic_filter.top_k, ) - def _should_expand_mcp_tools(self, tools: list[Any]) -> bool: + def _should_expand_mcp_tools(self, tools: Iterable[Mapping[str, object]]) -> bool: """ Check if tools contain MCP references with server_url="litellm_proxy". @@ -82,9 +90,9 @@ class SemanticToolFilterHook(CustomLogger): async def _expand_mcp_tools( self, - tools: list[Any], + tools: Iterable[Mapping[str, object]], user_api_key_dict: "UserAPIKeyAuth", - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Expand MCP references to actual tool definitions. @@ -111,7 +119,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Convert Pydantic models to dicts for compatibility - openai_tools_as_dicts: Final = [] + openai_tools_as_dicts: Final[list[dict[str, object]]] = [] for tool in openai_tools: if hasattr(tool, "model_dump"): tool_dict = tool.model_dump(exclude_none=True) @@ -141,8 +149,8 @@ class SemanticToolFilterHook(CustomLogger): async def _filter_expanded_tools( self, data: dict, - expanded_tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + expanded_tools: list[dict[str, object]], + ) -> list[dict[str, object]]: """ Apply the semantic filter to expanded MCP tool definitions. @@ -159,7 +167,7 @@ class SemanticToolFilterHook(CustomLogger): return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools) - def _selected_tool_names(self, filtered_tools: list[dict[str, Any]]) -> list[str]: + def _selected_tool_names(self, filtered_tools: Sequence[object]) -> list[str]: """Names of the semantically selected tools, as produced by the MCP expansion.""" names: Final = (self.filter._extract_tool_info(tool)[0] for tool in filtered_tools) return [name for name in names if name] @@ -217,10 +225,10 @@ class SemanticToolFilterHook(CustomLogger): def _emit_filter_metadata( self, data: dict, - mcp_tools: list[object], - filtered_mcp_tools: list[object], - native_tools: list[object], - filtered_tools: list[object], + mcp_tools: Sequence[object], + filtered_mcp_tools: Sequence[object], + native_tools: Sequence[object], + filtered_tools: Sequence[object], ) -> None: """ Emit response-header metadata when MCP tools were filtered. @@ -252,10 +260,10 @@ class SemanticToolFilterHook(CustomLogger): def _emit_filter_metadata_safe( self, data: dict, - mcp_tools: list[object], - filtered_mcp_tools: list[object], - native_tools: list[object], - filtered_tools: list[object], + mcp_tools: Sequence[object], + filtered_mcp_tools: Sequence[object], + native_tools: Sequence[object], + filtered_tools: Sequence[object], ) -> None: """ Emit filter metadata without letting an emission failure abort the @@ -375,7 +383,7 @@ class SemanticToolFilterHook(CustomLogger): ) if mcp_tools: - filtered_mcp_tools = await self.filter.filter_tools( + filtered_mcp_tools: list[object] = await self.filter.filter_tools( query=user_query, available_tools=mcp_tools, ) @@ -419,9 +427,9 @@ class SemanticToolFilterHook(CustomLogger): self, data: dict, user_api_key_dict: "UserAPIKeyAuth", - response: Any, + response: object, request_headers: dict[str, str] | None = None, - litellm_call_info: dict[str, Any] | None = None, + litellm_call_info: dict[str, object] | None = None, ) -> dict[str, str] | None: """Add semantic filter stats and tool names to response headers.""" from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH @@ -446,7 +454,7 @@ class SemanticToolFilterHook(CustomLogger): return headers - def _get_tool_names_csv(self, tools: list[Any]) -> str: + def _get_tool_names_csv(self, tools: Sequence[object]) -> str: """Extract tool names and return as CSV string.""" if not tools: return "" @@ -461,7 +469,7 @@ class SemanticToolFilterHook(CustomLogger): @staticmethod async def initialize_from_config( - config: dict[str, Any] | None, + config: SemanticToolFilterConfig | None, llm_router: Optional["Router"], ) -> Optional["SemanticToolFilterHook"]: """ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 064b53e07b7..da54c8d6de5 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4,7 +4,8 @@ import json import re import time from collections import OrderedDict -from collections.abc import Mapping, MutableMapping +from collections.abc import Mapping, MutableMapping, Sequence +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -52,7 +53,7 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance -_SPECIAL_HEADERS_CACHE: Final = frozenset(v.value.lower() for v in SpecialHeaders._member_map_.values()) +_SPECIAL_HEADERS_CACHE: Final = frozenset(str(v.value).lower() for v in SpecialHeaders) _REDACTED_HEADER_VALUE: Final = "***REDACTED***" _CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names() | frozenset( @@ -123,7 +124,7 @@ def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: _ANTHROPIC_SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]+$") -def _sanitize_for_log(value: Any) -> str: +def _sanitize_for_log(value: object) -> str: """ Basic log sanitization helper to reduce log-injection risk. @@ -161,7 +162,7 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig - from litellm.types.proxy.policy_engine import PolicyMatchContext + from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext ProxyConfig = _ProxyConfig else: @@ -318,7 +319,7 @@ _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_overr _URL_DESTINATION_REQUEST_FIELDS: Final = ("model", "file_id") -def _reject_url_valued_destinations(data: dict[str, Any]) -> None: +def _reject_url_valued_destinations(data: dict[str, object]) -> None: """Reject URL-valued ``model``/``file_id`` unless admin-allowlisted. Some providers (HuggingFace, Oobabooga, Gemini files) accept a URL in the @@ -377,7 +378,7 @@ def _invalid_metadata_type_error(field: str, value: object) -> ProxyException: ) -def _normalized_metadata_object(field: str, value: object) -> Mapping[str, Any]: +def _normalized_metadata_object(field: str, value: object) -> Mapping[str, object]: """Return ``value`` as a metadata object or raise a 400 like OpenAI does. A JSON string that parses to an object is accepted because multipart/form-data @@ -392,6 +393,23 @@ def _normalized_metadata_object(field: str, value: object) -> Mapping[str, Any]: raise _invalid_metadata_type_error(field=field, value=value) +def _normalized_metadata_slot( + request_data: MutableMapping[str, object], metadata_variable_name: str +) -> dict[str, object]: + """Return the request's metadata slot as a dict, normalising it in place first. + + Metadata can arrive as a JSON string (multipart/form-data, ``extra_body``). Parsing it here keeps + existing entries alive through a merge instead of silently overwriting them with an empty dict. + """ + raw: Final = request_data.get(metadata_variable_name) + if isinstance(raw, dict): + return raw + parsed: Final = safe_json_loads(raw) if isinstance(raw, str) else None + normalized: Final[dict[str, object]] = parsed if isinstance(parsed, dict) else {} + request_data[metadata_variable_name] = normalized + return normalized + + def _strip_untrusted_request_header_controls( headers: Any, *, @@ -407,7 +425,7 @@ def _strip_untrusted_request_header_controls( headers.pop(header_name, None) -def _is_false_like(value: Any) -> bool: +def _is_false_like(value: object) -> bool: if isinstance(value, bool): return value is False if isinstance(value, str): @@ -452,7 +470,7 @@ def _key_or_team_allows_client_pricing_override( ) -def _strip_client_message_redaction_opt_out(data: dict[str, Any]) -> None: +def _strip_client_message_redaction_opt_out(data: dict[str, object]) -> None: stripped: Final[list[str]] = [] if "turn_off_message_logging" in data and _is_false_like(data["turn_off_message_logging"]): stripped.append("turn_off_message_logging") @@ -503,7 +521,7 @@ def _strip_client_callback_credentials( ) -def _strip_client_pricing_overrides(data: dict[str, Any]) -> None: +def _strip_client_pricing_overrides(data: dict[str, object]) -> None: """Drop pricing overrides from the request body and any metadata variant. Skipped only when the calling key/team carries @@ -556,9 +574,9 @@ def _get_metadata_variable_name(request: Request) -> str: def _promoted_trace_control_fields( - requester_metadata: Mapping[str, Any], - litellm_metadata: Mapping[str, Any], -) -> tuple[tuple[str, Any], ...]: + requester_metadata: Mapping[str, object], + litellm_metadata: Mapping[str, object], +) -> tuple[tuple[str, object], ...]: """Return the caller's trace-control fields that ``litellm_metadata`` does not already set.""" return tuple( (key, value) @@ -1169,7 +1187,7 @@ class LiteLLMProxyRequestSetup: def add_litellm_data_for_backend_llm_call( *, headers: dict, - request_data: Mapping[str, Any], + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, general_settings: dict[str, Any] | None = None, ) -> LitellmDataForBackendLLMCall: @@ -1549,14 +1567,7 @@ class LiteLLMProxyRequestSetup: return _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata = request_data.get(_metadata_variable_name) - if isinstance(metadata, str): - parsed: Final = safe_json_loads(metadata) - metadata = parsed if isinstance(parsed, dict) else {} - request_data[_metadata_variable_name] = metadata - elif not isinstance(metadata, dict): - metadata = {} - request_data[_metadata_variable_name] = metadata + metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( @@ -1608,18 +1619,7 @@ class LiteLLMProxyRequestSetup: # from (litellm_metadata vs metadata) so the merged tags are visible # to _tag_max_budget_check. _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata = request_data.get(_metadata_variable_name) - # metadata can arrive as a JSON string (multipart/form-data, extra_body). - # Parse it so existing tags survive the merge — overwriting the string - # with {} would let a caller bypass _tag_max_budget_check on an - # over-budget body tag by also sending a within-budget header tag. - if isinstance(metadata, str): - parsed: Final = safe_json_loads(metadata) - metadata = parsed if isinstance(parsed, dict) else {} - request_data[_metadata_variable_name] = metadata - elif not isinstance(metadata, dict): - metadata = {} - request_data[_metadata_variable_name] = metadata + metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( @@ -1759,7 +1759,7 @@ async def add_litellm_data_to_request( # admin-injection strip below so the audit / spend-tracking consumers of # proxy_server_request["body"] see the cleaned metadata rather than # attacker-forged user_api_key_* fields. - _litellm_received_at: Final = getattr(request.state, "litellm_received_at", None) + _litellm_received_at: Final[datetime | None] = getattr(request.state, "litellm_received_at", None) arrival_time: Final = _litellm_received_at.timestamp() if _litellm_received_at is not None else time.time() data["proxy_server_request"] = { "url": str(request.url), @@ -2423,16 +2423,16 @@ def _resolve_provider_from_deployment( if deployment is None: continue - litellm_params = getattr(deployment, "litellm_params", None) + litellm_params: object = getattr(deployment, "litellm_params", None) if litellm_params is None: continue custom_provider = getattr(litellm_params, "custom_llm_provider", None) - if custom_provider: + if isinstance(custom_provider, str) and custom_provider: return custom_provider - deployment_model = getattr(litellm_params, "model", "") or "" - if "/" in deployment_model: + deployment_model = getattr(litellm_params, "model", "") + if isinstance(deployment_model, str) and "/" in deployment_model: return deployment_model.split("/", 1)[0] return None @@ -2855,8 +2855,8 @@ def _extract_policy_id(s: str) -> str | None: def _match_and_track_policies( data: dict, context: "PolicyMatchContext", - request_body_policies: Any, - policies_override: dict[str, Any] | None = None, + request_body_policies: Sequence[str], + policies_override: dict[str, "Policy"] | None = None, ) -> tuple[list[str], dict[str, str]]: """ Match policies via attachments and request body, track them in metadata. @@ -2914,7 +2914,7 @@ def _apply_resolved_guardrails_to_metadata( metadata_variable_name: str, context: "PolicyMatchContext", policy_names: list[str] | None = None, - policies: dict[str, Any] | None = None, + policies: dict[str, "Policy"] | None = None, ) -> None: """Apply resolved guardrails and pipelines to request metadata.""" from litellm._logging import verbose_proxy_logger @@ -3044,7 +3044,7 @@ async def add_guardrails_from_policy_engine( request_body_names.append(item) # Resolve policy versions by ID from in-memory cache (populated by sync job; no DB in hot path) - merged_policies: Final[dict[str, Any]] = dict(registry.get_all_policies()) + merged_policies: Final[dict[str, Policy]] = dict(registry.get_all_policies()) fetched_policy_names: Final[list[str]] = [] for policy_id in request_body_version_ids: result = registry.get_policy_by_id_for_request(policy_id=policy_id) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d1c08352919..24ba874dc97 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2467,7 +2467,7 @@ async def _validate_update_key_data( user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None, premium_user: bool, - prisma_client: Any, + prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" @@ -3700,7 +3700,7 @@ async def info_key_fn( except Exception: # if using pydantic v1 key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback - key_token_hash: Final = key_info.pop("token") + key_token_hash: Final[str | None] = key_info.pop("token") model_max_budget = key_info.get("model_max_budget") or {} budget_table: Final = key_info.get("litellm_budget_table") or {} @@ -5155,7 +5155,7 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio max_budget = key_in_db.max_budget if key_in_db.litellm_budget_table is not None: - budget_max_budget: Final = getattr(key_in_db.litellm_budget_table, "max_budget", None) + budget_max_budget: Final[float | None] = getattr(key_in_db.litellm_budget_table, "max_budget", None) if budget_max_budget is not None: if max_budget is None or budget_max_budget < max_budget: max_budget = budget_max_budget diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 108e6a7b47d..f58f3722741 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -13,7 +13,7 @@ import copy import json import os from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import Response, StreamingResponse @@ -90,7 +90,7 @@ class _ApplyPoliciesResultBase(TypedDict): class ApplyPoliciesResult(_ApplyPoliciesResultBase, total=False): """Result of apply_policies. agent_response set when agent_id provided.""" - agent_response: Any + agent_response: object class _ApplyPoliciesPerItemResultBase(TypedDict): @@ -103,7 +103,7 @@ class _ApplyPoliciesPerItemResultBase(TypedDict): class ApplyPoliciesPerItemResult(_ApplyPoliciesPerItemResultBase, total=False): """Result for one input when using inputs_list. agent_response set when agent_id provided.""" - agent_response: Any + agent_response: object class ApplyPoliciesListResult(TypedDict): @@ -295,8 +295,8 @@ async def test_policies_and_guardrails( from litellm.proxy.proxy_server import chat_completion, proxy_logging_obj from litellm.proxy.utils import handle_exception_on_proxy - def _serialize_chat_response(response: Any) -> Any: - if hasattr(response, "model_dump"): + def _serialize_chat_response(response: object) -> object: + if isinstance(response, BaseModel): return response.model_dump(exclude_unset=True) if isinstance(response, dict): return response @@ -306,7 +306,7 @@ async def test_policies_and_guardrails( inputs: GenericGuardrailAPIInputs, agent_id: str, user_api_key_dict: UserAPIKeyAuth, - ) -> Any: + ) -> object: body: Final = _chat_body_from_inputs(inputs, agent_id, data.request_data) req: Final = _request_with_json_body(body) resp: Final = Response() diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9a3bc82c6fa..2c817ed3143 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -12,7 +12,7 @@ import os import re from collections.abc import Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, Final, cast +from typing import TYPE_CHECKING, Annotated, Final, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -64,6 +64,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.utils import LlmProviders +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter @@ -112,7 +113,21 @@ def is_passthrough_request_streaming(request_body: object) -> bool: return bool(request_body.get("stream", False)) -def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, Any]: +def _optional_str(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _string_keyed_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +async def _json_request_body(request: Request) -> Mapping[str, object]: + return await request.json() + + +def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: """ Build the request metadata carrying key-level spend attribution and the pre-call budget reservation for a router-model passthrough request. @@ -201,7 +216,7 @@ async def llm_passthrough_factory_proxy_route( # anthropic is streaming when 'stream' = True is in the body if request.method == "POST": if "multipart/form-data" not in request.headers.get("content-type", ""): - _request_body = await request.json() + _request_body = await _json_request_body(request) else: _request_body = await get_form_data(request) @@ -374,7 +389,7 @@ async def vllm_proxy_route( endpoint=endpoint, request_query_params=request.query_params, request_headers=_safe_get_request_headers(request), - stream=request_body.get("stream", False), + stream=is_streaming_request, content=None, data=None, files=None, @@ -802,7 +817,7 @@ async def handle_bedrock_passthrough_router_model( # Use the common processing path (same as non-router models) # This ensures all metadata, hooks, and logging are properly initialized - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) data["model"] = model @@ -846,8 +861,8 @@ async def handle_bedrock_count_tokens( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - request_body: dict[str, Any], -) -> dict[str, Any]: + request_body: dict[str, object], +) -> dict[str, object]: """ Handle AWS Bedrock CountTokens API requests. @@ -864,7 +879,7 @@ async def handle_bedrock_count_tokens( handler: Final = BedrockCountTokensHandler() # Extract model from request body - model: Final = request_body.get("model") + model: Final = _optional_str(request_body.get("model")) if not model: raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) @@ -996,7 +1011,7 @@ async def bedrock_llm_proxy_route( "Bedrock passthrough: Using direct Bedrock model '%s' for endpoint '%s'", model, endpoint ) - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) data["method"] = request.method @@ -1095,7 +1110,7 @@ async def bedrock_proxy_route( headers: Final = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it try: - data: Final = await request.json() + data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) @@ -1186,7 +1201,7 @@ async def comprehend_medical_proxy_route( ) try: - data: Final = await request.json() + data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail=str(e)) @@ -1397,7 +1412,7 @@ async def assemblyai_proxy_route( is_streaming_request = False # assemblyai is streaming when 'stream' = True is in the body if request.method == "POST": - _request_body: Final = await request.json() + _request_body: Final = await _json_request_body(request) if _request_body.get("stream"): is_streaming_request = True @@ -1504,7 +1519,7 @@ async def azure_proxy_route( endpoint=endpoint, request_query_params=request.query_params, request_headers=_safe_get_request_headers(request), - stream=request_body.get("stream", False), + stream=is_streaming_request, content=None, data=None, files=None, @@ -1591,7 +1606,7 @@ async def azure_proxy_route( extra_headers = auth_credentials.get("headers") or {} - base_target_url = litellm_params.get("api_base") + base_target_url = _optional_str(litellm_params.get("api_base")) if base_target_url is None: raise Exception(f"API base not found for {part}") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( @@ -1712,7 +1727,7 @@ def get_vertex_pass_through_handler( def _override_vertex_params_from_router_credentials( - router_credentials: Any | None, + router_credentials: LiteLLM_ManagedVectorStore | None, vertex_project: str | None, vertex_location: str | None, ) -> tuple[str | None, str | None]: @@ -1732,14 +1747,14 @@ def _override_vertex_params_from_router_credentials( verbose_proxy_logger.debug("Using vector store credentials to override vertex project and location") - litellm_params: Final = router_credentials.get("litellm_params", {}) + litellm_params: Final = _string_keyed_mapping(router_credentials.get("litellm_params")) if not litellm_params: verbose_proxy_logger.warning("Vector store credentials found but litellm_params is empty") return vertex_project, vertex_location # Extract vertex_project and vertex_location from litellm_params - vector_store_project: Final = litellm_params.get("vertex_project") - vector_store_location: Final = litellm_params.get("vertex_location") + vector_store_project: Final = _optional_str(litellm_params.get("vertex_project")) + vector_store_location: Final = _optional_str(litellm_params.get("vertex_location")) if vector_store_project: verbose_proxy_logger.debug( @@ -1747,7 +1762,6 @@ def _override_vertex_params_from_router_credentials( vertex_project, vector_store_project, ) - vertex_project = vector_store_project else: verbose_proxy_logger.warning("Vector store credentials found but missing vertex_project in litellm_params") @@ -1757,11 +1771,10 @@ def _override_vertex_params_from_router_credentials( vertex_location, vector_store_location, ) - vertex_location = vector_store_location else: verbose_proxy_logger.warning("Vector store credentials found but missing vertex_location in litellm_params") - return vertex_project, vertex_location + return vector_store_project or vertex_project, vector_store_location or vertex_location _CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL: Final = ( @@ -1869,8 +1882,8 @@ def _forwarded_headers_for_credentialless_vertex_passthrough( async def _prepare_vertex_auth_headers( request: Request, - vertex_credentials: Any | None, - router_credentials: Any | None, + vertex_credentials: VertexPassThroughCredentials | None, + router_credentials: LiteLLM_ManagedVectorStore | None, vertex_project: str | None, vertex_location: str | None, base_target_url: str | None, @@ -1967,7 +1980,7 @@ async def _base_vertex_proxy_route( fastapi_response: Response, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, user_api_key_dict: UserAPIKeyAuth | None = None, - router_credentials: Any | None = None, + router_credentials: LiteLLM_ManagedVectorStore | None = None, ): """ Base function for Vertex AI passthrough routes. @@ -2851,7 +2864,7 @@ async def watsonx_proxy_route( is_streaming_request = False if request.method == "POST": if "multipart/form-data" not in request.headers.get("content-type", ""): - _request_body = await request.json() + _request_body = await _json_request_body(request) else: _request_body = await get_form_data(request) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3d60f4f5f3a..d48c61da5f0 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3146,6 +3146,14 @@ def _get_pass_through_endpoints_from_config() -> list[PassThroughGenericEndpoint return returned_endpoints +def _config_field_endpoints(response: ConfigFieldInfo) -> list[object] | None: + return response.field_value + + +def _request_app(request: Request) -> FastAPI: + return request.app + + async def _get_pass_through_endpoints_from_db( endpoint_id: str | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, @@ -3162,7 +3170,7 @@ async def _get_pass_through_endpoints_from_db( except Exception: return [] - pass_through_endpoint_data: Final[list | None] = response.field_value + pass_through_endpoint_data: Final = _config_field_endpoints(response) if pass_through_endpoint_data is None: return [] @@ -3325,7 +3333,7 @@ async def update_pass_through_endpoints( detail={"error": "No pass-through endpoints found"}, ) - pass_through_endpoint_data: Final[list | None] = response.field_value + pass_through_endpoint_data: Final[list | None] = _config_field_endpoints(response) if pass_through_endpoint_data is None: raise HTTPException( status_code=404, @@ -3396,7 +3404,7 @@ async def update_pass_through_endpoints( _custom_headers: dict | None = updated_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) - route_app: Final[FastAPI] = request.app + route_app: Final = _request_app(request) if updated_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( app=route_app, @@ -3488,7 +3496,7 @@ async def create_pass_through_endpoints( _custom_headers: dict | None = created_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) - route_app: Final[FastAPI] = request.app + route_app: Final = _request_app(request) if created_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( app=route_app, @@ -3556,7 +3564,7 @@ async def delete_pass_through_endpoints( response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Update field by removing endpoint - pass_through_endpoint_data: Final[list | None] = response.field_value + pass_through_endpoint_data: Final[list | None] = _config_field_endpoints(response) if response.field_value is None or pass_through_endpoint_data is None: raise HTTPException( status_code=400, diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 4d62f1d6d71..db574f859b3 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -7,12 +7,14 @@ Provides: """ import base64 +import json from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse, StreamingResponse +from starlette.datastructures import UploadFile import litellm from litellm._logging import verbose_proxy_logger @@ -45,6 +47,16 @@ if TYPE_CHECKING: router: Final = APIRouter() +def _as_string_keyed_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +def _response_attr(source: object, name: str) -> object: + return getattr(source, name, None) + + def _raise_vector_store_scan_depth_exceeded() -> None: raise HTTPException( status_code=400, @@ -53,8 +65,8 @@ def _raise_vector_store_scan_depth_exceeded() -> None: def _append_payload_to_scan_stack( - payload_stack: list[tuple[Any, int]], - value: Any, + payload_stack: list[tuple[object, int]], + value: object, next_depth: int, ) -> None: if isinstance(value, dict): @@ -117,7 +129,7 @@ async def _authorize_nested_vector_store_ids( def _build_file_metadata_entry( - response: Any, + response: object, file_data: tuple[str, bytes, str] | None = None, file_url: str | None = None, ) -> Mapping[str, str | int | None]: @@ -135,11 +147,11 @@ def _build_file_metadata_entry( from datetime import datetime, timezone # Extract file_id from response - file_id = None - if hasattr(response, "get"): - file_id = response.get("file_id") - elif hasattr(response, "file_id"): - file_id = response.file_id + mapping_response: Final = _as_string_keyed_mapping(response) + raw_file_id: Final = ( + mapping_response.get("file_id") if mapping_response is not None else _response_attr(response, "file_id") + ) + file_id: Final = raw_file_id if isinstance(raw_file_id, str) else None # Extract file information from file_data tuple filename = None @@ -152,7 +164,7 @@ def _build_file_metadata_entry( content_type = file_data[2] if len(file_data) > 2 else None # Build file metadata entry - file_entry: Final = { + file_entry: Final[dict[str, str | int | None]] = { "file_id": file_id, "filename": filename, "file_url": file_url, @@ -169,7 +181,7 @@ def _build_file_metadata_entry( async def _save_vector_store_to_db_from_rag_ingest( - response: Any, + response: object, ingest_options: Mapping[str, dict[str, str | None]], prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, @@ -197,10 +209,11 @@ async def _save_vector_store_to_db_from_rag_ingest( ) # Handle both dict and object responses - if hasattr(response, "get"): - vector_store_id = response.get("vector_store_id") + mapping_response: Final = _as_string_keyed_mapping(response) + if mapping_response is not None: + vector_store_id = mapping_response.get("vector_store_id") elif hasattr(response, "vector_store_id"): - vector_store_id = response.vector_store_id + vector_store_id = _response_attr(response, "vector_store_id") else: verbose_proxy_logger.warning("Unable to extract vector_store_id from response type: %s", type(response)) return @@ -266,14 +279,13 @@ async def _save_vector_store_to_db_from_rag_ingest( verbose_proxy_logger.info("Vector store %s already exists, appending file to metadata", vector_store_id) # Update existing vector store with new file - existing_metadata = existing_vector_store.vector_store_metadata or {} - if isinstance(existing_metadata, str): - import json + stored_metadata: Final = existing_vector_store.vector_store_metadata or {} + existing_metadata: dict[str, object] = ( + json.loads(stored_metadata) if isinstance(stored_metadata, str) else stored_metadata + ) - existing_metadata = json.loads(existing_metadata) - - ingested_files: Final = existing_metadata.get("ingested_files", []) - ingested_files.append(file_entry) + previous_files: Final = existing_metadata.get("ingested_files", []) + ingested_files: Final = [*previous_files, file_entry] if isinstance(previous_files, list) else [file_entry] existing_metadata["ingested_files"] = ingested_files # Update the vector store @@ -340,9 +352,9 @@ async def parse_rag_ingest_request( # Get file file_obj = form_data.get("file") - if file_obj is not None and hasattr(file_obj, "read"): + if isinstance(file_obj, UploadFile): file_content = await file_obj.read(MAX_UPLOAD_SIZE_BYTES + 1) - file_data = (file_obj.filename, file_content, file_obj.content_type) + file_data = (file_obj.filename or "", file_content, file_obj.content_type or "") # Parse JSON from 'request' form field (contains full request body as JSON) request_json_str: Final[str | bytes | None] = form_data.get("request") diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 020698dabd9..d32d6ab8861 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,7 +10,7 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Final, TypedDict, cast from fastapi import Request, Response @@ -38,19 +38,55 @@ class _StreamOutputItem(TypedDict, total=False): content: ReadOnly[Sequence[_StreamContentPart | None]] +class _StreamResponsePayload(TypedDict, total=False): + status: ReadOnly[str] + error: ReadOnly[dict[str, object] | None] + usage: ReadOnly[dict[str, object] | None] + reasoning: ReadOnly[dict[str, object] | None] + tool_choice: ReadOnly[object] + tools: ReadOnly[list[object] | None] + model: ReadOnly[str | None] + instructions: ReadOnly[str | None] + temperature: ReadOnly[float | None] + top_p: ReadOnly[float | None] + max_output_tokens: ReadOnly[int | None] + previous_response_id: ReadOnly[str | None] + text: ReadOnly[dict[str, object] | None] + truncation: ReadOnly[str | None] + parallel_tool_calls: ReadOnly[bool | None] + user: ReadOnly[str | None] + store: ReadOnly[bool | None] + incomplete_details: ReadOnly[dict[str, object] | None] + output: ReadOnly[Sequence[_StreamOutputItem]] + + +class _StreamEvent(TypedDict, total=False): + type: ReadOnly[str] + item: ReadOnly[_StreamOutputItem] + item_id: ReadOnly[str] + part: ReadOnly[_StreamContentPart] + content_index: ReadOnly[int] + delta: ReadOnly[str] + response: ReadOnly[_StreamResponsePayload] + + +def _parse_stream_event(serialized_event: str) -> _StreamEvent: + return json.loads(serialized_event) + + async def background_streaming_task( polling_id: str, - data, + data: dict[str, object], polling_handler: ResponsePollingHandler, request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - general_settings, + general_settings: dict[str, object], llm_router: "Router | None", proxy_config: "ProxyConfig", proxy_logging_obj: "ProxyLogging", - select_data_generator, - user_model, + select_data_generator: Callable[..., object] | None, + user_model: str | None, user_temperature: float | None, user_request_timeout: float | None, user_max_tokens: int | None, @@ -180,7 +216,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event = _parse_stream_event(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 7cad3f0a022..5f43785e57c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import Any, Final, NoReturn, cast +from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast from fastapi import HTTPException, status @@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +from litellm.types.router import DeploymentTypedDict @dataclass @@ -637,7 +638,7 @@ def _get_budget_limit_counters( for window in budget_limits: window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") - max_budget = window_dict.get("max_budget") + max_budget = _to_float(window_dict.get("max_budget")) if not budget_duration or max_budget is None or max_budget <= 0: continue window_start = get_budget_window_start(window_dict) @@ -663,18 +664,20 @@ def _get_budget_limit_counters( return counters -def _coerce_window(window: Any) -> dict: - if isinstance(window, dict): +def _coerce_window(window: object) -> Mapping[str, object]: + if isinstance(window, Mapping): return window if isinstance(window, str): try: - parsed: Final = json.loads(window) - return parsed if isinstance(parsed, dict) else {} + parsed: Final[object] = json.loads(window) except Exception: return {} - if hasattr(window, "model_dump"): - return window.model_dump() - return {} + return parsed if isinstance(parsed, Mapping) else {} + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return {} + dumped: Final[object] = model_dump() + return dumped if isinstance(dumped, Mapping) else {} async def _reserve_counter( @@ -891,7 +894,7 @@ def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float return default_reserved_cost -def get_budget_window_start(window: Any) -> datetime | None: +def get_budget_window_start(window: object) -> datetime | None: window_dict: Final = _coerce_window(window) budget_duration: Final = window_dict.get("budget_duration") if budget_duration is None: @@ -909,7 +912,7 @@ def get_budget_window_start(window: Any) -> datetime | None: return reset_at - timedelta(seconds=duration_seconds) -def _coerce_datetime(value: Any) -> datetime | None: +def _coerce_datetime(value: object) -> datetime | None: if value is None: return None if isinstance(value, datetime): @@ -1183,11 +1186,11 @@ def _get_model_cost_infos( def _deployment_tiered_pricing_table( - deployment: dict[str, Any], + deployment: DeploymentTypedDict, llm_router: Router, -) -> list[dict] | None: - model_id: Final = deployment.get("model_info", {}).get("id") - backend_model: Final = deployment.get("litellm_params", {}).get("model") +) -> Sequence[Mapping[str, object]] | None: + model_id: Final = _get_value(_get_value(deployment, "model_info"), "id") + backend_model: Final = _get_value(_get_value(deployment, "litellm_params"), "model") if not isinstance(model_id, str) or not isinstance(backend_model, str): return None deployment_model_info: Final = llm_router.get_deployment_model_info(model_id=model_id, model_name=backend_model) @@ -1352,7 +1355,7 @@ def _estimate_output_tokens( return min(requested, model_ceiling) -def _count_text_tokens(model: str, text: Any) -> int: +def _count_text_tokens(model: str, text: object) -> int: if text is None: return 0 @@ -1392,8 +1395,8 @@ def _is_input_only_route(route: str) -> bool: ) -def _to_float(value: Any) -> float | None: - if value is None: +def _to_float(value: object) -> float | None: + if not isinstance(value, (SupportsFloat, SupportsIndex, str, bytes, bytearray)): return None try: return float(value) @@ -1401,8 +1404,8 @@ def _to_float(value: Any) -> float | None: return None -def _to_int(value: Any) -> int | None: - if value is None: +def _to_int(value: object) -> int | None: + if not isinstance(value, (SupportsInt, SupportsIndex, str, bytes, bytearray)): return None try: return int(value) @@ -1410,7 +1413,7 @@ def _to_int(value: Any) -> int | None: return None -def _get_value(obj: Any, key: str) -> Any: - if isinstance(obj, dict): +def _get_value(obj: object, key: str) -> object: + if isinstance(obj, Mapping): return obj.get(key) return getattr(obj, key, None) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 52261d2c305..603271abd72 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1,9 +1,10 @@ import os import re import secrets +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel @@ -187,7 +188,28 @@ def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | return resolved_id -def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> dict: +_MISSING_ATTRIBUTE: Final = object() + + +def _attribute_or_missing(source: object, name: str) -> object: + return getattr(source, name, _MISSING_ATTRIBUTE) + + +@runtime_checkable +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + +def _dumped_usage_info(usage_info: object) -> object: + if isinstance(usage_info, _ModelDumpable): + return usage_info.model_dump() + instance_dict: Final = _attribute_or_missing(usage_info, "__dict__") + if instance_dict is not _MISSING_ATTRIBUTE: + return instance_dict + return usage_info + + +def _extract_usage_for_ocr_call(response_obj: object, response_obj_dict: dict) -> dict: """ Extract usage information for OCR/AOCR calls. @@ -208,12 +230,10 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d usage_info = response_obj_dict.get("usage_info") # Try to extract usage_info from object attributes if not found in dict - if not usage_info and hasattr(response_obj, "usage_info"): - usage_info = response_obj.usage_info - if hasattr(usage_info, "model_dump"): - usage_info = usage_info.model_dump() - elif hasattr(usage_info, "__dict__"): - usage_info = vars(usage_info) + if not usage_info: + attribute_usage_info: Final = _attribute_or_missing(response_obj, "usage_info") + if attribute_usage_info is not _MISSING_ATTRIBUTE: + usage_info = _dumped_usage_info(attribute_usage_info) # For OCR, we track pages instead of tokens if usage_info is not None: @@ -549,6 +569,14 @@ def _ensure_datetime_utc(timestamp: datetime) -> datetime: return timestamp +async def _query_raw_rows( + prisma_client: PrismaClient, + sql_query: str, + *args: object, +) -> Sequence[Mapping[str, object]] | None: + return await prisma_client.db.query_raw(sql_query, *args) + + async def get_spend_by_team( start_date: dt, end_date: dt, @@ -610,7 +638,7 @@ async def get_spend_by_team( group_by_day; """ - db_response: Final = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id) + db_response: Final = await _query_raw_rows(prisma_client, sql_query, start_date, end_date, team_id) if db_response is None: return [] @@ -685,7 +713,7 @@ async def get_spend_by_team_and_customer( group_by_day; """ - db_response: Final = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id, customer_id) + db_response: Final = await _query_raw_rows(prisma_client, sql_query, start_date, end_date, team_id, customer_id) if db_response is None: return [] @@ -740,7 +768,7 @@ def _sanitize_request_body_for_spend_logs_payload( return {} visited.add(obj_id) - def _sanitize_value(value: Any) -> Any: + def _sanitize_value(value: object) -> object: if isinstance(value, dict): return _sanitize_request_body_for_spend_logs_payload(value, visited, max_string_length_prompt_in_db) elif isinstance(value, list): @@ -1035,7 +1063,7 @@ def _sanitize_error_information_for_spend_logs( return cast(StandardLoggingPayloadErrorInformation, sanitized) -def _convert_to_json_serializable_dict(obj: Any, visited: set | None = None, max_depth: int = 20) -> Any: +def _convert_to_json_serializable_dict(obj: object, visited: set[int] | None = None, max_depth: int = 20) -> object: """ Convert object to JSON-serializable dict, handling Pydantic models safely. @@ -1089,6 +1117,13 @@ def _convert_to_json_serializable_dict(obj: Any, visited: set | None = None, max visited.remove(obj_id) +def _convert_mapping_to_json_serializable(obj: Mapping[str, object]) -> dict[str, object]: + converted: Final = _convert_to_json_serializable_dict(obj) + if isinstance(converted, dict): + return converted + return dict(obj) + + def _get_proxy_server_request_for_spend_logs_payload( metadata: dict, litellm_params: dict, @@ -1125,7 +1160,7 @@ def _get_proxy_server_request_for_spend_logs_payload( # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): - _request_body = _convert_to_json_serializable_dict(_request_body) + _request_body = _convert_mapping_to_json_serializable(_request_body) perform_redaction(model_call_details=_request_body, result=None) _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) @@ -1170,7 +1205,7 @@ def _get_response_for_spend_logs_payload( if payload is None: return "{}" if _should_store_prompts_and_responses_in_spend_logs(): - response_obj: Any = payload.get("response") + response_obj: object = payload.get("response") if response_obj is None: return "{}" diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index a1eb7ed06eb..52258602581 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -3,10 +3,11 @@ import asyncio import json import os from collections import Counter -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import ( - Any, Final, + NamedTuple, Protocol, cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read ) @@ -15,6 +16,7 @@ from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile from pydantic import ConfigDict, JsonValue, ValidationError, create_model from pydantic.fields import FieldInfo +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -44,6 +46,31 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router: Final = APIRouter() +JsonSchemaItems: Final = TypedDict( + "JsonSchemaItems", + {"$ref": ReadOnly[str], "enum": ReadOnly[Sequence[JsonValue]]}, + total=False, +) + + +class JsonSchemaNode(TypedDict, total=False): + type: ReadOnly[str] + description: ReadOnly[str] + enum: ReadOnly[Sequence[JsonValue]] + anyOf: ReadOnly[Sequence["JsonSchemaNode"]] + items: ReadOnly["JsonSchemaItems"] + properties: ReadOnly[Mapping[str, "JsonSchemaNode"]] + + +_EMPTY_SCHEMA_DEFS: Final[Mapping[str, "JsonSchemaNode"]] = MappingProxyType({}) + + +class JsonSchemaPropertyEntry(TypedDict): + description: ReadOnly[str] + type: ReadOnly[str] + items: NotRequired[ReadOnly["JsonSchemaItems"]] + + class _SsoSettingsMappingRow(Protocol): @property def sso_settings(self) -> Mapping[str, object] | None: ... @@ -157,10 +184,10 @@ class UIThemeConfig(BaseModel): class SettingsResponse(BaseModel): """Base response model for settings with values and schema information""" - values: dict[str, Any] + values: dict[str, object] """The current configuration values""" - field_schema: dict[str, Any] + field_schema: dict[str, object] """Schema information including descriptions and property types for UI display""" @@ -548,6 +575,62 @@ async def delete_allowed_ip( return {"message": f"IP {ip_address.ip} deleted successfully", "status": "success"} +def _resolve_non_null_variant(field_info: JsonSchemaNode) -> JsonSchemaNode: + """Pydantic v2 renders Optional fields as ``anyOf: [actual_type, null]``.""" + if "anyOf" not in field_info: + return field_info + return next((variant for variant in field_info["anyOf"] if variant.get("type") != "null"), field_info) + + +def _schema_items_entry(resolved: JsonSchemaNode, defs: Mapping[str, JsonSchemaNode]) -> "JsonSchemaItems | None": + """Items info (including enum values) for array fields, so the UI can render a multi-select dropdown.""" + if "items" not in resolved: + return None + items: Final = resolved["items"] + if "$ref" not in items: + return items + ref_def: Final = defs.get(items["$ref"].split("/")[-1]) + if ref_def is None or "enum" not in ref_def: + return None + enum_items: Final[JsonSchemaItems] = {"enum": ref_def["enum"]} + return enum_items + + +def _schema_property_entry(field_info: JsonSchemaNode, defs: Mapping[str, JsonSchemaNode]) -> JsonSchemaPropertyEntry: + resolved: Final = _resolve_non_null_variant(field_info) + items_entry: Final = _schema_items_entry(resolved, defs) + description: Final = field_info.get("description", "") + type_name: Final = resolved.get("type", "string") + if items_entry is None: + entry: Final[JsonSchemaPropertyEntry] = {"description": description, "type": type_name} + return entry + entry_with_items: Final[JsonSchemaPropertyEntry] = { + "description": description, + "type": type_name, + "items": items_entry, + } + return entry_with_items + + +class _RootSchema(NamedTuple): + description: str + properties: Mapping[str, JsonSchemaNode] + nested_defs: Mapping[str, JsonSchemaNode] + defs: Mapping[str, JsonSchemaNode] + + +def _root_schema(settings_class: type[BaseModel]) -> _RootSchema: + from pydantic import TypeAdapter + + raw_schema: Final = TypeAdapter(settings_class).json_schema(by_alias=True) + return _RootSchema( + description=raw_schema.get("description", ""), + properties=raw_schema["properties"], + nested_defs=raw_schema.get("definitions", _EMPTY_SCHEMA_DEFS), + defs=raw_schema["$defs"] if "$defs" in raw_schema else raw_schema.get("definitions", _EMPTY_SCHEMA_DEFS), + ) + + async def _get_settings_with_schema( settings_key: str, settings_class: type[BaseModel], @@ -561,69 +644,43 @@ async def _get_settings_with_schema( settings_class: The Pydantic class to use for schema config: The config dictionary """ - from pydantic import TypeAdapter - litellm_settings: Final = config.get("litellm_settings", {}) or {} settings_data: Final = litellm_settings.get(settings_key, {}) or {} # Create the settings object settings: Final = settings_class(**(settings_data)) # Get the schema - schema: Final = TypeAdapter(settings_class).json_schema(by_alias=True) + root_schema: Final = _root_schema(settings_class) # Convert to dict for response settings_dict: Final = settings.model_dump() # Add descriptions to the response - result: Final = { - "values": settings_dict, - "field_schema": { - "description": schema.get("description", ""), - "properties": {}, - }, + schema_properties_out: Final[Mapping[str, JsonSchemaPropertyEntry]] = { + field_name: _schema_property_entry(field_info, root_schema.defs) + for field_name, field_info in root_schema.properties.items() } - # Add property descriptions - defs: Final = schema.get("$defs", schema.get("definitions", {})) - for field_name, field_info in schema["properties"].items(): - # For Optional fields, Pydantic v2 uses anyOf with [actual_type, null]. - # Resolve the non-null variant to get the real type and items. - resolved = field_info - if "anyOf" in field_info: - for variant in field_info["anyOf"]: - if variant.get("type") != "null": - resolved = variant - break - - prop_entry: dict = { - "description": field_info.get("description", ""), - "type": resolved.get("type", "string"), - } - # Pass through items info (including enum values) for array fields - # so the UI can render a multi-select dropdown - if "items" in resolved: - items = resolved["items"] - # Resolve $ref to enum definitions if needed - if "$ref" in items: - ref_name = items["$ref"].split("/")[-1] - ref_def = defs.get(ref_name, {}) - if "enum" in ref_def: - prop_entry["items"] = {"enum": ref_def["enum"]} - else: - prop_entry["items"] = items - result["field_schema"]["properties"][field_name] = prop_entry - # Add nested object descriptions - for def_name, def_schema in schema.get("definitions", {}).items(): - result["field_schema"][def_name] = { + nested_defs_out: Final[Mapping[str, Mapping[str, object]]] = { + def_name: { "description": def_schema.get("description", ""), "properties": { prop_name: {"description": prop_info.get("description", "")} for prop_name, prop_info in def_schema.get("properties", {}).items() }, } + for def_name, def_schema in root_schema.nested_defs.items() + } - return result + return { + "values": settings_dict, + "field_schema": { + "description": root_schema.description, + "properties": schema_properties_out, + **nested_defs_out, + }, + } @router.get( @@ -930,32 +987,29 @@ async def get_sso_settings(): resolved: Final = resolve_sso_config(sso_db_settings, os.environ) # Get the schema for UI display - from pydantic import TypeAdapter - - schema: Final = TypeAdapter(SSOConfig).json_schema(by_alias=True) + root_schema: Final = _root_schema(SSOConfig) # Convert to dict for response, masking OAuth client secrets so plaintext # is never sent to the UI. sso_dict: Final = mask_sensitive_keys(resolved.config.model_dump(), set(SSO_SECRET_FIELDS)) # Add descriptions to the response - result: Final = { - "values": sso_dict, - "provenance": resolved.provenance, - "field_schema": { - "description": schema.get("description", ""), - "properties": {}, - }, - } - - # Add property descriptions - for field_name, field_info in schema["properties"].items(): - result["field_schema"]["properties"][field_name] = { + schema_properties_out: Final[Mapping[str, Mapping[str, str]]] = { + field_name: { "description": field_info.get("description", ""), "type": field_info.get("type", "string"), } + for field_name, field_info in root_schema.properties.items() + } - return result + return { + "values": sso_dict, + "provenance": resolved.provenance, + "field_schema": { + "description": root_schema.description, + "properties": schema_properties_out, + }, + } @router.patch( @@ -1309,7 +1363,7 @@ UI_SETTINGS_CACHE_KEY: Final = "ui_settings:settings_dict" UI_SETTINGS_CACHE_TTL: Final = 600 # 10 minutes -async def get_ui_settings_cached() -> dict[str, Any]: +async def get_ui_settings_cached() -> dict[str, JsonValue]: """ Return the persisted UI settings dict, using DualCache for reads. diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 368fd481e63..c1e09a7937f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -90,6 +90,16 @@ def _is_json_array(value: object) -> TypeIs[list[object]]: # guard-ok: trivial return isinstance(value, list) +def _optional_str(value: object) -> str | None: + """Keep a JSON payload entry only when it is a string, since the wire format is caller-controlled.""" + return value if isinstance(value, str) else None + + +def _json_array_or_empty(value: object) -> Sequence[object]: + """Narrow a JSON payload entry that the caller iterates, tolerating a missing or malformed value.""" + return value if _is_json_array(value) else () + + def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verifies every value is str return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) @@ -301,7 +311,7 @@ class BaseResponsesAPIStreamingIterator: ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - _item: Final = getattr(openai_responses_api_chunk, "item", None) + _item: Final[object] = getattr(openai_responses_api_chunk, "item", None) if _item is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_item, @@ -309,7 +319,7 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: - _annotation: Final = getattr(openai_responses_api_chunk, "annotation", None) + _annotation: Final[object] = getattr(openai_responses_api_chunk, "annotation", None) if _annotation is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_annotation, @@ -1081,7 +1091,7 @@ class _HasModelDumpJson(Protocol): def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... -def _dump_response_object(obj: object) -> dict[str, Any]: +def _dump_response_object(obj: object) -> dict[str, object]: if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): @@ -1254,7 +1264,7 @@ def _build_synthetic_response_events( ) if item_type == "message": - content_parts: Sequence[object] = output_item_payload.get("content", []) or [] + content_parts: Sequence[object] = _json_array_or_empty(output_item_payload.get("content")) for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( @@ -1302,7 +1312,7 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - summaries: Sequence[object] = output_item_payload.get("summary", []) or [] + summaries: Sequence[object] = _json_array_or_empty(output_item_payload.get("summary")) for summary_index, summary in enumerate(summaries): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") @@ -2018,7 +2028,7 @@ class ManagedResponsesWebSocketHandler: model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: dict[str, object] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, @@ -2031,9 +2041,9 @@ class ManagedResponsesWebSocketHandler: self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: dict[str, Any] = litellm_metadata or {} - self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( - "deployment_model_name" + self.litellm_metadata: dict[str, object] = litellm_metadata or {} + self.model_group: str | None = _optional_str( + self.litellm_metadata.get("model_group") or self.litellm_metadata.get("deployment_model_name") ) self.api_key = api_key self.api_base = api_base @@ -2055,7 +2065,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> str | None: + def _serialize_chunk(chunk: object) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: if isinstance(chunk, _HasModelDumpJson): @@ -2246,7 +2256,7 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: + def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, object]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} @@ -2462,12 +2472,12 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model: Final[str | None] = call_kwargs.pop("model", None) + requested_model: Final[str | None] = _optional_str(call_kwargs.pop("model", None)) model: Final[str] = ( self.model if requested_model is None or requested_model == self.model_group else requested_model ) - previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) + previous_response_id: Final[str | None] = _optional_str(call_kwargs.pop("previous_response_id", None)) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history diff --git a/litellm/types/guardrail_base_init.py b/litellm/types/guardrail_base_init.py new file mode 100644 index 00000000000..9174e8d840f --- /dev/null +++ b/litellm/types/guardrail_base_init.py @@ -0,0 +1,24 @@ +"""Typed view of the scalar keyword payload guardrails forward to ``CustomGuardrail.__init__``. + +Guardrail subclasses collect their base-class options in ``**kwargs`` and splat them into +``super().__init__``. Declaring the payload's shape here lets the checker resolve each +forwarded argument to its real parameter type instead of ``Any``. +""" + +from typing_extensions import ReadOnly, TypedDict + + +class GuardrailBaseInitKwargs(TypedDict, total=False): + guardrail_name: ReadOnly[str | None] + default_on: ReadOnly[bool] + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] From 8e687f10047ce7f0454204b94ed08e7bac2151b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:04:44 -0700 Subject: [PATCH 022/126] test(embeddings): move legacy intercepts to the wire for the omitted-format path The omitted-format path deliberately no longer dispatches through embeddings.create, so four legacy tests now intercept at the transport or client.post instead. Also adds a bypass error-path unit test, rewords a stale comment and a README scope note, and ratchets the lint budgets down. --- basedpyright-code-budget.json | 2 +- litellm/llms/hosted_vllm/embedding/README.md | 2 +- litellm/utils.py | 8 +-- ruff-strict-budget.json | 8 +-- test-quality-budget.json | 4 +- .../test_litellm_proxy_provider.py | 71 +++++++++++-------- tests/llm_translation/test_nvidia_nim.py | 48 ++++++++----- tests/local_testing/test_exceptions.py | 17 +++-- tests/local_testing/test_router.py | 4 +- ...penai_embedding_encoding_format_default.py | 19 +++++ type-discipline-budget.json | 2 +- 11 files changed, 117 insertions(+), 68 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 962d1266fd7..ba57fd1278b 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -105,7 +105,7 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38721 + "limit": 38720 }, "reportUnknownParameterType": { "limit": 19778 diff --git a/litellm/llms/hosted_vllm/embedding/README.md b/litellm/llms/hosted_vllm/embedding/README.md index 32b7ea5c560..50474aabdeb 100644 --- a/litellm/llms/hosted_vllm/embedding/README.md +++ b/litellm/llms/hosted_vllm/embedding/README.md @@ -4,7 +4,7 @@ VLLM is a superset of OpenAI's `embedding` endpoint. ## `encoding_format` -For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request: +For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request. `hosted_vllm/...` models use a separate handler that never adds the field on its own, so this resolution applies to the `openai/...`-style routes only: 1. Explicit value on the embedding call (`encoding_format=...`). 2. Model config (`litellm_params.encoding_format` on the proxy `model_list` entry). diff --git a/litellm/utils.py b/litellm/utils.py index fa2226dbf2c..66a07d2e2db 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3542,10 +3542,10 @@ def get_optional_params_embeddings( non_default_params=non_default_params, optional_params={}, kwargs=kwargs ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini": - # OpenAI SDKs (and litellm's own client) send encoding_format="float" - # by default; float lists are exactly what the vertex API returns, so - # the param is a no-op — don't reject the provider default. Other - # values (e.g. "base64") stay on the unsupported-param path below. + # OpenAI SDKs send encoding_format="float" by default; float lists are + # exactly what the vertex API returns, so the param is a no-op — don't + # reject the provider default. Other values (e.g. "base64") stay on + # the unsupported-param path below. if non_default_params.get("encoding_format") == "float": non_default_params.pop("encoding_format") supported_params = get_supported_openai_params( diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index c60988eccc0..73c69a732fe 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2003 + "limit": 2001 }, "ANN202": { "limit": 845 @@ -240,13 +240,13 @@ "limit": 96 }, "TRY201": { - "limit": 405 + "limit": 403 }, "TRY203": { - "limit": 113 + "limit": 111 }, "TRY300": { - "limit": 857 + "limit": 855 }, "UP028": { "limit": 2 diff --git a/test-quality-budget.json b/test-quality-budget.json index ee33eb581d6..d834c581609 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 742 + "limit": 741 }, "TQ003": { "limit": 62 @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11139 + "limit": 11135 } } diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index 1cb805bf9ba..8630259877d 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -5,6 +5,7 @@ from io import BytesIO from unittest.mock import AsyncMock +import httpx import litellm from litellm import completion, embedding import pytest @@ -92,44 +93,54 @@ async def test_litellm_gateway_from_sdk_embedding(is_async): litellm.set_verbose = True litellm._turn_on_debug() + captured_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "my-vllm-model", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + if is_async: from openai import AsyncOpenAI - openai_client = AsyncOpenAI(api_key="fake-key") - mock_method = AsyncMock() - patch_target = openai_client.embeddings.create + openai_client = AsyncOpenAI( + api_key="fake-key", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + response = await litellm.aembedding( + model="litellm_proxy/my-vllm-model", + input="Hello world", + client=openai_client, + api_base="my-custom-api-base", + ) else: from openai import OpenAI - openai_client = OpenAI(api_key="fake-key") - mock_method = MagicMock() - patch_target = openai_client.embeddings.create + openai_client = OpenAI( + api_key="fake-key", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + response = litellm.embedding( + model="litellm_proxy/my-vllm-model", + input="Hello world", + client=openai_client, + api_base="my-custom-api-base", + ) - with patch.object(patch_target.__self__, patch_target.__name__, new=mock_method): - try: - if is_async: - await litellm.aembedding( - model="litellm_proxy/my-vllm-model", - input="Hello world", - client=openai_client, - api_base="my-custom-api-base", - ) - else: - litellm.embedding( - model="litellm_proxy/my-vllm-model", - input="Hello world", - client=openai_client, - api_base="my-custom-api-base", - ) - except Exception as e: - print(e) + request_body = captured_bodies[0] + print("Request body - {}".format(request_body)) - mock_method.assert_called_once() - - print("Call KWARGS - {}".format(mock_method.call_args.kwargs)) - - assert "Hello world" == mock_method.call_args.kwargs["input"] - assert "my-vllm-model" == mock_method.call_args.kwargs["model"] + assert "Hello world" == request_body["input"] + assert "my-vllm-model" == request_body["model"] + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] @pytest.mark.parametrize("is_async", [False, True]) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 7ee4f347f72..d5942e674d0 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -63,27 +63,39 @@ def test_embedding_nvidia_nim(): litellm.set_verbose = True from openai import OpenAI + captured_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "nvidia/nv-embedqa-e5-v5", + "usage": {"prompt_tokens": 6, "total_tokens": 6}, + }, + ) + client = OpenAI( api_key="fake-api-key", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), ) - with patch.object(client.embeddings.with_raw_response, "create") as mock_client: - try: - litellm.embedding( - model="nvidia_nim/nvidia/nv-embedqa-e5-v5", - input="What is the meaning of life?", - input_type="passage", - dimensions=1024, - client=client, - ) - except Exception as e: - print(e) - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - print("request_body: ", request_body) - assert request_body["input"] == "What is the meaning of life?" - assert request_body["model"] == "nvidia/nv-embedqa-e5-v5" - assert request_body["extra_body"]["input_type"] == "passage" - assert request_body["dimensions"] == 1024 + response = litellm.embedding( + model="nvidia_nim/nvidia/nv-embedqa-e5-v5", + input="What is the meaning of life?", + input_type="passage", + dimensions=1024, + client=client, + ) + request_body = captured_bodies[0] + print("request_body: ", request_body) + assert request_body["input"] == "What is the meaning of life?" + assert request_body["model"] == "nvidia/nv-embedqa-e5-v5" + assert request_body["input_type"] == "passage" + assert request_body["dimensions"] == 1024 + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] def test_chat_completion_nvidia_nim_with_tools(): diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 8370046446d..e6392cda406 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -5,7 +5,7 @@ import traceback from typing import Any import httpx -from openai import AsyncOpenAI, AuthenticationError, BadRequestError, OpenAIError, RateLimitError +from openai import AsyncAzureOpenAI, AsyncOpenAI, AuthenticationError, AzureOpenAI, BadRequestError, OpenAIError, RateLimitError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -895,7 +895,12 @@ def _pre_call_utils( ): if call_type == "embedding": data["input"] = "Hello world!" - mapped_target: Any = client.embeddings.with_raw_response + if isinstance(client, (AzureOpenAI, AsyncAzureOpenAI)): + mapped_target: Any = client.embeddings.with_raw_response + patched_attr = "create" + else: + mapped_target = client + patched_attr = "post" if sync_mode: original_function = litellm.embedding else: @@ -905,6 +910,7 @@ def _pre_call_utils( if streaming is True: data["stream"] = True mapped_target = client.chat.completions.with_raw_response # type: ignore + patched_attr = "create" if sync_mode: original_function = litellm.completion else: @@ -914,12 +920,13 @@ def _pre_call_utils( if streaming is True: data["stream"] = True mapped_target = client.completions.with_raw_response # type: ignore + patched_attr = "create" if sync_mode: original_function = litellm.text_completion else: original_function = litellm.atext_completion - return data, original_function, mapped_target + return data, original_function, mapped_target, patched_attr def _pre_call_utils_httpx( @@ -1003,7 +1010,7 @@ async def test_exception_with_headers(sync_mode, provider, model, call_type, str ) data = {"model": model} - data, original_function, mapped_target = _pre_call_utils( + data, original_function, mapped_target, patched_attr = _pre_call_utils( call_type=call_type, data=data, client=openai_client, @@ -1049,7 +1056,7 @@ async def test_exception_with_headers(sync_mode, provider, model, call_type, str with patch.object( mapped_target, - "create", + patched_attr, side_effect=_return_exception, ): new_retry_after_mock_client = MagicMock(return_value=-1) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 370c43f8f44..c714bb4f9a7 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -2032,8 +2032,8 @@ def test_router_dynamic_cooldown_correct_retry_after_time(): raise exception with patch.object( - openai_client.embeddings.with_raw_response, - "create", + openai_client, + "post", side_effect=_return_exception, ): new_retry_after_mock_client = MagicMock(return_value=-1) diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/test_litellm/test_openai_embedding_encoding_format_default.py index 9842bf30585..7a42eaf0f0a 100644 --- a/tests/test_litellm/test_openai_embedding_encoding_format_default.py +++ b/tests/test_litellm/test_openai_embedding_encoding_format_default.py @@ -100,3 +100,22 @@ async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( request_body: Final = json.loads(mock_route.calls.last.request.read()) assert "encoding_format" not in request_body assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_omitted_encoding_format_maps_provider_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 429, + headers={"retry-after": "42", "x-should-retry": "false"}, + json={"error": {"message": "rate limited", "type": "rate_limit_error"}}, + ) + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", max_retries=0 + ) + + assert int(exc_info.value.litellm_response_headers["retry-after"]) == 42 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3f1a7defe7..67ebb2b3b17 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16564 + "limit": 16562 }, "LIT011": { "limit": 5577 From 2a88384e4ec2d2f183fb69f223ed376b2eb24991 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:15:32 -0700 Subject: [PATCH 023/126] style(utils): drop an em-dash from the vertex encoding_format comment --- litellm/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 66a07d2e2db..c4250dd0e4b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3543,9 +3543,9 @@ def get_optional_params_embeddings( ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini": # OpenAI SDKs send encoding_format="float" by default; float lists are - # exactly what the vertex API returns, so the param is a no-op — don't - # reject the provider default. Other values (e.g. "base64") stay on - # the unsupported-param path below. + # exactly what the vertex API returns, so the param is a no-op and the + # provider default is not rejected. Other values (e.g. "base64") stay + # on the unsupported-param path below. if non_default_params.get("encoding_format") == "float": non_default_params.pop("encoding_format") supported_params = get_supported_openai_params( From 39d81380ba34a6d3d519ce4faef24046ef47b62e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 19:25:48 +0000 Subject: [PATCH 024/126] chore(lint): fix post-merge type regressions and ratchet lint budgets --- basedpyright-code-budget.json | 20 +++++++++---------- litellm/caching/valkey_semantic_cache.py | 5 +++-- .../arize/arize_phoenix_prompt_manager.py | 6 ++++-- .../context_management/editors/compact.py | 15 +++++++++----- litellm/llms/vertex_ai/common_utils.py | 7 ++++--- .../guardrail_hooks/bedrock_guardrails.py | 4 +--- litellm/proxy/litellm_pre_call_utils.py | 6 ++++-- .../key_management_endpoints.py | 5 ++++- ruff-strict-budget.json | 16 +++++++-------- type-discipline-budget.json | 8 ++++---- 10 files changed, 52 insertions(+), 40 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef88ae574fb..00c79c1e418 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 17270 + "limit": 16279 }, "reportArgumentType": { - "limit": 2539 + "limit": 2530 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5486 + "limit": 5063 }, "reportFunctionMemberAccess": { "limit": 7 @@ -42,7 +42,7 @@ "limit": 12 }, "reportIndexIssue": { - "limit": 35 + "limit": 30 }, "reportInvalidTypeForm": { "limit": 34 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5658 + "limit": 5642 }, "reportMissingTypeArgument": { - "limit": 15425 + "limit": 15404 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,19 +105,19 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38721 + "limit": 38622 }, "reportUnknownParameterType": { - "limit": 19778 + "limit": 19748 }, "reportUnknownVariableType": { - "limit": 30290 + "limit": 30210 }, "reportUnnecessaryCast": { "limit": 117 }, "reportUnnecessaryComparison": { - "limit": 697 + "limit": 696 }, "reportUnnecessaryContains": { "limit": 5 diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index b63b2e0dc10..ec91651ed33 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -19,7 +19,7 @@ import hashlib import os from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any, Final, Protocol +from typing import Any, Final, Protocol, cast from redis import Redis from redis.asyncio import Redis as AsyncRedis @@ -294,7 +294,8 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose("No prompt provided for semantic caching") return - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + metadata: Final = cast("dict[str, object] | None", kwargs.get("metadata")) # cast-ok: untyped kwargs + embedding: Final = await self._get_async_embedding(prompt, metadata=metadata) await self._ensure_index_async(len(embedding)) doc_key: Final = self._doc_key(key) diff --git a/litellm/integrations/arize/arize_phoenix_prompt_manager.py b/litellm/integrations/arize/arize_phoenix_prompt_manager.py index 0c616e845f8..0c9e868c146 100644 --- a/litellm/integrations/arize/arize_phoenix_prompt_manager.py +++ b/litellm/integrations/arize/arize_phoenix_prompt_manager.py @@ -4,7 +4,7 @@ Fetches prompt versions from Arize Phoenix and provides workspace-based access c """ from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Any, Final, cast from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -203,7 +203,9 @@ class ArizePhoenixTemplateManager: # Combine rendered content final_content = " ".join(rendered_content_parts) - rendered_messages.append({"role": role, "content": final_content}) + rendered_messages.append( + cast("AllMessageValues", {"role": role, "content": final_content}) # cast-ok: Phoenix roles are OpenAI + ) return rendered_messages diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 4551ff5213f..a45bfb93640 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -36,6 +36,7 @@ from litellm.types.llms.anthropic import ( CompactionBlock, UsageIteration, ) +from litellm.types.llms.openai import AllMessageValues if TYPE_CHECKING: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -867,7 +868,7 @@ def _append_text_to_content(content: object, extra_text: str) -> object: class _SummaryCallUserKwarg(TypedDict, total=False): - user: ReadOnly[object] + user: ReadOnly[str] class _SummaryCallRegionKwarg(TypedDict, total=False): @@ -876,11 +877,11 @@ class _SummaryCallRegionKwarg(TypedDict, total=False): class _SummaryCallKwargs(TypedDict): model: ReadOnly[str] - messages: ReadOnly[list[dict[str, object]]] + messages: ReadOnly[list[AllMessageValues]] max_tokens: ReadOnly[int] timeout: ReadOnly[float] litellm_metadata: ReadOnly[Mapping[str, object]] - user: NotRequired[ReadOnly[object]] + user: NotRequired[ReadOnly[str]] allowed_model_region: NotRequired[ReadOnly[str]] @@ -927,11 +928,15 @@ async def _call_summary_model( end_user_id: Final = metadata.get("user_api_key_end_user_id") call_kwargs: Final[_SummaryCallKwargs] = { "model": summary_model, - "messages": summary_messages, + "messages": cast("list[AllMessageValues]", summary_messages), # cast-ok: built as OpenAI chat messages "max_tokens": max_tokens, "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, "litellm_metadata": metadata, - **(_SummaryCallUserKwarg(user=end_user_id) if end_user_id else _SummaryCallUserKwarg()), + **( + _SummaryCallUserKwarg(user=end_user_id) + if isinstance(end_user_id, str) and end_user_id + else _SummaryCallUserKwarg() + ), **( _SummaryCallRegionKwarg(allowed_model_region=allowed_model_region) if allowed_model_region is not None diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 48649cf3105..1de316835a1 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,7 +1,7 @@ import re from copy import deepcopy from enum import Enum -from typing import Any, Final, Literal, get_type_hints +from typing import Any, Final, Literal, cast, get_type_hints import httpx @@ -726,8 +726,9 @@ def set_schema_property_ordering(schema: dict[str, object], depth: int = 0) -> d schema["propertyOrdering"] = [k for k, v in schema["properties"].items()] for k, v in schema["properties"].items(): set_schema_property_ordering(v, depth + 1) - if "items" in schema: - set_schema_property_ordering(schema["items"], depth + 1) + items: Final = schema.get("items") + if isinstance(items, dict): + set_schema_property_ordering(cast("dict[str, object]", items), depth + 1) # cast-ok: JSON Schema child return schema diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 29fcafa40fa..4a3c8b0628d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -364,9 +364,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_request["content"] = bedrock_request_content return bedrock_request - def _build_response_content_items( - self, response: object, has_grounding: bool - ) -> list[BedrockContentItem]: + def _build_response_content_items(self, response: object, has_grounding: bool) -> list[BedrockContentItem]: """Build content item(s) from the model response. When the request supplied grounding, the response is qualified ``guard_content`` so Bedrock can score it. """ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index efa7cb04315..b49541f63ed 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -7,7 +7,7 @@ from collections import OrderedDict from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast from fastapi import HTTPException, Request from pydantic import ValidationError as PydanticValidationError @@ -1343,6 +1343,8 @@ class LiteLLMProxyRequestSetup: def get_sanitized_user_information_from_key( user_api_key_dict: UserAPIKeyAuth, ) -> StandardLoggingUserAPIKeyMetadata: + stripped_metadata: Final = strip_callback_config(user_api_key_dict.metadata) + auth_metadata: Final = cast("dict[str, str] | None", stripped_metadata) # cast-ok: metadata is free-form JSON user_api_key_logged_metadata: Final = StandardLoggingUserAPIKeyMetadata( user_api_key_hash=user_api_key_dict.api_key, # just the hashed token user_api_key_alias=user_api_key_dict.key_alias, @@ -1365,7 +1367,7 @@ class LiteLLMProxyRequestSetup: user_api_key_budget_reset_at=( user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None ), - user_api_key_auth_metadata=strip_callback_config(user_api_key_dict.metadata), + user_api_key_auth_metadata=auth_metadata, ) return user_api_key_logged_metadata diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e1157b75107..f0f684b0fa1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2476,6 +2476,9 @@ async def _validate_update_key_data( user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "Database not connected"}) + # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) validate_budget_duration(data.budget_duration) @@ -2594,7 +2597,7 @@ async def _validate_update_key_data( # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key: Final = getattr(existing_key_row, "team_id", None) is not None can_skip_admin_check: Final = (caller_is_creator or _key_is_team_key) and not _is_budget_change - if (not _is_proxy_admin) and prisma_client is not None and not can_skip_admin_check: + if (not _is_proxy_admin) and not can_skip_admin_check: hashed_key: Final = existing_key_row.token await _check_key_admin_access( user_api_key_dict=user_api_key_dict, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index c60988eccc0..33877524fb5 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 3012 + "limit": 3004 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 827 + "limit": 825 }, "ANN201": { "limit": 2003 }, "ANN202": { - "limit": 845 + "limit": 843 }, "ANN204": { - "limit": 702 + "limit": 700 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 655 + "limit": 517 }, "ASYNC230": { "limit": 11 @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 175 + "limit": 174 }, "RUF012": { "limit": 239 @@ -198,7 +198,7 @@ "limit": 58 }, "SIM102": { - "limit": 315 + "limit": 313 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1117 + "limit": 1105 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3f1a7defe7..be3c9589015 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22705 + "limit": 22655 }, "LIT002": { - "limit": 26854 + "limit": 26830 }, "LIT003": { "limit": 269 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16564 + "limit": 16546 }, "LIT011": { - "limit": 5577 + "limit": 5558 }, "LIT012": { "limit": 4508 From 4e2574ce08701e89d6f61a5df8557e48d3c5d4bc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:44:43 -0700 Subject: [PATCH 025/126] fix(proxy): match /v1/audio/speech content-type to the returned audio format --- .../litellm_core_utils/audio_utils/utils.py | 30 ++++++++++++++- litellm/proxy/proxy_server.py | 15 +++----- .../audio_utils/test_utils.py | 30 +++++++++++++++ .../proxy/proxy_server/test_routes_audio.py | 38 ++++++++++++++++++- .../test_audio_speech_prometheus_hooks.py | 2 + 5 files changed, 104 insertions(+), 11 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 3b3775a8fe6..4d75d5dc8f5 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -7,7 +7,12 @@ import os from dataclasses import dataclass from typing import Final -from litellm.types.files import get_file_mime_type_from_extension +from litellm.types.files import ( + AUDIO_FILE_TYPES, + FILE_EXTENSIONS, + FILE_MIME_TYPES, + get_file_mime_type_from_extension, +) from litellm.types.utils import FileTypes @@ -323,3 +328,26 @@ def calculate_request_duration(file: FileTypes) -> float | None: except Exception: # Silently fail if duration extraction fails return None + + +DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg" + + +def _speech_media_type_for_response_format(response_format: str) -> str | None: + file_type: Final = next( + (candidate for candidate, extensions in FILE_EXTENSIONS.items() if response_format.lower() in extensions), + None, + ) + if file_type is None or file_type not in AUDIO_FILE_TYPES: + return None + return FILE_MIME_TYPES[file_type] + + +def resolve_speech_media_type(upstream_content_type: str | None, response_format: str | None) -> str: + upstream_media_type: Final = (upstream_content_type or "").split(";", 1)[0].strip().lower() + if upstream_media_type.startswith("audio/"): + return upstream_media_type + requested_media_type: Final = ( + None if response_format is None else _speech_media_type_for_response_format(response_format) + ) + return requested_media_type or DEFAULT_SPEECH_MEDIA_TYPE diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af0aa9743bc..96ff4f0daef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -263,6 +263,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs, @@ -10877,15 +10878,11 @@ async def audio_speech( if callback_headers: custom_headers.update(callback_headers) - # Determine media type based on model type - media_type = "audio/mpeg" # Default for OpenAI TTS - request_model: Final = data.get("model", "") - if request_model: - request_model_lower: Final = request_model.lower() - if "gemini" in request_model_lower and ( - "tts" in request_model_lower or "preview-tts" in request_model_lower - ): - media_type = "audio/wav" # Gemini TTS returns WAV format after conversion + requested_format: Final = data.get("response_format") + media_type: Final = resolve_speech_media_type( + upstream_content_type=response.response.headers.get("content-type"), + response_format=requested_format if isinstance(requested_format, str) else None, + ) return StreamingResponse( _audio_speech_chunk_generator(response), diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py b/tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py new file mode 100644 index 00000000000..87207588ba5 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py @@ -0,0 +1,30 @@ +import pytest + +from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type + + +@pytest.mark.parametrize( + ("upstream_content_type", "response_format", "expected"), + [ + ("audio/wav", None, "audio/wav"), + ("AUDIO/WAV", None, "audio/wav"), + ("audio/flac; charset=binary", "mp3", "audio/flac"), + ("application/json", "flac", "audio/flac"), + ("application/octet-stream", "pcm", "audio/pcm"), + (None, "wav", "audio/wav"), + (None, "WAV", "audio/wav"), + (None, "opus", "audio/opus"), + (None, "aac", "audio/aac"), + (None, "mp3", "audio/mpeg"), + (None, "mp4", "audio/mpeg"), + (None, "bogus", "audio/mpeg"), + (None, None, "audio/mpeg"), + ("", None, "audio/mpeg"), + ], +) +def test_resolve_speech_media_type(upstream_content_type, response_format, expected): + resolved = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=response_format, + ) + assert resolved == expected diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py index 74542a3eaf6..b99affc2ac3 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -12,13 +12,15 @@ from __future__ import annotations import io from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from litellm.proxy import proxy_server @pytest.fixture -def patched_speech(monkeypatch): +def patched_speech(monkeypatch, request): + upstream_content_type = getattr(request, "param", "audio/mpeg") monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) monkeypatch.setattr( proxy_server, @@ -37,6 +39,11 @@ def patched_speech(monkeypatch): monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data) class _FakeBinaryResp: + response = httpx.Response( + status_code=200, + headers={} if upstream_content_type is None else {"content-type": upstream_content_type}, + ) + async def aiter_bytes(self, chunk_size: int = 8192): async def _gen(): yield b"\x00\x01\x02" @@ -152,6 +159,35 @@ def test_audio_speech_happy_path(client, auth_as, patched_speech, path): } +@pytest.mark.parametrize( + ("patched_speech", "response_format", "expected_content_type"), + [ + ("audio/wav", "wav", "audio/wav"), + ("audio/flac", "flac", "audio/flac"), + ("audio/pcm", "pcm", "audio/pcm"), + ("audio/wav", "mp3", "audio/wav"), + ("application/json", "flac", "audio/flac"), + (None, "wav", "audio/wav"), + (None, None, "audio/mpeg"), + ], + indirect=["patched_speech"], +) +def test_audio_speech_content_type_matches_audio_format( + client, auth_as, patched_speech, response_format, expected_content_type +): + """Regression for LIT-6482: /v1/audio/speech mislabeled wav/flac/pcm as audio/mpeg.""" + payload = { + "model": "tts-1", + "input": "Hi", + "voice": "alloy", + **({} if response_format is None else {"response_format": response_format}), + } + with auth_as(): + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 200 + assert response.headers.get("content-type", "").split(";")[0] == expected_content_type + + @pytest.mark.parametrize("path", ["/v1/audio/speech", "/audio/speech"]) def test_audio_speech_error(client, auth_as, patched_speech_error, path): """Pins ``POST /v1/audio/speech`` and ``POST /audio/speech`` (error).""" diff --git a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py index 99f6f3a9b72..959cb2b1e89 100644 --- a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py +++ b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py @@ -2,6 +2,7 @@ import asyncio import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi.testclient import TestClient @@ -29,6 +30,7 @@ def _make_mock_tts_response(): inner = MagicMock() inner.aiter_bytes = _aiter_bytes inner._hidden_params = {} + inner.response = httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}) async def _resolver(): return inner From c251d6d609e5a85ae8df28411304de9dc7d84b06 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:11:51 -0700 Subject: [PATCH 026/126] fix(vertex_ai): label TTS audio bytes with their real content-type --- .../litellm_core_utils/audio_utils/utils.py | 22 ++++++++ .../text_to_speech/transformation.py | 8 ++- .../audio_utils/test_utils.py | 30 ---------- .../litellm_core_utils/test_audio_utils.py | 56 +++++++++++++++++++ .../text_to_speech/test_transformation.py | 43 ++++++++++++++ 5 files changed, 126 insertions(+), 33 deletions(-) delete mode 100644 tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 4d75d5dc8f5..89222bf8107 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -11,6 +11,7 @@ from litellm.types.files import ( AUDIO_FILE_TYPES, FILE_EXTENSIONS, FILE_MIME_TYPES, + FileType, get_file_mime_type_from_extension, ) from litellm.types.utils import FileTypes @@ -351,3 +352,24 @@ def resolve_speech_media_type(upstream_content_type: str | None, response_format None if response_format is None else _speech_media_type_for_response_format(response_format) ) return requested_media_type or DEFAULT_SPEECH_MEDIA_TYPE + + +_OGG_OPUS_HEAD_WINDOW: Final = 64 +_MPEG_FRAME_SYNC_MASK: Final = 0xE0 +_MPEG_FRAME_LAYER_MASK: Final = 0x06 + + +def speech_media_type_from_audio_bytes(audio: bytes) -> str | None: + if audio[:4] == b"RIFF" and audio[8:12] == b"WAVE": + return FILE_MIME_TYPES[FileType.WAV] + if audio[:4] == b"fLaC": + return FILE_MIME_TYPES[FileType.FLAC] + if audio[:4] == b"OggS": + is_opus: Final = b"OpusHead" in audio[:_OGG_OPUS_HEAD_WINDOW] + return FILE_MIME_TYPES[FileType.OPUS if is_opus else FileType.OGG] + if audio[:3] == b"ID3": + return FILE_MIME_TYPES[FileType.MP3] + if len(audio) < 2 or audio[0] != 0xFF or (audio[1] & _MPEG_FRAME_SYNC_MASK) != _MPEG_FRAME_SYNC_MASK: + return None + is_adts_aac: Final = (audio[1] & _MPEG_FRAME_LAYER_MASK) == 0 + return FILE_MIME_TYPES[FileType.AAC if is_adts_aac else FileType.MP3] diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index cf14ab88751..a5ff7eca021 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -11,6 +11,9 @@ from typing import TYPE_CHECKING, Any, Final, Union import httpx +from litellm.litellm_core_utils.audio_utils.utils import ( + speech_media_type_from_audio_bytes, +) from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -457,12 +460,11 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): if not response_content: raise ValueError("No audioContent in Vertex AI TTS response") - # Decode base64 to get binary content binary_data: Final = base64.b64decode(response_content) - - # Create an httpx.Response object with the binary data + media_type: Final = speech_media_type_from_audio_bytes(binary_data) response: Final = httpx.Response( status_code=200, + headers={} if media_type is None else {"content-type": media_type}, content=binary_data, ) diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py b/tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py deleted file mode 100644 index 87207588ba5..00000000000 --- a/tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py +++ /dev/null @@ -1,30 +0,0 @@ -import pytest - -from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type - - -@pytest.mark.parametrize( - ("upstream_content_type", "response_format", "expected"), - [ - ("audio/wav", None, "audio/wav"), - ("AUDIO/WAV", None, "audio/wav"), - ("audio/flac; charset=binary", "mp3", "audio/flac"), - ("application/json", "flac", "audio/flac"), - ("application/octet-stream", "pcm", "audio/pcm"), - (None, "wav", "audio/wav"), - (None, "WAV", "audio/wav"), - (None, "opus", "audio/opus"), - (None, "aac", "audio/aac"), - (None, "mp3", "audio/mpeg"), - (None, "mp4", "audio/mpeg"), - (None, "bogus", "audio/mpeg"), - (None, None, "audio/mpeg"), - ("", None, "audio/mpeg"), - ], -) -def test_resolve_speech_media_type(upstream_content_type, response_format, expected): - resolved = resolve_speech_media_type( - upstream_content_type=upstream_content_type, - response_format=response_format, - ) - assert resolved == expected diff --git a/tests/test_litellm/litellm_core_utils/test_audio_utils.py b/tests/test_litellm/litellm_core_utils/test_audio_utils.py index 0e8176fffce..693ff760e2d 100644 --- a/tests/test_litellm/litellm_core_utils/test_audio_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_audio_utils.py @@ -347,3 +347,59 @@ class TestNormalizeTranscriptionLanguageToBcp47: ) assert normalize_transcription_language_to_bcp47(language) == expected + + +class TestResolveSpeechMediaType: + @pytest.mark.parametrize( + ("upstream_content_type", "response_format", "expected"), + [ + ("audio/wav", None, "audio/wav"), + ("AUDIO/WAV", None, "audio/wav"), + ("audio/flac; charset=binary", "mp3", "audio/flac"), + ("application/json", "flac", "audio/flac"), + ("application/octet-stream", "pcm", "audio/pcm"), + (None, "wav", "audio/wav"), + (None, "WAV", "audio/wav"), + (None, "opus", "audio/opus"), + (None, "aac", "audio/aac"), + (None, "mp3", "audio/mpeg"), + (None, "mp4", "audio/mpeg"), + (None, "bogus", "audio/mpeg"), + (None, None, "audio/mpeg"), + ("", None, "audio/mpeg"), + ], + ) + def test_resolution(self, upstream_content_type, response_format, expected): + from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type + + resolved = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=response_format, + ) + assert resolved == expected + + +class TestSpeechMediaTypeFromAudioBytes: + @pytest.mark.parametrize( + ("audio", "expected"), + [ + (b"RIFF\x24\x00\x00\x00WAVEfmt ", "audio/wav"), + (b"fLaC\x00\x00\x00\x22", "audio/flac"), + (b"OggS" + b"\x00" * 24 + b"OpusHead", "audio/opus"), + (b"OggS" + b"\x00" * 24 + b"\x01vorbis", "audio/ogg"), + (b"ID3\x04\x00\x00\x00\x00\x00\x00", "audio/mpeg"), + (b"\xff\xfb\x90\x64", "audio/mpeg"), + (b"\xff\xf3\x80\x00", "audio/mpeg"), + (b"\xff\xf1\x50\x80", "audio/aac"), + (b"\xff\xf9\x50\x80", "audio/aac"), + (b"RIFF\x24\x00\x00\x00AVI LIST", None), + (b"\xff\x00\x00\x00", None), + (b"\x00\x01\x02\x03\x04\x05", None), + (b"\xff", None), + (b"", None), + ], + ) + def test_sniffing(self, audio, expected): + from litellm.litellm_core_utils.audio_utils.utils import speech_media_type_from_audio_bytes + + assert speech_media_type_from_audio_bytes(audio) == expected diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 05da22a73fd..fba337b5f2c 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -1,3 +1,4 @@ +import base64 from unittest.mock import MagicMock, Mock, patch import httpx @@ -126,6 +127,48 @@ class TestVertexAITextToSpeechConfig: assert voice_dict == voice_input +@pytest.mark.parametrize( + ("audio", "expected_content_type"), + [ + (b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00", "audio/wav"), + (b"\xff\xfb\x90\x64\x00\x00\x00\x00", "audio/mpeg"), + (b"OggS" + b"\x00" * 24 + b"OpusHead", "audio/opus"), + (b"fLaC\x00\x00\x00\x22", "audio/flac"), + ], +) +def test_transform_text_to_speech_response_labels_content_type(audio, expected_content_type): + raw_response = httpx.Response( + status_code=200, + json={"audioContent": base64.b64encode(audio).decode()}, + ) + + result = VertexAITextToSpeechConfig().transform_text_to_speech_response( + model="vertex_ai/chirp", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert result.response.headers["content-type"] == expected_content_type + assert result.response.content == audio + + +def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): + raw_pcm = b"\x00\x01\x02\x03\x04\x05\x06\x07" + raw_response = httpx.Response( + status_code=200, + json={"audioContent": base64.b64encode(raw_pcm).decode()}, + ) + + result = VertexAITextToSpeechConfig().transform_text_to_speech_response( + model="vertex_ai/chirp", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert "content-type" not in result.response.headers + assert result.response.content == raw_pcm + + @patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") From 4cd11e8b19405c192b380be2f6c9a1a61518d1f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:20:40 -0700 Subject: [PATCH 027/126] fix(audio_utils): validate MPEG frame headers before labeling sniffed audio --- .../litellm_core_utils/audio_utils/utils.py | 38 ++++++++++++++++--- .../litellm_core_utils/test_audio_utils.py | 6 +++ 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 89222bf8107..dab3e48f91a 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -355,8 +355,35 @@ def resolve_speech_media_type(upstream_content_type: str | None, response_format _OGG_OPUS_HEAD_WINDOW: Final = 64 -_MPEG_FRAME_SYNC_MASK: Final = 0xE0 -_MPEG_FRAME_LAYER_MASK: Final = 0x06 +_ADTS_SYNC_AND_LAYER_MASK: Final = 0xF6 +_ADTS_SYNC_AND_LAYER: Final = 0xF0 +_ADTS_SAMPLE_RATE_INDEX_LIMIT: Final = 13 +_MPEG_SYNC_MASK: Final = 0xE0 +_MPEG_LAYER_MASK: Final = 0x06 +_MPEG_RESERVED_VERSION: Final = 0x01 +_MPEG_INVALID_BITRATE_INDEX: Final = 0x0F +_MPEG_RESERVED_SAMPLE_RATE_INDEX: Final = 0x03 + + +def _adts_aac_frame_media_type(header: bytes) -> str | None: + sample_rate_index: Final = (header[2] >> 2) & 0x0F + return FILE_MIME_TYPES[FileType.AAC] if sample_rate_index < _ADTS_SAMPLE_RATE_INDEX_LIMIT else None + + +def _mpeg_audio_frame_media_type(header: bytes) -> str | None: + version: Final = (header[1] >> 3) & 0x03 + layer: Final = header[1] & _MPEG_LAYER_MASK + bitrate_index: Final = header[2] >> 4 + sample_rate_index: Final = (header[2] >> 2) & 0x03 + if ( + (header[1] & _MPEG_SYNC_MASK) != _MPEG_SYNC_MASK + or version == _MPEG_RESERVED_VERSION + or layer == 0 + or bitrate_index == _MPEG_INVALID_BITRATE_INDEX + or sample_rate_index == _MPEG_RESERVED_SAMPLE_RATE_INDEX + ): + return None + return FILE_MIME_TYPES[FileType.MP3] def speech_media_type_from_audio_bytes(audio: bytes) -> str | None: @@ -369,7 +396,8 @@ def speech_media_type_from_audio_bytes(audio: bytes) -> str | None: return FILE_MIME_TYPES[FileType.OPUS if is_opus else FileType.OGG] if audio[:3] == b"ID3": return FILE_MIME_TYPES[FileType.MP3] - if len(audio) < 2 or audio[0] != 0xFF or (audio[1] & _MPEG_FRAME_SYNC_MASK) != _MPEG_FRAME_SYNC_MASK: + if len(audio) < 3 or audio[0] != 0xFF: return None - is_adts_aac: Final = (audio[1] & _MPEG_FRAME_LAYER_MASK) == 0 - return FILE_MIME_TYPES[FileType.AAC if is_adts_aac else FileType.MP3] + if (audio[1] & _ADTS_SYNC_AND_LAYER_MASK) == _ADTS_SYNC_AND_LAYER: + return _adts_aac_frame_media_type(audio) + return _mpeg_audio_frame_media_type(audio) diff --git a/tests/test_litellm/litellm_core_utils/test_audio_utils.py b/tests/test_litellm/litellm_core_utils/test_audio_utils.py index 693ff760e2d..155f6680416 100644 --- a/tests/test_litellm/litellm_core_utils/test_audio_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_audio_utils.py @@ -393,8 +393,14 @@ class TestSpeechMediaTypeFromAudioBytes: (b"\xff\xf1\x50\x80", "audio/aac"), (b"\xff\xf9\x50\x80", "audio/aac"), (b"RIFF\x24\x00\x00\x00AVI LIST", None), + (b"\xff\xff\xff\xff\xff\xff", None), + (b"\xff\xfb\xf0\x00", None), + (b"\xff\xfb\x9c\x00", None), + (b"\xff\xeb\x90\x00", None), + (b"\xff\xf1\xf4\x80", None), (b"\xff\x00\x00\x00", None), (b"\x00\x01\x02\x03\x04\x05", None), + (b"\xff\xfb", None), (b"\xff", None), (b"", None), ], From d804b9d4fe4753f94913b28f169b74a01478b93c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:26:47 -0700 Subject: [PATCH 028/126] fix(vertex_ai): skip non-dict property values in set_schema_property_ordering The typed rewrite made the properties recursion call .get on every child, so a malformed schema with a string or list property value raised AttributeError where it previously passed through untouched. --- basedpyright-code-budget.json | 22 +++++++++---------- litellm/llms/vertex_ai/common_utils.py | 5 +++-- ruff-strict-budget.json | 18 +++++++-------- .../vertex_ai/test_vertex_ai_common_utils.py | 16 ++++++++++++++ type-discipline-budget.json | 8 +++---- 5 files changed, 43 insertions(+), 26 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 8634e9a5c32..d572a328926 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 16282 + "limit": 15294 }, "reportArgumentType": { - "limit": 2529 + "limit": 2520 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5062 + "limit": 4639 }, "reportFunctionMemberAccess": { "limit": 7 @@ -42,7 +42,7 @@ "limit": 12 }, "reportIndexIssue": { - "limit": 30 + "limit": 25 }, "reportInvalidTypeForm": { "limit": 34 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5642 + "limit": 5626 }, "reportMissingTypeArgument": { - "limit": 15404 + "limit": 15383 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,19 +105,19 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38621 + "limit": 38521 }, "reportUnknownParameterType": { - "limit": 19748 + "limit": 19718 }, "reportUnknownVariableType": { - "limit": 30210 + "limit": 30129 }, "reportUnnecessaryCast": { "limit": 117 }, "reportUnnecessaryComparison": { - "limit": 696 + "limit": 695 }, "reportUnnecessaryContains": { "limit": 5 @@ -141,6 +141,6 @@ "limit": 543 }, "reportUnusedVariable": { - "limit": 139 + "limit": 138 } } diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 1de316835a1..a36c920dda0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -724,8 +724,9 @@ def set_schema_property_ordering(schema: dict[str, object], depth: int = 0) -> d # retain propertyOrdering as an escape hatch if user already specifies it if "propertyOrdering" not in schema: schema["propertyOrdering"] = [k for k, v in schema["properties"].items()] - for k, v in schema["properties"].items(): - set_schema_property_ordering(v, depth + 1) + for v in schema["properties"].values(): + if isinstance(v, dict): + set_schema_property_ordering(cast("dict[str, object]", v), depth + 1) # cast-ok: JSON Schema child items: Final = schema.get("items") if isinstance(items, dict): set_schema_property_ordering(cast("dict[str, object]", items), depth + 1) # cast-ok: JSON Schema child diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 35bedee08a9..f479764f269 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 2996 + "limit": 2988 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 823 + "limit": 821 }, "ANN201": { "limit": 2003 }, "ANN202": { - "limit": 841 + "limit": 839 }, "ANN204": { - "limit": 698 + "limit": 696 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 378 + "limit": 240 }, "ASYNC230": { "limit": 11 @@ -117,7 +117,7 @@ "limit": 1 }, "PERF102": { - "limit": 23 + "limit": 22 }, "PERF401": { "limit": 12 @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 173 + "limit": 172 }, "RUF012": { "limit": 239 @@ -198,7 +198,7 @@ "limit": 58 }, "SIM102": { - "limit": 311 + "limit": 309 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1092 + "limit": 1080 }, "TRY002": { "limit": 524 diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index cc923f05831..d1d751989ea 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -195,6 +195,22 @@ def test_set_schema_property_ordering_with_excessive_nesting(): set_schema_property_ordering(schema) +def test_set_schema_property_ordering_skips_non_dict_property_values(): + """Non-dict property values must be skipped, not recursed into (they used to raise).""" + schema = { + "properties": { + "a": "hello", + "b": {"type": "string"}, + "c": ["x"], + "d": "a string mentioning items", + } + } + + result = set_schema_property_ordering(schema) + + assert result["propertyOrdering"] == ["a", "b", "c", "d"] + + def test_build_vertex_schema(): """Test build_vertex_schema with a sample schema""" from litellm.llms.vertex_ai.common_utils import _build_vertex_schema diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8c83bde8774..2518ff223ac 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22604 + "limit": 22554 }, "LIT002": { - "limit": 26806 + "limit": 26782 }, "LIT003": { "limit": 269 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16528 + "limit": 16510 }, "LIT011": { - "limit": 5539 + "limit": 5520 }, "LIT012": { "limit": 4506 From 608603ee63e4544a73a433464ee54ca124c83fa6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:36:34 -0700 Subject: [PATCH 029/126] fix(speech): honor pcm response_format for Gemini TTS and reject unsupported containers --- .../speech_to_completion_bridge/handler.py | 2 + .../transformation.py | 54 ++++++++++++------ .../test_transformation.py | 57 ++++++++++++++++++- 3 files changed, 96 insertions(+), 17 deletions(-) diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 9e949db625a..6c33621ec89 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -115,9 +115,11 @@ class SpeechToCompletionBridgeHandler: **request_data, ) + requested_response_format: Final = optional_params.get("response_format") if isinstance(result, ModelResponse): return self.transformation_handler.transform_response( model_response=result, + response_format=requested_response_format if isinstance(requested_response_format, str) else None, ) else: raise Exception(f"Unmapped response type. Got type: {type(result)}") diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 9b757ce86be..e2d3fadf852 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -21,6 +21,8 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None: GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16" +GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm" +GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT}) class ChatAudioParam(TypedDict): @@ -29,6 +31,26 @@ class ChatAudioParam(TypedDict): class SpeechToCompletionBridgeTransformationHandler: + def _validate_response_format( + self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object] + ) -> None: + if not self._is_gemini_tts_model(model): + return + response_format: Final = optional_params.get("response_format") + if not isinstance(response_format, str) or response_format in GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: + return + from litellm.exceptions import BadRequestError + + supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS)) + raise BadRequestError( + message=( + f"Gemini TTS only produces raw PCM16 audio, so response_format='{response_format}'" + f" is not supported. Supported response formats: {supported}." + ), + model=model, + llm_provider=custom_llm_provider, + ) + def _chat_completion_params(self, optional_params: Mapping[str, object]) -> Mapping[str, object]: return MappingProxyType( { @@ -67,6 +89,7 @@ class SpeechToCompletionBridgeTransformationHandler: litellm_logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, ) -> dict: + self._validate_response_format(model, custom_llm_provider, optional_params) user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": input} return_kwargs: Final = { "model": model, @@ -125,7 +148,14 @@ class SpeechToCompletionBridgeTransformationHandler: """Check if the model is a Gemini TTS model that returns PCM16 data.""" return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower()) - def transform_response(self, model_response: "ModelResponse") -> "HttpxBinaryResponseContent": + def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]: + if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT: + return decoded_audio, "audio/pcm" + return self._convert_pcm16_to_wav(decoded_audio), "audio/wav" + + def transform_response( + self, model_response: "ModelResponse", response_format: str | None + ) -> "HttpxBinaryResponseContent": import base64 import httpx @@ -136,23 +166,15 @@ class SpeechToCompletionBridgeTransformationHandler: audio_part: Final = cast(Choices, model_response.choices[0]).message.audio if audio_part is None: raise ValueError("No audio part found in the response") - audio_content: Final = audio_part.data + decoded_audio: Final = base64.b64decode(audio_part.data) - # Decode base64 to get binary content - binary_data = base64.b64decode(audio_content) - - # Check if this is a Gemini TTS model that returns raw PCM16 data model: Final = getattr(model_response, "model", "") - headers: Final = {} - if self._is_gemini_tts_model(model): - # Convert PCM16 to WAV format for proper audio file playback - binary_data = self._convert_pcm16_to_wav(binary_data) - headers["Content-Type"] = "audio/wav" - else: - headers["Content-Type"] = "audio/mpeg" - - # Create an httpx.Response object - response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers) + content, content_type = ( + self._gemini_tts_response_body(decoded_audio, response_format) + if self._is_gemini_tts_model(model) + else (decoded_audio, "audio/mpeg") + ) + response: Final = httpx.Response(status_code=200, content=content, headers={"Content-Type": content_type}) binary_response: Final = HttpxBinaryResponseContent(response) binary_response.set_response_cost(_completion_response_cost(model_response)) return binary_response diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py index c0c720bbaf6..953f028af3c 100644 --- a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py +++ b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py @@ -1,3 +1,4 @@ +import base64 from typing import Final from unittest.mock import MagicMock @@ -8,8 +9,17 @@ from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ( SpeechToCompletionBridgeTransformationHandler, ) +from litellm.types.utils import ChatCompletionAudioResponse, Choices, Message, ModelResponse GEMINI_TTS_MODEL: Final = "gemini-3.1-flash-tts-preview" +PCM_BYTES: Final = b"\x01\x02\x03\x04" * 6 + + +def _model_response(model: str, pcm: bytes) -> ModelResponse: + audio: Final = ChatCompletionAudioResponse( + data=base64.b64encode(pcm).decode(), expires_at=0, transcript="hello" + ) + return ModelResponse(model=model, choices=[Choices(message=Message(content=None, audio=audio))]) def _bridge_request(response_format: str | None) -> dict: @@ -28,7 +38,7 @@ def _bridge_request(response_format: str | None) -> dict: ) -@pytest.mark.parametrize("response_format", ["wav", "mp3", "pcm", None]) +@pytest.mark.parametrize("response_format", ["wav", "pcm", None]) def test_gemini_tts_request_keeps_speech_response_format_out_of_chat_params(response_format: str | None) -> None: request: Final = _bridge_request(response_format) @@ -60,3 +70,48 @@ def test_non_gemini_request_forwards_speech_response_format_as_audio_format() -> assert "response_format" not in request assert request["audio"] == {"voice": "alloy", "format": "wav"} + + +@pytest.mark.parametrize("response_format", ["mp3", "flac", "opus", "aac"]) +def test_gemini_tts_request_rejects_formats_gemini_cannot_produce(response_format: str) -> None: + with pytest.raises(litellm.BadRequestError) as excinfo: + _bridge_request(response_format) + + assert excinfo.value.status_code == 400 + assert response_format in str(excinfo.value) + assert "pcm" in str(excinfo.value) + assert "wav" in str(excinfo.value) + + +def test_gemini_tts_pcm_response_returns_raw_pcm_bytes() -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + response_format="pcm", + ) + + assert response.response.content == PCM_BYTES + assert response.response.headers["content-type"] == "audio/pcm" + + +@pytest.mark.parametrize("response_format", ["wav", None]) +def test_gemini_tts_wav_and_default_responses_wrap_pcm_in_wav(response_format: str | None) -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + response_format=response_format, + ) + + body: Final = response.response.content + assert body[:4] == b"RIFF" + assert body[8:12] == b"WAVE" + assert body[44:] == PCM_BYTES + assert response.response.headers["content-type"] == "audio/wav" + + +def test_non_gemini_response_keeps_original_bytes_and_mpeg_content_type() -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response("gpt-4o-audio-preview", PCM_BYTES), + response_format="mp3", + ) + + assert response.response.content == PCM_BYTES + assert response.response.headers["content-type"] == "audio/mpeg" From abfb6adc2b2e153be51a8f829bbe605cce60c12e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 17:13:26 -0700 Subject: [PATCH 030/126] refactor(proxy): bound the budget window seed by time instead of request ids The one-time seed for a budget window row subtracted the batch's own LiteLLM_SpendLogs rows by request_id, and request_id is the client's x-litellm-call-id whenever the response carries no id of its own. Carrying that set through the queue meant an unbounded, client-controlled aggregate that the commit-failure requeue kept alive across retries. Every log row at or after a batch's earliest start is owed by an increment that still reaches the row, so summing only rows before it needs nothing from the request. That drops request_ids end to end and closes the cross-pod double count the id list could not see. --- .../proxy/db/budget_window_spend_writer.py | 67 +++++------ litellm/proxy/db/db_spend_update_writer.py | 7 +- .../window_spend_update_queue.py | 20 +--- .../proxy/hooks/proxy_track_cost_callback.py | 3 +- litellm/proxy/proxy_server.py | 11 +- .../test_redis_update_buffer.py | 9 +- .../test_window_spend_update_queue.py | 73 ++---------- .../db/test_budget_window_spend_writer.py | 108 ++++++++---------- .../proxy/db/test_db_spend_update_writer.py | 70 ------------ .../hooks/test_proxy_track_cost_callback.py | 60 +--------- tests/test_litellm/proxy/test_proxy_server.py | 55 +-------- 11 files changed, 104 insertions(+), 379 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index f9188f95cfd..8b1ad0e24c7 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -7,17 +7,14 @@ instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold (issue #35766). Raw SQL rather than the Prisma upsert helper because the conditional roll cannot be expressed through the query builder. -Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, excluding -the requests whose increments are in the same batch so neither source counts -them twice. One gap survives that exclusion: without the Redis transaction -buffer every pod flushes its own increments, so a row seeded by one pod can -include spend logs whose increments are still queued on another pod, and those -increments are added again when that pod flushes. That is bounded by a single -flush interval, happens at most once per window row, and only ever over-counts: -the seed never omits spend, because every increment not yet in the row still -reaches it on its own pod's next flush. A row therefore lags real spend by at -most one flush interval of queued increments, the same lag the SpendLogs -aggregate it replaces (and every other spend column) already has. +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, summing +only rows that started before the batch being flushed so neither source counts +the same request twice. Anything at or after that cutoff is owed by an +increment that still reaches the row, on this pod's next flush or another +pod's, so a row lags real spend by at most one flush interval of queued +increments: the same lag the SpendLogs aggregate it replaces (and every other +spend column) already has. A request whose increment is lost before it flushes, +which today means the pod dying, is missed by both sources and stays missing. """ from collections.abc import Sequence @@ -69,13 +66,13 @@ _ROLL_WINDOW_SPEND_SQL: Final = ( _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" + "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" + "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( @@ -92,8 +89,8 @@ _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) class WindowSpendLogsAggregate(Protocol): - """Sums LiteLLM_SpendLogs for one entity since window_start, ignoring the - requests whose ids are handed in. + """Sums LiteLLM_SpendLogs for one entity between window_start and the + batch's earliest request. Injected so the flush can be exercised without a database and so the expensive aggregate stays swappable. @@ -105,21 +102,19 @@ class WindowSpendLogsAggregate(Protocol): entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Sequence[str], - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: ... -async def spend_logs_total_excluding( +async def spend_logs_total_before_batch( prisma_client: "PrismaClient", entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Sequence[str], - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: - """LiteLLM_SpendLogs spend for one entity since window_start, minus the - requests already accounted for by the increments being flushed. + """LiteLLM_SpendLogs spend for one entity since window_start, stopping + before the requests the increments being flushed already cover. The spend log writer drains its own queue on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so @@ -127,13 +122,11 @@ async def spend_logs_total_excluding( already in the table. Counting them in the seed and again in the increment is what made a fresh row land at twice the true spend. - The exclusion is bounded to rows that started at or after the batch's - earliest request. request_id can be chosen by the client - (x-litellm-call-id), so an unbounded exclusion would let a replayed old id - erase a historical row from the seed while its increment still lands. - Without a known start the batch's ids are not excluded at all: that can - only over-count once, which enforcement tolerates, whereas under-counting - is a budget bypass. + Every log row at or after the cutoff belongs to a request whose own + increment still reaches this row, on this pod's next flush or another pod's, + so bounding the sum by time needs nothing from the request itself. Without a + known start the whole window is summed: that can only over-count once, which + enforcement tolerates, whereas under-counting is a budget bypass. """ if entity_type == Litellm_EntityType.KEY.value: bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL @@ -143,13 +136,12 @@ async def spend_logs_total_excluding( return None rows: Final = ( await prisma_client.db.query_raw(unbounded_sql, entity_id, window_start) - if exclude_started_at is None or not exclude_request_ids + if batch_started_at is None else await prisma_client.db.query_raw( bounded_sql, entity_id, window_start, - tuple(exclude_request_ids), - _exclusion_lower_bound(exclude_started_at), + _exclusion_upper_bound(batch_started_at), ) ) if not rows: @@ -157,7 +149,7 @@ async def spend_logs_total_excluding( return float(rows[0].get("total") or 0.0) -def _exclusion_lower_bound(started_at: datetime) -> datetime: +def _exclusion_upper_bound(started_at: datetime) -> datetime: """LiteLLM_SpendLogs.startTime is TIMESTAMP(3); floor to the second so a millisecond rounding of the batch's own earliest row cannot slip under it.""" return to_naive_utc(started_at).replace(microsecond=0) @@ -194,8 +186,8 @@ async def _seed_base_for_missing_row( This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on every cold counter today, but here it runs once per window lifetime and off - the request path, and it excludes this batch's own requests so they are - counted by their increments alone. + the request path, and it stops before the queued increments so they are + counted once. """ if _primary_key(transaction) in existing_primary_keys: return 0.0 @@ -204,8 +196,7 @@ async def _seed_base_for_missing_row( entity_type=transaction["entity_type"], entity_id=transaction["entity_id"], window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), - exclude_request_ids=transaction["request_ids"], - exclude_started_at=_transaction_started_at(transaction), + batch_started_at=_transaction_started_at(transaction), ) return float(base or 0.0) @@ -241,7 +232,7 @@ def _upsert_params( async def commit_window_spend_updates( prisma_client: "PrismaClient", transactions: Sequence[WindowSpendTransaction], - spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding, + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_before_batch, ) -> None: """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 641c07914d9..b3fd2c3f22c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -215,11 +215,7 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ) -> str | None: - """Returns the LiteLLM_SpendLogs request_id this call was recorded - under, so the caller can tell the budget-window writer which log rows - its increments already cover. None when the payload could not be built. - """ + ) -> None: from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -310,7 +306,6 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") - return payload.get("request_id") except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py index 04dea66165e..43b069fd8c2 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -26,17 +26,11 @@ class WindowSpendTransaction(TypedDict): window_start is an ISO-8601 string rather than a datetime so the transaction survives the JSON round trip through the Redis buffer. - request_ids carries the LiteLLM_SpendLogs ids this spend came from. The - one-time seed for a window that has no row yet subtracts them from its - LiteLLM_SpendLogs aggregate, because the spend log writer flushes on its - own ~2s poll and will usually have persisted these rows before the window - queue flushes; without the exclusion the seed and the increment would each - count them. - - started_at is the earliest request start in the batch. The seed only - subtracts a request_id whose LiteLLM_SpendLogs.startTime is at or after it, - so a client that replays an old id through x-litellm-call-id cannot make the - seed drop the historical row that id already paid for. + started_at is the earliest request start in the batch. The one-time seed for + a window that has no row yet sums only LiteLLM_SpendLogs rows that started + before it, because the spend log writer flushes on its own ~2s poll and will + usually have persisted this batch's rows before the window queue flushes; + without the bound the seed and the increment would each count them. """ entity_type: ReadOnly[str] @@ -44,7 +38,6 @@ class WindowSpendTransaction(TypedDict): window_duration: ReadOnly[str] window_start: ReadOnly[str] spend: ReadOnly[float] - request_ids: ReadOnly[Sequence[str]] started_at: ReadOnly[str | None] @@ -72,7 +65,6 @@ def build_window_spend_transaction( window_duration: str, window_start: datetime, spend: float, - request_id: str | None = None, started_at: datetime | None = None, ) -> WindowSpendTransaction: return WindowSpendTransaction( @@ -81,7 +73,6 @@ def build_window_spend_transaction( window_duration=window_duration, window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), spend=spend, - request_ids=() if request_id is None else (request_id,), started_at=None if started_at is None else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"), @@ -101,7 +92,6 @@ def _merge_window_spend_transactions( window_duration=first["window_duration"], window_start=first["window_start"], spend=math.fsum(payload["spend"] for payload in payloads), - request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))), started_at=min(started_ats) if started_ats else None, ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index f02901f0e97..47aafda2337 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -587,7 +587,7 @@ async def _update_database_and_spend_counters( model_access_groups: Sequence[str] | None = None, ) -> None: try: - spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database( + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -623,7 +623,6 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, - request_id=spend_log_request_id, request_started_at=start_time, model_access_groups=model_access_groups, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 654fe4a3f2e..7c47eb1bc42 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2658,7 +2658,6 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, - request_id: str | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, ): @@ -2733,7 +2732,6 @@ async def increment_spend_counters( window_duration=duration, window_start=key_window_start, increment=cost, - request_id=request_id, request_started_at=request_started_at, ) @@ -2777,7 +2775,6 @@ async def increment_spend_counters( window_duration=duration, window_start=team_window_start, increment=cost, - request_id=request_id, request_started_at=request_started_at, ) @@ -3005,16 +3002,15 @@ async def _enqueue_window_spend_row_update( window_duration: str, window_start: datetime | None, increment: float, - request_id: str | None, request_started_at: datetime | None, ) -> None: """Queue this request's cost against the LiteLLM_BudgetWindowSpend row for the window, so enforcement can read a maintained total instead of aggregating LiteLLM_SpendLogs. - request_id is the LiteLLM_SpendLogs id this cost was recorded under and - request_started_at its startTime; the flush uses them to keep the one-time - seed from counting a request that its increment already covers. + request_started_at is this request's LiteLLM_SpendLogs startTime; the flush + stops the one-time seed there so a request its increment already covers is + not counted twice. Enqueued even when the cache increment was skipped for a reserved counter: the reservation only pre-charged the counter, and the row still owes the @@ -3035,7 +3031,6 @@ async def _enqueue_window_spend_row_update( window_duration=window_duration, window_start=window_start, spend=increment, - request_id=request_id, started_at=request_started_at, ) ) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 504654e103a..99ac1fe8b50 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -203,7 +203,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff "window_duration": "30d", "window_start": "2026-08-01T00:00:00.000000", "spend": 3.0, - "request_ids": ["req-1"], + "started_at": None, } ] ) @@ -233,13 +233,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff window_spend, ) = result - # Budget window spend from two pods is summed per window, not overwritten, - # and both pods' request ids reach the seed exclusion. + # Budget window spend from two pods is summed per window, not overwritten. assert window_spend is not None assert len(window_spend) == 1 assert window_spend[0]["spend"] == 6.0 assert window_spend[0]["entity_id"] == "hashed-token" - assert window_spend[0]["request_ids"] == ("req-1",) # Verify db spend was parsed correctly assert db_spend is not None @@ -326,7 +324,6 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd window_duration="30d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=3.0, - request_id="req-1", started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), ), ) @@ -500,7 +497,6 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up window_duration="30d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=1.25, - request_id="req-1", started_at=datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc), ) ) @@ -526,7 +522,6 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up "window_duration": "30d", "window_start": "2026-08-01T00:00:00.000000", "spend": 1.25, - "request_ids": ["req-1"], "started_at": "2026-08-10T12:00:00.000000", } ] diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py index b1ecda57afa..6632b1c8e35 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -19,7 +19,6 @@ def _txn( spend: float, duration: str = "30d", entity_type: str = "key", - request_id: str | None = None, started_at: datetime | None = None, ): return build_window_spend_transaction( @@ -28,7 +27,6 @@ def _txn( window_duration=duration, window_start=window_start, spend=spend, - request_id=request_id, started_at=started_at, ) @@ -38,13 +36,12 @@ def test_build_window_spend_transaction_stores_naive_utc_iso(): TIMESTAMP(3) column, so a non-UTC input must be converted, not truncated.""" non_utc = datetime(2026, 8, 1, 20, 0, tzinfo=timezone(timedelta(hours=-4))) - assert _txn("k1", non_utc, 1.0, request_id="req-1") == { + assert _txn("k1", non_utc, 1.0) == { "entity_type": "key", "entity_id": "k1", "window_duration": "30d", "window_start": "2026-08-02T00:00:00.000000", "spend": 1.0, - "request_ids": ("req-1",), "started_at": None, } @@ -59,19 +56,19 @@ def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso(): @pytest.mark.asyncio async def test_aggregation_keeps_the_earliest_started_at_of_the_batch(): - """The seed bounds its request-id exclusion at the batch's earliest start, - so a later start must never win the merge.""" + """The seed stops at the batch's earliest start, so a later start must never + win the merge: it would push the cutoff forward and count a request the + increments already cover.""" queue = WindowSpendUpdateQueue() earliest = datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-2", started_at=earliest + timedelta(seconds=5))) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1", started_at=earliest)) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-3")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, started_at=earliest + timedelta(seconds=5))) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, started_at=earliest)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() assert len(aggregated) == 1 assert aggregated[0]["started_at"] == "2026-08-10T12:00:00.000000" - assert aggregated[0]["request_ids"] == ("req-1", "req-2", "req-3") def test_to_naive_utc_leaves_naive_values_alone(): @@ -208,62 +205,12 @@ def test_aggregation_survives_the_redis_json_round_trip(): assert reloaded == aggregated -@pytest.mark.asyncio -async def test_aggregation_unions_the_request_ids_of_merged_increments(): - """The seed excludes exactly the requests its batch already covers, so every - merged increment's id has to survive aggregation.""" - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1")) - await queue.add_update(_txn("k1", WINDOW_A, 2.0, request_id="req-2")) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert len(aggregated) == 1 - assert aggregated[0]["request_ids"] == ("req-1", "req-2") - - -@pytest.mark.asyncio -async def test_request_ids_stay_with_their_own_window(): - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) - await queue.add_update(_txn("k1", WINDOW_B, 2.0, request_id="req-b")) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert {payload["window_start"]: payload["request_ids"] for payload in aggregated} == { - "2026-08-01T00:00:00.000000": ("req-a",), - "2026-08-31T00:00:00.000000": ("req-b",), - } - - -@pytest.mark.asyncio -async def test_request_ids_are_deduplicated_and_ordered(): - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-b")) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert aggregated[0]["request_ids"] == ("req-a", "req-b") - - -@pytest.mark.asyncio -async def test_increment_without_a_request_id_carries_no_exclusion(): - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0)) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert aggregated[0]["request_ids"] == () - - -def test_request_ids_survive_the_redis_json_round_trip(): +def test_started_at_survives_the_redis_json_round_trip(): aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] + [(_txn("k1", WINDOW_A, 1.0, started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc)),)] ) reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) - assert reloaded[0]["request_ids"] == ("req-1",) + assert reloaded[0]["started_at"] == "2026-08-10T12:00:00.000000" assert reloaded[0]["spend"] == 1.0 diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py index a849317c930..74b71f63401 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -8,7 +8,7 @@ import pytest from litellm.proxy.db.budget_window_spend_writer import ( commit_window_spend_updates, roll_window_spend_row, - spend_logs_total_excluding, + spend_logs_total_before_batch, ) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, @@ -82,16 +82,14 @@ class _RecordingAggregate: entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Any, - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: self.calls.append( { "entity_type": entity_type, "entity_id": entity_id, "window_start": window_start, - "exclude_request_ids": tuple(exclude_request_ids), - "exclude_started_at": exclude_started_at, + "batch_started_at": batch_started_at, } ) return self.value @@ -99,8 +97,8 @@ class _RecordingAggregate: class _SpendLogsFake: """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, - honouring the exclusion exactly as the real aggregate's - NOT (request_id = ANY(...) AND startTime >= bound) does.""" + honouring the cutoff exactly as the real aggregate's + startTime < bound does.""" def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: self.rows = rows @@ -111,25 +109,22 @@ class _SpendLogsFake: entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Any, - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: - excluded = frozenset(exclude_request_ids) if exclude_started_at is not None else frozenset() return math.fsum( spend - for request_id, spend, started_at in self.rows - if not (request_id in excluded and started_at >= exclude_started_at) + for _request_id, spend, started_at in self.rows + if batch_started_at is None or started_at < batch_started_at ) -def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: +def _batch(spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: return { "entity_type": "key", "entity_id": "k1", "window_duration": "30d", "window_start": "2026-08-01T00:00:00.000000", "spend": spend, - "request_ids": request_ids, "started_at": None if started_at is None else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), @@ -345,7 +340,7 @@ async def test_unknown_entity_type_contributes_no_seed(): db = _FakeDB(existing_rows=[]) async def no_such_column( - prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at + prisma_client, entity_type, entity_id, window_start, batch_started_at ): return None @@ -363,7 +358,7 @@ async def test_unknown_entity_type_contributes_no_seed(): async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing(): db = _FakeDB(existing_rows=[]) - async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at): + async def unavailable(prisma_client, entity_type, entity_id, window_start, batch_started_at): return None await commit_window_spend_updates( @@ -399,18 +394,17 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio -async def test_seed_receives_the_batch_request_ids_and_earliest_start_to_exclude(): +async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): db = _FakeDB(existing_rows=[]) aggregate = _RecordingAggregate(value=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1", "req-2", "req-3"), 3.0),), + transactions=(_batch(3.0),), spend_logs_aggregate=aggregate, ) - assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") - assert aggregate.calls[0]["exclude_started_at"] == BATCH_STARTED_AT + assert aggregate.calls[0]["batch_started_at"] == BATCH_STARTED_AT @pytest.mark.asyncio @@ -420,11 +414,11 @@ async def test_seed_passes_no_start_bound_when_the_batch_has_none(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1",), 1.0, started_at=None),), + transactions=(_batch(1.0, started_at=None),), spend_logs_aggregate=aggregate, ) - assert aggregate.calls[0]["exclude_started_at"] is None + assert aggregate.calls[0]["batch_started_at"] is None @pytest.mark.asyncio @@ -444,7 +438,7 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + transactions=(_batch(0.000141),), spend_logs_aggregate=already_flushed, ) @@ -460,7 +454,7 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1",), 0.000047),), + transactions=(_batch(0.000047),), spend_logs_aggregate=spend_logs, ) @@ -469,17 +463,20 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): @pytest.mark.asyncio -async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed(): - """request_id can be chosen by the client via x-litellm-call-id. A request - that replays an id from before this batch writes no new LiteLLM_SpendLogs - row (the insert skips duplicates), so the seed must keep counting the - historical row that id belongs to; only its increment is new.""" +async def test_seed_skips_logs_from_requests_this_batch_never_saw(): + """A concurrent request on another pod can land its spend log before this + pod seeds the row. Its increment is still queued over there, so the cutoff + has to drop it from the seed even though this batch has no way to know its + id; counting it here and again on that pod's flush is the double count the + old id list could not catch.""" db = _FakeDB(existing_rows=[]) - spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),)) + spend_logs = _SpendLogsFake( + rows=(("older", 0.5, BEFORE_BATCH), ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1))), + ) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("replayed",), 0.000047),), + transactions=(_batch(0.000047),), spend_logs_aggregate=spend_logs, ) @@ -496,7 +493,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + transactions=(_batch(0.000141),), spend_logs_aggregate=nothing_flushed, ) @@ -509,57 +506,48 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): "entity_type, expected_column", [("key", "api_key = $1"), ("team", "team_id = $1")], ) -async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch_start_bound( - entity_type, expected_column -): +async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected_column): db = _FakeDB(existing_rows=[{"total": 1.25}]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type=entity_type, entity_id="e1", window_start=WINDOW_A, - exclude_request_ids=("req-1", "req-2"), - exclude_started_at=BATCH_STARTED_AT, + batch_started_at=BATCH_STARTED_AT, ) assert total == pytest.approx(1.25) ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized + assert "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" in normalized assert 'FROM "LiteLLM_SpendLogs"' in normalized # startTime is TIMESTAMP(3): the bound is floored to the second so the # batch's own earliest row cannot round under it. - assert params == ("e1", WINDOW_A, ("req-1", "req-2"), datetime(2026, 8, 10, 12, 0, 0)) - # The ids are bound, never spliced into the statement. - assert "req-1" not in query + assert params == ("e1", WINDOW_A, datetime(2026, 8, 10, 12, 0, 0)) + # Nothing the caller supplied reaches the statement text. + assert "e1" not in query @pytest.mark.asyncio -@pytest.mark.parametrize( - "exclude_request_ids, exclude_started_at", - [(("req-1",), None), ((), BATCH_STARTED_AT)], -) -async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_bound( - exclude_request_ids, exclude_started_at -): - """Ids without a start bound would reopen the replayed-id hole, so the - seed counts everything instead; at worst that over-counts one batch.""" +async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): + """A batch with no known start cannot place the cutoff, so the seed counts + everything; at worst that over-counts one batch, which enforcement + tolerates, where under-counting is a budget bypass.""" db = _FakeDB(existing_rows=[{"total": 1.25}]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="e1", window_start=WINDOW_A, - exclude_request_ids=exclude_request_ids, - exclude_started_at=exclude_started_at, + batch_started_at=None, ) assert total == pytest.approx(1.25) ((query, params),) = db.query_raw_calls - assert "request_id" not in query + assert '"startTime" <' not in query assert params == ("e1", WINDOW_A) @@ -567,13 +555,12 @@ async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_boun async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type="user", entity_id="u1", window_start=WINDOW_A, - exclude_request_ids=(), - exclude_started_at=None, + batch_started_at=None, ) assert total is None @@ -584,13 +571,12 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="k-unknown", window_start=WINDOW_A, - exclude_request_ids=(), - exclude_started_at=None, + batch_started_at=None, ) assert total == 0.0 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index d28cf8c9c6a..11ef911de3e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2581,7 +2581,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ window_duration="30d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=0.5, - request_id="req-1", ) await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() @@ -2611,7 +2610,6 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): window_duration="7d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=2.0, - request_id="req-1", ), ) mock_redis_update_buffer = AsyncMock() @@ -2638,74 +2636,6 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): db_writer.pod_lock_manager.release_lock.assert_awaited_once() -@pytest.mark.asyncio -async def test_update_database_returns_the_spend_log_request_id(): - """The budget-window seed excludes the log rows its increments already - cover, so the caller needs the id this call was recorded under. It cannot - be re-derived: cache hits append time.time() to the id.""" - db_writer = DBSpendUpdateWriter() - db_writer._insert_spend_log_to_db = AsyncMock() - db_writer._enqueue_tool_usage_transaction = AsyncMock() - - with ( - patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam - "litellm.proxy.proxy_server", - disable_spend_logs=False, - prisma_client=MagicMock(), - litellm_proxy_budget_name="test-budget", - ) - ): - request_id = await db_writer.update_database( - token="test-token", - user_id="test-user", - end_user_id=None, - team_id="test-team", - org_id=None, - kwargs={"model": "gpt-4", "custom_llm_provider": "openai", "litellm_call_id": "call-xyz"}, - completion_response=MagicMock(), - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.1, - ) - await asyncio.sleep(0) - - assert request_id is not None - # Same id the spend log row was queued under. - assert request_id == db_writer._insert_spend_log_to_db.call_args[1]["payload"]["request_id"] - - -@pytest.mark.asyncio -async def test_update_database_returns_none_when_the_payload_cannot_be_built(): - db_writer = DBSpendUpdateWriter() - - with ( - patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam - "litellm.proxy.proxy_server", - disable_spend_logs=False, - prisma_client=MagicMock(), - litellm_proxy_budget_name="test-budget", - ), - patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam - "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", - side_effect=Exception("payload boom"), - ), - ): - request_id = await db_writer.update_database( - token="test-token", - user_id="test-user", - end_user_id=None, - team_id="test-team", - org_id=None, - kwargs={"model": "gpt-4"}, - completion_response=MagicMock(), - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.1, - ) - - assert request_id is None - - @pytest.mark.asyncio async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at(): """Spend flushes must leave settings_updated_at alone, or it decays into diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 2a540f4f522..8043a1aca3f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -567,9 +567,7 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re @pytest.mark.asyncio async def test_update_database_and_spend_counters_updates_counters_after_db_update(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - return_value="chatcmpl-abc123" - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} start_time = datetime.now() @@ -602,7 +600,6 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], - request_id="chatcmpl-abc123", request_started_at=start_time, model_access_groups=("premium",), ) @@ -1884,61 +1881,6 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( ) -@pytest.mark.asyncio -async def test_update_database_and_spend_counters_forwards_the_spend_log_request_id(): - """The budget-window flush excludes the log rows its increments already - cover. That only works if the id update_database recorded the row under is - handed to the counter update, so this seam is load-bearing.""" - proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - return_value="chatcmpl-abc123" - ) - increment_spend_counters = AsyncMock() - - await _update_database_and_spend_counters( - proxy_logging_obj=proxy_logging_obj, - increment_spend_counters=increment_spend_counters, - user_api_key="test_api_key", - user_id="test_user_id", - end_user_id=None, - team_id="test_team_id", - org_id="test_org_id", - kwargs={}, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.2, - budget_reservation=None, - ) - - assert increment_spend_counters.await_args.kwargs["request_id"] == "chatcmpl-abc123" - - -@pytest.mark.asyncio -async def test_update_database_and_spend_counters_forwards_a_missing_request_id_as_none(): - proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=None) - increment_spend_counters = AsyncMock() - - await _update_database_and_spend_counters( - proxy_logging_obj=proxy_logging_obj, - increment_spend_counters=increment_spend_counters, - user_api_key="test_api_key", - user_id="test_user_id", - end_user_id=None, - team_id="test_team_id", - org_id="test_org_id", - kwargs={}, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.2, - budget_reservation=None, - ) - - assert increment_spend_counters.await_args.kwargs["request_id"] is None - - class _FakeDeploymentLookup: """Deployment lookup returning the access groups each deployment declares.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1de3ed6e56d..43280258153 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11399,35 +11399,10 @@ async def test_no_window_spend_row_enqueued_without_budget_limits(): assert enqueued == [] -@pytest.mark.asyncio -async def test_window_spend_row_carries_the_spend_log_request_id(): - """The flush excludes these ids from its one-time seed, so the id threaded - here has to be the same one the LiteLLM_SpendLogs row was written under.""" - from litellm.proxy.proxy_server import increment_spend_counters - - reset_at = datetime.now(timezone.utc) + timedelta(days=10) - key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] - - with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", - team_id=None, - user_id=None, - response_cost=0.25, - request_id="chatcmpl-abc123", - ) - enqueued = await _drain(queue) - - assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",) - - @pytest.mark.asyncio async def test_window_spend_row_carries_the_request_start_time(): - """The seed only excludes a batch id whose LiteLLM_SpendLogs.startTime is at - or after this, so it must be the same start the spend log was written with.""" + """The seed sums LiteLLM_SpendLogs only up to this point, so it must be the + same start the spend log row was written with.""" from litellm.proxy.proxy_server import increment_spend_counters reset_at = datetime.now(timezone.utc) + timedelta(days=10) @@ -11442,7 +11417,6 @@ async def test_window_spend_row_carries_the_request_start_time(): team_id=None, user_id=None, response_cost=0.25, - request_id="chatcmpl-abc123", request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), ) enqueued = await _drain(queue) @@ -11451,26 +11425,7 @@ async def test_window_spend_row_carries_the_request_start_time(): @pytest.mark.asyncio -async def test_window_spend_row_without_a_request_id_excludes_nothing(): - from litellm.proxy.proxy_server import increment_spend_counters - - reset_at = datetime.now(timezone.utc) + timedelta(days=10) - key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] - - with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) - enqueued = await _drain(queue) - - assert enqueued[0]["request_ids"] == () - - -@pytest.mark.asyncio -async def test_team_window_spend_row_carries_the_request_id(): +async def test_team_window_spend_row_carries_the_request_start_time(): from litellm.proxy.proxy_server import increment_spend_counters reset_at = datetime.now(timezone.utc) + timedelta(days=3) @@ -11485,11 +11440,11 @@ async def test_team_window_spend_row_carries_the_request_id(): team_id="team-1", user_id=None, response_cost=1.5, - request_id="chatcmpl-team", + request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), ) enqueued = await _drain(queue) - assert enqueued[0]["request_ids"] == ("chatcmpl-team",) + assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" def _mock_startup_prisma_client(health_check_error=None, connect_error=None): From b67b44bdaaa747d6d40b24c2b9583b9caf332a38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:10:14 -0700 Subject: [PATCH 031/126] fix(proxy): map audio_speech errors to their status codes instead of a blanket 500 --- litellm/proxy/proxy_server.py | 10 ++++++- .../proxy/proxy_server/test_routes_audio.py | 30 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ee1eb876b87..37034ea9a62 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10996,7 +10996,15 @@ async def audio_speech( ) verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - %s", e) verbose_proxy_logger.debug(traceback.format_exc()) - raise e + if isinstance(e, (ProxyException, HTTPException)): + raise e + raise ProxyException( + message=getattr(e, "message", f"{e}"), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + openai_code=getattr(e, "code", None), + code=getattr(e, "status_code", 500), + ) @router.post( diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py index b99affc2ac3..522a20cd34b 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -86,6 +86,24 @@ def patched_speech_error(monkeypatch): yield +@pytest.fixture +def patched_speech_provider_rejection(monkeypatch, patched_speech_error): + import litellm + + async def _raise(*args, **kwargs): + raise litellm.BadRequestError( + message=( + "Gemini TTS only produces raw PCM16 audio, so response_format='mp3' is not supported." + " Supported response formats: pcm, wav." + ), + model="gemini-3.1-flash-tts-preview", + llm_provider="gemini", + ) + + monkeypatch.setattr(proxy_server, "route_request", _raise) + yield + + @pytest.fixture def patched_transcription(monkeypatch): router = MagicMock() @@ -198,6 +216,18 @@ def test_audio_speech_error(client, auth_as, patched_speech_error, path): assert len(response.content) > 0 +def test_audio_speech_bad_request_maps_to_400(client, auth_as, patched_speech_provider_rejection): + """Regression for LIT-6501: a BadRequestError from the speech path surfaced as a generic 500.""" + payload = {"model": "gemini-tts", "input": "Hi", "voice": "Kore", "response_format": "mp3"} + with auth_as(): + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 400 + error = response.json()["error"] + assert "response_format='mp3'" in error["message"] + assert "pcm" in error["message"] + assert "wav" in error["message"] + + @pytest.mark.parametrize("path", ["/v1/audio/transcriptions", "/audio/transcriptions"]) def test_audio_transcription_happy_path(client, auth_as, patched_transcription, path): """Pins ``POST /v1/audio/transcriptions`` / ``POST /audio/transcriptions`` (happy).""" From 35a375e26f99f6246fd2e68e4c002532fef111a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:10:15 -0700 Subject: [PATCH 032/126] fix(speech): stop vertex gemini tts from dropping response_format in cloud tts param mapping --- litellm/utils.py | 4 ++++ tests/test_litellm/test_utils.py | 20 ++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/litellm/utils.py b/litellm/utils.py index 5e9e115ed54..e7028301dab 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9409,6 +9409,10 @@ class ProviderConfigManager: return RunwayMLTextToSpeechConfig() elif litellm.LlmProviders.VERTEX_AI == provider: + if "gemini" in model: + # Gemini TTS uses the speech_to_completion bridge, and Google Cloud TTS param + # mapping would drop response_format before the bridge sees it (LIT-6501) + return None from litellm.llms.vertex_ai.text_to_speech.transformation import ( VertexAITextToSpeechConfig, ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 6524353aa48..935e4c61535 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1416,6 +1416,26 @@ def test_get_provider_rerank_config(): assert isinstance(config, HostedVLLMRerankConfig) +def test_get_provider_text_to_speech_config_vertex_gemini_skips_cloud_tts(): + """Regression for LIT-6501: mapping vertex Gemini TTS params through Google Cloud TTS + dropped response_format before the speech_to_completion bridge could honor it.""" + from litellm.llms.vertex_ai.text_to_speech.transformation import VertexAITextToSpeechConfig + from litellm.utils import LlmProviders + + assert ( + ProviderConfigManager.get_provider_text_to_speech_config( + model="gemini-2.5-flash-preview-tts", provider=LlmProviders.VERTEX_AI + ) + is None + ) + assert isinstance( + ProviderConfigManager.get_provider_text_to_speech_config( + model="en-US-Studio-O", provider=LlmProviders.VERTEX_AI + ), + VertexAITextToSpeechConfig, + ) + + # Models that should be skipped during testing OLD_PROVIDERS = ["aleph_alpha", "palm"] SKIP_MODELS = [ From 5c7e6b80c9e4274b9582e7ffc898f91b224f5351 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 10:18:15 +0000 Subject: [PATCH 033/126] test: isolate global MCP registry and pin savings tests to bundled cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/conftest.py | 31 +++++++++++++++++++ .../proxy/spend_tracking/test_savings.py | 2 ++ 2 files changed, 33 insertions(+) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index b477bf3f406..c559e023c47 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -2,6 +2,37 @@ import os import pytest +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, +) + + +@pytest.fixture(autouse=True) +def _hermetic_mcp_server_registry(): + """Snapshot and restore the global manager's server-registry state around every test. + + ``global_mcp_server_manager`` is a module-global singleton, and many tests in this + package seed ``registry``/``config_mcp_servers`` (or clear them) without cleaning up. + In a shared CI shard the leaked entries poison later tests in the same worker, e.g. + the ``all_proxy_servers`` sentinel expansion in ``auth/`` suddenly sees a bridge + server registered by a discovery test, so the outcome depends on xdist scheduling. + Restoring the state here makes ordering irrelevant. + """ + saved_registry = dict(global_mcp_server_manager.registry) + saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers) + saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping) + saved_oauth_slots = global_mcp_server_manager._oauth_discovery_slots + try: + yield + finally: + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry.update(saved_registry) + global_mcp_server_manager.config_mcp_servers.clear() + global_mcp_server_manager.config_mcp_servers.update(saved_config_servers) + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.update(saved_tool_mapping) + global_mcp_server_manager._oauth_discovery_slots = saved_oauth_slots + @pytest.fixture(autouse=True) def _hermetic_server_root_path(): diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index c3297ee6ae9..7dd18587df3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -14,6 +14,8 @@ from litellm.proxy.spend_tracking.savings import ( from litellm.router import Router from litellm.types.utils import Usage +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + def _anthropic_costs(model: str) -> tuple[float, float]: info = litellm.get_model_info(model=model, custom_llm_provider="anthropic") From 99a6dd02af50e6f822cf57d438a75097232bde15 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:46:11 -0700 Subject: [PATCH 034/126] fix(proxy): narrow audio_speech response before reading upstream content-type --- litellm/proxy/proxy_server.py | 5 ++++- .../proxy/proxy_server/test_routes_audio.py | 21 +++++++------------ 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 70057201af8..c996735ddbb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11068,8 +11068,11 @@ async def audio_speech( custom_headers.update(callback_headers) requested_format: Final = data.get("response_format") + upstream_content_type: Final = ( + response.response.headers.get("content-type") if isinstance(response, HttpxBinaryResponseContent) else None + ) media_type: Final = resolve_speech_media_type( - upstream_content_type=response.response.headers.get("content-type"), + upstream_content_type=upstream_content_type, response_format=requested_format if isinstance(requested_format, str) else None, ) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py index 522a20cd34b..de76c7257cf 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -16,6 +16,7 @@ import httpx import pytest from litellm.proxy import proxy_server +from litellm.types.llms.openai import HttpxBinaryResponseContent @pytest.fixture @@ -38,20 +39,14 @@ def patched_speech(monkeypatch, request): monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data) - class _FakeBinaryResp: - response = httpx.Response( - status_code=200, - headers={} if upstream_content_type is None else {"content-type": upstream_content_type}, - ) - - async def aiter_bytes(self, chunk_size: int = 8192): - async def _gen(): - yield b"\x00\x01\x02" - - return _gen() - async def _llm_call(): - return _FakeBinaryResp() + return HttpxBinaryResponseContent( + httpx.Response( + status_code=200, + headers={} if upstream_content_type is None else {"content-type": upstream_content_type}, + content=b"\x00\x01\x02", + ) + ) async def _fake_route_request(*args, **kwargs): return _llm_call() From dd031f1036eaff49c428da4a133e904b7b3a702f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:42:07 +0000 Subject: [PATCH 035/126] fix(ci): parse paginated gh api output without splitting on unicode line breaks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/close_duplicate_issues.py | 30 +++++++----- .../test_github_close_duplicate_issues.py | 46 +++++++++++++++++++ 2 files changed, 64 insertions(+), 12 deletions(-) create mode 100644 tests/test_litellm/test_github_close_duplicate_issues.py diff --git a/.github/scripts/close_duplicate_issues.py b/.github/scripts/close_duplicate_issues.py index ec522af4f88..4c837c06418 100755 --- a/.github/scripts/close_duplicate_issues.py +++ b/.github/scripts/close_duplicate_issues.py @@ -39,6 +39,23 @@ def gh(*args: str) -> str: return result.stdout +def parse_concatenated_json(raw: str) -> list[dict]: + """Parse the concatenated JSON documents that `gh api --paginate` emits.""" + decoder = json.JSONDecoder() + issues: list[dict] = [] + idx = 0 + while idx < len(raw): + if raw[idx].isspace(): + idx += 1 + continue + parsed, idx = decoder.raw_decode(raw, idx) + if isinstance(parsed, list): + issues.extend(parsed) + else: + issues.append(parsed) + return issues + + def fetch_open_issues(repo: str | None) -> list[dict]: """Fetch all open issues (excluding PRs) via gh api --paginate.""" if repo: @@ -49,18 +66,7 @@ def fetch_open_issues(repo: str | None) -> list[dict]: endpoint = "repos/{owner}/{repo}/issues?state=open&per_page=100&sort=created&direction=asc" cmd = ["api", "--paginate", endpoint] - raw = gh(*cmd) - # gh --paginate concatenates JSON arrays, so we may get multiple arrays - issues = [] - for line in raw.strip().splitlines(): - line = line.strip() - if not line: - continue - parsed = json.loads(line) - if isinstance(parsed, list): - issues.extend(parsed) - else: - issues.append(parsed) + issues = parse_concatenated_json(gh(*cmd)) # Filter out pull requests (they also appear in the issues endpoint) return [i for i in issues if "pull_request" not in i] diff --git a/tests/test_litellm/test_github_close_duplicate_issues.py b/tests/test_litellm/test_github_close_duplicate_issues.py new file mode 100644 index 00000000000..0ee9b3f096d --- /dev/null +++ b/tests/test_litellm/test_github_close_duplicate_issues.py @@ -0,0 +1,46 @@ +"""Unit tests for `.github/scripts/close_duplicate_issues.py`.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / ".github" + / "scripts" + / "close_duplicate_issues.py" +) + + +@pytest.fixture(scope="module") +def dedupe_module(): + spec = importlib.util.spec_from_file_location("close_duplicate_issues", SCRIPT_PATH) + assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" + module = importlib.util.module_from_spec(spec) + sys.modules["close_duplicate_issues"] = module + spec.loader.exec_module(module) + return module + + +def test_parse_concatenated_json_joins_paginated_arrays(dedupe_module): + page_one = json.dumps([{"number": 1, "title": "a"}, {"number": 2, "title": "b"}]) + page_two = json.dumps([{"number": 3, "title": "c"}]) + issues = dedupe_module.parse_concatenated_json(page_one + page_two) + assert [i["number"] for i in issues] == [1, 2, 3] + + +@pytest.mark.parametrize("separator", ["\u2028", "\u2029", "\x85"]) +def test_parse_concatenated_json_survives_unicode_line_breaks_in_bodies( + dedupe_module, separator +): + body = f"first{separator}second" + page_one = json.dumps([{"number": 1, "title": "a", "body": body}]) + page_two = json.dumps([{"number": 2, "title": "b", "body": "plain"}]) + issues = dedupe_module.parse_concatenated_json(page_one + page_two) + assert [i["number"] for i in issues] == [1, 2] + assert issues[0]["body"] == body From 797848dd8243fe10c4bbb58126c97b373ecb87ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:18:59 -0700 Subject: [PATCH 036/126] fix(cost): bill OCR annotation pages via annotation_cost_per_page --- litellm/cost_calculator.py | 18 ++++--- litellm/llms/base_llm/ocr/transformation.py | 1 + .../llms/mistral/ocr/test_mistral_ocr_cost.py | 54 +++++++++++++++++++ 3 files changed, 66 insertions(+), 7 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3adc1c25dfd..d3ad6124751 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1910,12 +1910,15 @@ def ocr_cost( if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: float | None = None - if model_info is not None: - ocr_cost_per_page = model_info.get("ocr_cost_per_page") + ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None + annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None + annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed - if pages_processed is None: + annotation_pages: Final = response.usage_info.pages_processed_annotation or 0 + has_billable_annotation_pages: Final = annotation_rate is not None and annotation_pages > 0 + + if pages_processed is None and not has_billable_annotation_pages: if cost_per_credit is not None or ocr_cost_per_page is None: # Surface missing usage data instead of silently under-reporting # cost. The previous behavior raised ValueError; we now return 0.0 @@ -1931,7 +1934,7 @@ def ocr_cost( return 0.0, 0.0 raise ValueError("OCR response pages_processed is None") - if ocr_cost_per_page is None: + if ocr_cost_per_page is None and not has_billable_annotation_pages: # No per-page pricing configured. Either the model is on credit-based # pricing (and credits weren't returned, so the credit branch above did # not match) or the model has no OCR pricing entry at all. Surface a @@ -1947,8 +1950,9 @@ def ocr_cost( ) return 0.0, 0.0 - total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed - return total_ocr_processing_cost, 0.0 + ocr_pages_cost: Final = (ocr_cost_per_page or 0.0) * (pages_processed or 0) + annotation_pages_cost: Final = (annotation_rate or 0.0) * annotation_pages + return ocr_pages_cost + annotation_pages_cost, 0.0 def vector_store_search_cost( diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index d1c77186ea8..3b302837032 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -75,6 +75,7 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" pages_processed: int | None = None + pages_processed_annotation: int | None = None credits: float | None = None doc_size_bytes: int | None = None diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 890df597933..a0e1616d4b2 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -24,6 +24,9 @@ OCR3_MODEL = "mistral/mistral-ocr-2512" OCR3_COST_PER_PAGE = 0.002 OCR3_ANNOTATION_COST_PER_PAGE = 0.003 +AZURE_DOC_AI_MODEL = "azure_ai/mistral-document-ai-2512" +AZURE_DOC_AI_COST_PER_PAGE = 0.003 + def _ocr_response(model: str, pages_processed: int) -> OCRResponse: return OCRResponse( @@ -33,6 +36,14 @@ def _ocr_response(model: str, pages_processed: int) -> OCRResponse: ) +def _annotated_ocr_response(model: str, pages_processed: int | None, annotation_pages: int) -> OCRResponse: + return OCRResponse( + pages=[], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed, pages_processed_annotation=annotation_pages), + ) + + @pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) def test_model_info_ocr4_price(model: str) -> None: info = litellm.get_model_info(model=f"mistral/{model}", custom_llm_provider="mistral") @@ -79,3 +90,46 @@ def test_ocr3_cost_scales_with_pages(local_model_cost_map, pages_processed: int) call_type="ocr", ) assert cost == pytest.approx(OCR3_COST_PER_PAGE * pages_processed) + + +def test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", 2, 3), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(2 * OCR3_COST_PER_PAGE + 3 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_ocr3_bills_annotation_only_response(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", 0, 3), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(3 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_ocr3_bills_annotation_pages_when_pages_processed_missing(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", None, 4), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(4 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate(local_model_cost_map) -> None: + info = litellm.get_model_info(model=AZURE_DOC_AI_MODEL, custom_llm_provider="azure_ai") + assert info.get("annotation_cost_per_page") is None + assert info["ocr_cost_per_page"] == AZURE_DOC_AI_COST_PER_PAGE + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-document-ai-2512", 0, 1), + model=AZURE_DOC_AI_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + ) + assert cost == pytest.approx(AZURE_DOC_AI_COST_PER_PAGE) From 7b4b92f54f87fd52835aeeeb90a087d111d3483a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:37:26 +0000 Subject: [PATCH 037/126] fix(registry): update veo 3.1 pricing with resolution tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 26 ++++++++++++------- model_prices_and_context_window.json | 26 ++++++++++++------- 2 files changed, 32 insertions(+), 20 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7071eaa0807..cc963a80d83 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23777,8 +23777,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23792,7 +23794,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23820,8 +23823,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23835,7 +23840,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -43369,8 +43375,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43399,8 +43405,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7071eaa0807..cc963a80d83 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23777,8 +23777,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23792,7 +23794,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23820,8 +23823,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23835,7 +23840,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -43369,8 +43375,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43399,8 +43405,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], From 9125a5b7a0af11892dda9260160efcfa52125619 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 21:39:34 +0000 Subject: [PATCH 038/126] fix(responses): json-encode object tool call arguments in the chat completions bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../custom_tools.py | 17 +++++++ .../streaming_iterator.py | 9 ++-- .../transformation.py | 11 +++-- .../test_litellm_completion_responses.py | 49 +++++++++++++++++++ .../test_streaming_iterator_transformation.py | 34 +++++++++++++ 5 files changed, 111 insertions(+), 9 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index cccae06c74b..bd6abd3f45f 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -45,6 +45,23 @@ def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: return tool_name in custom_tool_names +def serialize_tool_call_arguments(raw_arguments: object, default: str = "") -> str: + """Render tool call arguments as the JSON string tool-call schemas require. + + Arguments normally arrive already JSON-encoded, but clients and providers + also send the decoded object. ``str()`` on a dict yields a Python repr with + single quotes, which every downstream JSON parser rejects with errors like + "Expecting ',' delimiter". + """ + if raw_arguments is None or raw_arguments == "": + return default + if isinstance(raw_arguments, str): + return raw_arguments + if isinstance(raw_arguments, (dict, list, tuple, bool, int, float)): + return json.dumps(raw_arguments) + return str(raw_arguments) + + def unwrap_custom_tool_arguments(arguments: str) -> str: """Extract the raw content string from JSON-wrapped arguments. diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..b2edf2bf9ed 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, + serialize_tool_call_arguments, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -213,10 +214,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_args_delta = "" if isinstance(fn, dict): fn_name = str(fn.get("name") or "") - fn_args_delta = str(fn.get("arguments") or "") + fn_args_delta = serialize_tool_call_arguments(fn.get("arguments")) else: fn_name = str(getattr(fn, "name", "") or "") - fn_args_delta = str(getattr(fn, "arguments", "") or "") + fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) output_index = self._get_or_assign_tool_output_index(call_id) @@ -284,10 +285,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_args = "" if isinstance(fn, dict): fn_name = str(fn.get("name") or "") - fn_args = str(fn.get("arguments") or "") + fn_args = serialize_tool_call_arguments(fn.get("arguments")) else: fn_name = str(getattr(fn, "name", "") or "") - fn_args = str(getattr(fn, "arguments", "") or "") + fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) # Track if this is a new tool call that wasn't streamed diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f39df38d069..3ca01f2eb74 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -93,6 +93,7 @@ from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, is_custom_tool_call, + serialize_tool_call_arguments, unwrap_custom_tool_arguments, validated_allowed_callers, ) @@ -1010,7 +1011,7 @@ class LiteLLMCompletionResponsesConfig: type=cast(Literal["function"], tool_use_type), function=ChatCompletionToolCallFunctionChunk( name=str(function.get("name", "")), - arguments=str(function.get("arguments", "{}")), + arguments=serialize_tool_call_arguments(function.get("arguments"), "{}"), ), index=index, ) @@ -1539,7 +1540,7 @@ class LiteLLMCompletionResponsesConfig: type=cast(Literal["function"], _tool_use_definition.get("type") or "function"), function=ChatCompletionToolCallFunctionChunk( name=function.get("name") or "", - arguments=str(function.get("arguments") or ""), + arguments=serialize_tool_call_arguments(function.get("arguments")), ), index=0, ) @@ -1589,7 +1590,7 @@ class LiteLLMCompletionResponsesConfig: type="function", function=ChatCompletionToolCallFunctionChunk( name=f"{namespace}__{raw_name}" if qualify else raw_name, - arguments=str(raw_arguments or ""), + arguments=serialize_tool_call_arguments(raw_arguments), ), index=0, ) @@ -2022,7 +2023,7 @@ class LiteLLMCompletionResponsesConfig: function_definition = tool.function tool_name = function_definition.name or "" tool_id = tool.id or "" - tool_arguments = function_definition.get("arguments") or "" + tool_arguments = serialize_tool_call_arguments(function_definition.get("arguments")) # Check if this is a custom tool if is_custom_tool_call(tool_name, custom_tool_names): @@ -2557,7 +2558,7 @@ class LiteLLMCompletionResponsesConfig: type="function", function=Function( name=tool_call.get("name") or "", - arguments=tool_call.get("arguments") or "", + arguments=serialize_tool_call_arguments(tool_call.get("arguments")), ), ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b96d2eb5322..16e099f0404 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -966,6 +966,55 @@ class TestFunctionCallTransformation: assert result[0]["tool_calls"][0]["function"]["arguments"] == "{}" + def test_function_call_transformation_json_encodes_object_arguments(self): + """A decoded arguments object must be JSON-encoded, not str()'d. + + Clients and providers sometimes send `arguments` as an object rather + than a JSON string; `str()` on a dict produces a Python repr with + single quotes, which downstream JSON parsers reject with errors like + "Expecting ',' delimiter". + """ + function_call_item = { + "type": "function_call", + "name": "shell", + "arguments": {"command": "ls", "timeout": 30, "flags": ["-l", "-a"]}, + "call_id": "call_123", + "id": "call_123", + "status": "completed", + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( + function_call=function_call_item + ) + + arguments = result[0].get("tool_calls", [])[0].get("function", {}).get("arguments") + assert json.loads(arguments) == {"command": "ls", "timeout": 30, "flags": ["-l", "-a"]} + assert "'" not in arguments + + def test_create_tool_call_chunk_json_encodes_object_arguments(self): + """Cached tool_call definitions with object arguments stay valid JSON.""" + chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + tool_use_definition={ + "id": "call_456", + "type": "function", + "function": {"name": "shell", "arguments": {"command": "ls"}}, + }, + tool_call_id="call_456", + index=0, + ) + + assert json.loads(chunk["function"]["arguments"]) == {"command": "ls"} + + def test_create_tool_call_chunk_keeps_empty_arguments_default(self): + """Missing arguments still fall back to an empty JSON object.""" + chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + tool_use_definition={"id": "call_789", "type": "function", "function": {"name": "shell"}}, + tool_call_id="call_789", + index=0, + ) + + assert chunk["function"]["arguments"] == "{}" + def test_complete_input_transformation_with_function_calls(self): """Test the complete transformation with the exact input from the issue""" test_input = [ diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 823f656ddc5..01148f627f1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -10,6 +10,7 @@ before response.completed, and that every event of a bridged stream carries the spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ +import json from unittest.mock import AsyncMock, MagicMock import pytest @@ -523,3 +524,36 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids assert len(set(response_ids)) == 1 assert response_ids[0].startswith("resp_") + + +def test_object_tool_call_arguments_stream_as_valid_json(): + """A provider that sends decoded object arguments must still stream valid JSON. + + `str()` on a dict yields a Python repr with single quotes, which clients + parsing function_call_arguments reject with errors like + "Expecting ',' delimiter". + """ + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_obj", + "type": "function", + "function": {"name": "shell", "arguments": {"command": "ls", "flags": ["-l"]}}, + } + ] + ) + + streamed_arguments = "".join( + evt.delta + for evt in iterator._pending_tool_events + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ) + + assert json.loads(streamed_arguments) == {"command": "ls", "flags": ["-l"]} From ade21da9d075a9431178594163791eee5ab61e2d Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 21:54:05 +0000 Subject: [PATCH 039/126] refactor(responses): simplify tool call argument serializer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_completion_transformation/custom_tools.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index bd6abd3f45f..90491739bb0 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -53,13 +53,11 @@ def serialize_tool_call_arguments(raw_arguments: object, default: str = "") -> s single quotes, which every downstream JSON parser rejects with errors like "Expecting ',' delimiter". """ - if raw_arguments is None or raw_arguments == "": - return default if isinstance(raw_arguments, str): - return raw_arguments - if isinstance(raw_arguments, (dict, list, tuple, bool, int, float)): - return json.dumps(raw_arguments) - return str(raw_arguments) + return raw_arguments or default + if raw_arguments is None: + return default + return json.dumps(raw_arguments, default=str) def unwrap_custom_tool_arguments(arguments: str) -> str: From c344c7a66beb3ab184182d227eae9fc4c346283d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 19:11:04 +0000 Subject: [PATCH 040/126] fix(registry): add zai/glm-5.2, together Qwen3.8-Flash, cerebras/gemma-4-31b, elevenlabs/scribe_v2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 50 +++++++++++++++++++ model_prices_and_context_window.json | 50 +++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc963a80d83..b7927767eeb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -55046,5 +55046,55 @@ "max_tokens": 40960, "mode": "embedding", "source": "https://docs.fireworks.ai/serverless/pricing" + }, + "zai/glm-5.2": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.z.ai/guides/overview/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3.8-Flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "cerebras/gemma-4-31b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 131072, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/gemma-4-31b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "elevenlabs/scribe_v2": { + "input_cost_per_second": 6.11e-05, + "litellm_provider": "elevenlabs", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc963a80d83..b7927767eeb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -55046,5 +55046,55 @@ "max_tokens": 40960, "mode": "embedding", "source": "https://docs.fireworks.ai/serverless/pricing" + }, + "zai/glm-5.2": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.z.ai/guides/overview/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3.8-Flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "cerebras/gemma-4-31b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 131072, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/gemma-4-31b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "elevenlabs/scribe_v2": { + "input_cost_per_second": 6.11e-05, + "litellm_provider": "elevenlabs", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } From 9c577c6045b01c0002e68d584da69b8bad996ed4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 14:04:15 -0700 Subject: [PATCH 041/126] test(e2e): assert user-observable behavior instead of DOM structure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The UI e2e suite had a class of assertions that pin how the dashboard is built rather than what it does, so an ordinary refactor turns them red without any user-visible change. Geometry. The auto-router template select had two tests made of pixel arithmetic plus a data-side="bottom" check, which is Base UI's own positioner signal. The regression they guard (#38554) is a popup opening on top of the control that spawned it, so both cases collapse to one invariant: the options never cover the trigger. It now runs at both viewport heights and reads the popup as role=listbox. The models header test compared the tabs and refresh centers within 2px, which a padding change flips; it now asserts the two share a row. Structure. The logs drawer test walked xpath=../../.. from a text node and read collapsed state off chevron icon classes. SectionHeader now renders a real disclosure button with aria-expanded, and its two copy buttons carry distinct names instead of both being "Copy". Sidebar group toggles expose aria-expanded too, so the migration spec can ask for a collapsed group by state rather than by nesting depth. Positional lookups. keyRow.locator("button").first(), row.locator("td") .first() and getByTestId(grid).locator("div").first() all named a position where they meant an action; they now name the control. Table scoping moves from "table tbody" to role=row. Timing. Nine waitForTimeout calls are gone. Every assertion that followed them already retried to its own timeout, so the sleeps only slowed the run down. Both files under tests/users/ were wrapped in test.skip("...", () => {}), which registers one skipped test and never runs the body, so the four tests inside had never executed and were written against a UI that has since changed (the search placeholder is "Search by email…", the ID filters moved into a drawer, pagination is labelled "Go to previous page"). Rewritten against the current surface: the suite goes from 104 collected tests to 107. Left in place deliberately: the chip and dialog-footer data-slot selectors, because the accessible names they work around live in components/ui/, which is shadcn CLI-managed and not hand-edited. --- .../tests/internal-user/internalUser.spec.ts | 7 +- .../internal-user/internalUserNoTeam.spec.ts | 7 +- .../internalUserWithTeams.spec.ts | 9 +- .../internal-viewer/internalViewer.spec.ts | 4 +- tests/e2e/ui/tests/logs/logs.spec.ts | 38 +++--- tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts | 2 +- .../ui/tests/migration/migratedPages.spec.ts | 10 +- .../e2e/ui/tests/modelsPage/addModel.spec.ts | 29 ++--- .../autoRouterTemplateSelect.spec.ts | 53 +++------ .../tests/modelsPage/responsiveHeader.spec.ts | 9 +- tests/e2e/ui/tests/proxy-admin/keys.spec.ts | 15 ++- .../e2e/ui/tests/team-admin/teamAdmin.spec.ts | 2 +- tests/e2e/ui/tests/usage/usagePage.spec.ts | 7 +- tests/e2e/ui/tests/users/searchUsers.spec.ts | 110 ++++++------------ .../ui/tests/users/viewInternalUsers.spec.ts | 59 +++------- .../src/components/leftnav.test.tsx | 14 +++ .../src/components/leftnav.tsx | 1 + .../LogDetailsDrawer/SectionHeader.test.tsx | 18 +++ .../LogDetailsDrawer/SectionHeader.tsx | 73 +++++++----- 19 files changed, 198 insertions(+), 269 deletions(-) diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index b8424b06115..26e34dd2fe5 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -22,8 +22,7 @@ test.describe("Internal User", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - const dropdown = page.locator('[data-slot="combobox-content"]:visible'); - await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({ timeout: 5_000 }); + await expect(page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }).first()).toBeVisible({ timeout: 5_000 }); }); test("Team info page omits the Settings tab for non-admin members", async ({ page }) => { @@ -43,12 +42,12 @@ test.describe("Internal User", () => { // Anchor on the user's own seeded key so the absence check below cannot // pass vacuously against an empty table. - await expect(page.locator("table tbody").getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()).toBeVisible({ + await expect(page.getByRole("row").filter({ hasText: E2E_INTERNAL_USER_KEY_ALIAS }).first()).toBeVisible({ timeout: 10_000, }); // The litellm-dashboard team is the proxy's internal bookkeeping team — // its keys must never leak into an internal user's Virtual Keys table. - await expect(page.locator("table tbody").getByText("litellm-dashboard")).toHaveCount(0); + await expect(page.getByRole("row").filter({ hasText: "litellm-dashboard" })).toHaveCount(0); }); }); diff --git a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts index c44305187f1..653e096b713 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts @@ -30,16 +30,13 @@ test.describe("Internal User with no team memberships", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); - await expect(dropdown).toBeVisible({ timeout: 5_000 }); - // Wait for the settled-empty state, not a transient one. The dropdown shows // "Loading teams…" while teams load and only swaps in "No teams found" once // the request resolves with nothing (team_dropdown.tsx passes both copies to // PaginatedSearchSelect). Asserting on it means a regression where teams DO // load for this user fails here instead of racing a one-shot count() against // an in-flight request. - await expect(dropdown.getByText("No teams found")).toBeVisible({ timeout: 10_000 }); - await expect(dropdown.getByRole("option")).toHaveCount(0); + await expect(page.getByText("No teams found")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("option")).toHaveCount(0); }); }); diff --git a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts index 68319154554..49e27a36673 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts @@ -21,13 +21,10 @@ test.describe("Internal User with team memberships", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); - await expect(dropdown).toBeVisible({ timeout: 5_000 }); - // Both seeded memberships render, and nothing else does — proving the // dropdown is scoped to the user's teams rather than empty or unfiltered. - await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS, { exact: true })).toBeVisible({ timeout: 10_000 }); - await expect(dropdown.getByText(E2E_TEAM_ORG_ALIAS, { exact: true })).toBeVisible(); - await expect(dropdown.getByRole("option")).toHaveCount(2); + await expect(page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS })).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("option", { name: E2E_TEAM_ORG_ALIAS })).toBeVisible(); + await expect(page.getByRole("option")).toHaveCount(2); }); }); diff --git a/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts index 4de86c46398..dd40976341d 100644 --- a/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts +++ b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts @@ -59,9 +59,9 @@ test.describe("Internal Viewer", () => { await expect(page.getByRole("button", { name: /Create New Key/i })).toHaveCount(0); // Open the viewer's own key info page - const keyRow = page.locator("tr", { hasText: E2E_VIEWER_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_VIEWER_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_VIEWER_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); // None of the destructive / mutating actions should render diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 56a3d0f0109..610a88c6cd0 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -11,12 +11,11 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, sendChatCompletion, waitForSpendLog } const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; -/** - * Walking up from the label is the only stable handle: the header carries no role, test id or class, - * and its copy button is icon-only with a hover-only tooltip. - */ -const sectionHeader = (drawer: Locator, label: "Input" | "Output"): Locator => - drawer.getByText(label, { exact: true }).locator("xpath=../../.."); +const sectionToggle = (drawer: Locator, label: "Input" | "Output"): Locator => + drawer.getByRole("button", { name: new RegExp(`^${label}\\b`) }); + +const sectionCopy = (drawer: Locator, label: "Input" | "Output"): Locator => + drawer.getByRole("button", { name: `Copy ${label.toLowerCase()}` }); /** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */ const requestLogsRows = (page: PlaywrightPage): Locator => @@ -95,14 +94,14 @@ test.describe("Logs page", () => { await expect(drawer).toBeVisible({ timeout: 20_000 }); // Copy request: the Input card's copy button puts the prompt on the clipboard. - await sectionHeader(drawer, "Input").getByRole("button").click(); + await sectionCopy(drawer, "Input").click(); await expect(page.getByText("Input copied")).toBeVisible({ timeout: 10_000, }); expect(await page.evaluate(() => navigator.clipboard.readText())).toContain(prompt); // Copy response: the Output card's copy button puts the completion on it. - await sectionHeader(drawer, "Output").getByRole("button").click(); + await sectionCopy(drawer, "Output").click(); await expect(page.getByText("Output copied")).toBeVisible({ timeout: 10_000, }); @@ -125,24 +124,15 @@ test.describe("Logs page", () => { timeout: 20_000, }); - // The body collapses via `max-height: 0; overflow: hidden`, which zeroes its own bounding - // box, so the wrapper reads as hidden while the clipped text node inside it does not. - const header = sectionHeader(drawer, "Input"); - const body = header.locator("xpath=following-sibling::div[1]"); - await expect(header.locator(".lucide-chevron-up")).toBeVisible(); - await expect(body).toBeVisible(); + const toggle = sectionToggle(drawer, "Input"); + await expect(toggle).toHaveAttribute("aria-expanded", "true"); + await expect(drawer.getByText(prompt, { exact: false })).toBeVisible(); - await header.click(); - await expect(header.locator(".lucide-chevron-down")).toBeVisible({ - timeout: 10_000, - }); - await expect(body).toBeHidden({ timeout: 10_000 }); + await toggle.click(); + await expect(toggle).toHaveAttribute("aria-expanded", "false", { timeout: 10_000 }); - await header.click(); - await expect(header.locator(".lucide-chevron-up")).toBeVisible({ - timeout: 10_000, - }); - await expect(body).toBeVisible({ timeout: 10_000 }); + await toggle.click(); + await expect(toggle).toHaveAttribute("aria-expanded", "true", { timeout: 10_000 }); await expect(drawer.getByText(prompt, { exact: false })).toBeVisible({ timeout: 10_000, }); diff --git a/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts b/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts index 46799c8a18f..aa7cdf82498 100644 --- a/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts +++ b/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts @@ -73,7 +73,7 @@ test.describe("MCP Servers - edit and delete", () => { test("Deleting a server removes it", async ({ page }) => { expect(await findServerByName(page, serverName), `created server ${serverName} exists`).toBeTruthy(); - const card = page.getByTestId("mcp-servers-grid").locator("div").filter({ hasText: serverName }).first(); + const card = page.getByTestId("mcp-servers-grid").getByRole("button", { name: serverName }); await card.getByRole("button", { name: "Server actions" }).click(); await page.getByRole("menuitem", { name: "Delete" }).click(); diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 3ad4b217d08..473d0b795f1 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -36,16 +36,10 @@ async function expectRendered(page: Page) { async function clickSidebar(page: Page, segment: string) { const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - // A collapsed group is a menu item with a group-toggle button but no - // rendered submenu yet; clicking the toggle expands it. - const collapsedGroup = sidebar(page) - .locator( - '[data-slot="sidebar-menu-item"]:has(> [data-slot="sidebar-menu-button"]):not(:has(> [data-slot="sidebar-menu-sub"])) > [data-slot="sidebar-menu-button"]', - ) - .first(); + const collapsedGroup = sidebar(page).getByRole("button", { expanded: false }).first(); if (!(await collapsedGroup.isVisible().catch(() => false))) break; await collapsedGroup.click(); - await page.waitForTimeout(250); + await expect(collapsedGroup).toHaveAttribute("aria-expanded", "true"); } await link.click(); } diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index dad716b4c83..84d1c01b452 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -188,7 +188,7 @@ test.describe("Add Model", () => { await expect(resultsModal).toBeHidden({ timeout: 5_000 }); const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); }); expect(created.model_name, "the model is created under the name that was typed").toBe(publicName); expect(created.litellm_params?.api_base, "the api base survives the form").toBe(MOCK_LLM_BASE); @@ -254,7 +254,7 @@ test.describe("Add Model", () => { // Click Add Model button by its text const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); }); // The form sends custom_llm_provider separately from the name, so both halves have to arrive. expect(created.model_name, "the selected model is what goes on the wire").toBe("claude-haiku-4-5"); @@ -267,11 +267,9 @@ test.describe("Add Model", () => { // Navigate to All Models tab await page.getByRole("tab", { name: "All Models" }).click(); await page.waitForLoadState("networkidle"); - await page.waitForTimeout(2000); // Search for the model we just added await page.getByPlaceholder("Search model names").fill("claude-haiku-4-5"); - await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { @@ -279,8 +277,9 @@ test.describe("Add Model", () => { }); // Verify the model name appears in the table body - const tableBody = page.locator("table tbody"); - await expect(tableBody.getByText("claude-haiku-4-5").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: "claude-haiku-4-5" })).not.toHaveCount(0, { + timeout: 15_000, + }); // A row proves the name is there, not what the deployment routes to. const stored = await findDeploymentByName(page, "claude-haiku-4-5"); @@ -333,11 +332,11 @@ test.describe("Add Model", () => { const teamDropdown = page.getByTestId("team-dropdown").getByRole("combobox"); await expect(teamDropdown).toBeVisible({ timeout: 5_000 }); await teamDropdown.click(); - const teamOption = page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ID).first(); + const teamOption = page.getByRole("option", { name: E2E_TEAM_CRUD_ID }).first(); await expect(teamOption).toBeVisible({ timeout: 5_000 }); await teamOption.click(); - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); // Scope to the toast container so a stale toast can't satisfy this. await expect(page.locator("[data-sonner-toast]").getByText("created successfully").last()).toBeVisible({ @@ -347,12 +346,9 @@ test.describe("Add Model", () => { // The Models table renders team-scoped models with the team id in the row. await page.getByRole("tab", { name: "All Models" }).click(); await page.waitForLoadState("networkidle"); - // networkidle fires before the table finishes re-rendering. - await page.waitForTimeout(2000); await page.getByPlaceholder("Search model names").fill("cohere"); - await page.waitForTimeout(1000); - + // 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, @@ -361,7 +357,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 - .locator("table tbody tr") + .getByRole("row") .filter({ hasText: "cohere/" }) .filter({ hasText: E2E_TEAM_CRUD_ID }); await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); @@ -387,7 +383,7 @@ test.describe("Add Model", () => { // Click Add Model button by its text const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); }); // A wildcard with the star stripped becomes a plain "cohere" deployment that matches nothing. expect(created.model_name, "the wildcard route goes on the wire intact").toBe("cohere/*"); @@ -398,11 +394,9 @@ test.describe("Add Model", () => { // Navigate to All Models tab await page.getByRole("tab", { name: "All Models" }).click(); await page.waitForLoadState("networkidle"); - await page.waitForTimeout(2000); // Search for the wildcard model await page.getByPlaceholder("Search model names").fill("cohere"); - await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { @@ -410,8 +404,7 @@ test.describe("Add Model", () => { }); // Verify the wildcard model appears in the table body (wildcard models show as "cohere/*") - const tableBody = page.locator("table tbody"); - await expect(tableBody.getByText("cohere/").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: "cohere/" })).not.toHaveCount(0, { timeout: 15_000 }); // "cohere/" in the table also matches a plain cohere deployment; require the wildcard exactly. const stored = await findDeploymentByName(page, "cohere/*"); diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts index 1d080ec82b8..fe168267a54 100644 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts @@ -17,49 +17,32 @@ async function openTemplateSelect(page: PlaywrightPage) { return trigger; } -function pollPixelsBelowTrigger(trigger: Locator, popup: Locator) { +function pollOptionsCoverTrigger(trigger: Locator, options: Locator) { return expect.poll(async () => { const triggerBox = await trigger.boundingBox(); - const popupBox = await popup.boundingBox(); - if (!triggerBox || !popupBox) return null; - return popupBox.y - (triggerBox.y + triggerBox.height); - }); -} - -function pollPopupOverlapsTrigger(trigger: Locator, popup: Locator) { - return expect.poll(async () => { - const triggerBox = await trigger.boundingBox(); - const popupBox = await popup.boundingBox(); - if (!triggerBox || !popupBox) return null; - return popupBox.y < triggerBox.y + triggerBox.height && popupBox.y + popupBox.height > triggerBox.y; + const optionsBox = await options.boundingBox(); + if (!triggerBox || !optionsBox) return null; + return optionsBox.y < triggerBox.y + triggerBox.height && optionsBox.y + optionsBox.height > triggerBox.y; }); } test.describe("Auto Router template select anchoring", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - test("opens the options below the trigger rather than over it", async ({ page }) => { - await page.setViewportSize({ width: 1280, height: 900 }); - const trigger = await openTemplateSelect(page); + for (const { room, height } of [ + { room: "with room below it", height: 900 }, + { room: "with no room below it", height: 560 }, + ]) { + test(`keeps the trigger uncovered when the options open ${room}`, async ({ page }) => { + await page.setViewportSize({ width: 1280, height }); + const trigger = await openTemplateSelect(page); + await trigger.scrollIntoViewIfNeeded(); - await trigger.click(); - const popup = page.locator('[data-slot="select-content"]'); - await expect(popup).toBeVisible(); + await trigger.click(); + const options = page.getByRole("listbox"); + await expect(options).toBeVisible(); - // Item-aligned mode reports "none" and puts the active item over the trigger. - await expect(popup).toHaveAttribute("data-side", "bottom"); - await pollPixelsBelowTrigger(trigger, popup).toBeGreaterThanOrEqual(0); - }); - - test("flips above the trigger instead of covering it when there is no room below", async ({ page }) => { - await page.setViewportSize({ width: 1280, height: 560 }); - const trigger = await openTemplateSelect(page); - await trigger.scrollIntoViewIfNeeded(); - - await trigger.click(); - const popup = page.locator('[data-slot="select-content"]'); - await expect(popup).toBeVisible(); - - await pollPopupOverlapsTrigger(trigger, popup).toBe(false); - }); + await pollOptionsCoverTrigger(trigger, options).toBe(false); + }); + } }); diff --git a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts index 6ad1ccb8451..366c371208f 100644 --- a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts @@ -7,9 +7,7 @@ test.describe("Models and Endpoints responsive header", () => { viewport: { width: 900, height: 720 }, }); - test("keeps the refresh action on the same row as the tabs", async ({ - page, - }) => { + test("keeps the refresh action on the same row as the tabs", async ({ page }) => { await page.goto("/ui"); await page .getByRole("complementary") @@ -26,8 +24,7 @@ test.describe("Models and Endpoints responsive header", () => { expect(tabsBox).not.toBeNull(); expect(refreshBox).not.toBeNull(); - const tabsCenterY = tabsBox!.y + tabsBox!.height / 2; - const refreshCenterY = refreshBox!.y + refreshBox!.height / 2; - expect(Math.abs(tabsCenterY - refreshCenterY)).toBeLessThanOrEqual(2); + const sharesARow = refreshBox!.y < tabsBox!.y + tabsBox!.height && refreshBox!.y + refreshBox!.height > tabsBox!.y; + expect(sharesARow, "refresh wrapped onto its own row below the tabs").toBe(true); }); }); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index 0c38641dcc7..f5bee68f245 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -43,7 +43,7 @@ test.describe("Proxy Admin - Keys", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); + await page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }).first().click(); // Select models — the popup is portaled to the body, so scope options to the page. await page.getByRole("combobox", { name: "Select models" }).click(); @@ -74,10 +74,9 @@ test.describe("Proxy Admin - Keys", () => { const before = await findKeyByAlias(page, E2E_REGENERATE_KEY_ALIAS); expect(before?.token, `seeded key ${E2E_REGENERATE_KEY_ALIAS} has a token`).toBeTruthy(); - // Key IDs are rendered as buttons in the table - const keyRow = page.locator("tr", { hasText: E2E_REGENERATE_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_REGENERATE_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_REGENERATE_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); @@ -109,9 +108,9 @@ test.describe("Proxy Admin - Keys", () => { const before = await findKeyByAlias(page, E2E_UPDATE_LIMITS_KEY_ALIAS); expect(before, `seeded key ${E2E_UPDATE_LIMITS_KEY_ALIAS} exists`).toBeTruthy(); - const keyRow = page.locator("tr", { hasText: E2E_UPDATE_LIMITS_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_UPDATE_LIMITS_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_UPDATE_LIMITS_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); @@ -147,9 +146,9 @@ test.describe("Proxy Admin - Keys", () => { await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); - const keyRow = page.locator("tr", { hasText: E2E_DELETE_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_DELETE_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_DELETE_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index f93cca75347..5116cb5df19 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -142,7 +142,7 @@ test.describe("Team Admin", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); + await page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }).first().click(); // Models — pick "All Team Models". The popup is portaled to the body, so // scope the option lookup to the page. diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index 8fa59beb905..f61ab018e1b 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -51,20 +51,19 @@ test.describe("Usage page", () => { const card = await openUsage(page); // Table view (the default): the key is listed by its alias. - const row = card.locator("tbody tr").filter({ hasText: alias }); + const row = card.getByRole("row").filter({ hasText: alias }); await expect(row, `${alias} missing from Top Virtual Keys`).toHaveCount(1, { timeout: 30_000, }); // Chart view swaps the table out for the bar chart, and back. await card.getByText("Chart View", { exact: true }).click(); - await expect(card.locator("tbody tr")).toHaveCount(0, { timeout: 10_000 }); + await expect(card.getByRole("table")).toHaveCount(0, { timeout: 10_000 }); await card.getByText("Table View", { exact: true }).click(); await expect(row).toHaveCount(1, { timeout: 10_000 }); - // Clicking the Key ID cell fetches key info and opens the detail panel. // The alias is already in the row behind the modal, so match the panel's own controls. - await row.locator("td").first().click(); + await row.getByRole("button", { name: token }).click(); const keyInfo = page.getByRole("tab", { name: "Overview", exact: true }); await expect(keyInfo, "key info panel did not open").toBeVisible({ timeout: 20_000, diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index e87218b5a5e..5e7e3e35b91 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -1,91 +1,51 @@ -import { test, expect, Page } from "@playwright/test"; +import { test, expect, Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; -test.skip("Internal Users Search", () => { +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; + +const userRows = (page: PlaywrightPage) => page.getByRole("row").filter({ has: page.getByRole("cell") }); + +async function goToInternalUsers(page: PlaywrightPage) { + await navigateToPage(page, Page.Users); + await expect(page.getByRole("columnheader", { name: "User ID" })).toBeVisible({ timeout: 30_000 }); + await expect(userRows(page)).not.toHaveCount(0, { timeout: 30_000 }); +} + +test.describe("Internal Users Search", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - async function goToInternalUsers(page: Page) { - await page.goto("/ui"); - - const tab = page.getByRole("menuitem", { name: "Internal User" }); - await expect(tab).toBeVisible(); - await tab.click(); - - await expect(page.locator("tbody tr").first()).toBeVisible(); - await expect(page.locator('[data-slot="skeleton"]')).toHaveCount(0); - } - - test("can search users by email", async ({ page }) => { + test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => { await goToInternalUsers(page); - const rows = page.locator("tbody tr"); - const searchInput = page.getByPlaceholder("Search by email..."); + const search = page.getByPlaceholder("Search by email…"); + await expect(search).toBeVisible(); - await expect(searchInput).toBeVisible(); + await search.fill("noteam@"); + await expect(userRows(page)).toHaveCount(1, { timeout: 30_000 }); + await expect(userRows(page).first()).toContainText("noteam@test.local"); - // Ensure initial data is loaded - const initialCount = await rows.count(); - expect(initialCount).toBeGreaterThan(0); - - // 🔹 Apply filter + wait for backend response - await Promise.all([ - page.waitForResponse( - (res) => - res.url().includes("/user/list") && - res.url().includes("user_email=test%40") && // encoded "test@" - res.status() === 200, - ), - searchInput.fill("test@"), - ]); - await page.waitForTimeout(5000); - const filteredCount = await rows.count(); - await expect(filteredCount).toBeLessThan(initialCount); - - // 🔹 Clear filter + wait for unfiltered request - await Promise.all([ - page.waitForResponse( - (res) => res.url().includes("/user/list") && !res.url().includes("user_email=") && res.status() === 200, - ), - searchInput.clear(), - ]); - - const resetCount = await rows.count(); - await expect(resetCount).toBe(initialCount); + await search.clear(); + await expect(userRows(page).filter({ hasText: "admin@test.local" })).not.toHaveCount(0, { timeout: 30_000 }); }); - test("can filter users by user ID and SSO ID", async ({ page }) => { + test("filters the table down to one user by user ID", async ({ page }) => { await goToInternalUsers(page); - const rows = page.locator("tbody tr"); - // Ensure initial data is loaded - const initialCount = await rows.count(); - expect(initialCount).toBeGreaterThan(0); + await page.getByRole("button", { name: "Filters" }).click(); + await page.getByTestId("users-filter-user-id").fill("e2e-internal-noteam"); + await page.getByTestId("filter-drawer-apply").click(); - const filtersButton = page.getByRole("button", { - name: "Filters", - exact: true, - }); - await filtersButton.click(); + await expect(userRows(page)).toHaveCount(1, { timeout: 30_000 }); + await expect(userRows(page).first()).toContainText("noteam@test.local"); + }); - const userIdInput = page.getByPlaceholder("Filter by User ID"); - const ssoIdInput = page.getByPlaceholder("Filter by SSO ID"); - await Promise.all([ - page.waitForResponse( - (res) => res.url().includes("/user/list") && res.url().includes("user_ids=user") && res.status() === 200, - ), - userIdInput.fill("user"), - ]); + test("shows no users when the SSO ID matches nobody", async ({ page }) => { + await goToInternalUsers(page); - await Promise.all([ - page.waitForResponse( - (res) => - res.url().includes("/user/list") && - res.url().includes("user_ids=user") && - res.url().includes("sso_user_ids=sso") && - res.status() === 200, - ), - ssoIdInput.fill("sso"), - ]); - const combinedFilteredCount = await rows.count(); - await expect(combinedFilteredCount).toBeLessThan(initialCount); + await page.getByRole("button", { name: "Filters" }).click(); + await page.getByTestId("users-filter-sso-id").fill("e2e-sso-id-that-matches-nobody"); + await page.getByTestId("filter-drawer-apply").click(); + + await expect(userRows(page)).toHaveCount(0, { timeout: 30_000 }); }); }); diff --git a/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts index 614191372d0..b46fb4d112a 100644 --- a/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts +++ b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts @@ -1,54 +1,29 @@ -import { test, expect, Page } from "@playwright/test"; +import { test, expect, Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; -test.skip("Internal Users Page", () => { +async function goToInternalUsers(page: PlaywrightPage) { + await navigateToPage(page, Page.Users); + await expect(page.getByRole("columnheader", { name: "User ID" })).toBeVisible({ timeout: 30_000 }); + await expect(userRows(page)).not.toHaveCount(0, { timeout: 30_000 }); +} + +const userRows = (page: PlaywrightPage) => page.getByRole("row").filter({ has: page.getByRole("cell") }); + +test.describe("Internal Users Page", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - async function goToInternalUsers(page: Page) { - await page.goto("/ui"); - - const internalUserTab = page.getByRole("menuitem", { name: "Internal User" }); - await expect(internalUserTab).toBeVisible(); - await internalUserTab.click(); - - const firstRow = page.locator("tbody tr").first(); - await expect(firstRow).toBeVisible(); - await expect(page.locator('[data-slot="skeleton"]')).toHaveCount(0); - } - - test("renders internal users table correctly", async ({ page }) => { + test("lists the seeded users under the identifying columns", async ({ page }) => { await goToInternalUsers(page); - const rows = page.locator("tbody tr"); - const rowCount = await rows.count(); - expect(rowCount).toBeGreaterThan(0); - - const userIdHeader = page.getByRole("columnheader", { name: "User ID" }); - await expect(userIdHeader).toBeVisible(); - - const virtualKeysHeader = page.getByRole("columnheader", { name: "Virtual Keys" }); - await expect(virtualKeysHeader).toBeVisible(); + await expect(page.getByRole("columnheader", { name: "User ID" })).toBeVisible(); + await expect(page.getByRole("columnheader", { name: "Virtual Keys" })).toBeVisible(); }); - test("pagination controls work correctly", async ({ page }) => { + test("cannot page backwards off the first page", async ({ page }) => { await goToInternalUsers(page); - const paginationInfo = page.locator(".text-sm.text-gray-700"); - const prevButton = page.getByRole("button", { name: "Previous" }); - const nextButton = page.getByRole("button", { name: "Next" }); - - const infoText = (await paginationInfo.textContent()) || ""; - - // On first page, Previous should be disabled - if (infoText.includes("1 -")) { - await expect(prevButton).toBeDisabled(); - } - - await page.waitForTimeout(1000); - // Check if there are more pages - const hasMorePages = infoText.includes("of") && !infoText.endsWith("25 of 25"); - if (hasMorePages) { - await expect(nextButton).toBeEnabled(); - } + await expect(page.getByRole("button", { name: "Go to previous page" })).toBeDisabled(); }); }); diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index b0bb0e8a5b5..c3e1f924d09 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -213,6 +213,20 @@ describe("Sidebar (leftnav)", () => { expect(screen.getByText("Search Tools")).toBeInTheDocument(); }); }); + it("reports whether a nested tab is expanded", async () => { + renderWithProviders(); + + const toggle = screen.getByText("Tools").closest("button")!; + expect(toggle).toHaveAttribute("aria-expanded", "false"); + + act(() => { + fireEvent.click(toggle); + }); + await waitFor(() => { + expect(toggle).toHaveAttribute("aria-expanded", "true"); + }); + }); + it("keeps Router Settings as a single Settings child", () => { // Router Settings is admin-only, so getAvailablePages() filters it out entirely and the // page_utils duplicate-key guard cannot see it. Walk menuGroups directly, otherwise a diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 42389facac2..51ba36348e1 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -570,6 +570,7 @@ const Sidebar_: React.FC = ({ toggleGroup(item.key)} title={collapsed ? labelText(item) : undefined} > diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx index 5aee6b33ec5..6cd9f476628 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx @@ -53,6 +53,24 @@ describe("SectionHeader", () => { expect(onToggleCollapse).toHaveBeenCalledTimes(1); }); + it("reports its collapsed state to assistive technology", () => { + const { rerender } = render( + , + ); + + expect(screen.getByRole("button", { name: /^Input/ })).toHaveAttribute("aria-expanded", "true"); + + rerender(); + + expect(screen.getByRole("button", { name: /^Input/ })).toHaveAttribute("aria-expanded", "false"); + }); + + it("names each copy button for the section it belongs to", () => { + render(); + + expect(screen.getByRole("button", { name: "Copy output" })).toBeInTheDocument(); + }); + it("stays inert when no toggle handler is given", async () => { const onCopy = vi.fn(); render(); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx index 93e9953b2ee..6a18aaf8642 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx @@ -17,6 +17,8 @@ interface SectionHeaderProps { turnCount?: number; } +const SUMMARY_CLASSES = "flex flex-1 items-center gap-4"; + export function SectionHeader({ type, tokens, @@ -26,42 +28,53 @@ export function SectionHeader({ onToggleCollapse, turnCount, }: SectionHeaderProps) { + const summary = ( + <> + {onToggleCollapse && + (isCollapsed ? ( + + ) : ( + + ))} + +
+ {type === "input" ? ( + + ) : ( + ✨ + )} + {type === "input" ? "Input" : "Output"} +
+ + {tokens !== undefined && Tokens: {tokens.toLocaleString()}} + + {cost !== undefined && Cost: ${cost.toFixed(6)}} + + {turnCount !== undefined && turnCount > 0 && ( + Turns: {turnCount} + )} + + ); + return (
-
- {onToggleCollapse && - (isCollapsed ? ( - - ) : ( - - ))} - -
- {type === "input" ? ( - - ) : ( - ✨ - )} - {type === "input" ? "Input" : "Output"} -
- - {tokens !== undefined && ( - Tokens: {tokens.toLocaleString()} - )} - - {cost !== undefined && Cost: ${cost.toFixed(6)}} - - {turnCount !== undefined && turnCount > 0 && ( - Turns: {turnCount} - )} -
+ {onToggleCollapse ? ( + + ) : ( +
{summary}
+ )} { e.stopPropagation(); onCopy(); From 78e1c658b4828ac5595d1bdabb259d873691148d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 14:45:11 -0700 Subject: [PATCH 042/126] fix(e2e): assert sidebar expansion without a self-resolving locator The migration smoke waited on `getByRole("button", { expanded: false })` after clicking it. Playwright re-resolves that locator on every retry, so once the clicked group flipped to expanded it matched the next collapsed group instead, and the assertion could never pass. Count the remaining collapsed groups and wait for that count to drop by one. --- tests/e2e/ui/tests/migration/migratedPages.spec.ts | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 473d0b795f1..547330190bd 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -35,11 +35,12 @@ async function expectRendered(page: Page) { */ async function clickSidebar(page: Page, segment: string) { const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); + const collapsedGroups = sidebar(page).getByRole("button", { expanded: false }); for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - const collapsedGroup = sidebar(page).getByRole("button", { expanded: false }).first(); - if (!(await collapsedGroup.isVisible().catch(() => false))) break; - await collapsedGroup.click(); - await expect(collapsedGroup).toHaveAttribute("aria-expanded", "true"); + const stillCollapsed = await collapsedGroups.count(); + if (stillCollapsed === 0) break; + await collapsedGroups.first().click(); + await expect(collapsedGroups).toHaveCount(stillCollapsed - 1); } await link.click(); } From be5997f3664f7359204fe58f3111a93f24471aee Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:48:16 -0700 Subject: [PATCH 043/126] feat: add Azure AI DeepSeek V4 Flash 0731 pricing --- model_prices_and_context_window.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7071eaa0807..6b4a9818554 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9959,6 +9959,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 1.9e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.1e-07, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/embed-v-4-0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", From cc258b5473932c939903d589604f83f2ca260469 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 14:54:11 -0700 Subject: [PATCH 044/126] test(e2e): keep the placement guarantees the geometry rewrites dropped The consolidated popup test only asserted the options never cover the trigger, so opening above the trigger with room below it, the regression PR #38554 fixed, would have passed. Split it back into a below-trigger case and a cramped-viewport case. The header test accepted a single pixel of vertical intersection; require the refresh control's centre to sit within the tab row instead. --- .../autoRouterTemplateSelect.spec.ts | 58 +++++++++++++------ .../tests/modelsPage/responsiveHeader.spec.ts | 3 +- 2 files changed, 42 insertions(+), 19 deletions(-) diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts index fe168267a54..7fc20104f20 100644 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts @@ -17,32 +17,54 @@ async function openTemplateSelect(page: PlaywrightPage) { return trigger; } +async function boxes(trigger: Locator, options: Locator) { + const triggerBox = await trigger.boundingBox(); + const optionsBox = await options.boundingBox(); + return triggerBox && optionsBox ? { triggerBox, optionsBox } : null; +} + +function pollOptionsOpenBelowTrigger(trigger: Locator, options: Locator) { + return expect.poll(async () => { + const box = await boxes(trigger, options); + return box && box.optionsBox.y >= box.triggerBox.y + box.triggerBox.height; + }); +} + function pollOptionsCoverTrigger(trigger: Locator, options: Locator) { return expect.poll(async () => { - const triggerBox = await trigger.boundingBox(); - const optionsBox = await options.boundingBox(); - if (!triggerBox || !optionsBox) return null; - return optionsBox.y < triggerBox.y + triggerBox.height && optionsBox.y + optionsBox.height > triggerBox.y; + const box = await boxes(trigger, options); + return ( + box && + box.optionsBox.y < box.triggerBox.y + box.triggerBox.height && + box.optionsBox.y + box.optionsBox.height > box.triggerBox.y + ); }); } test.describe("Auto Router template select anchoring", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - for (const { room, height } of [ - { room: "with room below it", height: 900 }, - { room: "with no room below it", height: 560 }, - ]) { - test(`keeps the trigger uncovered when the options open ${room}`, async ({ page }) => { - await page.setViewportSize({ width: 1280, height }); - const trigger = await openTemplateSelect(page); - await trigger.scrollIntoViewIfNeeded(); + test("opens the options below the trigger when there is room below it", async ({ page }) => { + await page.setViewportSize({ width: 1280, height: 900 }); + const trigger = await openTemplateSelect(page); + await trigger.scrollIntoViewIfNeeded(); - await trigger.click(); - const options = page.getByRole("listbox"); - await expect(options).toBeVisible(); + await trigger.click(); + const options = page.getByRole("listbox"); + await expect(options).toBeVisible(); - await pollOptionsCoverTrigger(trigger, options).toBe(false); - }); - } + await pollOptionsOpenBelowTrigger(trigger, options).toBe(true); + }); + + test("keeps the trigger uncovered when the options open with no room below it", async ({ page }) => { + await page.setViewportSize({ width: 1280, height: 560 }); + const trigger = await openTemplateSelect(page); + await trigger.scrollIntoViewIfNeeded(); + + await trigger.click(); + const options = page.getByRole("listbox"); + await expect(options).toBeVisible(); + + await pollOptionsCoverTrigger(trigger, options).toBe(false); + }); }); diff --git a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts index 366c371208f..aabdf18d427 100644 --- a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts @@ -24,7 +24,8 @@ test.describe("Models and Endpoints responsive header", () => { expect(tabsBox).not.toBeNull(); expect(refreshBox).not.toBeNull(); - const sharesARow = refreshBox!.y < tabsBox!.y + tabsBox!.height && refreshBox!.y + refreshBox!.height > tabsBox!.y; + const refreshCenterY = refreshBox!.y + refreshBox!.height / 2; + const sharesARow = refreshCenterY > tabsBox!.y && refreshCenterY < tabsBox!.y + tabsBox!.height; expect(sharesARow, "refresh wrapped onto its own row below the tabs").toBe(true); }); }); From 1a80b7ae252e017b937dfb9cf665cc04152ce4ec Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 15:00:39 -0700 Subject: [PATCH 045/126] fix: sync Azure AI model backup registry --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7071eaa0807..6b4a9818554 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9959,6 +9959,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 1.9e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.1e-07, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/embed-v-4-0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", From fdc259077ea48225403bc16c0dbb964f4298eda2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:05:54 -0700 Subject: [PATCH 046/126] test(e2e/ui): automate 8 manual QA checklist flows Adds Playwright coverage for the RC checklist items an audit marked automatable today: Playground to Logs hand-off, public Agent/MCP hub tabs, team models in the Playground dropdown via a team key, Add Model with a stored credential, internal user team key creation, a second admin account, team model deletion, and Presidio guardrail CRUD without a live sidecar. Seeds e2e-team-keygen with the /key/generate member permission so the internal user key flow avoids the team-list cache lag --- tests/e2e/ui/constants.ts | 3 + tests/e2e/ui/fixtures/seed.sql | 14 ++- tests/e2e/ui/helpers/traffic.ts | 28 ++++++ .../ui/tests/guardrails/guardrails.spec.ts | 83 ++++++++++++++++ .../tests/internal-user/internalUser.spec.ts | 55 +++++++++++ .../internalUserWithTeams.spec.ts | 16 +++- tests/e2e/ui/tests/logs/logs.spec.ts | 26 ++++- tests/e2e/ui/tests/modelHub/modelHub.spec.ts | 81 +++++++++++++++- .../e2e/ui/tests/modelsPage/addModel.spec.ts | 81 ++++++++++++++++ .../tests/modelsPage/deleteTeamModel.spec.ts | 70 ++++++++++++++ .../ui/tests/proxy-admin/secondAdmin.spec.ts | 94 +++++++++++++++++++ .../e2e/ui/tests/team-admin/teamAdmin.spec.ts | 88 +++++++++++++++++ 12 files changed, 631 insertions(+), 8 deletions(-) create mode 100644 tests/e2e/ui/tests/guardrails/guardrails.spec.ts create mode 100644 tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts create mode 100644 tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts index 9d918736262..bb33c90ddf3 100644 --- a/tests/e2e/ui/constants.ts +++ b/tests/e2e/ui/constants.ts @@ -29,6 +29,7 @@ export const E2E_PROXY_ADMIN_USER_ID = "e2e-proxy-admin"; export const E2E_PROXY_ADMIN_EMAIL = "admin@test.local"; export const E2E_INTERNAL_USER_ID = "e2e-internal-user"; export const E2E_INTERNAL_USER_EMAIL = "internal@test.local"; +export const E2E_TEAM_ADMIN_USER_ID = "e2e-team-admin"; // Key aliases for seeded test keys (match seed.sql) export const E2E_UPDATE_LIMITS_KEY_ALIAS = "e2eUpdateLimitsKey"; @@ -46,3 +47,5 @@ export const E2E_TEAM_ORG_ID = "e2e-team-org"; export const E2E_TEAM_ORG_ALIAS = "E2E Team In Org"; export const E2E_TEAM_NO_ADMIN_ID = "e2e-team-no-admin"; export const E2E_TEAM_NO_ADMIN_ALIAS = "E2E Team No Admin"; +export const E2E_TEAM_KEYGEN_ID = "e2e-team-keygen"; +export const E2E_TEAM_KEYGEN_ALIAS = "E2E Team Keygen"; diff --git a/tests/e2e/ui/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql index a1218633cdb..e77b4a16b3d 100644 --- a/tests/e2e/ui/fixtures/seed.sql +++ b/tests/e2e/ui/fixtures/seed.sql @@ -29,7 +29,7 @@ INSERT INTO "LiteLLM_UserTable" ("user_id", "user_email", "user_role", "teams", VALUES ('e2e-proxy-admin', 'admin@test.local', 'proxy_admin', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-admin-viewer', 'adminviewer@test.local', 'proxy_admin_viewer', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), + ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org","e2e-team-keygen"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), @@ -63,6 +63,17 @@ INSERT INTO "LiteLLM_TeamTable" ( '[{"role":"user","user_id":"e2e-invitable-user"}]'::jsonb, '{}'::jsonb, '{"fake-openai-gpt-4"}', 0.0, '{}'::jsonb, '{}'::jsonb, false); +INSERT INTO "LiteLLM_TeamTable" ( + "team_id", "team_alias", "organization_id", "admins", "members", + "members_with_roles", "metadata", "models", "spend", "model_spend", "model_max_budget", "blocked", + "team_member_permissions" +) VALUES + ('e2e-team-keygen', 'E2E Team Keygen', NULL, + '{}', '{"e2e-internal-user"}', + '[{"role":"user","user_id":"e2e-internal-user"}]'::jsonb, + '{}'::jsonb, '{"fake-openai-gpt-4"}', 0.0, '{}'::jsonb, '{}'::jsonb, false, + '{"/key/generate"}'); + -- 6. Team Memberships (only user_id, team_id, spend — no created_at/updated_at) INSERT INTO "LiteLLM_TeamMembership" ("user_id", "team_id", "spend") VALUES @@ -72,6 +83,7 @@ VALUES ('e2e-removable-member', 'e2e-team-crud', 0.0), ('e2e-team-admin', 'e2e-team-delete', 0.0), ('e2e-internal-user', 'e2e-team-org', 0.0), + ('e2e-internal-user', 'e2e-team-keygen', 0.0), ('e2e-invitable-user', 'e2e-team-no-admin', 0.0); -- 7. Verification Tokens (API Keys) diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index a2fc9463c94..ebd3c9a417f 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -84,6 +84,34 @@ export async function waitForSpendLog( throw new Error(`spend log for request ${requestId} never appeared (last /spend/logs status ${lastStatus})`); } +export async function waitForSpendLogByPrompt( + request: APIRequestContext, + prompt: string, + timeoutMs = 60_000, +): Promise { + const deadline = Date.now() + timeoutMs; + let lastStatus = 0; + while (Date.now() < deadline) { + const res = await request.get(`${rootPath()}/spend/logs`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + lastStatus = res.status(); + if (res.ok()) { + const rows: { request_id?: string; messages?: unknown; proxy_server_request?: unknown }[] = await res.json(); + const row = (Array.isArray(rows) ? rows : []).find( + (candidate) => + JSON.stringify(candidate.messages ?? "").includes(prompt) || + JSON.stringify(candidate.proxy_server_request ?? "").includes(prompt), + ); + if (row?.request_id) { + return row.request_id; + } + } + await new Promise((r) => setTimeout(r, 2_000)); + } + throw new Error(`no spend log row carrying prompt ${prompt} appeared (last /spend/logs status ${lastStatus})`); +} + const isoDay = (d: Date): string => d.toISOString().slice(0, 10); /** diff --git a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts new file mode 100644 index 00000000000..a3ce6a73075 --- /dev/null +++ b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts @@ -0,0 +1,83 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_NO_ADMIN_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; + +test.describe("Guardrails", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Create a Presidio guardrail, see it in team settings, and delete it", async ({ page }) => { + const guardrailName = `e2e-presidio-${Date.now()}`; + + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: /Add New Guardrail/i }).click(); + await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click(); + + const dialog = page.getByRole("dialog", { name: "Create guardrail" }); + await expect(dialog).toBeVisible({ timeout: 10_000 }); + + await dialog.getByLabel("Guardrail Name").fill(guardrailName); + + const providerSelect = dialog.getByRole("combobox", { name: "Guardrail Provider" }); + await providerSelect.click(); + await providerSelect.fill("Presidio"); + await page.getByRole("option", { name: "Presidio PII" }).click(); + + await dialog.getByLabel("Mode", { exact: true }).click(); + await page.keyboard.type("pre_call"); + await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 }); + await page.keyboard.press("Enter"); + await expect(dialog.locator('[data-slot="combobox-chip"]').filter({ hasText: "pre_call" })).toBeVisible({ + timeout: 5_000, + }); + await dialog.getByText("Create guardrail", { exact: true }).click(); + + await dialog.getByLabel("presidio_analyzer_api_base").fill("http://127.0.0.1:9999"); + await expect(dialog.getByLabel("presidio_analyzer_api_base")).toHaveValue("http://127.0.0.1:9999"); + await dialog.getByLabel("presidio_anonymizer_api_base").fill("http://127.0.0.1:9999"); + await expect(dialog.getByLabel("presidio_anonymizer_api_base")).toHaveValue("http://127.0.0.1:9999"); + + await dialog.getByRole("button", { name: "Next" }).click(); + await expect(dialog.getByText("Configure PII Protection")).toBeVisible({ timeout: 10_000 }); + await dialog.getByRole("button", { name: "Select All & Mask" }).click(); + + await dialog.getByRole("button", { name: "Create Guardrail" }).click(); + await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 }); + + const row = page.locator("table tbody tr").filter({ hasText: guardrailName }); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, E2E_TEAM_NO_ADMIN_ID); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + const guardrailsSelect = page.getByRole("combobox", { name: "Select guardrails" }); + await expect(guardrailsSelect).toBeVisible({ timeout: 10_000 }); + await guardrailsSelect.click(); + await guardrailsSelect.fill(guardrailName); + await expect(page.getByRole("option", { name: guardrailName })).toBeVisible({ timeout: 10_000 }); + await page.keyboard.press("Escape"); + + await navigateToPage(page, Page.Guardrails); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + await row.getByRole("button", { name: "Open guardrail actions" }).click(); + await page.getByRole("menuitem", { name: "Delete" }).click(); + + const deleteModal = page.getByRole("dialog", { name: "Delete Guardrail" }); + await expect(deleteModal).toBeVisible({ timeout: 5_000 }); + await deleteModal.getByRole("button", { name: "Delete", exact: true }).click(); + + await expect(page.getByText(`Guardrail "${guardrailName}" deleted successfully`)).toBeVisible({ + timeout: 10_000, + }); + await expect(row).toHaveCount(0, { timeout: 15_000 }); + + await page.reload(); + await expect(page.getByRole("button", { name: /Add New Guardrail/i })).toBeVisible({ timeout: 20_000 }); + await expect(page.locator("table tbody tr").filter({ hasText: guardrailName })).toHaveCount(0); + }); +}); diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index b8424b06115..6c893e469c6 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -3,10 +3,13 @@ import { E2E_INTERNAL_USER_KEY_ALIAS, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID, + E2E_TEAM_KEYGEN_ALIAS, INTERNAL_USER_STORAGE_PATH, } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage, clickTeamId } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; +import { keySourceSelect, onlyVisible, openPlayground, selectModel, sendMessage } from "../../helpers/playground"; test.describe("Internal User", () => { test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); @@ -38,6 +41,58 @@ test.describe("Internal User", () => { await expect(page.getByRole("tab", { name: "Members" })).not.toBeVisible(); }); + test("Internal user creates a team key and uses it in the Playground", async ({ page, request }) => { + const suffix = Date.now(); + const auth = { Authorization: `Bearer ${masterKey()}` }; + + let apiKey = ""; + try { + await navigateToPage(page, Page.ApiKeys); + + await page.getByRole("button", { name: /Create New Key/i }).click(); + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + await expect(page.getByRole("radio", { name: "You", exact: true })).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("radio", { name: "Another User" })).toHaveCount(0); + + const keyName = `e2e-internal-team-key-${suffix}`; + await page.getByLabel(/Key Name/).fill(keyName); + + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); + await teamSelect.click(); + await page.keyboard.type(E2E_TEAM_KEYGEN_ALIAS); + await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_KEYGEN_ALIAS).first().click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Team Models", exact: true }).click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + await page.keyboard.press("Escape"); + + await openPlayground(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + + const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(apiKey); + + await selectModel(page, CHAT_MODEL_A); + await sendMessage(page, `internal user team key ping ${keyName}`); + + await expect(page.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 60_000 }); + } finally { + if (apiKey) { + await request.post("/key/delete", { headers: auth, data: { keys: [apiKey] } }); + } + } + }); + test("Virtual Keys page does not surface litellm-dashboard team keys", async ({ page }) => { await navigateToPage(page, Page.ApiKeys); diff --git a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts index 68319154554..23e9ed78d9c 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts @@ -1,12 +1,17 @@ import { test, expect } from "@playwright/test"; -import { INTERNAL_USER_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_ORG_ALIAS } from "../../constants"; +import { + INTERNAL_USER_STORAGE_PATH, + E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_KEYGEN_ALIAS, + E2E_TEAM_ORG_ALIAS, +} from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; /** * Differential partner to internalUserNoTeam.spec.ts: the seeded - * e2e-internal-user belongs to exactly two teams, so the Create Key dropdown - * must list both. Without this, the no-team spec's "zero options" assertion + * e2e-internal-user belongs to exactly three teams, so the Create Key dropdown + * must list all of them. Without this, the no-team spec's "zero options" assertion * would still pass against a bug that empties the dropdown for everyone. */ test.describe("Internal User with team memberships", () => { @@ -24,10 +29,11 @@ test.describe("Internal User with team memberships", () => { const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); await expect(dropdown).toBeVisible({ timeout: 5_000 }); - // Both seeded memberships render, and nothing else does — proving the + // All seeded memberships render, and nothing else does — proving the // dropdown is scoped to the user's teams rather than empty or unfiltered. await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS, { exact: true })).toBeVisible({ timeout: 10_000 }); await expect(dropdown.getByText(E2E_TEAM_ORG_ALIAS, { exact: true })).toBeVisible(); - await expect(dropdown.getByRole("option")).toHaveCount(2); + await expect(dropdown.getByText(E2E_TEAM_KEYGEN_ALIAS, { exact: true })).toBeVisible(); + await expect(dropdown.getByRole("option")).toHaveCount(3); }); }); diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 56a3d0f0109..46dfbf47478 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -2,7 +2,14 @@ import { test, expect, type Locator, type Page as PlaywrightPage } from "@playwr import { ADMIN_STORAGE_PATH } from "../../constants"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; -import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; +import { + CHAT_MODEL_A, + MOCK_RESPONSE_TEXT, + sendChatCompletion, + waitForSpendLog, + waitForSpendLogByPrompt, +} from "../../helpers/traffic"; +import { openPlayground, selectModel, sendMessage } from "../../helpers/playground"; /** * Anchored to traffic this spec generates itself, with a unique prompt and end user per run, so it @@ -47,6 +54,23 @@ test.describe("Logs page", () => { permissions: ["clipboard-read", "clipboard-write"], }); + test("a chat sent from the Playground lands in Logs with its content", async ({ page, request }) => { + const prompt = `logs-playground-prompt-${uniqueSuffix()}`; + await openPlayground(page); + await selectModel(page, CHAT_MODEL_A); + await sendMessage(page, prompt); + await expect(page.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 60_000 }); + + const requestId = await waitForSpendLogByPrompt(request, prompt); + + const row = await openLogsForRequest(page, requestId); + await row.click(); + const drawer = page.getByRole("dialog").first(); + await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 }); + await expect(drawer.getByText(prompt, { exact: false }).first()).toBeVisible({ timeout: 20_000 }); + await expect(drawer.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 20_000 }); + }); + test("a served request expands to its request and response", async ({ page, request }) => { const prompt = `logs-detail-prompt-${uniqueSuffix()}`; const requestId = await sendChatCompletion(request, { diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 16ec94c1dc8..7fe3894d75d 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -1,7 +1,8 @@ -import { test, expect } from "@playwright/test"; +import { test, expect, type APIRequestContext } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; +import { masterKey } from "../../helpers/traffic"; test.describe("AI Hub (internal admin view)", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -77,4 +78,82 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { // agents/MCP servers exist, so we don't assert on them in a fresh CI run. await expect(page.getByRole("tab", { name: "Model Hub" })).toBeVisible({ timeout: 10_000 }); }); + + test("Agent Hub and MCP Hub tabs render their public entries", async ({ page, request }) => { + const suffix = `${Date.now()}`; + const agentName = `e2e-public-agent-${suffix}`; + const mcpServerName = `e2e_public_mcp_${suffix}`; + const auth = { Authorization: `Bearer ${masterKey()}` }; + + const seedPublicEntries = async (api: APIRequestContext): Promise<{ agentId: string; serverId: string }> => { + const agentRes = await api.post("/v1/agents", { + headers: auth, + data: { + agent_name: agentName, + agent_card_params: { + name: agentName, + description: "E2E public agent", + version: "1.0.0", + url: "http://127.0.0.1:9999/", + capabilities: {}, + skills: [], + defaultInputModes: ["text"], + defaultOutputModes: ["text"], + }, + }, + }); + expect(agentRes.ok(), `agent create failed (${agentRes.status()}): ${await agentRes.text()}`).toBe(true); + const agentId = (await agentRes.json()).agent_id as string; + + const serverRes = await api.post("/v1/mcp/server", { + headers: auth, + data: { + server_name: mcpServerName, + url: "http://127.0.0.1:9999/mcp", + transport: "http", + description: "E2E public MCP server", + }, + }); + expect(serverRes.ok(), `mcp server create failed (${serverRes.status()}): ${await serverRes.text()}`).toBe(true); + const serverId = (await serverRes.json()).server_id as string; + + const agentPublicRes = await api.post("/v1/agents/make_public", { + headers: auth, + data: { agent_ids: [agentId] }, + }); + expect(agentPublicRes.ok(), `agents make_public failed: ${await agentPublicRes.text()}`).toBe(true); + const mcpPublicRes = await api.post("/v1/mcp/make_public", { + headers: auth, + data: { mcp_server_ids: [serverId] }, + }); + expect(mcpPublicRes.ok(), `mcp make_public failed: ${await mcpPublicRes.text()}`).toBe(true); + + return { agentId, serverId }; + }; + + const { agentId, serverId } = await seedPublicEntries(request); + try { + await page.goto(`/ui/model_hub_table?key=${masterKey()}`); + await dismissFeedbackPopup(page); + + const agentHubTab = page.getByRole("tab", { name: "Agent Hub" }); + await expect(agentHubTab).toBeVisible({ timeout: 15_000 }); + await agentHubTab.click(); + await expect(page.getByText("Available Agents")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("row").filter({ hasText: agentName })).toHaveCount(1, { timeout: 10_000 }); + await expect(page.getByText("E2E public agent").first()).toBeVisible(); + + const mcpHubTab = page.getByRole("tab", { name: "MCP Hub" }); + await expect(mcpHubTab).toBeVisible(); + await mcpHubTab.click(); + await expect(page.getByText("Available MCP Servers")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("row").filter({ hasText: mcpServerName })).toHaveCount(1, { timeout: 10_000 }); + await expect(page.getByText("E2E public MCP server").first()).toBeVisible(); + } finally { + await request.post("/v1/agents/make_public", { headers: auth, data: { agent_ids: [] } }); + await request.post("/v1/mcp/make_public", { headers: auth, data: { mcp_server_ids: [] } }); + await request.delete(`/v1/agents/${agentId}`, { headers: auth }); + await request.delete(`/v1/mcp/server/${serverId}`, { headers: auth }); + } + }); }); diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index dad716b4c83..f8ca02b36d0 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -212,6 +212,87 @@ test.describe("Add Model", () => { .toBe(true); }); + test("Add a model with a stored credential, pass Test Connect, and serve traffic", async ({ page, request }) => { + const masterKey = users[Role.ProxyAdmin].password; + const auth = { Authorization: `Bearer ${masterKey}` }; + const credentialName = `e2e-cred-reuse-${Date.now()}`; + const createCred = await page.request.post("/credentials", { + headers: auth, + data: { + credential_name: credentialName, + credential_values: { api_key: "fake-key", api_base: MOCK_LLM_BASE }, + credential_info: { custom_llm_provider: "openai" }, + }, + }); + expect(createCred.ok(), `POST /credentials failed (${createCred.status()}): ${await createCred.text()}`).toBe(true); + + try { + await navigateToPage(page, Page.Models); + await page.getByRole("tab", { name: "Add Model" }).click(); + + await selectProvider(page, "OpenAI-Compatible Endpoints (Together AI, etc.)"); + + const publicName = `e2e-cred-model-${Date.now()}`; + uiAddedModelName = publicName; + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); + await page.keyboard.press("Escape"); + await page.getByPlaceholder("Enter custom model name").fill(publicName); + + const credentialSelect = page.getByRole("combobox", { name: "Existing Credentials" }); + await credentialSelect.click(); + await credentialSelect.fill(credentialName); + await page.getByRole("option", { name: credentialName, exact: true }).click(); + + await expect(page.locator("#api_key")).toHaveCount(0); + await expect(page.locator("#api_base")).toHaveCount(0); + + await page.getByRole("button", { name: "Test Connect" }).click(); + await expect(page.getByText("Connection Test Results")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByTestId("connection-success-msg")).toBeVisible({ timeout: 30_000 }); + + const resultsModal = page.getByRole("dialog", { name: "Connection Test Results" }); + await resultsModal.locator('[data-slot="dialog-footer"]').getByRole("button", { name: "Close" }).click(); + await expect(resultsModal).toBeHidden({ timeout: 5_000 }); + + const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { + await page.getByRole("button", { name: "Add Model" }).last().click(); + }); + expect(created.litellm_params?.litellm_credential_name, "the picked credential goes on the wire").toBe( + credentialName, + ); + expect(created.litellm_params?.api_key, "no raw api key goes on the wire").toBeUndefined(); + + await expect(page.getByText("created successfully")).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => { + try { + await sendChatCompletion(request, { model: publicName, prompt: `hello via ${credentialName}` }); + return true; + } catch { + return false; + } + }, + { + message: `model ${publicName} added with a stored credential never served a request`, + timeout: 30_000, + }, + ) + .toBe(true); + } finally { + const stored = uiAddedModelName ? await findDeploymentByName(page, uiAddedModelName) : undefined; + const id = stored?.model_info?.id; + if (id) { + await page.request.post("/model/delete", { headers: auth, data: { id } }); + uiAddedModelName = ""; + } + await page.request.delete(`/credentials/${credentialName}`, { headers: auth }); + } + }); + test("Test connection with bad credentials shows failure", async ({ page }) => { await navigateToPage(page, Page.Models); await page.getByRole("tab", { name: "Add Model" }).click(); diff --git a/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts b/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts new file mode 100644 index 00000000000..dca7d9f006b --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts @@ -0,0 +1,70 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; + +async function findDeploymentByName(page: PlaywrightPage, modelName: string): Promise | undefined> { + const body = await readBack<{ data: Record[] }>(page, "/v2/model/info"); + return body.data.find((row) => row.model_name === modelName); +} + +test.describe("Delete team model", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Delete a team-scoped model and verify it leaves the team's model list", async ({ page }) => { + const modelName = `e2e-team-model-delete-${Date.now()}`; + const createResponse = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: modelName, + litellm_params: { + model: "openai/fake-gpt-4", + api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`, + api_key: "fake-key", + }, + model_info: { team_id: E2E_TEAM_CRUD_ID }, + }, + }); + expect(createResponse.ok(), `/model/new failed: ${createResponse.status()} ${await createResponse.text()}`).toBe( + true, + ); + + await expect + .poll(async () => (await findDeploymentByName(page, modelName)) !== undefined, { + message: `deployment ${modelName} never appeared in /v2/model/info after create`, + timeout: 30_000, + }) + .toBe(true); + + await navigateToPage(page, Page.Models); + await page.getByPlaceholder("Search model names").fill(modelName); + + const row = page.locator("table tbody tr").filter({ hasText: modelName }); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + await expect(row.getByText(E2E_TEAM_CRUD_ID)).toBeVisible({ timeout: 10_000 }); + + await row.getByRole("button", { name: "Delete model" }).click(); + + const modal = page.getByRole("dialog", { name: "Delete Model" }); + await expect(modal).toBeVisible({ timeout: 5_000 }); + await expect(modal.getByText(modelName).first()).toBeVisible(); + await modal.getByRole("button", { name: "Delete", exact: true }).click(); + + await expect(page.getByText("Model deleted successfully").first()).toBeVisible({ timeout: 10_000 }); + await expect(row).toHaveCount(0, { timeout: 15_000 }); + + await expect + .poll(async () => await findDeploymentByName(page, modelName), { + message: `deployment ${modelName} still readable from /v2/model/info after delete`, + timeout: 15_000, + }) + .toBeUndefined(); + + await page.reload(); + await page.getByPlaceholder("Search model names").fill(modelName); + await expect(page.getByText("No models found").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.locator("table tbody tr").filter({ hasText: modelName })).toHaveCount(0); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts new file mode 100644 index 00000000000..2dd30060d2d --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts @@ -0,0 +1,94 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +test.describe("Second proxy admin", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("an invited admin can log in, mint a key, and call a model with it", async ({ page, browser, request }) => { + const suffix = Date.now(); + const email = `second-admin-${suffix}@test.local`; + const password = "e2e-second-admin-password"; + const auth = { Authorization: `Bearer ${masterKey()}` }; + + const adminContext = await browser.newContext({ storageState: ADMIN_STORAGE_PATH }); + let userId = ""; + try { + const adminPage = await adminContext.newPage(); + await navigateToPage(adminPage, Page.Users); + await dismissFeedbackPopup(adminPage); + + await adminPage.getByRole("button", { name: "+ Invite User", exact: true }).click(); + const dialog = adminPage.getByRole("dialog", { name: "Invite User" }); + await expect(dialog).toBeVisible({ timeout: 5_000 }); + + await dialog.getByLabel("User Email").fill(email); + + await dialog.getByLabel(/Global Proxy Role/).click(); + await adminPage.getByRole("option", { name: /Admin \(All Permissions\)/ }).click(); + + const createdResponse = adminPage.waitForResponse( + (res) => res.url().includes("/user/new") && res.request().method() === "POST", + ); + await dialog.getByRole("button", { name: "Invite User" }).click(); + const createdBody = await (await createdResponse).json(); + userId = (createdBody.data?.user_id ?? createdBody.user_id) as string; + expect(userId, "created user id from /user/new").toBeTruthy(); + + await expect(adminPage.getByText("API user Created").first()).toBeVisible({ timeout: 10_000 }); + } finally { + await adminContext.close(); + } + + try { + const passwordRes = await request.post("/user/update", { + headers: auth, + data: { user_email: email, password }, + }); + expect(passwordRes.ok(), `setting password failed (${passwordRes.status()}): ${await passwordRes.text()}`).toBe( + true, + ); + + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page.getByPlaceholder("Enter your password").fill(password); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + + await navigateToPage(page, Page.ApiKeys); + await page.getByRole("button", { name: /Create New Key/i }).click(); + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + await page.getByLabel(/Key Name/).fill(`e2e-second-admin-key-${suffix}`); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Proxy Models", exact: true }).click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + await page.keyboard.press("Escape"); + + const response = await page.request.post("/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}` }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: `second admin ping ${suffix}` }], + }, + }); + expect(response.status()).toBe(200); + const body = await response.json(); + expect(body.choices?.[0]?.message?.content).toBe(MOCK_RESPONSE_TEXT); + } finally { + if (userId) { + await request.post("/user/delete", { headers: auth, data: { user_ids: [userId] } }); + } + } + }); +}); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index f93cca75347..c23fa2c994c 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -1,6 +1,7 @@ import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; import { E2E_INTERNAL_USER_KEY_ALIAS, + E2E_TEAM_ADMIN_USER_ID, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID, TEAM_ADMIN_STORAGE_PATH, @@ -8,6 +9,8 @@ import { import { Page } from "../../fixtures/pages"; import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; +import { keySourceSelect, modelSelect, onlyVisible, openPlayground } from "../../helpers/playground"; /** * Every identifier a roster is addressable by. Which of user_id / user_email is populated depends on @@ -128,6 +131,91 @@ test.describe("Team Admin", () => { .not.toContain("e2e-removable-member"); }); + test("Team admin sees all team models in the Playground model dropdown", async ({ page, request }) => { + const suffix = Date.now(); + const teamModelName = `e2e-team-dropdown-model-${suffix}`; + const auth = { Authorization: `Bearer ${masterKey()}` }; + + const teamRes = await request.post("/team/new", { + headers: auth, + data: { + team_alias: `e2e-playground-team-${suffix}`, + models: [CHAT_MODEL_A], + members_with_roles: [{ role: "admin", user_id: E2E_TEAM_ADMIN_USER_ID }], + }, + }); + expect(teamRes.ok(), `team create failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + const teamId = (await teamRes.json()).team_id as string; + + let modelId = ""; + let teamKey = ""; + try { + const modelRes = await request.post("/model/new", { + headers: auth, + data: { + model_name: teamModelName, + litellm_params: { + model: "openai/fake-gpt-4", + api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`, + api_key: "fake-key", + }, + model_info: { team_id: teamId }, + }, + }); + expect(modelRes.ok(), `model create failed (${modelRes.status()}): ${await modelRes.text()}`).toBe(true); + modelId = (await modelRes.json()).model_info?.id as string; + + const keyRes = await request.post("/key/generate", { headers: auth, data: { team_id: teamId } }); + expect(keyRes.ok(), `key generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); + teamKey = (await keyRes.json()).key as string; + + await expect + .poll( + async () => { + const res = await request.get("/model_group/info", { + headers: { Authorization: `Bearer ${teamKey}` }, + }); + if (!res.ok()) return false; + const body: { data?: { model_group?: string }[] } = await res.json(); + return (body.data ?? []).some((group) => group.model_group === teamModelName); + }, + { + message: `model group ${teamModelName} never became visible to the team key`, + timeout: 30_000, + }, + ) + .toBe(true); + + await openPlayground(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + + const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(teamKey); + + const select = modelSelect(page); + await select.click(); + await select.fill(teamModelName); + await expect(onlyVisible(page.getByRole("option", { name: teamModelName }))).toBeVisible({ + timeout: 15_000, + }); + + await select.fill(CHAT_MODEL_A); + await expect(onlyVisible(page.getByRole("option", { name: CHAT_MODEL_A }))).toBeVisible({ + timeout: 15_000, + }); + } finally { + if (teamKey) { + await request.post("/key/delete", { headers: auth, data: { keys: [teamKey] } }); + } + if (modelId) { + await request.post("/model/delete", { headers: auth, data: { id: modelId } }); + } + await request.post("/team/delete", { headers: auth, data: { team_ids: [teamId] } }); + } + }); + test("Team admin can create a team key with All Team Models", async ({ page }) => { await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); From abd8beec018eb8faf4e40530f3063b9ed6cb15fb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:28:34 -0700 Subject: [PATCH 047/126] fix(e2e): measure the clipped popup and assert the table's empty state Two assertions were checking the wrong thing. The anchoring tests read getByRole("listbox"), which resolves to SelectPrimitive.List; that sits at full content height inside the popup that clips and scrolls it, so the box overlapped the trigger even when nothing visible did. Measure the popup. The SSO-ID search expected zero rows, but DataTable renders a "No results" message row when a filter matches nothing, so the count is one. Assert the empty state the user actually sees. --- .../modelsPage/autoRouterTemplateSelect.spec.ts | 12 ++++++------ tests/e2e/ui/tests/users/searchUsers.spec.ts | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts index 7fc20104f20..51df50a2e68 100644 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts @@ -23,6 +23,8 @@ async function boxes(trigger: Locator, options: Locator) { return triggerBox && optionsBox ? { triggerBox, optionsBox } : null; } +const clippedPopup = (page: PlaywrightPage) => page.locator('[data-slot="select-content"]'); + function pollOptionsOpenBelowTrigger(trigger: Locator, options: Locator) { return expect.poll(async () => { const box = await boxes(trigger, options); @@ -50,10 +52,9 @@ test.describe("Auto Router template select anchoring", () => { await trigger.scrollIntoViewIfNeeded(); await trigger.click(); - const options = page.getByRole("listbox"); - await expect(options).toBeVisible(); + await expect(page.getByRole("listbox")).toBeVisible(); - await pollOptionsOpenBelowTrigger(trigger, options).toBe(true); + await pollOptionsOpenBelowTrigger(trigger, clippedPopup(page)).toBe(true); }); test("keeps the trigger uncovered when the options open with no room below it", async ({ page }) => { @@ -62,9 +63,8 @@ test.describe("Auto Router template select anchoring", () => { await trigger.scrollIntoViewIfNeeded(); await trigger.click(); - const options = page.getByRole("listbox"); - await expect(options).toBeVisible(); + await expect(page.getByRole("listbox")).toBeVisible(); - await pollOptionsCoverTrigger(trigger, options).toBe(false); + await pollOptionsCoverTrigger(trigger, clippedPopup(page)).toBe(false); }); }); diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index 5e7e3e35b91..ee1a3f18f69 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -46,6 +46,6 @@ test.describe("Internal Users Search", () => { await page.getByTestId("users-filter-sso-id").fill("e2e-sso-id-that-matches-nobody"); await page.getByTestId("filter-drawer-apply").click(); - await expect(userRows(page)).toHaveCount(0, { timeout: 30_000 }); + await expect(page.getByRole("row").filter({ hasText: "No results" })).toBeVisible({ timeout: 30_000 }); }); }); From bab347a28e651c7de780145dd99763124b8c6d1f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 15:37:56 -0700 Subject: [PATCH 048/126] test(gcs_pubsub): expect router_metadata key in spend logs fixture --- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 1838fb16e91..fc73aa554d4 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"router_metadata\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, From 5eff708d0fd53fc474627b674ffe143f494d3aa1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 15:49:03 -0700 Subject: [PATCH 049/126] fix(proxy): keep persisted spend another pod has not incremented in the window seed The batch start told the seed which LiteLLM_SpendLogs rows were its own, but using it as a hard cutoff also dropped rows another pod had already persisted. Those rows are only repaid by that pod's own increment, so if it died first the window row stayed permanently under the recorded spend. The seed now reads both sums in one scan and takes off this batch's own spend, flooring at the pre-batch total for the case where its log rows have not landed yet. Redis payloads keep an empty request_ids so a leader from before the field was dropped can still merge what it pops during a rolling deploy. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- .../proxy/db/budget_window_spend_writer.py | 95 +++++++++------ .../redis_update_buffer.py | 10 +- .../window_spend_update_queue.py | 37 +++++- .../test_redis_update_buffer.py | 34 ++++++ .../db/test_budget_window_spend_writer.py | 108 ++++++++++-------- 5 files changed, 197 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 8b1ad0e24c7..8cf2f737063 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -7,17 +7,15 @@ instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold (issue #35766). Raw SQL rather than the Prisma upsert helper because the conditional roll cannot be expressed through the query builder. -Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, summing -only rows that started before the batch being flushed so neither source counts -the same request twice. Anything at or after that cutoff is owed by an -increment that still reaches the row, on this pod's next flush or another -pod's, so a row lags real spend by at most one flush interval of queued -increments: the same lag the SpendLogs aggregate it replaces (and every other -spend column) already has. A request whose increment is lost before it flushes, -which today means the pod dying, is missed by both sources and stays missing. +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once and takes +off what the increments being flushed will add, so neither source counts the +same request twice. A row therefore lags real spend by at most one flush +interval of increments queued elsewhere: the same lag the SpendLogs aggregate +it replaces (and every other spend column) already has. """ from collections.abc import Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Final, Protocol @@ -64,32 +62,45 @@ _ROLL_WINDOW_SPEND_SQL: Final = ( ) _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' - "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" + "SELECT COALESCE(SUM(spend), 0.0) AS total, " + "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' - "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" + "SELECT COALESCE(SUM(spend), 0.0) AS total, " + "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) +@dataclass(frozen=True, slots=True) +class WindowSeedTotals: + """The two sums a seed needs: everything persisted for the window, and the + part of it that predates the batch being flushed.""" + + total: float + before_batch: float + + class WindowSpendLogsAggregate(Protocol): - """Sums LiteLLM_SpendLogs for one entity between window_start and the + """Sums LiteLLM_SpendLogs for one entity since window_start, split at the batch's earliest request. Injected so the flush can be exercised without a database and so the @@ -103,18 +114,18 @@ class WindowSpendLogsAggregate(Protocol): entity_id: str, window_start: datetime, batch_started_at: datetime | None, - ) -> float | None: ... + ) -> WindowSeedTotals | None: ... -async def spend_logs_total_before_batch( +async def spend_logs_seed_totals( prisma_client: "PrismaClient", entity_type: str, entity_id: str, window_start: datetime, batch_started_at: datetime | None, -) -> float | None: - """LiteLLM_SpendLogs spend for one entity since window_start, stopping - before the requests the increments being flushed already cover. +) -> WindowSeedTotals | None: + """LiteLLM_SpendLogs spend for one entity since window_start, both in full + and up to the start of the batch being flushed, in one scan. The spend log writer drains its own queue on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so @@ -122,11 +133,12 @@ async def spend_logs_total_before_batch( already in the table. Counting them in the seed and again in the increment is what made a fresh row land at twice the true spend. - Every log row at or after the cutoff belongs to a request whose own - increment still reaches this row, on this pod's next flush or another pod's, - so bounding the sum by time needs nothing from the request itself. Without a - known start the whole window is summed: that can only over-count once, which - enforcement tolerates, whereas under-counting is a budget bypass. + Both halves are needed because neither is safe alone: the full sum + double-counts this batch, and the sum before the batch drops spend another + pod has already persisted but not yet incremented. _seed_base picks between + them. Without a known batch start the two are the same sum, so the seed + counts everything: that can only over-count once, which enforcement + tolerates, whereas under-counting is a budget bypass. """ if entity_type == Litellm_EntityType.KEY.value: bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL @@ -145,8 +157,11 @@ async def spend_logs_total_before_batch( ) ) if not rows: - return 0.0 - return float(rows[0].get("total") or 0.0) + return WindowSeedTotals(total=0.0, before_batch=0.0) + return WindowSeedTotals( + total=float(rows[0].get("total") or 0.0), + before_batch=float(rows[0].get("before_batch") or 0.0), + ) def _exclusion_upper_bound(started_at: datetime) -> datetime: @@ -186,19 +201,33 @@ async def _seed_base_for_missing_row( This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on every cold counter today, but here it runs once per window lifetime and off - the request path, and it stops before the queued increments so they are + the request path, and it discounts the queued increments so they are counted once. """ if _primary_key(transaction) in existing_primary_keys: return 0.0 - base: Final = await spend_logs_aggregate( + totals: Final = await spend_logs_aggregate( prisma_client=prisma_client, entity_type=transaction["entity_type"], entity_id=transaction["entity_id"], window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), batch_started_at=_transaction_started_at(transaction), ) - return float(base or 0.0) + if totals is None: + return 0.0 + return _seed_base(totals=totals, batch_spend=transaction["spend"]) + + +def _seed_base(totals: WindowSeedTotals, batch_spend: float) -> float: + """What the window already held before the increments about to be applied. + + Subtracting the batch's own spend from the full sum keeps every other + request in the seed, including the ones another pod persisted and has not + incremented yet, which a plain cutoff would drop for good if that pod died. + When this batch's own log rows have not landed yet the subtraction takes + spend that was never counted, so the sum before the batch is the floor. + """ + return max(totals.total - batch_spend, totals.before_batch) def _transaction_started_at(transaction: WindowSpendTransaction) -> datetime | None: @@ -232,7 +261,7 @@ def _upsert_params( async def commit_window_spend_updates( prisma_client: "PrismaClient", transactions: Sequence[WindowSpendTransaction], - spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_before_batch, + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_seed_totals, ) -> None: """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 4dd23270bf8..c06f2e04aca 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -46,6 +46,7 @@ from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdate from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, + to_wire_payload, ) from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( @@ -298,7 +299,7 @@ class RedisUpdateBuffer: ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), ( - window_spend_update_transactions, + tuple(map(to_wire_payload, window_spend_update_transactions)), REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, ), @@ -484,7 +485,12 @@ class RedisUpdateBuffer: (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), - (window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY), + ( + None + if window_spend_update_transactions is None + else tuple(map(to_wire_payload, window_spend_update_transactions)), + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ), ) rpush_list: Final = tuple( diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py index 43b069fd8c2..372a6666c02 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -27,10 +27,11 @@ class WindowSpendTransaction(TypedDict): transaction survives the JSON round trip through the Redis buffer. started_at is the earliest request start in the batch. The one-time seed for - a window that has no row yet sums only LiteLLM_SpendLogs rows that started - before it, because the spend log writer flushes on its own ~2s poll and will - usually have persisted this batch's rows before the window queue flushes; - without the bound the seed and the increment would each count them. + a window that has no row yet uses it to tell this batch's own + LiteLLM_SpendLogs rows from everything else, because the spend log writer + flushes on its own ~2s poll and will usually have persisted this batch's + rows before the window queue flushes; without that split the seed and the + increment would each count them. """ entity_type: ReadOnly[str] @@ -41,6 +42,34 @@ class WindowSpendTransaction(TypedDict): started_at: ReadOnly[str | None] +class WindowSpendWirePayload(WindowSpendTransaction): + """How an increment is encoded in the shared Redis buffer. + + request_ids is dead weight here: workers built before this field was + dropped index it while merging whatever they pop, and the pop is + destructive, so a leader still running one of those during a rolling deploy + would raise on a payload without the key and lose those increments. It is + always empty, which only makes such a leader seed without exclusions. + + TODO: remove once no supported version reads it, i.e. one release after the + field stopped being written. + """ + + request_ids: ReadOnly[Sequence[str]] + + +def to_wire_payload(transaction: WindowSpendTransaction) -> WindowSpendWirePayload: + return WindowSpendWirePayload( + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_duration=transaction["window_duration"], + window_start=transaction["window_start"], + spend=transaction["spend"], + started_at=transaction.get("started_at"), + request_ids=(), + ) + + def to_naive_utc(value: datetime) -> datetime: """LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC.""" if value.tzinfo is None: diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 99ac1fe8b50..8f3508fc4e9 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -523,10 +523,44 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up "window_start": "2026-08-01T00:00:00.000000", "spend": 1.25, "started_at": "2026-08-10T12:00:00.000000", + "request_ids": [], } ] +@pytest.mark.asyncio +async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_update_buffer, mock_redis_cache): + """A leader from before the field was dropped indexes request_ids while + merging what it popped, and the pop is destructive, so a payload without + the key would cost a rolling deploy those increments.""" + from datetime import datetime, timezone + + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=1.25, + ) + ) + + await redis_update_buffer.restore_transactions_to_redis( + window_spend_update_transactions=await window_queue.flush_and_get_aggregated_window_spend_transactions(), + ) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + restored = json.loads(rpush_list[0]["values"][0]) + assert [payload["request_ids"] for payload in restored] == [[]] + + @pytest.mark.asyncio async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure( redis_update_buffer, mock_redis_cache diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py index 74b71f63401..130f0c56ccf 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -6,9 +6,10 @@ from typing import Any import pytest from litellm.proxy.db.budget_window_spend_writer import ( + WindowSeedTotals, commit_window_spend_updates, roll_window_spend_row, - spend_logs_total_before_batch, + spend_logs_seed_totals, ) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, @@ -70,10 +71,15 @@ class _FakePrismaClient: class _RecordingAggregate: - """Stands in for the LiteLLM_SpendLogs seed aggregate.""" + """Stands in for the LiteLLM_SpendLogs seed aggregate. before_batch + defaults to the full total, the state where none of this batch's own log + rows have been persisted yet.""" - def __init__(self, value: float = 5.0) -> None: - self.value = value + def __init__(self, total: float = 5.0, before_batch: float | None = None) -> None: + self.totals = WindowSeedTotals( + total=total, + before_batch=total if before_batch is None else before_batch, + ) self.calls: list[dict[str, Any]] = [] async def __call__( @@ -83,7 +89,7 @@ class _RecordingAggregate: entity_id: str, window_start: datetime, batch_started_at: datetime | None, - ) -> float | None: + ) -> WindowSeedTotals | None: self.calls.append( { "entity_type": entity_type, @@ -92,13 +98,13 @@ class _RecordingAggregate: "batch_started_at": batch_started_at, } ) - return self.value + return self.totals class _SpendLogsFake: """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, - honouring the cutoff exactly as the real aggregate's - startTime < bound does.""" + splitting them at the batch start exactly as the real aggregate's + SUM(...) FILTER (WHERE startTime < bound) does.""" def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: self.rows = rows @@ -110,11 +116,14 @@ class _SpendLogsFake: entity_id: str, window_start: datetime, batch_started_at: datetime | None, - ) -> float | None: - return math.fsum( - spend - for _request_id, spend, started_at in self.rows - if batch_started_at is None or started_at < batch_started_at + ) -> WindowSeedTotals | None: + return WindowSeedTotals( + total=math.fsum(spend for _request_id, spend, _started_at in self.rows), + before_batch=math.fsum( + spend + for _request_id, spend, started_at in self.rows + if batch_started_at is None or started_at < batch_started_at + ), ) @@ -151,7 +160,7 @@ async def test_missing_row_is_seeded_from_spend_logs_once(): existed, so a brand new primary key inserts the SpendLogs total plus this increment.""" db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=5.0) + aggregate = _RecordingAggregate(total=5.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -178,7 +187,7 @@ async def test_existing_row_is_never_reseeded(): """The seed is a full LiteLLM_SpendLogs scan; running it for a row that is already maintained would both cost a scan and double count.""" db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) - aggregate = _RecordingAggregate(value=5.0) + aggregate = _RecordingAggregate(total=5.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -195,7 +204,7 @@ async def test_existing_row_is_never_reseeded(): @pytest.mark.asyncio async def test_seed_runs_only_for_the_primary_keys_that_are_missing(): db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) - aggregate = _RecordingAggregate(value=5.0) + aggregate = _RecordingAggregate(total=5.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -218,7 +227,7 @@ async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): """The conflict arm adds the increment alone so two pods that both seed the same new window cannot add the SpendLogs base twice.""" db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=9.0) + aggregate = _RecordingAggregate(total=9.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -255,7 +264,7 @@ async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one @pytest.mark.asyncio async def test_upsert_never_interpolates_values_into_the_sql(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -273,7 +282,7 @@ async def test_upserts_are_ordered_by_primary_key_then_window_start(): """Cross-pod lock ordering, plus an older window must be applied before the roll that supersedes it or the roll would be undone.""" db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -301,7 +310,7 @@ async def test_upserts_are_ordered_by_primary_key_then_window_start(): @pytest.mark.asyncio async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -339,9 +348,7 @@ async def test_unknown_entity_type_contributes_no_seed(): anything else starts from its increment alone.""" db = _FakeDB(existing_rows=[]) - async def no_such_column( - prisma_client, entity_type, entity_id, window_start, batch_started_at - ): + async def no_such_column(prisma_client, entity_type, entity_id, window_start, batch_started_at): return None await commit_window_spend_updates( @@ -396,7 +403,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -410,7 +417,7 @@ async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): @pytest.mark.asyncio async def test_seed_passes_no_start_bound_when_the_batch_has_none(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -463,15 +470,19 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): @pytest.mark.asyncio -async def test_seed_skips_logs_from_requests_this_batch_never_saw(): - """A concurrent request on another pod can land its spend log before this - pod seeds the row. Its increment is still queued over there, so the cutoff - has to drop it from the seed even though this batch has no way to know its - id; counting it here and again on that pod's flush is the double count the - old id list could not catch.""" +async def test_seed_keeps_spend_another_pod_persisted_after_this_batch_started(): + """A concurrent request on another pod can land its spend log after this + batch started but before this pod seeds the row. Dropping it on a plain + time cutoff would lose that spend for the rest of the window if that pod + died before flushing its increment, so the seed takes off only this batch's + own spend and keeps everything else.""" db = _FakeDB(existing_rows=[]) spend_logs = _SpendLogsFake( - rows=(("older", 0.5, BEFORE_BATCH), ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1))), + rows=( + ("older", 0.5, BEFORE_BATCH), + ("mine", 0.000047, BATCH_STARTED_AT), + ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1)), + ), ) await commit_window_spend_updates( @@ -481,7 +492,7 @@ async def test_seed_skips_logs_from_requests_this_batch_never_saw(): ) ((_, params),) = db.batcher.calls - assert params[INSERT_SPEND] == pytest.approx(0.500047) + assert params[INSERT_SPEND] == pytest.approx(0.750047) @pytest.mark.asyncio @@ -506,10 +517,10 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): "entity_type, expected_column", [("key", "api_key = $1"), ("team", "team_id = $1")], ) -async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected_column): - db = _FakeDB(existing_rows=[{"total": 1.25}]) +async def test_seed_aggregate_sql_splits_the_window_at_the_batch_start(entity_type, expected_column): + db = _FakeDB(existing_rows=[{"total": 1.25, "before_batch": 0.75}]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type=entity_type, entity_id="e1", @@ -517,11 +528,11 @@ async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected batch_started_at=BATCH_STARTED_AT, ) - assert total == pytest.approx(1.25) + assert totals == WindowSeedTotals(total=1.25, before_batch=0.75) ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" in normalized + assert "FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" in normalized assert 'FROM "LiteLLM_SpendLogs"' in normalized # startTime is TIMESTAMP(3): the bound is floored to the second so the # batch's own earliest row cannot round under it. @@ -532,12 +543,13 @@ async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected @pytest.mark.asyncio async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): - """A batch with no known start cannot place the cutoff, so the seed counts - everything; at worst that over-counts one batch, which enforcement - tolerates, where under-counting is a budget bypass.""" - db = _FakeDB(existing_rows=[{"total": 1.25}]) + """A batch with no known start cannot place the split, so both halves are + the same sum and the seed counts everything; at worst that over-counts one + batch, which enforcement tolerates, where under-counting is a budget + bypass.""" + db = _FakeDB(existing_rows=[{"total": 1.25, "before_batch": 1.25}]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="e1", @@ -545,7 +557,7 @@ async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): batch_started_at=None, ) - assert total == pytest.approx(1.25) + assert totals == WindowSeedTotals(total=1.25, before_batch=1.25) ((query, params),) = db.query_raw_calls assert '"startTime" <' not in query assert params == ("e1", WINDOW_A) @@ -555,7 +567,7 @@ async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type="user", entity_id="u1", @@ -563,7 +575,7 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs batch_started_at=None, ) - assert total is None + assert totals is None assert db.query_raw_calls == [] @@ -571,7 +583,7 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="k-unknown", @@ -579,4 +591,4 @@ async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): batch_started_at=None, ) - assert total == 0.0 + assert totals == WindowSeedTotals(total=0.0, before_batch=0.0) From b5c156a7d5b0ec911c34519a0846890ce6489566 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 15:49:49 -0700 Subject: [PATCH 050/126] style(proxy): drop the explicit return None update_database no longer needs Reverting the function to -> None left two bare `return None` statements that RET501 rejects now that None is the only value it can return. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- litellm/proxy/db/db_spend_update_writer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b3fd2c3f22c..202a95ba29b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -232,7 +232,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return None + return if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -318,7 +318,7 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) - return None + return async def _enqueue_tool_usage_transaction( self, From 8ab132b8beba9a656321c090c4873f06e385dfdf Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:54:20 -0700 Subject: [PATCH 051/126] fix(key_management): allow non-admin key_type preset transitions on /key/update A non-admin switching an existing key's type between the safe preset buckets (llm_api_routes, info_routes, and empty = full access) got a 403 from the allowed_routes admin gate, because /key/update, unlike /key/generate and /key/regenerate, had no carve-out for preset-derived values. Skip the gate only when both the incoming and the stored allowed_routes consist entirely of safe presets, so clearing an admin-set custom route restriction still requires proxy admin. --- .../key_management_endpoints.py | 30 +++++-- .../test_key_management_endpoints.py | 88 +++++++++++++++++-- 2 files changed, 106 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 52a192f8537..3c3a135ef72 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -733,6 +733,22 @@ def _check_allowed_routes_caller_permission( ) +def _is_safe_preset_route_transition( + incoming_allowed_routes: list | None, + existing_allowed_routes: list | None, +) -> bool: + """ + True when every route on BOTH sides is a safe `key_type` preset bucket + (empty = full access, which non-admins already get from a default + `/key/generate`). Requiring the existing side too keeps an owner from + clearing an admin-set custom route restriction (LIT-4139). + """ + return all( + route in _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS + for route in (*(incoming_allowed_routes or []), *(existing_allowed_routes or [])) + ) + + def _check_permissions_caller_permission( data: GenerateRequestBase, user_api_key_dict: UserAPIKeyAuth, @@ -2533,11 +2549,15 @@ async def _validate_update_key_data( _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - _check_allowed_routes_caller_permission( - allowed_routes=data.allowed_routes, - user_api_key_dict=user_api_key_dict, - allowed_routes_was_provided="allowed_routes" in data.model_fields_set, - ) + if not _is_safe_preset_route_transition( + incoming_allowed_routes=data.allowed_routes, + existing_allowed_routes=existing_key_row.allowed_routes, + ): + _check_allowed_routes_caller_permission( + allowed_routes=data.allowed_routes, + user_api_key_dict=user_api_key_dict, + allowed_routes_was_provided="allowed_routes" in data.model_fields_set, + ) _check_passthrough_routes_caller_permission( data=data, user_api_key_dict=user_api_key_dict, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 42e56ceabd3..0349e43da76 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -11033,6 +11033,79 @@ class TestLIT1884KeyUpdateValidation: ) +class TestLIT4891SafePresetKeyTypeTransition: + def _make_existing_key(self, allowed_routes): + row = MagicMock() + row.user_id = "internal-user-123" + row.created_by = "internal-user-123" + row.token = "hashed_token" + row.team_id = None + row.max_budget = None + row.spend = 0.0 + row.organization_id = None + row.project_id = None + row.allowed_routes = allowed_routes + return row + + def _make_auth(self): + return UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + async def _run_update(self, data, existing_key_row): + await _validate_update_key_data( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=self._make_auth(), + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + + @pytest.mark.asyncio + async def test_non_admin_owner_can_clear_safe_preset_to_full_access(self): + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + + @pytest.mark.asyncio + async def test_non_admin_owner_can_switch_full_access_to_safe_preset(self): + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["llm_api_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=[]), + ) + + @pytest.mark.asyncio + async def test_non_admin_owner_can_switch_between_safe_presets(self): + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["info_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_clear_custom_route_restriction(self): + with pytest.raises(HTTPException) as exc_info: + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["/chat/completions"]), + ) + assert exc_info.value.status_code == 403 + assert "Only proxy admins can set" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_non_admin_cannot_set_non_preset_routes(self): + with pytest.raises(HTTPException) as exc_info: + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["management_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + assert exc_info.value.status_code == 403 + assert "Only proxy admins can set" in str(exc_info.value.detail) + + class TestKeyOwnerPrivilegeEscalation: """ Policy: @@ -12007,9 +12080,10 @@ class TestAllowedRoutesCallerPermission: @pytest.mark.asyncio async def test_non_admin_update_key_explicit_empty_allowed_routes_rejected(self): - """`update_key_fn` rejects a non-admin when `allowed_routes` is - present as `[]` in the request body. The value matches the model - default but `model_fields_set` distinguishes the two.""" + """`update_key_fn` rejects a non-admin clearing a custom (non-preset) + route restriction with an explicit `[]` in the request body. The value + matches the model default but `model_fields_set` distinguishes the + two. Clearing from a safe preset is allowed (LIT-4891).""" from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, ) @@ -12032,7 +12106,7 @@ class TestAllowedRoutesCallerPermission: patch( "litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key", new_callable=AsyncMock, - return_value=MagicMock(), + return_value=MagicMock(allowed_routes=["/chat/completions"]), ), ): with pytest.raises(ProxyException) as exc_info: @@ -12047,8 +12121,8 @@ class TestAllowedRoutesCallerPermission: @pytest.mark.asyncio async def test_non_admin_update_key_explicit_null_allowed_routes_rejected(self): - """`update_key_fn` rejects a non-admin when `allowed_routes` is - present as `null` in the request body.""" + """`update_key_fn` rejects a non-admin clearing a custom (non-preset) + route restriction with an explicit `null` in the request body.""" from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, ) @@ -12071,7 +12145,7 @@ class TestAllowedRoutesCallerPermission: patch( "litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key", new_callable=AsyncMock, - return_value=MagicMock(), + return_value=MagicMock(allowed_routes=["/chat/completions"]), ), ): with pytest.raises(ProxyException) as exc_info: From 859bd01ddab619367671a39ae39d46aa3d349d3c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:57:49 -0700 Subject: [PATCH 052/126] fix(e2e): assert the users table's own empty-state copy UsersTable overrides DataTable's default noDataMessage with its own EmptyState, so the row reads "No users found" rather than "No results". Assert that, and pair it with the seeded user being absent so the check cannot pass while the filter silently does nothing. --- tests/e2e/ui/tests/users/searchUsers.spec.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index ee1a3f18f69..fa8f32764e8 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -46,6 +46,7 @@ test.describe("Internal Users Search", () => { await page.getByTestId("users-filter-sso-id").fill("e2e-sso-id-that-matches-nobody"); await page.getByTestId("filter-drawer-apply").click(); - await expect(page.getByRole("row").filter({ hasText: "No results" })).toBeVisible({ timeout: 30_000 }); + await expect(page.getByText("No users found")).toBeVisible({ timeout: 30_000 }); + await expect(userRows(page).filter({ hasText: "noteam@test.local" })).toHaveCount(0); }); }); From 92b46538e1f83455a038f5ac047e6ca522db8479 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:38:44 -0700 Subject: [PATCH 053/126] fix(policy_engine): restore request guardrails list after pipeline allow --- .../proxy/policy_engine/pipeline_executor.py | 52 +++++++-- .../policy_engine/test_pipeline_executor.py | 100 ++++++++++++++++++ 2 files changed, 142 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 50cb813c6fa..4be0f556ed7 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,6 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time +from collections.abc import Sequence from typing import Any, Final, Literal import litellm @@ -114,11 +115,7 @@ class PipelineExecutor: # Handle terminal actions if action == "allow": - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": return PipelineExecutionResult( @@ -138,11 +135,7 @@ class PipelineExecutor: # action == "next" → continue to next step # Ran out of steps without a terminal action → default allow - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) @staticmethod async def _run_step( @@ -251,6 +244,45 @@ class PipelineExecutor: return None +def _allow_result( + step_results: Sequence[PipelineStepResult], + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> PipelineExecutionResult: + """Build the terminal-allow result, propagating pipeline modifications without the per-step guardrail override.""" + restored: Final = _restore_request_guardrails(working_data, request_data) + return PipelineExecutionResult( + terminal_action="allow", + step_results=list(step_results), # mutable-ok: PipelineExecutionResult field is a list + modified_data=restored if restored != request_data else None, + ) + + +def _restore_request_guardrails( + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> dict: # mutable-ok: merged back into the request dict, which downstream code mutates + """ + Restore the request's own metadata["guardrails"] activation list. + + _run_step overrides it to [step.guardrail] so should_run_guardrail() allows each + step; letting that override escape via modified_data permanently drops every + independently activated guardrail from later lifecycle stages (post_call, etc.). + """ + working_metadata: Final = working_data.get("metadata") + if not isinstance(working_metadata, dict): + return working_data + request_metadata: Final = request_data.get("metadata") + original_guardrails: Final = request_metadata.get("guardrails") if isinstance(request_metadata, dict) else None + stripped: Final = {k: v for k, v in working_metadata.items() if k != "guardrails"} # mutable-ok: request dict + if original_guardrails is not None: + restored: Final = {**stripped, "guardrails": original_guardrails} # mutable-ok: request dict + return {**working_data, "metadata": restored} # mutable-ok: request dict + if not stripped and not isinstance(request_metadata, dict): + return {k: v for k, v in working_data.items() if k != "metadata"} # mutable-ok: request dict + return {**working_data, "metadata": stripped} # mutable-ok: request dict + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 054a5af4148..4fcb7d22588 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -749,6 +749,106 @@ async def test_single_step_pipeline_allow(monkeypatch): assert guard.calls == 1 +@pytest.mark.asyncio +async def test_allow_restores_independent_guardrails_list(monkeypatch): + """ + Request activates an independent guardrail; an unrelated pipeline runs and allows. + Expected: no modified_data escapes, so the request's guardrails list survives + and the independent guardrail still runs at later lifecycle stages (post_call). + Regression: LIT-6587 (pipeline clobbered the list with its last step's guardrail). + """ + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = { + "messages": [{"role": "user", "content": "clean content"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert pipeline_guard.calls == 1 + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert propagated["metadata"]["guardrails"] == ["independent-output-guard"] + assert data["metadata"]["guardrails"] == ["independent-output-guard"] + + +@pytest.mark.asyncio +async def test_allow_does_not_leak_guardrails_into_bare_request(monkeypatch): + """A request without metadata must not gain a metadata.guardrails list from the pipeline.""" + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = {"messages": [{"role": "user", "content": "clean content"}]} + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert "guardrails" not in propagated.get("metadata", {}) + assert "metadata" not in data + + +@pytest.mark.asyncio +async def test_data_forwarding_keeps_changes_and_restores_guardrails_list(monkeypatch): + """A pass_data pipeline's modifications propagate while the request's guardrails list is restored.""" + pii_guard = PiiMaskingGuardrail(guardrail_name="pii-masker") + content_guard = ContentCheckGuardrail(guardrail_name="content-check") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="pii-masker", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="content-check", on_fail="block", on_pass="allow"), + ], + ) + + monkeypatch.setattr(litellm, "callbacks", [pii_guard, content_guard]) + + data = { + "messages": [{"role": "user", "content": "Hello John Smith"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="pii-then-safety", + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]" + assert result.modified_data["metadata"]["guardrails"] == ["independent-output-guard"] + + @pytest.mark.asyncio async def test_step_results_include_duration(monkeypatch): """Step results should include timing information.""" From e60111438395c2ab92774417c1fc53f3610a91c1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:15:54 -0700 Subject: [PATCH 054/126] fix(bedrock): mask signed request headers in guardrail debug log --- .../guardrail_hooks/bedrock_guardrails.py | 6 ++- .../test_bedrock_guardrails.py | 49 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index a6635ea0776..84d3bd6071d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -32,6 +32,9 @@ 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.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 +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -1172,11 +1175,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): aws_region_name=aws_region_name, api_key=api_key, ) + headers_dict: Final = dict(prepared_request.headers) # mutable-ok: the masking helper requires a dict verbose_proxy_logger.debug( "Bedrock AI request body: %s, url %s, headers: %s", bedrock_request_data, prepared_request.url, - prepared_request.headers, + _get_masked_values(headers_dict), ) httpx_response: Final = await self._sign_and_post( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 235a5c0c09b..bcda1b8b61d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5590,3 +5590,52 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca assert payload["error"]["message"] == "Violated guardrail policy" assert payload["error"]["code"] == "400" assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" + + +@pytest.mark.asyncio +async def test_apply_guardrail_debug_log_masks_signed_request_headers(): + import logging + + from litellm._logging import verbose_proxy_logger + + session_token = "FakeSessionTokenValueThatMustNeverAppearInLogs1234567890" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_access_key_id="ASIAFAKEACCESSKEYID1", + aws_secret_access_key="fakeSecretAccessKeyForSigning", + aws_session_token=session_token, + aws_region_name="us-east-1", + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"action": "NONE", "outputs": []} + + captured_records: list[logging.LogRecord] = [] + + class _RecordingHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + captured_records.append(record) + + handler = _RecordingHandler(level=logging.DEBUG) + previous_level = verbose_proxy_logger.level + verbose_proxy_logger.addHandler(handler) + verbose_proxy_logger.setLevel(logging.DEBUG) + try: + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_response + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={}, + ) + finally: + verbose_proxy_logger.removeHandler(handler) + verbose_proxy_logger.setLevel(previous_level) + + rendered_messages = [record.getMessage() for record in captured_records] + header_lines = [message for message in rendered_messages if "headers:" in message] + assert header_lines, "expected the signed-request debug line to be logged" + assert any("X-Amz-Security-Token" in message for message in header_lines) + assert all(session_token not in message for message in rendered_messages) From 849269d52d9f3b9d566627a6756b3eb0c2f16672 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:40:18 -0700 Subject: [PATCH 055/126] build(rust): configure native extension profiles --- .github/actions/cache-cargo-build/action.yml | 23 ++++++++++---------- litellm-rust/Cargo.toml | 9 ++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 3 ++- pyproject.toml | 5 ++++- 4 files changed, 26 insertions(+), 14 deletions(-) diff --git a/.github/actions/cache-cargo-build/action.yml b/.github/actions/cache-cargo-build/action.yml index 36c6c790b84..0ccb58c012c 100644 --- a/.github/actions/cache-cargo-build/action.yml +++ b/.github/actions/cache-cargo-build/action.yml @@ -4,17 +4,16 @@ description: >- so only the first job on a given Cargo.lock compiles the bridge from scratch. litellm builds through maturin, which compiles litellm-rust/crates/python-bridge - in release mode before it can produce a wheel. `uv sync` therefore pays a full - build in every job that installs the workspace: measured at 2m40s per unit shard - on 2026-08-21, more than the whole unit tier spends running tests. Nothing caught - it, because the uv cache holds wheels uv downloads rather than wheels it builds, - and a path dependency whose source moves every commit could never hit that cache - anyway. Cargo rebuilds only what changed when its target directory survives, so a - warm job pays for the bridge crate alone. + in the dev profile for editable installs. `uv sync` therefore pays a full build + in every job that installs the workspace. Nothing caught it, because the uv cache + holds wheels uv downloads rather than wheels it builds, and a path dependency + whose source moves every commit could never hit that cache anyway. Cargo rebuilds + only what changed when its target directory survives, so a warm job pays for the + bridge crate alone. - The key namespace is separate from test-rust.yml's. Both cache the same directory, - but that workflow fills it with debug and clippy artifacts, which a release build - cannot reuse, and a shared key would let whichever ran first deny the other a save. + The key namespace is separate from test-rust.yml's check and release caches. They + cache the same directory for different workloads, and a shared key would let + whichever ran first deny the others a save. runs: using: composite @@ -26,6 +25,6 @@ runs: ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-cargo-dev-${{ hashFiles('litellm-rust/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-cargo-release- + ${{ runner.os }}-cargo-dev- diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 481ea3f8f66..c17a0605fc7 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -30,3 +30,12 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" + +[profile.release] +opt-level = 3 +lto = "thin" +codegen-units = 1 +panic = "unwind" +debug = false +incremental = false +strip = "symbols" diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 0c4a753f762..d461a483ae0 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,8 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["extension-module"] +default = ["abi3"] +abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] [dependencies] diff --git a/pyproject.toml b/pyproject.toml index 34c1fec1c11..23e96e475f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -262,7 +262,7 @@ healthcheck = [ ] [build-system] -requires = ["maturin==1.9.4"] +requires = ["maturin==1.15.0"] build-backend = "maturin" [tool.maturin] @@ -270,6 +270,9 @@ manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml" module-name = "litellm.rust_bridge._native" python-source = "." bindings = "pyo3" +features = ["extension-module"] +profile = "release" +editable-profile = "dev" include = ["litellm/proxy/_experimental/out/**"] exclude = [ "litellm/proxy/enterprise", From 9b774d3dcff85654431f30c075ca1985321c6885 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:50:07 -0700 Subject: [PATCH 056/126] fix(ci): isolate editable Cargo cache namespace --- .github/actions/cache-cargo-build/action.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/actions/cache-cargo-build/action.yml b/.github/actions/cache-cargo-build/action.yml index 0ccb58c012c..c3b8ce22c68 100644 --- a/.github/actions/cache-cargo-build/action.yml +++ b/.github/actions/cache-cargo-build/action.yml @@ -25,6 +25,6 @@ runs: ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-dev-${{ hashFiles('litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-maturin-dev-${{ hashFiles('litellm-rust/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-cargo-dev- + ${{ runner.os }}-maturin-dev- From 76cfa6339b13c5437bda88d617367ecb7c2ffa22 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:28:12 -0700 Subject: [PATCH 057/126] test: give the mocked prepared request real headers for the masked debug log --- .../test_litellm/proxy/guardrails/test_guardrail_endpoints.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 9b2117b7647..9511732fd50 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -917,7 +917,9 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): "Content-Type": "application/json", "Authorization": "Bearer test-api-key-789", } - mock_request_instance.prepare.return_value = Mock() + mock_request_instance.prepare.return_value = Mock( + headers=mock_request_instance.headers + ) mock_aws_request.return_value = mock_request_instance await guardrail_hook.make_bedrock_api_request( From e34f43328c8f6e0bd0df10747f2453ac0433b684 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 31 Aug 2026 17:33:12 -0700 Subject: [PATCH 058/126] fix(proxy): ship psycopg so partitioned SpendLogs detection actually runs (#38994) ProxyExtrasDBManager.spend_logs_is_partitioned() (#38452) silently returns False when psycopg can't be imported, and psycopg was never added to the extra_proxy install, so every production image lacks it. Schema reconciliation then generates the unfiltered primary-key rewrite against a genuinely partitioned LiteLLM_SpendLogs and Postgres rejects it, exactly the failure the fix was meant to prevent. Ships psycopg via extra_proxy and logs a warning when it's still missing instead of failing silently. --- .../litellm_proxy_extras/utils.py | 7 ++++++ pyproject.toml | 5 +++++ .../test_litellm_proxy_extras_utils.py | 22 +++++++++++++++++++ uv.lock | 4 ++++ 4 files changed, 38 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b2dc0a52c8f..b8032dd0d28 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -512,6 +512,13 @@ class ProxyExtrasDBManager: try: import psycopg except ImportError: + logger.warning( + "psycopg is not installed; skipping the LiteLLM_SpendLogs " + "partition check. If this table is partitioned (see " + "db_scripts/partition_spend_logs.sql), schema reconciliation " + "will try to rewrite its primary key and fail. Install the " + "litellm[extra_proxy] extra, which now includes psycopg." + ) return False cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) diff --git a/pyproject.toml b/pyproject.toml index 34c1fec1c11..96dcf1121bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,11 @@ cli = [ ] extra_proxy = [ "prisma>=0.11.0,<1.0", + # Used by ProxyExtrasDBManager.spend_logs_is_partitioned() to detect a + # partitioned LiteLLM_SpendLogs and keep schema reconciliation from + # fighting its composite primary key. + "psycopg>=3.2,<4.0", + "psycopg-binary>=3.2,<4.0", "azure-identity>=1.25.2,<2.0", "azure-keyvault-secrets>=4.10.0,<5.0", # Not in PyPI proxy extra. diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 498d0cb4723..b3d457707b8 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -681,3 +681,25 @@ class TestSpendLogsPartitionDetectionSchemaScope: def test_only_partitioned_relations_match(self, monkeypatch): query, _ = self._detect(monkeypatch, "postgresql://u:p@localhost:5432/db") assert "pg_partitioned_table" in query + + +class TestSpendLogsPartitionDetectionMissingPsycopg: + """psycopg ships in the `extra_proxy` install, but a stripped-down image + can still lack it. When it does, detection must fail closed to False + (never crash the migration path) and say so loudly, because a silent + False here is what let a genuinely partitioned LiteLLM_SpendLogs hit the + unfiltered primary-key rewrite in production.""" + + def test_missing_psycopg_returns_false(self, monkeypatch): + monkeypatch.setitem(sys.modules, "psycopg", None) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") + assert ProxyExtrasDBManager.spend_logs_is_partitioned() is False + + def test_missing_psycopg_logs_a_warning(self, monkeypatch, caplog): + monkeypatch.setitem(sys.modules, "psycopg", None) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + ProxyExtrasDBManager.spend_logs_is_partitioned() + assert any( + "psycopg is not installed" in record.message for record in caplog.records + ) diff --git a/uv.lock b/uv.lock index 8ef72116466..7df1a35b28a 100644 --- a/uv.lock +++ b/uv.lock @@ -4306,6 +4306,8 @@ extra-proxy = [ { name = "google-cloud-iam" }, { name = "google-cloud-kms" }, { name = "prisma" }, + { name = "psycopg" }, + { name = "psycopg-binary" }, { name = "redisvl" }, { name = "resend" }, ] @@ -4544,6 +4546,8 @@ requires-dist = [ { name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" }, { name = "prisma", marker = "extra == 'extra-proxy'", specifier = ">=0.11.0,<1.0" }, { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, + { name = "psycopg", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, + { name = "psycopg-binary", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "pydantic", specifier = ">=2.10.0,<3.0.0" }, { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, From 33cc9c1c48d7d8551828d0137a97300833294582 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:55:53 -0700 Subject: [PATCH 059/126] fix(ui): keep litellm_credential_name from LiteLLM Params JSON when no credential is selected (#39005) * fix(ui): keep litellm_credential_name from LiteLLM Params JSON when no credential is selected Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): drop null litellm_credential_name from AddModelPanel payload fixture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../panels/AddModelPanel.integration.test.tsx | 1 - .../handle_add_model_submit.test.tsx | 24 +++++++++++++++++++ .../add_model/handle_add_model_submit.tsx | 5 +++- 3 files changed, 28 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx index 19e1e3aa8bd..efba26734ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx @@ -87,7 +87,6 @@ const alwaysMounted = { api_key: undefined, api_base: undefined, custom_llm_provider: "openai", - litellm_credential_name: null, model: "gpt-4o", }; diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx index 7ef09d34924..2aa8f93b72b 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx @@ -73,4 +73,28 @@ describe("prepareModelAddRequest", () => { expect(deployment.litellmParamsObj.litellm_credential_name).toBe("selected-credential"); expect(deployment.litellmParamsObj.timeout).toBe(5); }); + + it("keeps litellm_credential_name from LiteLLM Params JSON when no credential is selected", async () => { + const formValues = { + model_mappings: [ + { + public_name: "Public Model", + litellm_model: "litellm/public", + }, + ], + model_name: "custom-model-name", + litellm_extra_params: JSON.stringify({ + litellm_credential_name: "from-json", + timeout: 5, + }), + litellm_credential_name: null, + }; + + const deployments = await prepareModelAddRequest({ ...formValues }, "token", null); + + expect(deployments).toHaveLength(1); + const [deployment] = deployments!; + expect(deployment.litellmParamsObj.litellm_credential_name).toBe("from-json"); + expect(deployment.litellmParamsObj.timeout).toBe(5); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx index bb2f78fa84e..b4ae51d5423 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx @@ -91,6 +91,9 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value === "") { continue; } + if (key === "litellm_credential_name" && value == null) { + continue; + } // Skip the custom_pricing and pricing_model fields as they're only used for UI control if (key === "custom_pricing" || key === "pricing_model" || key === "cache_control") { continue; @@ -124,7 +127,7 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value && value != undefined) { try { litellmExtraParams = JSON.parse(value); - if ("litellm_credential_name" in litellmExtraParams) { + if ("litellm_credential_name" in litellmExtraParams && formValues.litellm_credential_name) { delete litellmExtraParams.litellm_credential_name; } } catch (error) { From b473339ac02c9618299a03483c5f1dc4a118aed4 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 31 Aug 2026 18:00:29 -0700 Subject: [PATCH 060/126] =?UTF-8?q?Revert=20"fix(ui):=20keep=20litellm=5Fc?= =?UTF-8?q?redential=5Fname=20from=20LiteLLM=20Params=20JSON=20when=20n?= =?UTF-8?q?=E2=80=A6"=20(#39046)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 33cc9c1c48d7d8551828d0137a97300833294582. --- .../panels/AddModelPanel.integration.test.tsx | 1 + .../handle_add_model_submit.test.tsx | 24 ------------------- .../add_model/handle_add_model_submit.tsx | 5 +--- 3 files changed, 2 insertions(+), 28 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx index efba26734ff..19e1e3aa8bd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx @@ -87,6 +87,7 @@ const alwaysMounted = { api_key: undefined, api_base: undefined, custom_llm_provider: "openai", + litellm_credential_name: null, model: "gpt-4o", }; diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx index 2aa8f93b72b..7ef09d34924 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx @@ -73,28 +73,4 @@ describe("prepareModelAddRequest", () => { expect(deployment.litellmParamsObj.litellm_credential_name).toBe("selected-credential"); expect(deployment.litellmParamsObj.timeout).toBe(5); }); - - it("keeps litellm_credential_name from LiteLLM Params JSON when no credential is selected", async () => { - const formValues = { - model_mappings: [ - { - public_name: "Public Model", - litellm_model: "litellm/public", - }, - ], - model_name: "custom-model-name", - litellm_extra_params: JSON.stringify({ - litellm_credential_name: "from-json", - timeout: 5, - }), - litellm_credential_name: null, - }; - - const deployments = await prepareModelAddRequest({ ...formValues }, "token", null); - - expect(deployments).toHaveLength(1); - const [deployment] = deployments!; - expect(deployment.litellmParamsObj.litellm_credential_name).toBe("from-json"); - expect(deployment.litellmParamsObj.timeout).toBe(5); - }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx index b4ae51d5423..bb2f78fa84e 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx @@ -91,9 +91,6 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value === "") { continue; } - if (key === "litellm_credential_name" && value == null) { - continue; - } // Skip the custom_pricing and pricing_model fields as they're only used for UI control if (key === "custom_pricing" || key === "pricing_model" || key === "cache_control") { continue; @@ -127,7 +124,7 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value && value != undefined) { try { litellmExtraParams = JSON.parse(value); - if ("litellm_credential_name" in litellmExtraParams && formValues.litellm_credential_name) { + if ("litellm_credential_name" in litellmExtraParams) { delete litellmExtraParams.litellm_credential_name; } } catch (error) { From 4a163f1a6a05a64856caf1293ce33d5dd6ba4f8c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 18:15:18 -0700 Subject: [PATCH 061/126] test(e2e-ui): assert user-observable behavior instead of DOM structure in audit fixes Replace table tbody and data-slot locators with getByRole, restore prior public MCP hub entries instead of clearing the whitelist on cleanup, seed the public agent via the append-semantics per-agent route, and rework mutable cleanup state into const-scoped try/finally blocks --- .../ui/tests/guardrails/guardrails.spec.ts | 8 +- .../tests/internal-user/internalUser.spec.ts | 59 ++++++------- .../internalUserWithTeams.spec.ts | 8 -- tests/e2e/ui/tests/modelHub/modelHub.spec.ts | 27 +++--- .../tests/modelsPage/deleteTeamModel.spec.ts | 10 ++- .../ui/tests/proxy-admin/secondAdmin.spec.ts | 49 ++++++----- .../e2e/ui/tests/team-admin/teamAdmin.spec.ts | 86 +++++++++---------- 7 files changed, 123 insertions(+), 124 deletions(-) diff --git a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts index a3ce6a73075..77ff020510b 100644 --- a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts +++ b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts @@ -29,9 +29,7 @@ test.describe("Guardrails", () => { await page.keyboard.type("pre_call"); await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 }); await page.keyboard.press("Enter"); - await expect(dialog.locator('[data-slot="combobox-chip"]').filter({ hasText: "pre_call" })).toBeVisible({ - timeout: 5_000, - }); + await expect(dialog.getByText("pre_call", { exact: true })).toBeVisible({ timeout: 5_000 }); await dialog.getByText("Create guardrail", { exact: true }).click(); await dialog.getByLabel("presidio_analyzer_api_base").fill("http://127.0.0.1:9999"); @@ -46,7 +44,7 @@ test.describe("Guardrails", () => { await dialog.getByRole("button", { name: "Create Guardrail" }).click(); await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 }); - const row = page.locator("table tbody tr").filter({ hasText: guardrailName }); + const row = page.getByRole("row").filter({ hasText: guardrailName }); await expect(row).toHaveCount(1, { timeout: 15_000 }); await navigateToPage(page, Page.Teams); @@ -78,6 +76,6 @@ test.describe("Guardrails", () => { await page.reload(); await expect(page.getByRole("button", { name: /Add New Guardrail/i })).toBeVisible({ timeout: 20_000 }); - await expect(page.locator("table tbody tr").filter({ hasText: guardrailName })).toHaveCount(0); + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(0); }); }); diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index 79733b3289b..f392c5104da 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -44,35 +44,34 @@ test.describe("Internal User", () => { const suffix = Date.now(); const auth = { Authorization: `Bearer ${masterKey()}` }; - let apiKey = ""; + await navigateToPage(page, Page.ApiKeys); + + await page.getByRole("button", { name: /Create New Key/i }).click(); + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + await expect(page.getByRole("radio", { name: "You", exact: true })).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("radio", { name: "Another User" })).toHaveCount(0); + + const keyName = `e2e-internal-team-key-${suffix}`; + await page.getByLabel(/Key Name/).fill(keyName); + + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); + await teamSelect.click(); + await page.keyboard.type(E2E_TEAM_KEYGEN_ALIAS); + await page.getByRole("option", { name: E2E_TEAM_KEYGEN_ALIAS }).first().click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Team Models", exact: true }).click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + await page.keyboard.press("Escape"); + try { - await navigateToPage(page, Page.ApiKeys); - - await page.getByRole("button", { name: /Create New Key/i }).click(); - await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); - - await expect(page.getByRole("radio", { name: "You", exact: true })).toBeVisible({ timeout: 10_000 }); - await expect(page.getByRole("radio", { name: "Another User" })).toHaveCount(0); - - const keyName = `e2e-internal-team-key-${suffix}`; - await page.getByLabel(/Key Name/).fill(keyName); - - const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); - await teamSelect.click(); - await page.keyboard.type(E2E_TEAM_KEYGEN_ALIAS); - await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_KEYGEN_ALIAS).first().click(); - - await page.getByRole("combobox", { name: "Select models" }).click(); - await page.getByRole("option", { name: "All Team Models", exact: true }).click(); - await page.keyboard.press("Escape"); - - await page.getByRole("button", { name: "Create Key", exact: true }).click(); - - await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); - apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); - expect(apiKey).toMatch(/^sk-/); - await page.keyboard.press("Escape"); - await openPlayground(page); await keySourceSelect(page).click(); await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); @@ -86,9 +85,7 @@ test.describe("Internal User", () => { await expect(page.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 60_000 }); } finally { - if (apiKey) { - await request.post("/key/delete", { headers: auth, data: { keys: [apiKey] } }); - } + await request.post("/key/delete", { headers: auth, data: { keys: [apiKey] } }); } }); diff --git a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts index d4c636e0541..62681e9ceb5 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts @@ -8,12 +8,6 @@ import { import { Page } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; -/** - * Differential partner to internalUserNoTeam.spec.ts: the seeded - * e2e-internal-user belongs to exactly three teams, so the Create Key dropdown - * must list all of them. Without this, the no-team spec's "zero options" assertion - * would still pass against a bug that empties the dropdown for everyone. - */ test.describe("Internal User with team memberships", () => { test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); @@ -26,8 +20,6 @@ test.describe("Internal User with team memberships", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - // All seeded memberships render, and nothing else does — proving the - // dropdown is scoped to the user's teams rather than empty or unfiltered. await expect(page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS })).toBeVisible({ timeout: 10_000 }); await expect(page.getByRole("option", { name: E2E_TEAM_ORG_ALIAS })).toBeVisible(); await expect(page.getByRole("option", { name: E2E_TEAM_KEYGEN_ALIAS })).toBeVisible(); diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 7fe3894d75d..6877fc9c48d 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -85,7 +85,17 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { const mcpServerName = `e2e_public_mcp_${suffix}`; const auth = { Authorization: `Bearer ${masterKey()}` }; - const seedPublicEntries = async (api: APIRequestContext): Promise<{ agentId: string; serverId: string }> => { + const publicMcpServerIds = async (api: APIRequestContext): Promise => { + const res = await api.get("/public/mcp_hub"); + expect(res.ok(), `public mcp_hub read failed (${res.status()}): ${await res.text()}`).toBe(true); + const servers: { server_id: string }[] = await res.json(); + return servers.map((server) => server.server_id); + }; + + const seedPublicEntries = async ( + api: APIRequestContext, + priorMcpIds: string[], + ): Promise<{ agentId: string; serverId: string }> => { const agentRes = await api.post("/v1/agents", { headers: auth, data: { @@ -117,21 +127,19 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { expect(serverRes.ok(), `mcp server create failed (${serverRes.status()}): ${await serverRes.text()}`).toBe(true); const serverId = (await serverRes.json()).server_id as string; - const agentPublicRes = await api.post("/v1/agents/make_public", { - headers: auth, - data: { agent_ids: [agentId] }, - }); - expect(agentPublicRes.ok(), `agents make_public failed: ${await agentPublicRes.text()}`).toBe(true); + const agentPublicRes = await api.post(`/v1/agents/${agentId}/make_public`, { headers: auth }); + expect(agentPublicRes.ok(), `agent make_public failed: ${await agentPublicRes.text()}`).toBe(true); const mcpPublicRes = await api.post("/v1/mcp/make_public", { headers: auth, - data: { mcp_server_ids: [serverId] }, + data: { mcp_server_ids: [...priorMcpIds, serverId] }, }); expect(mcpPublicRes.ok(), `mcp make_public failed: ${await mcpPublicRes.text()}`).toBe(true); return { agentId, serverId }; }; - const { agentId, serverId } = await seedPublicEntries(request); + const priorMcpIds = await publicMcpServerIds(request); + const { agentId, serverId } = await seedPublicEntries(request, priorMcpIds); try { await page.goto(`/ui/model_hub_table?key=${masterKey()}`); await dismissFeedbackPopup(page); @@ -150,8 +158,7 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { await expect(page.getByRole("row").filter({ hasText: mcpServerName })).toHaveCount(1, { timeout: 10_000 }); await expect(page.getByText("E2E public MCP server").first()).toBeVisible(); } finally { - await request.post("/v1/agents/make_public", { headers: auth, data: { agent_ids: [] } }); - await request.post("/v1/mcp/make_public", { headers: auth, data: { mcp_server_ids: [] } }); + await request.post("/v1/mcp/make_public", { headers: auth, data: { mcp_server_ids: priorMcpIds } }); await request.delete(`/v1/agents/${agentId}`, { headers: auth }); await request.delete(`/v1/mcp/server/${serverId}`, { headers: auth }); } diff --git a/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts b/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts index dca7d9f006b..96abd9833c0 100644 --- a/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts @@ -5,8 +5,10 @@ import { navigateToPage } from "../../helpers/navigation"; import { readBack } from "../../helpers/roundTrip"; import { masterKey } from "../../helpers/traffic"; -async function findDeploymentByName(page: PlaywrightPage, modelName: string): Promise | undefined> { - const body = await readBack<{ data: Record[] }>(page, "/v2/model/info"); +type DeploymentRow = { model_name?: string }; + +async function findDeploymentByName(page: PlaywrightPage, modelName: string): Promise { + const body = await readBack<{ data: DeploymentRow[] }>(page, "/v2/model/info"); return body.data.find((row) => row.model_name === modelName); } @@ -41,7 +43,7 @@ test.describe("Delete team model", () => { await navigateToPage(page, Page.Models); await page.getByPlaceholder("Search model names").fill(modelName); - const row = page.locator("table tbody tr").filter({ hasText: modelName }); + const row = page.getByRole("row").filter({ hasText: modelName }); await expect(row).toHaveCount(1, { timeout: 15_000 }); await expect(row.getByText(E2E_TEAM_CRUD_ID)).toBeVisible({ timeout: 10_000 }); @@ -65,6 +67,6 @@ test.describe("Delete team model", () => { await page.reload(); await page.getByPlaceholder("Search model names").fill(modelName); await expect(page.getByText("No models found").first()).toBeVisible({ timeout: 15_000 }); - await expect(page.locator("table tbody tr").filter({ hasText: modelName })).toHaveCount(0); + await expect(page.getByRole("row").filter({ hasText: modelName })).toHaveCount(0); }); }); diff --git a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts index 2dd30060d2d..5a8bc84cc13 100644 --- a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts @@ -13,35 +13,38 @@ test.describe("Second proxy admin", () => { const password = "e2e-second-admin-password"; const auth = { Authorization: `Bearer ${masterKey()}` }; - const adminContext = await browser.newContext({ storageState: ADMIN_STORAGE_PATH }); - let userId = ""; - try { - const adminPage = await adminContext.newPage(); - await navigateToPage(adminPage, Page.Users); - await dismissFeedbackPopup(adminPage); + const inviteAdminUser = async (): Promise => { + const adminContext = await browser.newContext({ storageState: ADMIN_STORAGE_PATH }); + try { + const adminPage = await adminContext.newPage(); + await navigateToPage(adminPage, Page.Users); + await dismissFeedbackPopup(adminPage); - await adminPage.getByRole("button", { name: "+ Invite User", exact: true }).click(); - const dialog = adminPage.getByRole("dialog", { name: "Invite User" }); - await expect(dialog).toBeVisible({ timeout: 5_000 }); + await adminPage.getByRole("button", { name: "+ Invite User", exact: true }).click(); + const dialog = adminPage.getByRole("dialog", { name: "Invite User" }); + await expect(dialog).toBeVisible({ timeout: 5_000 }); - await dialog.getByLabel("User Email").fill(email); + await dialog.getByLabel("User Email").fill(email); - await dialog.getByLabel(/Global Proxy Role/).click(); - await adminPage.getByRole("option", { name: /Admin \(All Permissions\)/ }).click(); + await dialog.getByLabel(/Global Proxy Role/).click(); + await adminPage.getByRole("option", { name: /Admin \(All Permissions\)/ }).click(); - const createdResponse = adminPage.waitForResponse( - (res) => res.url().includes("/user/new") && res.request().method() === "POST", - ); - await dialog.getByRole("button", { name: "Invite User" }).click(); - const createdBody = await (await createdResponse).json(); - userId = (createdBody.data?.user_id ?? createdBody.user_id) as string; - expect(userId, "created user id from /user/new").toBeTruthy(); + const createdResponse = adminPage.waitForResponse( + (res) => res.url().includes("/user/new") && res.request().method() === "POST", + ); + await dialog.getByRole("button", { name: "Invite User" }).click(); + const createdBody = await (await createdResponse).json(); + const createdUserId = (createdBody.data?.user_id ?? createdBody.user_id) as string; + expect(createdUserId, "created user id from /user/new").toBeTruthy(); - await expect(adminPage.getByText("API user Created").first()).toBeVisible({ timeout: 10_000 }); - } finally { - await adminContext.close(); - } + await expect(adminPage.getByText("API user Created").first()).toBeVisible({ timeout: 10_000 }); + return createdUserId; + } finally { + await adminContext.close(); + } + }; + const userId = await inviteAdminUser(); try { const passwordRes = await request.post("/user/update", { headers: auth, diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index d902959de4c..26a6fa50b4b 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -147,8 +147,6 @@ test.describe("Team Admin", () => { expect(teamRes.ok(), `team create failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); const teamId = (await teamRes.json()).team_id as string; - let modelId = ""; - let teamKey = ""; try { const modelRes = await request.post("/model/new", { headers: auth, @@ -163,55 +161,57 @@ test.describe("Team Admin", () => { }, }); expect(modelRes.ok(), `model create failed (${modelRes.status()}): ${await modelRes.text()}`).toBe(true); - modelId = (await modelRes.json()).model_info?.id as string; + const modelId = (await modelRes.json()).model_info?.id as string; - const keyRes = await request.post("/key/generate", { headers: auth, data: { team_id: teamId } }); - expect(keyRes.ok(), `key generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); - teamKey = (await keyRes.json()).key as string; + try { + const keyRes = await request.post("/key/generate", { headers: auth, data: { team_id: teamId } }); + expect(keyRes.ok(), `key generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); + const teamKey = (await keyRes.json()).key as string; - await expect - .poll( - async () => { - const res = await request.get("/model_group/info", { - headers: { Authorization: `Bearer ${teamKey}` }, - }); - if (!res.ok()) return false; - const body: { data?: { model_group?: string }[] } = await res.json(); - return (body.data ?? []).some((group) => group.model_group === teamModelName); - }, - { - message: `model group ${teamModelName} never became visible to the team key`, - timeout: 30_000, - }, - ) - .toBe(true); + try { + await expect + .poll( + async () => { + const res = await request.get("/model_group/info", { + headers: { Authorization: `Bearer ${teamKey}` }, + }); + if (!res.ok()) return false; + const body: { data?: { model_group?: string }[] } = await res.json(); + return (body.data ?? []).some((group) => group.model_group === teamModelName); + }, + { + message: `model group ${teamModelName} never became visible to the team key`, + timeout: 30_000, + }, + ) + .toBe(true); - await openPlayground(page); - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + await openPlayground(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(teamKey); + const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(teamKey); - const select = modelSelect(page); - await select.click(); - await select.fill(teamModelName); - await expect(onlyVisible(page.getByRole("option", { name: teamModelName }))).toBeVisible({ - timeout: 15_000, - }); + const select = modelSelect(page); + await select.click(); + await select.fill(teamModelName); + await expect(onlyVisible(page.getByRole("option", { name: teamModelName }))).toBeVisible({ + timeout: 15_000, + }); - await select.fill(CHAT_MODEL_A); - await expect(onlyVisible(page.getByRole("option", { name: CHAT_MODEL_A }))).toBeVisible({ - timeout: 15_000, - }); - } finally { - if (teamKey) { - await request.post("/key/delete", { headers: auth, data: { keys: [teamKey] } }); - } - if (modelId) { + await select.fill(CHAT_MODEL_A); + await expect(onlyVisible(page.getByRole("option", { name: CHAT_MODEL_A }))).toBeVisible({ + timeout: 15_000, + }); + } finally { + await request.post("/key/delete", { headers: auth, data: { keys: [teamKey] } }); + } + } finally { await request.post("/model/delete", { headers: auth, data: { id: modelId } }); } + } finally { await request.post("/team/delete", { headers: auth, data: { team_ids: [teamId] } }); } }); From 3fadcd71553dcf02c76e465a17af34cca715e7ef Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 31 Aug 2026 18:19:40 -0700 Subject: [PATCH 062/126] fix(auth): quiet malformed virtual key rejections to stdout (#38838) * fix(auth): quiet malformed virtual key rejections to stdout Reduce noisy invalid-api-key error logs by classifying malformed virtual keys and routing their rejections to stdout as WARNING instead of stderr as ERROR. Suppressible via LITELLM_LOG=ERROR or log_client_error_tracebacks=true. Changes: - auth_utils: is_invalid_virtual_key_error() classifier and marker functions - auth_exception_handler: log invalid keys as WARNING to child logger before identity seeding and callbacks, escalate non-401 transforms to ERROR - user_api_key_auth: websocket early-raise WebSocketException(1008) to avoid double-logging at HTTP layer - _logging: child logger verbose_proxy_stdout_logger with no handler/level; LevelRoutingStreamHandler routes its WARNING records to stdout; handler setLevel in _turn_on_json() closes JSON config handler level leak - test_auth_exception_handler: new test case verifying malformed-key logs at WARNING with marker retention through transformations Fixes LIT-5362 * fix(auth): classify malformed-key 401 by raise-site marker, not message text Review round 1 (Greptile P2, veria Low): - Move the marker attribute name to litellm/constants.py per the shared sentinel convention - Stamp the marker on the malformed-key 401 where it is raised and classify only by it. Message text is caller-influenceable on other 401s (vector store ids, organization ids are interpolated into their messages), so a phrase match would let a request body demote an authorization failure to the quiet log path - Regression test: a 401 carrying the phrase but not the marker stays at ERROR on stderr --- litellm/_logging.py | 15 ++- litellm/constants.py | 6 + litellm/proxy/auth/auth_exception_handler.py | 103 +++++++++++------- litellm/proxy/auth/auth_utils.py | 38 +++++++ litellm/proxy/auth/user_api_key_auth.py | 14 ++- .../proxy/auth/test_auth_exception_handler.py | 36 +++++- 6 files changed, 162 insertions(+), 50 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index fbb35b72be2..9435562f890 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -264,13 +264,17 @@ def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str: class LevelRoutingStreamHandler(logging.StreamHandler): - """Writes records below WARNING to stdout and WARNING and above to stderr. + """Writes records below WARNING and invalid-key warnings to stdout, others to stderr. Collectors that derive severity from the stream report every stderr line as an error. + Invalid-key warnings route to stdout so LITELLM_LOG=ERROR can suppress them. """ def emit(self, record: logging.LogRecord) -> None: - preferred: Final = sys.stdout if record.levelno < logging.WARNING else sys.stderr + is_stdout_record: Final = record.levelno < logging.WARNING or ( + record.levelno == logging.WARNING and record.name == verbose_proxy_stdout_logger.name + ) + preferred: Final = sys.stdout if is_stdout_record else sys.stderr if preferred is None or getattr(preferred, "closed", False): self.stream = sys.stderr # rebind-ok: fall back to the pre-fix stream rather than raising per record else: @@ -508,6 +512,9 @@ else: handler.setFormatter(formatter) verbose_proxy_logger = logging.getLogger("LiteLLM Proxy") +# Malformed virtual key rejections log through this child; LevelRoutingStreamHandler +# writes its WARNING records to stdout. It has no handler or level of its own. +verbose_proxy_stdout_logger: Final = verbose_proxy_logger.getChild("stdout") verbose_router_logger = logging.getLogger("LiteLLM Router") verbose_logger = logging.getLogger("LiteLLM") @@ -520,6 +527,7 @@ verbose_logger.addHandler(handler) # handlers (JSON mode, uvicorn log config, a host app's root handler). verbose_router_logger.addFilter(_stdout_truncation_filter) verbose_proxy_logger.addFilter(_stdout_truncation_filter) +verbose_proxy_stdout_logger.addFilter(_stdout_truncation_filter) verbose_logger.addFilter(_stdout_truncation_filter) @@ -683,6 +691,7 @@ def _turn_on_json(): - Adds a JSON formatter to all loggers """ handler: Final = LevelRoutingStreamHandler() + handler.setLevel(numeric_level) handler.setFormatter(JsonFormatter()) _initialize_loggers_with_handler(handler) # Set up exception handlers @@ -700,12 +709,14 @@ def _disable_debugging(): verbose_logger.disabled = True verbose_router_logger.disabled = True verbose_proxy_logger.disabled = True + verbose_proxy_stdout_logger.disabled = True def _enable_debugging(): verbose_logger.disabled = False verbose_router_logger.disabled = False verbose_proxy_logger.disabled = False + verbose_proxy_stdout_logger.disabled = False def print_verbose(print_statement): diff --git a/litellm/constants.py b/litellm/constants.py index 1c8a6dc5874..c482ab0e39a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1427,6 +1427,12 @@ DEFAULT_SOFT_BUDGET: Final = float( ) # by default all litellm proxy keys have a soft budget of 50.0 # makes it clear this is a rate limit error for a litellm virtual key RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash" +# Prefix of the 401 raised when a submitted virtual key is not shaped like one. +INVALID_VIRTUAL_KEY_ERROR_MESSAGE: Final = "LiteLLM Virtual Key expected" +# Attribute stamped on that 401 at its raise site so log routing recognises it by +# provenance. Message text is caller-influenceable on other 401s, so it must not +# be used to classify. +INVALID_VIRTUAL_KEY_ERROR_MARKER: Final = "_litellm_invalid_virtual_key_error" # Python garbage collection threshold configuration # Format: "gen0,gen1,gen2" e.g., "1000,50,50" diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index a42187b3a44..64878a480a7 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -2,13 +2,14 @@ Handles Authentication Errors """ +import logging from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException, Request, status import litellm -from litellm._logging import verbose_proxy_logger +from litellm._logging import verbose_proxy_logger, verbose_proxy_stdout_logger from litellm.constants import EMPTY_MAPPING from litellm.integrations.otel.runtime import seed_request_identity from litellm.litellm_core_utils.core_helpers import is_expected_client_error @@ -18,7 +19,11 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_utils import _get_request_ip_address +from litellm.proxy.auth.auth_utils import ( + _get_request_ip_address, + is_invalid_virtual_key_error, + mark_invalid_virtual_key_error, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -36,6 +41,41 @@ else: Span = Any +def _as_proxy_exception(e: Exception) -> ProxyException: + """Convert an authentication failure into the ProxyException the client receives.""" + if isinstance(e, litellm.BudgetExceededError): + return ProxyException( + message=e.message, + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), + ) + if isinstance(e, HTTPException): + return ProxyException( + message=getattr(e, "detail", f"Authentication Error({e})"), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED), + ) + if isinstance(e, ProxyException): + return e + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + return ProxyException( + message=( + "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." + ), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) + return ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=status.HTTP_401_UNAUTHORIZED, + ) + + def _with_requester_ip_address(request_data: dict[str, object], requester_ip: str | None) -> dict[str, object]: """Auth gate rejections are raised before `add_litellm_data_to_request` records the caller IP, so their failure logs would otherwise carry no IP nor key/user identity.""" @@ -110,16 +150,21 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True, ) - log_fn: Final = ( - verbose_proxy_logger.error - if is_expected_client_error(e) and not litellm.log_client_error_tracebacks - else verbose_proxy_logger.exception - ) - log_fn( + + # Log authentication failures before identity seeding and callbacks, so the log + # survives a raising callback pipeline. Classify and route malformed virtual-key + # rejections to WARNING on stdout (suppressible via LITELLM_LOG=ERROR). + log_extra: Final = {"requester_ip": requester_ip} + is_invalid_virtual_key: Final = is_invalid_virtual_key_error(e) + is_quiet_log: Final = is_invalid_virtual_key and not litellm.log_client_error_tracebacks + logger: Final = verbose_proxy_stdout_logger if is_quiet_log else verbose_proxy_logger + logger.log( + logging.WARNING if is_quiet_log else logging.ERROR, "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", e, requester_ip, - extra={"requester_ip": requester_ip}, + exc_info=True if litellm.log_client_error_tracebacks or not is_expected_client_error(e) else None, + extra=log_extra, ) # Log this exception to OTEL, Datadog etc. Reuse the identity resolved @@ -167,35 +212,13 @@ class UserAPIKeyAuthExceptionHandler: if transformed_exception is not None: e = transformed_exception - if isinstance(e, litellm.BudgetExceededError): - raise ProxyException( - message=e.message, - type=ProxyErrorTypes.budget_exceeded, - param=None, - code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), + final_exception: Final = mark_invalid_virtual_key_error(_as_proxy_exception(e), is_invalid_virtual_key) + # If a quiet-logged malformed-key transform yields non-401, escalate to ERROR + if is_quiet_log and str(final_exception.code) != str(status.HTTP_401_UNAUTHORIZED): + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", + final_exception, + requester_ip, + extra=log_extra, ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED), - ) - elif isinstance(e, ProxyException): - raise e - if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): - raise ProxyException( - message=( - "Service Unavailable, the authentication database is " - "temporarily unreachable. Please retry shortly." - ), - type=ProxyErrorTypes.no_db_connection, - param="None", - code=status.HTTP_503_SERVICE_UNAVAILABLE, - ) - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_401_UNAUTHORIZED, - ) + raise final_exception diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9b1a6ba5aa7..83ee10b8108 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -15,6 +15,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, EMPTY_MAPPING, + INVALID_VIRTUAL_KEY_ERROR_MARKER, MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS, ) @@ -34,6 +35,43 @@ from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS from litellm.types.utils import CustomPricingLiteLLMParams +def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: + """True when an authentication error rejects a malformed virtual key. + + Classifies only by the marker stamped where that 401 is raised. Message + content is never inspected: other 401s interpolate caller-supplied values + (vector store ids, organization ids) into their messages, so a phrase + match would let a request body demote an authorization failure to the + quiet log path. + """ + if not isinstance(exception, (HTTPException, ProxyException)): + return False + + code: Final[object] = getattr(exception, "code", None) + status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None) + if str(status_code) != str(status.HTTP_401_UNAUTHORIZED): + return False + + return getattr(exception, INVALID_VIRTUAL_KEY_ERROR_MARKER, False) is True + + +def mark_invalid_virtual_key_error(exception: ProxyException, is_invalid_virtual_key: bool) -> ProxyException: + """Return an independently marked malformed-key exception after callback transformations.""" + if not is_invalid_virtual_key or str(exception.code) != str(status.HTTP_401_UNAUTHORIZED): + return exception + marked_exception: Final = ProxyException( + message=exception.message, + type=exception.type, + param=exception.param, + code=exception.code, + headers=exception.headers.copy(), + openai_code=None if exception.openai_code is None else str(exception.openai_code), + provider_specific_fields=exception.provider_specific_fields, + ) + setattr(marked_exception, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) + return marked_exception + + def _get_request_ip_address(request: Request, use_x_forwarded_for: bool | None = False) -> str | None: client_ip = None if use_x_forwarded_for is True and "x-forwarded-for" in request.headers: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e92d090a2fb..5fb6dad0cd7 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -19,12 +19,15 @@ import fastapi import orjson from fastapi import HTTPException, Request, WebSocket, status from fastapi.security.api_key import APIKeyHeader +from starlette.exceptions import WebSocketException import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, + INVALID_VIRTUAL_KEY_ERROR_MARKER, + INVALID_VIRTUAL_KEY_ERROR_MESSAGE, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, ) @@ -65,6 +68,7 @@ from litellm.proxy.auth.auth_utils import ( get_model_from_request, get_request_route, get_request_route_template, + is_invalid_virtual_key_error, iter_request_fallback_targets, normalize_request_route, pre_db_read_auth_checks, @@ -539,6 +543,8 @@ async def user_api_key_auth_websocket(websocket: WebSocket): try: return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") except Exception as e: + if is_invalid_virtual_key_error(e): + raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) verbose_proxy_logger.exception(e) await websocket.close(code=status.WS_1008_POLICY_VIOLATION) raise HTTPException(status_code=403, detail=str(e)) @@ -1867,13 +1873,17 @@ async def _user_api_key_auth_builder( _masked_key: Final = f"{api_key[:4]}****{api_key[-4:]}" if len(api_key) > 8 else "****" if not api_key.startswith("sk-"): _hint = _JWT_AUTH_DISABLED_HINT if not enable_jwt_auth and JWTHandler.is_jwt(token=api_key) else "" - raise HTTPException( + _malformed_key_error = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=( - f"LiteLLM Virtual Key expected. Received={_masked_key}, " + f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={_masked_key}, " f"expected to start with 'sk-'.{_hint}" ), ) # prevent token hashes from being used + # Stamp provenance here so log routing classifies this 401 by + # where it was raised, never by its message text. + setattr(_malformed_key_error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) + raise _malformed_key_error else: verbose_logger.warning( "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 90b3b29d919..90be51cfa5b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -26,6 +26,7 @@ from prisma.errors import ( from litellm._logging import verbose_proxy_logger +from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER from litellm.exceptions import BudgetExceededError from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler @@ -703,23 +704,43 @@ async def test_auth_failure_ip_stamp_does_not_mutate_callers_request_data(): assert request_data == {"model": "gpt-4o"} +def _marked_malformed_key_error() -> HTTPException: + """Build the malformed-key 401 as its raise site does: marker stamped on it.""" + error = HTTPException(status_code=401, detail="LiteLLM Virtual Key expected. Received=test") + setattr(error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) + return error + + @pytest.mark.asyncio @pytest.mark.parametrize( - "auth_error,expect_traceback", + "auth_error,expect_traceback,expect_level", [ pytest.param( ProxyException( message="Authentication Error", type=ProxyErrorTypes.auth_error, param=None, code=401 ), False, + "ERROR", id="expected_401_no_traceback", ), - pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error_keeps_traceback"), + pytest.param(ValueError("unexpected internal error"), True, "ERROR", id="unexpected_error_keeps_traceback"), + pytest.param( + _marked_malformed_key_error(), + False, + "WARNING", + id="malformed_virtual_key_warning_no_traceback", + ), + pytest.param( + HTTPException(status_code=401, detail="LiteLLM Virtual Key expected. Received=test"), + False, + "ERROR", + id="phrase_without_marker_stays_loud", + ), ], ) -async def test_handle_authentication_error_traceback_only_for_unexpected_errors(auth_error, expect_traceback, caplog): +async def test_handle_authentication_error_traceback_only_for_unexpected_errors(auth_error, expect_traceback, expect_level, caplog): """Regression for LIT-6043: expected 4xx auth rejections must not format a - traceback via logger.exception; unexpected errors must keep it.""" + traceback via logger.exception; malformed virtual keys log at WARNING.""" handler = UserAPIKeyAuthExceptionHandler() with ( @@ -740,8 +761,8 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( try: try: raise auth_error - except (ProxyException, ValueError) as caught: - with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException): + except (ProxyException, ValueError, HTTPException) as caught: + with caplog.at_level(expect_level, logger="LiteLLM Proxy"), pytest.raises((ProxyException, HTTPException)): await handler._handle_authentication_error( caught, MagicMock(), @@ -756,3 +777,6 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] assert len(records) == 1 assert (records[0].exc_info is not None) is expect_traceback + assert records[0].levelname == expect_level + expected_logger_name = "LiteLLM Proxy.stdout" if expect_level == "WARNING" else "LiteLLM Proxy" + assert records[0].name == expected_logger_name From f7accc4e29da707c2e1797c2217b6d1c420356f8 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 19:33:04 -0700 Subject: [PATCH 063/126] test(e2e): drop the two mgmt registry cells no shared-proxy test can cover mgmt.cache_settings.update.happy_path and mgmt.config_override.hashicorp_vault.happy_path were the last two uncovered Management/UI cells, and neither can be covered against the shared proxy the e2e suites run on. Both routes reconfigure the whole process rather than a resource the test owns. /cache/settings persists whatever it receives into a row that outranks the YAML cache_params and is re-applied on a timer, so a partial write downgrades a TLS cluster to a plaintext standalone node and every later Redis call hangs. That is what took out 60 of 72 tests on 2026-07-25 and got the original test removed in PR #34664. /config_overrides/hashicorp_vault has the same shape: a POST sets the HCP_VAULT_* env vars, swaps litellm.secret_manager_client process-wide, and writes a row the config-reload poll re-applies, so every os.environ/ lookup on the pod resolves against the test's Vault until the DELETE lands. Its constructor also never dials Vault, so a POST to a bogus address still returns 200 and a smoke test built on it would pass for the wrong reason. Keeping rows we have decided not to cover only inflates the denominator, so drop them and record the reasoning where someone would go to write the test. Filing the isolated-proxy harness they both need separately; the cells come back with it. Management/UI goes 75/77 to 75/75, headline 402/544 to 402/542. --- tests/e2e/coverage_registry/mgmt.yaml | 2 -- tests/e2e/management/test_config_misc_endpoints_e2e.py | 10 +++++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 1e6de0c3d6a..d571fb36546 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -64,14 +64,12 @@ - {id: mgmt.budget.list_v1.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "management_v1/budgets.py:129", rationale: "Budget enumeration the Budgets page can page, sort and filter"} - {id: mgmt.budget.list_v1.admin_only, module: mgmt, tier: P1, surface: api, assertions: [admin_only], source: "management_v1/budgets.py:129", rationale: "A caller without admin view is refused, not served an empty page"} - {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"} -- {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke). Deliberately uncovered: the previous test read the live settings and wrote them back, which proves nothing (identical values in, so a no-op POST still passes) while being able to break the deployment. /cache/settings persists what it receives and that row outranks YAML cache_params, re-applied on a timer, so a write that omits ssl or redis_startup_nodes turns a TLS cluster into a plaintext standalone node and every later Redis call hangs. That took out 60 of 72 tests on 2026-07-25. GET cannot round-trip it either: it resolves the stored row overlaid with REDIS_* env and never reads YAML, so on a fresh deploy it cannot see YAML ssl to echo back. A safe test needs an isolated proxy, or LIT-4816 fixed so a partial write cannot downgrade transport. Do not re-add a read-then-write-back test against a shared proxy."} - {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"} - {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"} - {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"} - {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"} - {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"} - {id: mgmt.fallback_management.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "fallback_management_endpoints.py", rationale: "Fallback config (smoke)"} -- {id: mgmt.config_override.hashicorp_vault.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "config_override_endpoints.py", rationale: "Vault integration (smoke)"} - {id: mgmt.workflow.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "workflow_management_endpoints.py", rationale: "Workflow tracking (smoke)"} - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} diff --git a/tests/e2e/management/test_config_misc_endpoints_e2e.py b/tests/e2e/management/test_config_misc_endpoints_e2e.py index 195732c0201..099ffa4b3bd 100644 --- a/tests/e2e/management/test_config_misc_endpoints_e2e.py +++ b/tests/e2e/management/test_config_misc_endpoints_e2e.py @@ -7,9 +7,13 @@ so a read-back reflects the change. Router settings, which mutate global proxy state, are exercised with a benign, self-restoring change so a shared proxy is left as it was found. -Cache settings are deliberately not covered here; see the rationale on -mgmt.cache_settings.update.happy_path in coverage_registry/mgmt.yaml before adding -a test for that route. +Cache settings and the Vault config override are deliberately not covered here. +Both routes reconfigure the whole proxy: /cache/settings persists what it receives +into a row that outranks the YAML cache_params and is re-applied on a timer, and +/config_overrides/hashicorp_vault swaps the process-wide secret manager. Neither can +be exercised safely against the shared proxy the suites run on, so they need an +isolated proxy before a test lands. Do not add a read-then-write-back test for +either one. """ from __future__ import annotations From ccd76dac505035885b34e54d325ea9bfe9a6718d Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 31 Aug 2026 19:38:39 -0700 Subject: [PATCH 064/126] fix(proxy): wire team-level logging callbacks into passthrough endpoints (#38979) * fix(proxy): wire team-level logging callbacks into passthrough endpoints LIT-5152: passthrough routes now wire dynamic team-level callbacks (success_callback, failure_callback, callback_vars) into Logging constructor, mirroring the add_litellm_data_to_request behavior. Three hardening fixes: 1. Catch TypeError/AttributeError in _get_validated_callback_metadata when team logging metadata has wrong shape (e.g., logging list instead of dict), preventing HTTP 500 on passthrough routes with malformed config. 2. Wrap websocket passthrough logging initialization in try/except, since the socket is already accepted at that point; errors after accept() yield abrupt close (1006/1011) rather than clean HTTP error response. 3. Handle malformed deprecated callback_settings gracefully with try/except. 4. Wrap HTTP passthrough callback resolution in try/except to prevent 500 on malformed team metadata (backward-compatibility fix). Changes: - pass_through_endpoints.py: wire dynamic callbacks in HTTP+WS paths, handle malformed metadata gracefully with try/except fallbacks - litellm_pre_call_utils.py: expand exception handling in validators - test file: regression test for happy-path team callback wiring * refactor(proxy): share passthrough team-callback resolution and cover its fail-open path Collapse the duplicated callback wiring on the HTTP and websocket passthrough paths into one helper that returns a frozen wiring value, log resolution failures at error level so a broken logging config stays visible, and add regression tests for malformed team metadata and an operational lookup failure. Reverts the _get_validated_callback_metadata except widening: it changed behavior for normal LLM routes, which is outside this ticket's scope. * fix(proxy): keep passthrough alive when team callback vars hold env references The deprecated team_metadata.callback_settings branch builds TeamCallbackMetadata directly, skipping the AddTeamCallback validation that strips os.environ/ references from the newer logging list. Stamping those vars onto the Logging object made its constructor raise, so a team on the legacy shape got HTTP 500 on every passthrough call. Validate the resolved vars inside the fail-open boundary instead, so the request goes through with dynamic callbacks skipped and the reason logged. * fix(proxy): lint violations in team callback wiring helper * style: format lint --- .../pass_through_endpoints.py | 95 ++++++++++- .../test_pass_through_endpoints.py | 161 ++++++++++++++++++ 2 files changed, 251 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 09d3dedaafa..ddd27f0fc2a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,9 +6,10 @@ import posixpath import traceback from base64 import b64encode from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from itertools import groupby -from typing import Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -47,6 +48,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) +from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_no_callback_env_reference from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -78,7 +80,10 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + _get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above +) from litellm.proxy.utils import normalize_route_for_root_path from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str @@ -90,7 +95,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, PassthroughStandardLoggingPayload, ) -from litellm.types.utils import Usage +from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging @@ -99,6 +104,9 @@ from .upstream_usage_headers import ( apply_upstream_reported_usage, ) +if TYPE_CHECKING: + from litellm.proxy.proxy_server import ProxyConfig + router: Final = APIRouter() pass_through_endpoint_logging: Final = PassThroughEndpointLogging() @@ -752,6 +760,67 @@ def _build_passthrough_failure_request_payload( return request_payload +@dataclass(frozen=True, slots=True) +class _TeamCallbackWiring: + success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg + failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg + logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg + + +def _resolve_team_callback_wiring( + user_api_key_dict: UserAPIKeyAuth, + proxy_config: "ProxyConfig", + route_description: str, +) -> _TeamCallbackWiring: + """Resolve key/team dynamic logging callbacks for a passthrough request. + + Mirrors add_litellm_data_to_request: callback_vars are unpacked top-level + (read by initialize_standard_callback_dynamic_params) and also stamped on + the proxy-owned trusted-vars field (read by get_trusted_callback_params). + + Fails open: a callback resolution or validation error is logged at error + level and the request proceeds without dynamic callbacks, since a broken + logging config must not fail the customer's upstream call (and the + websocket is already accepted by the time this runs on that path). The + env-reference check runs here because the deprecated callback_settings + branch skips AddTeamCallback validation, and Logging.__init__ would + otherwise reject the vars mid-request. + """ + try: + callback_settings_obj: Final = _get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, proxy_config=proxy_config + ) + if callback_settings_obj and callback_settings_obj.callback_vars: + for ( + item + ) in callback_settings_obj.callback_vars.items(): # rebind-ok: dict.items iteration for env-ref validation + validate_no_callback_env_reference(item[0], item[1], source="key/team callback metadata") + except Exception: # noqa: BLE001 - a broken logging config must never fail the passthrough request + verbose_proxy_logger.exception( + "%s: failed to resolve team logging callbacks, continuing without them", + route_description, + ) + return _TeamCallbackWiring() + if callback_settings_obj is None: + return _TeamCallbackWiring() + callback_vars: Final = callback_settings_obj.callback_vars + success_callbacks: Final = callback_settings_obj.success_callback + failure_callbacks: Final = callback_settings_obj.failure_callback + logging_kwargs: Final = ( + None + if not callback_vars + else { # mutable-ok: Logging arg + **callback_vars, + TRUSTED_CALLBACK_VARS_FIELD: callback_vars, + } + ) + return _TeamCallbackWiring( + success_callbacks=None if success_callbacks is None else [*success_callbacks], # mutable-ok: Logging arg + failure_callbacks=None if failure_callbacks is None else [*failure_callbacks], # mutable-ok: Logging arg + logging_kwargs=logging_kwargs, + ) + + async def _log_passthrough_upstream_failure( response: httpx.Response, user_api_key_dict: UserAPIKeyAuth, @@ -845,7 +914,7 @@ async def pass_through_request( from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( PassthroughGuardrailHandler, ) - from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj ######################################################### # Initialize variables @@ -930,6 +999,11 @@ async def pass_through_request( # read e.g. ``chat gpt-4o`` instead of ``chat unknown``. passthrough_model: Final = (_parsed_body.get("model") if isinstance(_parsed_body, dict) else None) or "unknown" start_time: Final = datetime.now() + team_callbacks: Final = _resolve_team_callback_wiring( + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_description="pass_through_endpoint", + ) logging_obj = Logging( model=passthrough_model, messages=[{"role": "user", "content": safe_dumps(_parsed_body)}], @@ -938,6 +1012,9 @@ async def pass_through_request( start_time=start_time, litellm_call_id=litellm_call_id, function_id="1245", + dynamic_success_callbacks=team_callbacks.success_callbacks, + dynamic_failure_callbacks=team_callbacks.failure_callbacks, + kwargs=team_callbacks.logging_kwargs, ) # Store passthrough guardrails config on logging_obj for field targeting @@ -2022,7 +2099,7 @@ async def websocket_passthrough_request( setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream """ from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) @@ -2055,6 +2132,11 @@ async def websocket_passthrough_request( upstream_headers[header_name] = header_value # Initialize logging object similar to HTTP passthrough + team_callbacks: Final = _resolve_team_callback_wiring( + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_description="websocket_passthrough", + ) logging_obj: Final = Logging( model="unknown", messages=[{"role": "user", "content": "WebSocket connection"}], @@ -2063,6 +2145,9 @@ async def websocket_passthrough_request( start_time=start_time, litellm_call_id=litellm_call_id, function_id="websocket_passthrough", + dynamic_success_callbacks=team_callbacks.success_callbacks, + dynamic_failure_callbacks=team_callbacks.failure_callbacks, + kwargs=team_callbacks.logging_kwargs, ) # Create passthrough logging payload diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a3f56adb86f..f5ae0fe5977 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5462,3 +5462,164 @@ def test_the_marker_check_distinguishes_the_two_route_kinds(): builtin = MagicMock(spec=Request) builtin.scope = {"endpoint": llm_passthrough_endpoints.anthropic_proxy_route} assert request_dispatched_to_pass_through_endpoint(builtin) is False + + +async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: UserAPIKeyAuth) -> tuple[int, object]: + import litellm + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"ok": True}) + + real_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": resolve_pass_through_request_timeout(None)}, + ) + cache_dict = litellm.in_memory_llm_clients_cache.cache_dict + cache_key = next((key for key, cached in cache_dict.items() if cached is real_handler), None) + assert cache_key is not None + cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))) + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + mock_request.body = AsyncMock(return_value=b'{"model": "gemini-2.0-flash"}') + + captured_data: dict = {} + + async def capture_pre_call_hook(user_api_key_dict, data, call_type): + captured_data.update(data) + return data + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=capture_pre_call_hook) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=None) + + try: + with patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ): + response = await pass_through_request( + request=mock_request, + target="https://upstream.example.test/v1/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + finally: + cache_dict[cache_key] = real_handler + + return response.status_code, captured_data.get("litellm_logging_obj") + + +@pytest.mark.asyncio +async def test_pass_through_request_wires_team_callbacks(): + """LIT-5152 regression: pass_through_request must resolve team-level logging + callbacks from key/team metadata and wire them into the Logging object, the + same way add_litellm_data_to_request does for normal LLM routes.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={ + "logging": [ + { + "callback_name": "langfuse", + "callback_type": "success_and_failure", + "callback_vars": { + "langfuse_public_key": "pk_test", + "langfuse_secret_key": "sk_test", + "langfuse_host": "https://langfuse.example.test", + }, + } + ] + }, + ) + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(user_api_key_dict) + + assert status_code == 200 + assert logging_obj is not None + assert logging_obj.dynamic_success_callbacks, "team success callbacks not wired into Logging" + assert logging_obj.dynamic_failure_callbacks, "team failure callbacks not wired into Logging" + assert logging_obj.standard_callback_dynamic_params.get("langfuse_public_key") == "pk_test" + assert logging_obj.standard_callback_dynamic_params.get("langfuse_secret_key") == "sk_test" + assert logging_obj.standard_callback_dynamic_params.get("langfuse_host") == "https://langfuse.example.test" + assert ("langfuse_public_key", "pk_test") in logging_obj._trusted_callback_vars + + +@pytest.mark.asyncio +async def test_pass_through_request_survives_malformed_team_logging_metadata(): + """LIT-5152 fail-open: a malformed team ``logging`` value (here a non-iterable) + raises inside callback resolution; the passthrough request must still succeed, + just without dynamic callbacks.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={"logging": 5}, + ) + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(user_api_key_dict) + + assert status_code == 200 + assert logging_obj is not None + assert not logging_obj.dynamic_success_callbacks + assert not logging_obj.dynamic_failure_callbacks + + +@pytest.mark.asyncio +async def test_pass_through_request_survives_env_reference_in_deprecated_callback_settings(): + """LIT-5152 fail-open: the deprecated ``callback_settings`` team metadata skips + AddTeamCallback validation, so an ``os.environ/`` callback var would otherwise + blow up inside ``Logging.__init__`` and fail the request; the passthrough must + instead succeed without dynamic callbacks.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={ + "callback_settings": { + "success_callback": ["langfuse"], + "failure_callback": ["langfuse"], + "callback_vars": { + "langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY", + "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY", + "langfuse_host": "https://langfuse.example.test", + }, + } + }, + ) + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(user_api_key_dict) + + assert status_code == 200 + assert logging_obj is not None + assert not logging_obj.dynamic_success_callbacks + assert not logging_obj.dynamic_failure_callbacks + assert not logging_obj.standard_callback_dynamic_params.get("langfuse_public_key") + + +@pytest.mark.asyncio +async def test_resolve_team_callback_wiring_fails_open_on_operational_error(): + """LIT-5152 fail-open: an operational error while resolving callback metadata + (e.g. team config lookup hitting a dead secret manager) must not raise; the + request proceeds without dynamic callbacks and the error is logged.""" + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _resolve_team_callback_wiring, + ) + from litellm.proxy.proxy_server import ProxyConfig + + class RaisingTeamConfig(ProxyConfig): + def load_team_config(self, team_id: str) -> dict: + raise RuntimeError("secret manager unavailable") + + wiring = _resolve_team_callback_wiring( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", team_id="test-team"), + proxy_config=RaisingTeamConfig(), + route_description="pass_through_endpoint", + ) + + assert wiring.success_callbacks is None + assert wiring.failure_callbacks is None + assert wiring.logging_kwargs is None From 8d6d7f9ce9e46155e5b4a63ba894c6e6eb4cd3f9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 19:51:41 -0700 Subject: [PATCH 065/126] feat(complexity_router): opt-in modality-based capability routing for image requests (#39032) --- .../prompt_templates/common_utils.py | 35 ++ .../complexity_router/README.md | 22 ++ .../complexity_router/complexity_router.py | 215 +++++++++++- .../complexity_router/config.py | 12 + litellm/types/utils.py | 4 + litellm/utils.py | 24 +- ...ore_utils_prompt_templates_common_utils.py | 47 +++ .../router_strategy/test_complexity_router.py | 320 ++++++++++++++++++ tests/test_litellm/test_utils.py | 30 ++ .../RoutingDecisionCard.test.tsx | 6 + .../LogDetailsDrawer/RoutingDecisionCard.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 +- 12 files changed, 705 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 1c8f10d3307..bda365102af 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -205,6 +205,41 @@ def is_non_content_values_set(message: AllMessageValues) -> bool: return any(message.get(key, None) is not None for key in message if key not in ignore_keys) +_IMAGE_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_image", "image"}) +_IMAGE_SCAN_MAX_DEPTH: Final = 4 + + +def _content_parts_contain_image(parts: Sequence[object]) -> bool: + """Depth-bounded frontier walk over nested content lists, iterative because the repo bans + recursion; an Anthropic tool_result nests its image parts exactly one level down.""" + frontier = parts # rebind-ok: depth-bounded frontier walk + for _ in range(_IMAGE_SCAN_MAX_DEPTH): + if any(isinstance(part, Mapping) and part.get("type") in _IMAGE_CONTENT_PART_TYPES for part in frontier): + return True + frontier = tuple( # rebind-ok: depth-bounded frontier walk + nested + for part in frontier + if isinstance(part, Mapping) + for content in (part.get("content"),) + if isinstance(content, list) + for nested in content + ) + if not frontier: + return False + return False + + +def request_contains_image_content(messages: Sequence[Mapping[str, object]]) -> bool: + """Whether any message carries an image content part, across the dialects that reach + pre-routing hooks untranslated: chat-completions ``image_url``, Responses ``input_image``, + and Anthropic Messages ``image``, including images nested inside ``tool_result`` blocks.""" + return any( + isinstance(content, list) and _content_parts_contain_image(content) + for message in messages + for content in (message.get("content"),) + ) + + def _audio_or_image_in_message_content(message: AllMessageValues) -> bool: """ Checks if message content contains an image or audio diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 63ba760ff66..bc8df67cc28 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -154,6 +154,9 @@ model_list: # Fallback model if tier cannot be determined default_model: gpt-4o + + # Replace a routed model that cannot take image input (default: false) + modality_routing: true ``` ## Usage @@ -178,6 +181,25 @@ response = litellm.completion( ## Special Behaviors +### Modality-based capability routing + +The classifier reads text alone, so a request carrying an image can classify cheap and land on a +text-only model, which rejects it with a provider 400 no fallback catches. With +`modality_routing: true`, one gate inspects every decided placement: when the routed model is +explicitly declared `supports_vision: false` (deployment `model_info` first, the model cost map +otherwise; unmapped names stay routable, and a multi-deployment group must accept on every +deployment), the request is re-placed on the nearest HIGHER tier holding a capable model, with +routing plugins still applied to the re-pick, then on `default_model` (never on plugin routers +and never for a plan-floored decision), and otherwise rejected with a clear 400 naming the +router. The walk only ever goes up, so a plan-mode floor cannot be undercut; a router whose only +vision model sits below the decided tier gets the 400 and an actionable message instead. + +A same-tier re-pick keeps the decision's cause and adds `modality:image` to `signals`; a tier +change or default takeover records `cause: modality_escalation` with the displaced placement +(`modality_escalated_from:` or `modality_displaced_default_model`). Escalations are never +pinned by session affinity, and a KEPT session pin bypasses the gate entirely: a session pinned +to a text-only model keeps it even when an image arrives. + ### Heuristic-first chaining `classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 329da35eab3..be7653902a7 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -30,6 +30,7 @@ from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata +from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.types.utils import ( @@ -738,6 +739,10 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo size shrinks again the moment the client compacts: pinning the escalated tier would hold the session on the big-window model long after the oversized context that forced it is gone. The gate re-fires per request, so leaving these unpinned costs nothing but the classifier call. + + A modality escalation is transient the same way: it describes what this one call carries (an + image), not what the session's traffic looks like, and pinning it would hold every following + text turn on the vision-capable model the image forced. """ return decision is None or ( decision.get("cause") @@ -745,6 +750,7 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo "default_model_fallback", "plan_mode", "housekeeping", + "modality_escalation", ) and not decision.get("context_escalated") ) @@ -2274,6 +2280,175 @@ class ComplexityRouter(CustomLogger): return pinned_model return self.get_model_for_tier(escalated_tier) + def _model_accepts_image_input(self, model_name: str) -> bool: + """Whether a routed model or pool entry can serve an image request. + + Resolved through the deployments that would actually serve the name; a name with no + deployment on the router is served by the SDK directly and is checked against the model + cost map itself. Only an explicit supports_vision false excludes, a deployment-level + model_info override first and the map otherwise, so unmapped custom names stay routable. + + A multi-deployment group must accept on EVERY deployment: the router picks a deployment + inside the group after this gate runs, so a mixed group marked eligible could still hand + the image to its text-only member and fail with the exact 400 the gate exists to prevent. + """ + from litellm.utils import is_vision_explicitly_disabled + + def deployment_accepts(deployment: Mapping[str, Any]) -> bool: + declared: Final = (deployment.get("model_info") or EMPTY_MAPPING).get("supports_vision") + if declared is not None: + return declared is True + litellm_model: Final = (deployment.get("litellm_params") or EMPTY_MAPPING).get("model") or model_name + return not is_vision_explicitly_disabled(litellm_model) + + deployments: Final = self.litellm_router_instance.get_model_list(model_name=model_name) + if not deployments: + return not is_vision_explicitly_disabled(model_name) + return all(deployment_accepts(deployment) for deployment in deployments) + + def _modality_eligible_models(self) -> frozenset[str]: + """Every configured pool entry, plus default_model, that can serve an image request.""" + names: Final = frozenset(entry for pool in self._tier_pools().values() for entry in pool) | frozenset( + name for name in (self.config.default_model,) if name + ) + return frozenset(name for name in names if self._model_accepts_image_input(name)) + + async def _gate_response_modality( + self, + response: PreRoutingHookResponse, + messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick + resolved_messages: Sequence[Mapping[str, object]] | None, + request_kwargs: dict, # mutable-ok: same shape the hook receives + ) -> PreRoutingHookResponse: + """Replace a routed model that cannot accept this request's image input. + + The single modality owner, applied to the decided response at the hook's exits so every + routing path is covered uniformly. A KEPT session pin is exempt by design (its cause); + replacement picks and every other path are just responses. The re-placement walks + UPWARD-ONLY from the decision's tier (so a plan-mode floor can never be undercut), picks + through `_pick_model_for_tier` so routing plugins still apply, then falls to + default_model (never on plugin routers, and never on a plan-floored decision, since + default_model carries no tier guarantee), else raises the clear 400. The rewritten + decision keeps its cause on a same-tier repick and becomes modality_escalation when the + tier moved or default_model took over, with the displaced placement in signals. + """ + decision: Final = response.routing_decision + if ( + not self.config.modality_routing + or not resolved_messages + or response.model is None + or (decision is not None and decision.get("cause") == "session_affinity_pin") + or not request_contains_image_content(resolved_messages) + or self._model_accepts_image_input(response.model) + ): + return response + eligible: Final = self._modality_eligible_models() + names: Final = self.config.tier_names() + pools: Final = self._tier_pools() + decided: Final = decision.get("tier") if decision is not None else None + start: Final = names.index(decided) if isinstance(decided, str) and decided in names else 0 + capable: Final = next( + (name for name in names[start:] if any(entry in eligible for entry in pools.get(name, ()))), None + ) + if capable is not None: + new_tier: ComplexityTier | str | None = capable if self.config.has_custom_tiers else ComplexityTier(capable) + repick_messages: Final = list(resolved_messages) # mutable-ok: the pick's param is list-typed + new_model = await self._pick_model_for_tier( + new_tier, + messages, + repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them + request_kwargs, + allowed_models=tuple(entry for entry in pools.get(capable, ()) if entry in eligible), + ) + elif self._modality_default_model_usable(request_kwargs, resolved_messages, eligible): + new_tier = None + new_model = self._placed_default_model() + else: + import litellm + + raise litellm.BadRequestError( + message=( + f"Auto-router {self.model_name} received a request with image input, but no model " + f"at or above the decided tier accepts images and modality_routing is enabled. " + f"Tiers checked: {', '.join(names[start:])}. Add a vision-capable model to a tier, " + f"or set a vision-capable default_model, or remove the image content." + ), + model=self.model_name, + llm_provider="", + ) + self._restamp_adaptive_choice(request_kwargs, response.model, new_model) + same_tier: Final = capable is not None and decided == capable + base_cause: Final = (decision.get("cause") if decision is not None else None) or "default_fallback" + displaced_default: Final = decided is None and response.model == self.config.default_model + markers: Final = ( + "modality:image", + *((f"modality_escalated_from:{decided}",) if not same_tier and isinstance(decided, str) else ()), + *(("modality_displaced_default_model",) if not same_tier and displaced_default else ()), + ) + old_signals: Final = tuple(decision.get("signals") or ()) if decision is not None else () + new_decision: Final = self._build_routing_decision( + routed_model=new_model, + cause=base_cause if same_tier else "modality_escalation", + tier=new_tier, + score=decision.get("score") if decision is not None else None, + signals=(*old_signals, *markers), + matched_keyword=decision.get("matched_keyword") if decision is not None else None, + escalation_keyword=decision.get("escalation_keyword") if decision is not None else None, + escalated=bool(decision.get("escalated", False)) if decision is not None else False, + classifier_model=decision.get("classifier_model") if decision is not None else None, + classifier_cost=decision.get("classifier_cost") if decision is not None else None, + conversation_continuing=bool(decision.get("conversation_continuing", True)) + if decision is not None + else True, + tier_litellm_params=self._litellm_params_for_model(new_tier, new_model), + context_escalation_original_tier=( + decision.get("context_escalation_original_tier") if decision is not None else None + ), + ) + from litellm.types.router import PreRoutingHookResponse as HookResponse + + return HookResponse( + model=new_model, + messages=response.messages, + litellm_params=self._litellm_params_for_model(new_tier, new_model), + routing_decision=new_decision, + ) + + def _modality_default_model_usable( + self, + request_kwargs: Mapping[str, object], + resolved_messages: Sequence[Mapping[str, object]] | None, + eligible: frozenset[str], + ) -> bool: + """default_model may serve a gated request only when it is configured, plugin-free + (it is never checked against the plugin pipeline), capability-eligible, and the turn + carries no plan-mode sentinel. The sentinel is re-detected here rather than read off + the decision record, because the record only marks turns the floor RAISED; a sentinel + turn already at or above the floor keeps its ordinary cause, and default_model carries + no tier the floor could vouch for on any sentinel turn.""" + return ( + bool(self.config.default_model) + and not self.config.plugins + and self.config.default_model in eligible + and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None + ) + + def _placed_default_model(self) -> str: + """The default_model behind a usable-default verdict; the raise is the type-level + proof, not a reachable path.""" + model: Final = self.config.default_model + if model is None: + raise ValueError(f"Auto-router {self.model_name}: modality gate routed to an unset default_model") + return model + + @staticmethod + def _restamp_adaptive_choice(request_kwargs: Mapping[str, object], old_model: str, new_model: str) -> None: + """The adaptive feedback loop reads its chosen-model marker from request metadata; a + gate rewrite must move the marker with the model or rewards land on the displaced one.""" + metadata: Final = request_kwargs.get("metadata") + if isinstance(metadata, dict) and metadata.get("adaptive_router_chosen_model") == old_model: + metadata["adaptive_router_chosen_model"] = new_model + def _lexical_tier_override(self, user_message: str) -> KeywordOverride | None: """When keyword_tier_rules match literally, the most-severe matched tier wins. @@ -2655,25 +2830,30 @@ class ComplexityRouter(CustomLogger): session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) has_original_messages: Final = messages is not None and len(messages) > 0 return self._with_session_deployment_affinity( - PreRoutingHookResponse( - model=routed_model, - messages=messages if has_original_messages else None, - litellm_params=session_tier_litellm_params, - routing_decision=self._build_routing_decision( - routed_model=routed_model, - cause=cause, - tier=routed_pin_tier, - matched_keyword=pin_plan_sentinel if plan_floored else None, - escalation_keyword=pin_escalation_keyword, - escalated=escalated, - conversation_continuing=conversation_continuing, - tier_litellm_params=session_tier_litellm_params, - context_escalation_original_tier=pin_context_original_tier, + await self._gate_response_modality( + PreRoutingHookResponse( + model=routed_model, + messages=messages if has_original_messages else None, + litellm_params=session_tier_litellm_params, + routing_decision=self._build_routing_decision( + routed_model=routed_model, + cause=cause, + tier=routed_pin_tier, + matched_keyword=pin_plan_sentinel if plan_floored else None, + escalation_keyword=pin_escalation_keyword, + escalated=escalated, + conversation_continuing=conversation_continuing, + tier_litellm_params=session_tier_litellm_params, + context_escalation_original_tier=pin_context_original_tier, + ), ), + messages, + resolved_messages, + request_kwargs, ) ) - response: Final = await self._classify_and_route( + routed_response: Final = await self._classify_and_route( model=model, request_kwargs=request_kwargs, messages=messages, @@ -2682,6 +2862,11 @@ class ComplexityRouter(CustomLogger): conversation_continuing=conversation_continuing, resolved_messages=resolved_messages, ) + response: Final = ( + await self._gate_response_modality(routed_response, messages, resolved_messages, request_kwargs) + if routed_response is not None + else None + ) # Sentinel presence, not the plan_mode cause, gates the pin write: a plan-mode turn # classified at or above the floor keeps its ordinary cause, yet on an adaptive router # the hard floor constrained its pick, so pinning it would carry a plan-mode-shaped diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 3c5e8aafa18..70aeecb31c6 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -848,6 +848,18 @@ class ComplexityRouterConfig(BaseModel): "drift plus the response tokens." ), ) + modality_routing: bool = Field( + default=False, + description=( + "Route image-bearing requests only to models that can accept image input. The " + "classifier reads text alone, so an image request whose text classifies cheap " + "otherwise lands on a text-only model and fails with a provider 400. When enabled, " + "a routed model explicitly declared supports_vision false (deployment model_info " + "or the model cost map; unmapped names stay routable) is replaced by the nearest " + "HIGHER tier holding a capable model, then default_model, else a clear 400. A kept " + "session-affinity pin still wins even when an image arrives." + ), + ) # Semantic (embedding) matching for keyword_tier_rules instead of literal text matching semantic_keyword_matching: bool = Field( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a1b3523442b..55a32989b1c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2840,6 +2840,10 @@ RoutingDecisionCause = Literal[ # never called. The matched sentinel rides in matched_keyword. Distinct from the keyword causes, # which are operator-authored rules; these sentinels ship with the router. "housekeeping", + # modality_routing replaced the decided placement: the request carries an image and the + # routed model does not accept image input, so the nearest higher capable tier or + # default_model served instead. The displaced placement rides in signals. + "modality_escalation", "session_affinity_pin", "session_affinity_escalation", # classification_mode 'user_turn': the request is an agent loop's continuation turn (no new diff --git a/litellm/utils.py b/litellm/utils.py index d8b19a406e6..f5adca8f272 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2660,10 +2660,19 @@ def _is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, ``_supports_factory`` so caching, fallback, and normalisation improvements apply here automatically. """ + from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider + try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider - ) + declared: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared is not None: + model = model.removeprefix( + f"{declared}/" + ) # rebind-ok: mirrors get_llm_provider's split without its OAuth flow + custom_llm_provider = declared # rebind-ok: same + else: + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, custom_llm_provider=custom_llm_provider + ) model_info: Final = _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) val: Final = model_info.get(key) if val is False: @@ -2751,6 +2760,15 @@ def supports_computer_use(model: str, custom_llm_provider: str | None = None) -> ) +def is_vision_explicitly_disabled(model: str, custom_llm_provider: str | None = None) -> bool: + """True only when supports_vision is explicitly declared false for the model. + + The opt-out mirror of :func:`supports_vision`: a missing declaration reads as not + disabled, so unknown or newly added models stay eligible for image routing. + """ + return _is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision") + + def supports_vision(model: str, custom_llm_provider: str | None = None) -> bool: """ Check if the given model supports vision and return a boolean value. diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 772fbf98c57..9ab66d55f3f 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1433,3 +1433,50 @@ class TestFlattenTopLevelSchemaCombinators: flatten_top_level_schema_combinators(schema) assert schema == snapshot + + +class TestRequestContainsImageContent: + """One detector for every dialect that reaches pre-routing hooks untranslated.""" + + @pytest.mark.parametrize( + "part", + [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}}, + {"type": "input_image", "image_url": "data:image/png;base64,aGk="}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aGk="}}, + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [{"type": "image", "source": {"type": "base64", "data": "aGk="}}], + }, + ], + ) + def test_detects_every_image_dialect_including_tool_results(self, part): + from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content + + messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}, part]}] + assert request_contains_image_content(messages) is True + + @pytest.mark.parametrize( + "messages", + [ + [{"role": "user", "content": "plain string"}], + [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + [{"role": "user", "content": [{"type": "input_audio", "input_audio": {"data": "x"}}]}], + [{"role": "user", "content": [{"type": "tool_result", "content": [{"type": "text", "text": "ok"}]}]}], + [{"role": "user", "content": None}], + [], + ], + ) + def test_ignores_text_audio_and_degenerate_shapes(self, messages): + from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content + + assert request_contains_image_content(messages) is False + + def test_hostile_nesting_is_depth_bounded(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content + + nested: dict = {"type": "image", "source": {"type": "base64", "data": "aGk="}} + for _ in range(50): + nested = {"type": "tool_result", "content": [nested]} + assert request_contains_image_content([{"role": "user", "content": [nested]}]) is False diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 3f7844cffba..1ec8be88c9b 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10417,3 +10417,323 @@ class TestContextWindowEscalation: assert oversized["model_name"] == "big-model" assert small["model_name"] == "small-model" + + +IMG_PART = {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}} +PLAN_BODY = { + "messages": [{"role": "system", "content": [{"type": "text", "text": "Plan mode is active. Do not execute."}]}] +} + + +class TestModalityRouting: + """modality_routing: the response gate replaces a routed model that cannot take images.""" + + IMAGE_MESSAGE = [{"role": "user", "content": [{"type": "text", "text": "What color is this?"}, IMG_PART]}] + BASE_TIERS = {"SIMPLE": "text-cheap", "MEDIUM": "vision-mid", "COMPLEX": "vision-big"} + BASE_VISION = {"text-cheap": False, "vision-mid": True, "vision-big": True, "vision-default": True} + + @staticmethod + def _router(mock_router_instance, config, vision_by_model): + """vision_by_model: model name -> True/False (deployment model_info) or None (undeclared).""" + + def get_model_list(model_name=None): + if model_name not in vision_by_model: + return [] + declared = vision_by_model[model_name] + return [ + { + "model_name": model_name, + "litellm_params": {"model": f"openai/unmapped-{model_name}"}, + "model_info": {} if declared is None else {"supports_vision": declared}, + } + ] + + mock_router_instance.get_model_list = get_model_list + return ComplexityRouter( + model_name="modality-test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "config_extra, vision, send_image, expected_model, expect_marker", + [ + ({}, {"text-cheap": False}, True, "text-cheap", False), + ({"modality_routing": True}, {"text-cheap": False}, False, "text-cheap", False), + ({"modality_routing": True}, {"text-cheap": None}, True, "text-cheap", False), + ], + ids=["flag_off", "no_image", "undeclared_model_stays_routable"], + ) + async def test_gate_leaves_ungated_requests_untouched( + self, mock_router_instance, config_extra, vision, send_image, expected_model, expect_marker + ): + router = self._router(mock_router_instance, {"tiers": dict(self.BASE_TIERS), **config_extra}, vision) + request = self.IMAGE_MESSAGE if send_image else [{"role": "user", "content": "What color is the sky?"}] + result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=request) + assert result.model == expected_model + assert result.routing_decision["cause"] == "heuristic_scorer" + assert ("modality:image" in (result.routing_decision.get("signals") or ())) is expect_marker + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "part", + [ + IMG_PART, + {"type": "input_image", "image_url": "data:image/png;base64,aGk="}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aGk="}}, + {"type": "tool_result", "tool_use_id": "tu_1", "content": [dict(IMG_PART, type="image")]}, + ], + ids=["image_url", "input_image", "anthropic_image", "tool_result_nested"], + ) + async def test_every_image_dialect_escalates(self, mock_router_instance, part): + router = self._router( + mock_router_instance, {"tiers": dict(self.BASE_TIERS), "modality_routing": True}, dict(self.BASE_VISION) + ) + message = [{"role": "user", "content": [{"type": "text", "text": "What color is this?"}, part]}] + result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=message) + assert result.model == "vision-mid" + assert result.routing_decision["cause"] == "modality_escalation" + assert "modality_escalated_from:SIMPLE" in result.routing_decision["signals"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "path, expected_model, expected_cause", + [ + ("classifier_escalates", "vision-mid", "modality_escalation"), + ("same_tier_repick_keeps_cause", "vision-cheap", "heuristic_scorer"), + ("keyword_tier_escalates", "vision-mid", "modality_escalation"), + ("no_ask_capable_default_kept", "vision-default", "default_fallback"), + ("no_ask_text_default_displaced", "vision-mid", "modality_escalation"), + ("custom_tiers_walk", "premium-model", "modality_escalation"), + ("pin_kept_bypasses", "text-cheap", "session_affinity_pin"), + ("pin_replacement_gated", "vision-big", "modality_escalation"), + ("adaptive_pick_rewritten", "vision-mid", "modality_escalation"), + ], + ) + async def test_placements_across_decision_paths(self, mock_router_instance, path, expected_model, expected_cause): + config = {"tiers": dict(self.BASE_TIERS), "modality_routing": True} + vision = dict(self.BASE_VISION) + request_kwargs = {} + messages = self.IMAGE_MESSAGE + if path == "same_tier_repick_keeps_cause": + config["tiers"]["SIMPLE"] = ["text-cheap", "vision-cheap"] + vision["vision-cheap"] = True + with patch( # test-quality-ok: the mixed-pool repick is unreachable deterministically without pinning the first random pick + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + side_effect=lambda pool: sorted(pool)[0], + ): + router = self._router(mock_router_instance, config, vision) + result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=messages) + assert result.model == expected_model + assert result.routing_decision["cause"] == expected_cause + assert result.routing_decision["signals"][-1] == "modality:image" + return + if path == "keyword_tier_escalates": + config["keyword_tier_rules"] = [{"keywords": ["quick lookup"], "tier": "SIMPLE"}] + messages = [ + {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]} + ] + elif path == "no_ask_capable_default_kept": + config["default_model"] = "vision-default" + messages = [{"role": "user", "content": [IMG_PART]}] + elif path == "no_ask_text_default_displaced": + config["default_model"] = "text-default" + vision["text-default"] = False + messages = [{"role": "user", "content": [IMG_PART]}] + elif path == "custom_tiers_walk": + config = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "fallback_tier": "cheap", + "tier_definitions": [ + {"name": "cheap", "description": "trivial asks"}, + {"name": "premium", "description": "hard asks"}, + ], + "tiers": {"cheap": "cheap-model", "premium": "premium-model"}, + "keyword_tier_rules": [{"keywords": ["quick lookup"], "tier": "cheap"}], + "modality_routing": True, + } + vision = {"cheap-model": False, "premium-model": True} + messages = [ + {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]} + ] + elif path in ("pin_kept_bypasses", "pin_replacement_gated"): + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) + mock_router_instance.cache = cache + config["session_affinity"] = True + request_kwargs = {"metadata": {"session_id": "s1"}} + if path == "pin_replacement_gated": + config["tiers"]["MEDIUM"] = "text-mid" + vision["text-mid"] = False + messages = [ + {"role": "user", "content": [{"type": "text", "text": "LITELLM ESCALATE describe this"}, IMG_PART]} + ] + elif path == "adaptive_pick_rewritten": + config["adaptive"] = True + mock_router_instance.model_list = [] + mock_router_instance.model_name_to_deployment_indices = {} + router = self._router(mock_router_instance, config, vision) + result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages) + assert result.model == expected_model + assert result.routing_decision["cause"] == expected_cause + if path == "adaptive_pick_rewritten": + assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == expected_model + + @pytest.mark.asyncio + async def test_plan_floored_decision_never_falls_to_default_model(self, mock_router_instance): + """An upward-only walk cannot undercut the floor; default_model must not either.""" + config = { + "tiers": {"SIMPLE": "vision-cheap", "MEDIUM": "text-mid"}, + "default_model": "vision-default", + "plan_mode_min_tier": "MEDIUM", + "modality_routing": True, + } + vision = {"vision-cheap": True, "text-mid": False, "vision-default": True} + router = self._router(mock_router_instance, config, vision) + with pytest.raises(litellm.BadRequestError, match="no model"): + await router.async_pre_routing_hook( + model="m", + request_kwargs={"proxy_server_request": {"body": PLAN_BODY}}, + messages=[{"role": "user", "content": [{"type": "text", "text": "plan this"}, IMG_PART]}], + ) + + @pytest.mark.asyncio + async def test_at_floor_plan_turn_never_falls_to_default_model(self, mock_router_instance): + """A sentinel turn whose classified tier already satisfies the floor keeps its ordinary + cause, so the record carries no floor marker; the default arm must still refuse it.""" + config = { + "tiers": {"SIMPLE": "text-a", "MEDIUM": "text-b"}, + "default_model": "vision-default", + "plan_mode_min_tier": "SIMPLE", + "modality_routing": True, + } + vision = {"text-a": False, "text-b": False, "vision-default": True} + router = self._router(mock_router_instance, config, vision) + with pytest.raises(litellm.BadRequestError, match="no model"): + await router.async_pre_routing_hook( + model="m", + request_kwargs={"proxy_server_request": {"body": PLAN_BODY}}, + messages=[{"role": "user", "content": [{"type": "text", "text": "plan this"}, IMG_PART]}], + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "default_model, default_vision, expect_error", + [(None, None, True), ("text-default", False, True), ("vision-default", True, False)], + ids=["no_default", "text_only_default", "vision_default_serves"], + ) + async def test_no_capable_tier_above_uses_default_or_rejects( + self, mock_router_instance, default_model, default_vision, expect_error + ): + config = {"tiers": {"SIMPLE": "text-cheap", "COMPLEX": "text-big"}, "modality_routing": True} + vision = {"text-cheap": False, "text-big": False} + if default_model is not None: + config["default_model"] = default_model + vision[default_model] = default_vision + router = self._router(mock_router_instance, config, vision) + if expect_error: + with pytest.raises(litellm.BadRequestError, match="no model"): + await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE) + return + result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE) + assert result.model == "vision-default" + assert result.routing_decision["cause"] == "modality_escalation" + assert "modality_escalated_from:SIMPLE" in result.routing_decision["signals"] + + @pytest.mark.asyncio + async def test_mixed_deployment_group_is_treated_text_only(self, mock_router_instance): + def get_model_list(model_name=None): + declared = {"mixed-group": [True, False], "vision-big": [True]}.get(model_name) + if declared is None: + return [] + return [ + { + "model_name": model_name, + "litellm_params": {"model": f"openai/unmapped-{model_name}-{i}"}, + "model_info": {"supports_vision": accepts}, + } + for i, accepts in enumerate(declared) + ] + + mock_router_instance.get_model_list = get_model_list + router = ComplexityRouter( + model_name="modality-test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "mixed-group", "COMPLEX": "vision-big"}, + "modality_routing": True, + }, + ) + result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE) + assert result.model == "vision-big" + assert result.routing_decision["cause"] == "modality_escalation" + + @pytest.mark.asyncio + async def test_continuation_turn_screenshot_escalates_past_the_held_model(self, mock_router_instance): + """classification_mode user_turn replays the held model on continuation turns; a + continuation carrying a screenshot must still be re-placed when that model is text-only.""" + mock_router_instance.cache = DualCache() + config = { + "tiers": dict(self.BASE_TIERS), + "classification_mode": "user_turn", + "modality_routing": True, + } + router = self._router(mock_router_instance, config, dict(self.BASE_VISION)) + first = await router.async_pre_routing_hook( + model="m", + request_kwargs={"metadata": {"session_id": "cont-1"}}, + messages=[{"role": "user", "content": "hi there"}], + ) + assert first.model == "text-cheap" + continuation = [ + {"role": "user", "content": "hi there"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "tu_1", "name": "screenshot", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [{"type": "image", "source": {"type": "base64", "data": "aGk="}}], + } + ], + }, + ] + second = await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": "cont-1"}}, messages=continuation + ) + assert second.model == "vision-mid" + assert second.routing_decision["cause"] == "modality_escalation" + assert "modality_escalated_from:SIMPLE" in second.routing_decision["signals"] + + @pytest.mark.asyncio + async def test_rewrite_carries_the_context_escalation_record(self, mock_router_instance): + """A context-window escalation and a modality re-place are separate facts on one + record; rewriting for the image must not drop the sibling gate's fields.""" + from litellm.types.router import PreRoutingHookResponse + + router = self._router( + mock_router_instance, + {"tiers": dict(self.BASE_TIERS), "modality_routing": True}, + dict(self.BASE_VISION), + ) + decision = router._build_routing_decision( + routed_model="text-cheap", + cause="heuristic_scorer", + tier=ComplexityTier.SIMPLE, + context_escalation_original_tier=ComplexityTier.SIMPLE, + ) + response = PreRoutingHookResponse(model="text-cheap", messages=None, routing_decision=decision) + rewritten = await router._gate_response_modality(response, None, self.IMAGE_MESSAGE, {}) + assert rewritten.model == "vision-mid" + assert rewritten.routing_decision["cause"] == "modality_escalation" + assert rewritten.routing_decision["context_escalated"] is True + assert rewritten.routing_decision["context_escalation_original_tier"] == "SIMPLE" + + def test_modality_escalation_is_never_pinnable(self): + from litellm.router_strategy.complexity_router.complexity_router import _decision_is_pinnable + + assert _decision_is_pinnable({"cause": "modality_escalation"}) is False + assert _decision_is_pinnable({"cause": "heuristic_scorer"}) is True diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 6524353aa48..a7fa6ac5568 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -5765,3 +5765,33 @@ class TestHuggingFaceConfigFetch: assert _get_max_position_embeddings("some-org/some-model") == 512 request_timeout = hf_config_route.calls.last.request.extensions["timeout"] assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS + + +class TestIsVisionExplicitlyDisabled: + """github_copilot and chatgpt run an OAuth device flow inside get_llm_provider; the + explicit-disable lookup must adopt the declared prefix instead of resolving it, exactly + as _supports_factory does, or a capability check on a copilot deployment blocks routing + on a device-code prompt.""" + + @pytest.mark.parametrize("model", ["github_copilot/gpt-4o", "chatgpt/gpt-5"]) + def test_never_resolves_an_authenticating_prefix(self, model, monkeypatch): + from litellm.utils import is_vision_explicitly_disabled + + lookups: list = [] + + def _record(*args, **kwargs): + lookups.append((args, kwargs)) + raise RuntimeError("provider resolution must not run for an authenticating provider") + + monkeypatch.setattr(litellm, "get_llm_provider", _record) + + assert is_vision_explicitly_disabled(model) is False + assert lookups == [] + + def test_explicit_false_detected_and_absent_reads_enabled(self): + from litellm.utils import is_vision_explicitly_disabled + + assert ( + is_vision_explicitly_disabled("fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731") is True + ) + assert is_vision_explicitly_disabled("anthropic/claude-sonnet-4-5") is False diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx index e084fdf37e6..bf6a20a4af8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx @@ -186,6 +186,12 @@ describe("RoutingDecisionCard", () => { expect(screen.queryByText("housekeeping")).not.toBeInTheDocument(); }); + it("labels a modality escalation instead of showing the raw cause token", () => { + render(); + expect(screen.getByText("Escalated for image input")).toBeInTheDocument(); + expect(screen.queryByText("modality_escalation")).not.toBeInTheDocument(); + }); + it("shows the escalation keyword", () => { render( , diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index d2aa20901f5..aa1d45a859e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -90,6 +90,7 @@ const CONSTANT_CAUSE_LABELS: Record = { session_affinity_pin: "Pinned to session", session_affinity_escalation: "Escalated from session pin", user_turn_continuation: "Continuation turn, classifier skipped", + modality_escalation: "Escalated for image input", quality_tier: "Quality tier mapping", bandit: "Adaptive bandit", default_fallback: "Default model, no route matched", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c20545f6fb2..fb81191cfbf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34463,6 +34463,12 @@ export interface components { * @default 0.5 */ match_threshold: number; + /** + * Modality Routing + * @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, a routed model explicitly declared supports_vision false (deployment model_info or the model cost map; unmapped names stay routable) is replaced by the nearest HIGHER tier holding a capable model, then default_model, else a clear 400. A kept session-affinity pin still wins even when an image arrives. + * @default false + */ + modality_routing: boolean; /** * Plan Mode Min Tier * @description When set, requests carrying a coding-agent plan-mode sentinel (Claude Code plan mode, VS Code Copilot Plan mode, Copilot CLI's exit_plan_mode tool) are routed to at least this tier: the classified tier still wins when it is higher, and the floor also overrides a session-affinity pin to a lower tier for exactly the turns carrying the sentinel, without rewriting the pin -- the first turn after plan mode exits routes as if plan mode had never happened. Names a built-in tier, or with tier_definitions set, one of the defined tier names (list order is ascending severity, same as keyword_tier_rules). Unset disables detection entirely. The sentinels ride in client-injected prompt text, so a caller who pastes one can spend up to this tier's models -- never down, and never outside the configured pools. @@ -35620,7 +35626,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */ From 0565d33fa50b849c27edace7a413aec5d2611ed6 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 19:55:16 -0700 Subject: [PATCH 066/126] fix(ui): let the auto-router scoring tier list follow the theme (#39040) * fix(ui): let the auto-router scoring tier list follow the theme * test(ui): assert the tier list carries the muted-foreground token --- .../components/add_model/ClassificationMethodConfig.tsx | 2 +- .../components/add_model/ComplexityRouterConfig.test.tsx | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 2947e29319b..84f61c95eca 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -111,7 +111,7 @@ const HowClassificationWorks: React.FC<{ value: ComplexityRouterConfigValue }> = How Classification Works {scoringExplanation(value)} {scorerRuns && ranges && ( -
    +
    • {effectiveTierLabel("SIMPLE", value.tier_labels)}: Score < {ranges.simpleMedium}
      • diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 33ce1169c46..218c99e32c5 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -80,6 +80,15 @@ describe("ComplexityRouterConfig", () => { expect(screen.getByText(/Score > 0.60/)).toBeInTheDocument(); }); + it("leaves the score threshold list color to the theme instead of an inline style", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + const list = screen.getByText(/Score < 0.15/).closest("ul"); + expect(list).toBeInTheDocument(); + expect(list).toHaveClass("text-muted-foreground"); + expect(list?.style.color).toBe(""); + }); + it("should default to heuristic and hide classifier model/timeout fields", () => { renderWithProviders(); expect(screen.getByText("Advanced: Classification Method")).toBeInTheDocument(); From 502b3a2f794a6e806835e41ffa8e5fc951a7fda9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 20:02:08 -0700 Subject: [PATCH 067/126] feat(ui): auto-router controls for context-window escalation (#39054) * feat(ui): auto-router controls for context-window escalation Adds an Advanced: Context Window Escalation section to the auto-router form, both create and edit arms, with the toggle for enable_context_window_escalation and a clamped decimal input for context_window_escalation_buffer. An untouched control keeps both keys out of the payload so the router tracks the backend defaults; an explicit opt-out (false) survives the edit round-trip through the managed-keys projection and the hydrator, and preset prefill maps both keys straight through so a preset cannot silently drop them Resolves LIT-6601 * fix(ui): clearing the context-window buffer removes it from the payload Both review bots converged on the same defect: an emptied buffer field early-returned in commitBuffer, the draft was discarded on blur, and the stale number reappeared and stayed in the saved config, contradicting the copy that an empty field tracks the backend default. An empty commit now removes the key, which the managed-keys projection propagates as a real deletion on edit. Also trims the narrative comments the review flagged as restating behavior --- .../add_model/ComplexityRouterConfig.tsx | 13 ++++ .../ContextWindowEscalationConfig.tsx | 60 +++++++++++++++++ .../add_model/add_auto_router_tab.test.tsx | 65 +++++++++++++++++++ .../add_model/add_auto_router_tab.tsx | 2 + .../build_complexity_router_config.test.ts | 10 +++ .../build_complexity_router_config.ts | 12 ++++ ...d_updated_complexity_router_config.test.ts | 2 + .../edit_auto_router_modal.tsx | 14 ++++ .../src/lib/autorouter_presets.test.ts | 16 +++++ .../src/lib/autorouter_presets.ts | 2 + 10 files changed, 196 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 153afa0b586..111a7c9f10a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -29,6 +29,7 @@ import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; +import ContextWindowEscalationConfig from "./ContextWindowEscalationConfig"; import { Restricted, restrictedBy } from "./TierRestrictions"; import { type TierSetAction, applyTierSetAction, setFallbackTier } from "./tier_set_actions"; import { @@ -392,6 +393,13 @@ export interface ComplexityRouterConfigValue { tier_distance_penalty?: number; adaptive_eligible?: AdaptiveEligible; return_raw_model_name?: boolean; + /** + * Context-window escalation gate. Undefined means untouched, which keeps both keys out of the + * payload so the router tracks the backend defaults (enabled, 0.95 buffer); an explicit false + * is a real opt-out and must survive the edit round-trip. + */ + enable_context_window_escalation?: boolean; + context_window_escalation_buffer?: number; /** * Heuristic scorer knobs. Undefined means the operator never touched them, which keeps the key out of the * payload so the router tracks the backend defaults rather than freezing today's numbers. @@ -827,6 +835,11 @@ const ComplexityRouterConfig: React.FC = ({ ), }, + { + key: "context-window", + label: Advanced: Context Window Escalation, + children: , + }, { key: "response", label: Advanced: Response Format, diff --git a/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx new file mode 100644 index 00000000000..c0a65076d20 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx @@ -0,0 +1,60 @@ +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import React from "react"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; + +const ContextWindowEscalationConfig: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => { + const enabled = value.enable_context_window_escalation ?? true; + // A number input renders Number("0.") as "0", so a decimal cannot be typed without a local draft. + const [bufferDraft, setBufferDraft] = React.useState(null); + const commitBuffer = (raw: string) => { + setBufferDraft(null); + if (raw.trim() === "") { + onChange({ ...value, context_window_escalation_buffer: undefined }); + return; + } + const parsed = Number(raw); + if (!Number.isFinite(parsed)) return; + onChange({ ...value, context_window_escalation_buffer: Math.min(1, Math.max(0.01, parsed)) }); + }; + return ( + <> +
      + onChange({ ...value, enable_context_window_escalation: next })} + aria-label="Escalate oversized prompts to a tier that fits" + /> + Escalate oversized prompts to a tier that fits +
      + + When a prompt provably cannot fit the decided tier's context windows, route it to the lowest tier whose + window holds it instead of letting the provider reject it. Off means requests dispatch on complexity alone. + + {enabled && ( +
      + + setBufferDraft(event.target.value)} + onBlur={(event) => commitBuffer(event.target.value)} + /> + + Fraction of a model's window the counted prompt must fit within, above 0 up to 1. Empty tracks the + backend default of 0.95. + +
      + )} + + ); +}; + +export default ContextWindowEscalationConfig; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 01cb41bcb95..71b454dbb06 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -373,6 +373,71 @@ describe("AddAutoRouterTab", () => { }); }); + it("carries a context-window escalation opt-out through to the create payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-window-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Context Window Escalation")); + const toggle = await screen.findByRole("switch", { name: "Escalate oversized prompts to a tier that fits" }); + expect(toggle).toBeChecked(); + await user.click(toggle); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + enable_context_window_escalation: false, + }); + }); + + it("clamps the context-window buffer to 1 and keeps an untouched buffer out of the payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-buffer-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Context Window Escalation")); + const buffer = await screen.findByLabelText("Window fit buffer"); + fireEvent.change(buffer, { target: { value: "1.5" } }); + fireEvent.blur(buffer, { target: { value: "1.5" } }); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const config = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config; + expect(config).toMatchObject({ context_window_escalation_buffer: 1 }); + expect(config).not.toHaveProperty("enable_context_window_escalation"); + }); + + it("clearing the buffer removes it from the payload so the router tracks the backend default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-clear-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Context Window Escalation")); + const buffer = await screen.findByLabelText("Window fit buffer"); + fireEvent.change(buffer, { target: { value: "0.8" } }); + fireEvent.blur(buffer, { target: { value: "0.8" } }); + fireEvent.change(buffer, { target: { value: "" } }); + fireEvent.blur(buffer, { target: { value: "" } }); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).not.toHaveProperty( + "context_window_escalation_buffer", + ); + }); + // The scalar floor is the one scorer knob with no group dict behind it, so its wiring into the create // payload is only proven end to end. 0 is the case a truthy check would silently drop. it("carries a reasoning override floor of 0 through to the create payload", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 4e5e5e8d460..ed584a4882b 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -367,6 +367,8 @@ const AddAutoRouterTab: React.FC = ({ tokenThresholds: complexityRouterConfig.token_thresholds, dimensionWeights: complexityRouterConfig.dimension_weights, reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score, + enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation, + contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer, }; const submitRecommendedRouter = async (name: string) => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 84406a2093e..feddaaa0eac 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -59,6 +59,16 @@ describe("buildComplexityRouterConfig", () => { expect(config).toEqual(expected); }); + it("carries an explicit context-window escalation opt-out and buffer, false included", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + enableContextWindowEscalation: false, + contextWindowEscalationBuffer: 0.9, + }); + expect(config.enable_context_window_escalation).toBe(false); + expect(config.context_window_escalation_buffer).toBe(0.9); + }); + it("trims escalation keywords and drops blank entries", () => { const config = buildComplexityRouterConfig({ ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index ba500d116ce..9a14e956207 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -123,6 +123,8 @@ export interface BuildComplexityRouterConfigParams { dimensionWeights?: DimensionWeights; reasoningOverrideMinScore?: number; tierModelParams?: TierModelParamsByTier; + enableContextWindowEscalation?: boolean; + contextWindowEscalationBuffer?: number; } /** @@ -174,6 +176,8 @@ export interface ComplexityRouterConfigPayload { token_thresholds?: TokenThresholds; dimension_weights?: DimensionWeights; reasoning_override_min_score?: number; + enable_context_window_escalation?: boolean; + context_window_escalation_buffer?: number; tier_model_configs?: Record; } @@ -407,6 +411,8 @@ export const buildComplexityRouterConfig = ({ dimensionWeights, reasoningOverrideMinScore, tierModelParams, + enableContextWindowEscalation, + contextWindowEscalationBuffer, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const serializedTierModelConfigs = customTierSet ? serializeTierModelConfigs( @@ -463,6 +469,12 @@ export const buildComplexityRouterConfig = ({ adaptive_eligible: adaptiveEligible, }), ...(returnRawModelName && { return_raw_model_name: true }), + ...(enableContextWindowEscalation !== undefined && { + enable_context_window_escalation: enableContextWindowEscalation, + }), + ...(contextWindowEscalationBuffer !== undefined && { + context_window_escalation_buffer: contextWindowEscalationBuffer, + }), ...scorerKnobs, }; if (!customTierSet) return payload; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index a7bd4b8eab4..d8c2987ead5 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -487,6 +487,8 @@ describe("managed keys survive an untouched open-and-save", () => { token_thresholds: { simple: 20, complex: 500 }, dimension_weights: { tokenCount: 0.1 }, reasoning_override_min_score: 0.3, + enable_context_window_escalation: false, + context_window_escalation_buffer: 0.9, }; // tier_definitions, fallback_tier and classification_prompt cannot sit beside heuristic_first, which diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 425d5d51f06..4751c2e64b7 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -107,6 +107,8 @@ export interface StoredComplexityRouterConfig { tier_distance_penalty?: number; adaptive_eligible?: AdaptiveEligible; return_raw_model_name?: boolean; + enable_context_window_escalation?: unknown; + context_window_escalation_buffer?: unknown; } /** @@ -178,6 +180,14 @@ export const hydrateComplexityRouterConfig = ( tier_distance_penalty: parsedConfig.tier_distance_penalty, adaptive_eligible: parsedConfig.adaptive_eligible || "all", return_raw_model_name: parsedConfig.return_raw_model_name || false, + enable_context_window_escalation: + typeof parsedConfig.enable_context_window_escalation === "boolean" + ? parsedConfig.enable_context_window_escalation + : undefined, + context_window_escalation_buffer: + typeof parsedConfig.context_window_escalation_buffer === "number" + ? parsedConfig.context_window_escalation_buffer + : undefined, }; }; @@ -208,6 +218,8 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "token_thresholds", "dimension_weights", "reasoning_override_min_score", + "enable_context_window_escalation", + "context_window_escalation_buffer", ]); // Managed only when the caller passes the corresponding state. A caller that does not render @@ -307,6 +319,8 @@ export const buildUpdatedComplexityRouterConfig = ( dimensionWeights: value.dimension_weights, reasoningOverrideMinScore: value.reasoning_override_min_score, tierModelParams: value.tier_model_params, + enableContextWindowEscalation: value.enable_context_window_escalation, + contextWindowEscalationBuffer: value.context_window_escalation_buffer, }; const built = buildComplexityRouterConfig(builderParams); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 2de1ac11db2..f14d6279e32 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -573,6 +573,22 @@ describe("autorouter_presets", () => { expect(prefill.escalationKeywords).toEqual([]); }); + it("carries a preset's context-window escalation opt-out and buffer through the prefill", () => { + const prefill = buildPresetPrefill( + { + tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + classifier_type: "heuristic", + session_affinity: false, + deployment_affinity: true, + enable_context_window_escalation: false, + context_window_escalation_buffer: 0.9, + }, + groupsOnly(["gpt-5-nano"]), + ); + expect(prefill.complexityRouterConfig.enable_context_window_escalation).toBe(false); + expect(prefill.complexityRouterConfig.context_window_escalation_buffer).toBe(0.9); + }); + it("falls back to the defaults when a preset omits match_threshold and escalation_keywords", () => { const prefill = buildPresetPrefill( { diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index a35b868db5e..f96dd5ddb4c 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -295,6 +295,8 @@ export const buildPresetPrefill = ( tier_distance_penalty: config.tier_distance_penalty, adaptive_eligible: config.adaptive_eligible, return_raw_model_name: config.return_raw_model_name, + enable_context_window_escalation: config.enable_context_window_escalation, + context_window_escalation_buffer: config.context_window_escalation_buffer, }, customTechnicalKeywords: config.custom_technical_keywords ?? [], keywordTierRules: hydrateKeywordTierRules(config.keyword_tier_rules ?? []), From d9c43d5e17746b11414c8aa17755b1e11bec16f6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 20:07:28 -0700 Subject: [PATCH 068/126] fix(key management): read budget window usage from the window spend table Pass window_duration to get_current_spend so /key/info re-checks a stale-low counter against the LiteLLM_BudgetWindowSpend row instead of aggregating LiteLLM_SpendLogs, and reuse _budget_limit_windows for the stored-column coercion. Drop the /v2/key/info batch cap (a new 422 for callers that work today) and the unrelated CI timeout bump and soft_budget test --- .../key_management_endpoints.py | 103 +++--------- .../test_key_management_endpoints.py | 156 +++--------------- 2 files changed, 51 insertions(+), 208 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 03abcb2ff98..02bfa5c0ca8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3579,95 +3579,58 @@ async def _build_model_max_budget_usage( ) -def _budget_window_to_dict(window: object) -> Mapping[str, object] | None: - """Coerce a budget_limits entry to a dict; None when the entry is unusable.""" - if isinstance(window, dict): - return window - model_dump: Final = getattr(window, "model_dump", None) - if not callable(model_dump): +def _window_max_budget(window: Mapping[str, object]) -> float | None: + """A window's max_budget as a float; None when absent or unparseable.""" + value: Final = window.get("max_budget") + if not isinstance(value, (int, float, str)): return None try: - dumped: Final = model_dump() - except Exception: # noqa: BLE001 # model_dump implementations can raise arbitrary errors + return float(value) + except ValueError: return None - return dumped if isinstance(dumped, dict) else None - - -def _coerce_budget_limits(budget_limits: object) -> Sequence[object] | None: - """Coerce budget_limits to a sequence of windows, parsing JSON strings; None when unusable.""" - if isinstance(budget_limits, str): - try: - parsed: Final = json.loads(budget_limits) - except (TypeError, ValueError): - return None - return parsed if isinstance(parsed, list) else None - return budget_limits if isinstance(budget_limits, list) else None - - -def _parse_window_max_budget(value: object) -> float | None: - """Coerce a window's max_budget to float; None when absent or unparseable.""" - if isinstance(value, (int, float, str)): - try: - return float(value) - except (TypeError, ValueError): - return None - return None async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: str) -> Mapping[str, object]: """ - Return a copy of a budget window with current-window spend attached. + Copy of a budget window with current-window spend attached. - Per-window spend is not persisted in the DB; it lives in the cross-pod spend - counters (spend:key:{hashed_token}:window:{budget_duration}) that - _virtual_key_multi_budget_check enforces against, so we read the same - counters via get_current_spend. Passing max_budget + window_start makes the - read re-check against the authoritative spend-log aggregate when the counter - is stale-low (e.g. after a Redis flush), same as the enforcement path. + Reads the same cross-pod counter (spend:key:{hashed_token}:window:{budget_duration}) + that _virtual_key_multi_budget_check enforces against, passing the same + window_duration + window_start so a stale-low counter is re-checked against + the LiteLLM_BudgetWindowSpend row instead of a spend-log aggregate. """ from litellm.proxy.proxy_server import get_current_spend duration: Final = window.get("budget_duration") - if not duration: - return dict(window) # mutable-ok: per-window response copy, built once per window + if not isinstance(duration, str) or not duration: + return window spend: Final = await get_current_spend( counter_key=f"spend:key:{api_key_hash}:window:{duration}", fallback_spend=0.0, - max_budget=_parse_window_max_budget(window.get("max_budget")), + max_budget=_window_max_budget(window), window_entity_type="Key", window_entity_id=api_key_hash, + window_duration=duration, window_start=get_budget_window_start(window), ) return {**window, "current_spend": round(spend, 4)} # mutable-ok: per-window response copy, built once per window -async def _budget_limits_entry_with_usage(window: object, api_key_hash: str) -> object: - """Return the window as an enriched dict when dict-coercible; the original entry otherwise.""" - coerced: Final = _budget_window_to_dict(window) - if not coerced: - return window - return await _budget_window_with_usage(window=coerced, api_key_hash=api_key_hash) - - -async def _budget_limits_with_usage(budget_limits: object, api_key_hash: str) -> Sequence[object] | None: +async def _budget_limits_with_usage( + budget_limits: Sequence[object] | str | None, api_key_hash: str +) -> tuple[Mapping[str, object], ...] | None: """ - Return budget_limits as window dicts with current-window spend attached. - - None when budget_limits is not a usable (possibly JSON-encoded) list; the - caller keeps the original value then. Entries that are not dict-coercible - are preserved as-is. + budget_limits as window dicts with current-window spend attached; None when + the key has no windows so the caller keeps the stored value. """ - windows: Final = _coerce_budget_limits(budget_limits) - if windows is None: + windows: Final = _budget_limit_windows(budget_limits) + if not windows: return None - return [ # mutable-ok: entries are awaited, so they cannot be built inside a frozen wrapper - await _budget_limits_entry_with_usage(window=window, api_key_hash=api_key_hash) for window in windows - ] - - -# Caps per-request fan-out: each key with budget windows costs one spend-counter -# read (worst case a SpendLogs aggregation) per window. -MAX_KEY_INFO_KEYS_PER_REQUEST: Final = 100 + return tuple( + await asyncio.gather( + *(_budget_window_with_usage(window=window, api_key_hash=api_key_hash) for window in windows) + ) + ) @router.post( @@ -3712,18 +3675,6 @@ async def info_key_fn_v2( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail={"message": "Malformed request. No keys passed in."}, ) - requested_key_count: Final = len(data.keys or ()) + len(data.key_aliases or ()) - if requested_key_count > MAX_KEY_INFO_KEYS_PER_REQUEST: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ # mutable-ok: one-shot HTTPException payload matching the sibling detail dict above; never mutated after construction - "message": ( - f"Too many keys requested: {requested_key_count}. " - f"At most {MAX_KEY_INFO_KEYS_PER_REQUEST} keys and key_aliases combined per request." - ) - }, - ) - # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query: Final = list(data.keys) if data.keys else [] if data.key_aliases: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a2f248ca55c..425365c8560 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -539,57 +539,6 @@ async def test_key_generation_with_object_permission(monkeypatch): assert key_insert_calls[0]["data"].get("object_permission_id") == "objperm123" -@pytest.mark.asyncio -async def test_generate_key_with_soft_budget_creates_budget_row(monkeypatch): - """soft_budget on /key/generate must create a budget table row and link its budget_id to the key.""" - mock_prisma_client = AsyncMock() - mock_prisma_client.jsonify_object = lambda data: data - mock_prisma_client.db = MagicMock() - mock_budget_create = AsyncMock(return_value=MagicMock(budget_id="budget-soft-123")) - mock_prisma_client.db.litellm_budgettable = MagicMock() - mock_prisma_client.db.litellm_budgettable.create = mock_budget_create - - async def _insert_data_side_effect(*args, **kwargs): - if kwargs.get("table_name") == "user": - return MagicMock(models=[], spend=0) - return MagicMock( - token="hashed_token_soft", - litellm_budget_table=None, - object_permission=None, - ) - - mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles - from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.proxy.management_endpoints.key_management_endpoints import ( - generate_key_fn, - ) - - await generate_key_fn( - data=GenerateKeyRequest(soft_budget=5.0), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-1234", - user_id="admin-1", - ), - ) - - mock_budget_create.assert_awaited_once() - created_budget = mock_budget_create.call_args.kwargs["data"] - assert created_budget["soft_budget"] == 5.0 - assert created_budget["created_by"] == "admin-1" - - key_insert_calls = [ - call.kwargs - for call in mock_prisma_client.insert_data.call_args_list - if call.kwargs.get("table_name") == "key" - ] - assert len(key_insert_calls) == 1 - assert key_insert_calls[0]["data"].get("budget_id") == "budget-soft-123" - - @pytest.mark.asyncio async def test_generate_key_debug_log_never_contains_raw_token(monkeypatch, caplog): """Regression for LIT-4356: /key/generate must never emit the raw virtual key @@ -14267,6 +14216,7 @@ async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): assert call_kwargs["max_budget"] == 2.0 assert call_kwargs["window_entity_type"] == "Key" assert call_kwargs["window_entity_id"] == test_key_token + assert call_kwargs["window_duration"] == "1h" assert call_kwargs["window_start"] is not None @@ -14397,39 +14347,9 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): f"spend:key:{test_key_token}:window:1h", f"spend:key:{test_key_token}:window:1d", } - - -@pytest.mark.asyncio -async def test_info_key_fn_v2_rejects_oversized_batch(monkeypatch): - """/v2/key/info must reject over-cap batches before doing any DB work.""" - from unittest.mock import AsyncMock - - from litellm.proxy._types import KeyRequest, ProxyException - from litellm.proxy.management_endpoints.key_management_endpoints import ( - MAX_KEY_INFO_KEYS_PER_REQUEST, - info_key_fn_v2, - ) - - mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) - - user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin-batch-cap", - ) - - with pytest.raises(ProxyException) as exc_info: - await info_key_fn_v2( - data=KeyRequest( - keys=[f"hash-{i}" for i in range(MAX_KEY_INFO_KEYS_PER_REQUEST)], - key_aliases=["alias-over-cap"], - ), - user_api_key_dict=user_api_key_dict, - ) - - assert exc_info.value.code == "422" - mock_prisma_client.get_data.assert_not_awaited() + assert { + call.kwargs["window_duration"] for call in mock_get_current_spend.await_args_list + } == {"1h", "1d"} @pytest.mark.asyncio @@ -14452,7 +14372,7 @@ async def test_budget_limits_with_usage_json_string_input(monkeypatch): ) result = await _budget_limits_with_usage(budget_limits=raw, api_key_hash="hash-1") - assert result == [ + assert list(result) == [ { "budget_duration": "1h", "max_budget": 2.0, @@ -14464,8 +14384,8 @@ async def test_budget_limits_with_usage_json_string_input(monkeypatch): @pytest.mark.asyncio -async def test_budget_limits_with_usage_skips_unusable_inputs(monkeypatch): - """Invalid JSON strings, non-list values, and malformed windows are skipped.""" +async def test_budget_limits_with_usage_empty_windows_keep_stored_value(monkeypatch): + """A key with no windows (None, [], or "[]") returns None so /key/info keeps the stored value; no spend lookup runs.""" from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -14477,32 +14397,9 @@ async def test_budget_limits_with_usage_skips_unusable_inputs(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - # invalid JSON string and non-list values return None: callers keep the original - assert await _budget_limits_with_usage(budget_limits="{not json", api_key_hash="hash-1") is None - assert await _budget_limits_with_usage(budget_limits={"budget_duration": "1h"}, api_key_hash="hash-1") is None - - # windows that are falsy, missing budget_duration, or not dict-like - windows = [ - {}, - {"max_budget": 2.0}, - {"budget_duration": "1h", "max_budget": "not-a-number"}, - 42, - ] - result = await _budget_limits_with_usage(budget_limits=windows, api_key_hash="hash-1") - - # only the well-formed window (with unparseable max_budget coerced to None) - # triggers a spend lookup - mock_get_current_spend.assert_awaited_once() - call_kwargs = mock_get_current_spend.await_args.kwargs - assert call_kwargs["counter_key"] == "spend:key:hash-1:window:1h" - assert call_kwargs["max_budget"] is None - assert result is not None - assert result[0] == {} - assert result[1] == {"max_budget": 2.0} - assert result[2] == {"budget_duration": "1h", "max_budget": "not-a-number", "current_spend": 0.0} - assert result[3] == 42 - # input is not mutated - assert windows[2] == {"budget_duration": "1h", "max_budget": "not-a-number"} + for stored in (None, [], "[]"): + assert await _budget_limits_with_usage(budget_limits=stored, api_key_hash="hash-1") is None + mock_get_current_spend.assert_not_awaited() @pytest.mark.asyncio @@ -14523,17 +14420,19 @@ async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" ) - assert result == [{"budget_duration": "2d", "current_spend": 0.75}] + assert list(result) == [{"budget_duration": "2d", "current_spend": 0.75}] call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-no-max:window:2d" + assert call_kwargs["window_duration"] == "2d" assert call_kwargs["max_budget"] is None @pytest.mark.asyncio async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): - """Window objects with model_dump() are converted to dicts; failing windows pass through.""" - from unittest.mock import AsyncMock, MagicMock + """BudgetLimitEntry windows (the shape UserAPIKeyAuth carries) are dumped to dicts and annotated.""" + from unittest.mock import AsyncMock + from litellm.models.team import BudgetLimitEntry from litellm.proxy.management_endpoints.key_management_endpoints import ( _budget_limits_with_usage, ) @@ -14543,25 +14442,18 @@ async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - good_window = MagicMock() - good_window.model_dump.return_value = { - "budget_duration": "7d", - "max_budget": 10.0, - "reset_at": None, - } - bad_window = MagicMock() - bad_window.model_dump.side_effect = ValueError("boom") - result = await _budget_limits_with_usage( - budget_limits=[good_window, bad_window], api_key_hash="hash-2" + budget_limits=[BudgetLimitEntry(budget_duration="7d", max_budget=10.0)], + api_key_hash="hash-2", ) - # good window converted to dict and annotated; failing window left as-is - assert result is not None - assert isinstance(result[0], dict) - assert result[0]["current_spend"] == 1.0 - assert result[1] is bad_window - mock_get_current_spend.assert_awaited_once() + assert list(result) == [ + {"budget_duration": "7d", "max_budget": 10.0, "reset_at": None, "current_spend": 1.0} + ] + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == "spend:key:hash-2:window:7d" + assert call_kwargs["window_duration"] == "7d" + assert call_kwargs["max_budget"] == 10.0 @pytest.mark.asyncio From 46d073b26f9b0037bb824384edf851155c9aecd6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 20:48:42 -0700 Subject: [PATCH 069/126] refactor(key): report window spend under budget_limits_usage instead of inlining current_spend budget_limits now comes back exactly as stored on /key/info and /v2/key/info. The per-window usage moves to a sibling budget_limits_usage field keyed by budget_duration (current_spend, budget_limit, reset_at), mirroring model_max_budget_usage, so the stored shape that /key/update accepts never carries a computed field. --- .../key_management_endpoints.py | 52 +++++--- .../test_key_management_endpoints.py | 117 +++++++++--------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 +- 3 files changed, 97 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 02bfa5c0ca8..c43b6ddf06a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -20,6 +20,7 @@ import secrets import traceback from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast import fastapi @@ -3590,9 +3591,12 @@ def _window_max_budget(window: Mapping[str, object]) -> float | None: return None -async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: str) -> Mapping[str, object]: +async def _budget_window_usage( + window: Mapping[str, object], api_key_hash: str +) -> tuple[str, Mapping[str, object]] | None: """ - Copy of a budget window with current-window spend attached. + (budget_duration, usage entry) for one budget window; None when the window + has no budget_duration to key it by. Reads the same cross-pod counter (spend:key:{hashed_token}:window:{budget_duration}) that _virtual_key_multi_budget_check enforces against, passing the same @@ -3603,34 +3607,41 @@ async def _budget_window_with_usage(window: Mapping[str, object], api_key_hash: duration: Final = window.get("budget_duration") if not isinstance(duration, str) or not duration: - return window + return None + max_budget: Final = _window_max_budget(window) spend: Final = await get_current_spend( counter_key=f"spend:key:{api_key_hash}:window:{duration}", fallback_spend=0.0, - max_budget=_window_max_budget(window), + max_budget=max_budget, window_entity_type="Key", window_entity_id=api_key_hash, window_duration=duration, window_start=get_budget_window_start(window), ) - return {**window, "current_spend": round(spend, 4)} # mutable-ok: per-window response copy, built once per window + return duration, MappingProxyType( + { + "current_spend": round(spend, 4), + "budget_limit": max_budget, + "reset_at": window.get("reset_at"), + } + ) -async def _budget_limits_with_usage( +async def _build_budget_limits_usage( budget_limits: Sequence[object] | str | None, api_key_hash: str -) -> tuple[Mapping[str, object], ...] | None: +) -> Mapping[str, Mapping[str, object]] | None: """ - budget_limits as window dicts with current-window spend attached; None when - the key has no windows so the caller keeps the stored value. + Current-window spend per budget window, keyed by budget_duration, reported + next to the stored budget_limits (which is returned untouched). None when + the key has no windows, so the field only appears on keys that have them. """ windows: Final = _budget_limit_windows(budget_limits) if not windows: return None - return tuple( - await asyncio.gather( - *(_budget_window_with_usage(window=window, api_key_hash=api_key_hash) for window in windows) - ) + usages: Final = await asyncio.gather( + *(_budget_window_usage(window=window, api_key_hash=api_key_hash) for window in windows) ) + return MappingProxyType({duration: usage for duration, usage in (u for u in usages if u is not None)}) @router.post( @@ -3717,12 +3728,12 @@ async def info_key_fn_v2( user_api_key_cache=model_max_budget_limiter.dual_cache, ) if k_token_hash: - budget_limits_usage = await _budget_limits_with_usage( + budget_limits_usage = await _build_budget_limits_usage( budget_limits=k_dict.get("budget_limits"), api_key_hash=k_token_hash, ) if budget_limits_usage is not None: - k_dict["budget_limits"] = budget_limits_usage + k_dict["budget_limits_usage"] = budget_limits_usage filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3759,9 +3770,10 @@ async def info_key_fn( - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - model_max_budget_usage: dict | None - Current-window spend per model, present only when the key has per-model budgets - - budget_limits: list | None - Concurrent budget windows. Each entry includes - current_spend: spend accumulated in the window so far (read from the same cross-pod - spend counter the budget enforcement uses) + - budget_limits: list | None - Concurrent budget windows, exactly as stored + - budget_limits_usage: dict | None - Current-window spend per budget window, keyed by + budget_duration, present only when the key has budget windows (read from the same + cross-pod spend counter the budget enforcement uses) - models: list - Model_name's the key is allowed to call - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} @@ -3841,12 +3853,12 @@ async def info_key_fn( model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) - budget_limits_usage: Final = await _budget_limits_with_usage( + budget_limits_usage: Final = await _build_budget_limits_usage( budget_limits=key_info.get("budget_limits"), api_key_hash=key_token_hash, ) if budget_limits_usage is not None: - key_info["budget_limits"] = budget_limits_usage + key_info["budget_limits_usage"] = budget_limits_usage # Attach object_permission if object_permission_id is set key_info = await attach_object_permission_to_dict(key_info, prisma_client) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 425365c8560..f6e57607717 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -14143,11 +14143,11 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): @pytest.mark.asyncio -async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): +async def test_info_key_fn_reports_budget_limits_usage(monkeypatch): """ - /key/info should attach current_spend to each budget_limits window, read from - the same spend counter (spend:key:{token}:window:{duration}) that budget - enforcement uses. + /key/info reports current-window spend per budget window under budget_limits_usage, + keyed by budget_duration and read from the same counter enforcement uses, while + budget_limits itself comes back exactly as stored. """ from unittest.mock import AsyncMock, MagicMock @@ -14204,11 +14204,14 @@ async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): user_api_key_dict=user_api_key_dict, ) - windows = result["info"]["budget_limits"] - assert len(windows) == 1 - assert windows[0]["current_spend"] == 0.73 - assert windows[0]["max_budget"] == 2.0 - assert windows[0]["budget_duration"] == "1h" + assert result["info"]["budget_limits"] == budget_limits + assert result["info"]["budget_limits_usage"] == { + "1h": { + "current_spend": 0.73, + "budget_limit": 2.0, + "reset_at": "2026-08-15T18:00:00+00:00", + } + } mock_get_current_spend.assert_awaited_once() call_kwargs = mock_get_current_spend.await_args.kwargs @@ -14222,7 +14225,7 @@ async def test_info_key_fn_budget_limits_includes_current_spend(monkeypatch): @pytest.mark.asyncio async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): - """Keys without budget_limits should not trigger window spend lookups.""" + """Keys without budget windows get no budget_limits_usage field and trigger no spend lookup.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import LiteLLM_VerificationToken @@ -14272,12 +14275,13 @@ async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): ) assert result["info"]["budget_limits"] is None + assert "budget_limits_usage" not in result["info"] mock_get_current_spend.assert_not_awaited() @pytest.mark.asyncio -async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): - """/v2/key/info should attach current_spend to each budget_limits window.""" +async def test_info_key_fn_v2_reports_budget_limits_usage(monkeypatch): + """/v2/key/info reports budget_limits_usage per window and leaves budget_limits as stored.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken @@ -14286,6 +14290,18 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): ) test_key_token = "hashed_token_v2_window_test" + budget_limits = [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + }, + { + "reset_at": "2026-08-16T00:00:00+00:00", + "max_budget": 20.0, + "budget_duration": "1d", + }, + ] mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -14304,18 +14320,7 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): mock_key.team_id = None mock_key.model_dump.return_value = { "token": test_key_token, - "budget_limits": [ - { - "reset_at": "2026-08-15T18:00:00+00:00", - "max_budget": 2.0, - "budget_duration": "1h", - }, - { - "reset_at": "2026-08-16T00:00:00+00:00", - "max_budget": 20.0, - "budget_duration": "1d", - }, - ], + "budget_limits": [dict(w) for w in budget_limits], "user_id": "user-v2-w", "team_id": None, "litellm_budget_table": None, @@ -14335,10 +14340,19 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): ) assert len(result["info"]) == 1 - windows = result["info"][0]["budget_limits"] - assert len(windows) == 2 - assert windows[0]["current_spend"] == 1.25 - assert windows[1]["current_spend"] == 1.25 + assert result["info"][0]["budget_limits"] == budget_limits + assert result["info"][0]["budget_limits_usage"] == { + "1h": { + "current_spend": 1.25, + "budget_limit": 2.0, + "reset_at": "2026-08-15T18:00:00+00:00", + }, + "1d": { + "current_spend": 1.25, + "budget_limit": 20.0, + "reset_at": "2026-08-16T00:00:00+00:00", + }, + } assert mock_get_current_spend.await_count == 2 counter_keys = { call.kwargs["counter_key"] for call in mock_get_current_spend.await_args_list @@ -14353,13 +14367,13 @@ async def test_info_key_fn_v2_budget_limits_includes_current_spend(monkeypatch): @pytest.mark.asyncio -async def test_budget_limits_with_usage_json_string_input(monkeypatch): - """budget_limits stored as a JSON string should be parsed and annotated.""" +async def test_build_budget_limits_usage_json_string_input(monkeypatch): + """budget_limits stored as a JSON string is parsed and reported per window.""" import json as json_module from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=0.5) @@ -14370,26 +14384,19 @@ async def test_budget_limits_with_usage_json_string_input(monkeypatch): raw = json_module.dumps( [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] ) - result = await _budget_limits_with_usage(budget_limits=raw, api_key_hash="hash-1") + result = await _build_budget_limits_usage(budget_limits=raw, api_key_hash="hash-1") - assert list(result) == [ - { - "budget_duration": "1h", - "max_budget": 2.0, - "reset_at": None, - "current_spend": 0.5, - } - ] + assert result == {"1h": {"current_spend": 0.5, "budget_limit": 2.0, "reset_at": None}} mock_get_current_spend.assert_awaited_once() @pytest.mark.asyncio -async def test_budget_limits_with_usage_empty_windows_keep_stored_value(monkeypatch): - """A key with no windows (None, [], or "[]") returns None so /key/info keeps the stored value; no spend lookup runs.""" +async def test_build_budget_limits_usage_empty_windows_returns_none(monkeypatch): + """A key with no windows (None, [], or "[]") returns None so the field is left off; no spend lookup runs.""" from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=0.0) @@ -14398,17 +14405,17 @@ async def test_budget_limits_with_usage_empty_windows_keep_stored_value(monkeypa ) for stored in (None, [], "[]"): - assert await _budget_limits_with_usage(budget_limits=stored, api_key_hash="hash-1") is None + assert await _build_budget_limits_usage(budget_limits=stored, api_key_hash="hash-1") is None mock_get_current_spend.assert_not_awaited() @pytest.mark.asyncio -async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): - """A window with only budget_duration still gets current_spend, read without a budget ceiling.""" +async def test_build_budget_limits_usage_window_without_max_budget(monkeypatch): + """A window with only budget_duration still reports current_spend, read without a budget ceiling.""" from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=0.75) @@ -14416,11 +14423,11 @@ async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - result = await _budget_limits_with_usage( + result = await _build_budget_limits_usage( budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" ) - assert list(result) == [{"budget_duration": "2d", "current_spend": 0.75}] + assert result == {"2d": {"current_spend": 0.75, "budget_limit": None, "reset_at": None}} call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-no-max:window:2d" assert call_kwargs["window_duration"] == "2d" @@ -14428,13 +14435,13 @@ async def test_budget_limits_with_usage_window_without_max_budget(monkeypatch): @pytest.mark.asyncio -async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): - """BudgetLimitEntry windows (the shape UserAPIKeyAuth carries) are dumped to dicts and annotated.""" +async def test_build_budget_limits_usage_pydantic_windows(monkeypatch): + """BudgetLimitEntry windows (the shape UserAPIKeyAuth carries) are dumped to dicts and reported.""" from unittest.mock import AsyncMock from litellm.models.team import BudgetLimitEntry from litellm.proxy.management_endpoints.key_management_endpoints import ( - _budget_limits_with_usage, + _build_budget_limits_usage, ) mock_get_current_spend = AsyncMock(return_value=1.0) @@ -14442,14 +14449,12 @@ async def test_budget_limits_with_usage_pydantic_windows(monkeypatch): "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend ) - result = await _budget_limits_with_usage( + result = await _build_budget_limits_usage( budget_limits=[BudgetLimitEntry(budget_duration="7d", max_budget=10.0)], api_key_hash="hash-2", ) - assert list(result) == [ - {"budget_duration": "7d", "max_budget": 10.0, "reset_at": None, "current_spend": 1.0} - ] + assert result == {"7d": {"current_spend": 1.0, "budget_limit": 10.0, "reset_at": None}} call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-2:window:7d" assert call_kwargs["window_duration"] == "7d" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9b5adec516d..7b6f428c5a4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7710,9 +7710,10 @@ export interface paths { * - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - model_max_budget_usage: dict | None - Current-window spend per model, present only when * the key has per-model budgets - * - budget_limits: list | None - Concurrent budget windows. Each entry includes - * current_spend: spend accumulated in the window so far (read from the same cross-pod - * spend counter the budget enforcement uses) + * - budget_limits: list | None - Concurrent budget windows, exactly as stored + * - budget_limits_usage: dict | None - Current-window spend per budget window, keyed by + * budget_duration, present only when the key has budget windows (read from the same + * cross-pod spend counter the budget enforcement uses) * - models: list - Model_name's the key is allowed to call * - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits * - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} From d1320404feff822f822b812951a0782c30d6cbb5 Mon Sep 17 00:00:00 2001 From: Kolade Fajimi <107228310+koladefaj@users.noreply.github.com> Date: Tue, 1 Sep 2026 04:51:31 +0100 Subject: [PATCH 070/126] fix(redis): coerce env var string types and fix param discovery through decorator wrappers (#30644) * fix(redis): coerce env var string types and fix param discovery through decorator wrappers inspect.getfullargspec doesn't work on redis.Redis/redis.RedisCluster because their __init__ is wrapped by @deprecated_args, which replaces the explicit signature with *args/**kwargs internally. getfullargspec returns an empty arg list, so _get_redis_kwargs and _get_redis_cluster_kwargs silently dropped every real constructor parameter not in their hand-picked include_args set -- cluster_error_retry_attempts and connection_error_retry_attempts among them, so an operator's configured retry bound never reached the Redis Cluster client and it fell back to redis-py's own default instead. Rebased onto litellm_internal_staging, which had independently added _init_arg_names (MRO-walking, inspect.unwrap-based) for the same class of bug in _get_redis_url_kwargs. Reused that pattern (as _unwrapped_init_args, without the MRO walk: redis.Redis/RedisCluster declare every real parameter directly on their own __init__, and MRO-walking breaks the tests here that mock the class with autospec=True, since inspect.getmro needs a real __mro__) rather than introducing a second, differently-shaped fix for the same problem. _get_redis_cluster_kwargs now also honors its own client argument instead of ignoring it, so the async cluster client's own extra constructor kwargs (cluster_error_retry_attempts, connection_error_retry_attempts, decode_responses, ...) are no longer filtered out by introspecting the sync class regardless of which client is actually built. Also fixes environment variables and Helm --set values always arriving as strings: redis-py 8.x changed health_check_interval's arithmetic to require a real number, so a stringified value raised TypeError on every Redis operation instead of connecting. _coerce_redis_kwargs_types coerces to each parameter's declared type at the end of _get_redis_client_logic, with an explicit type table for max_connections/socket_timeout/socket_connect_timeout since redis-py 8.x changed the timeout defaults from None to int 5, which would otherwise make a fractional value fail int() and get dropped. Co-authored-by: mangabits <1457532+mangabits@users.noreply.github.com> * ci: verify redis-py client version compatibility across a version matrix * test(redis): assert an async-only cluster kwarg every matrix version declares connection_error_retry_attempts is on the async cluster constructor in redis-py 5.x only; 6.0 removed it in favor of retry. The 6.4.0, 7.4.1 and 8.0.1 legs were failing on that missing parameter name rather than on the behavior under test, while the allow-list itself was doing the right thing on all four versions. decode_responses is async-cluster-only on every version the matrix covers, so it stands in for the same property: the sync cluster class takes it through **kwargs and never names it in its signature. Reverting _get_redis_cluster_kwargs to ignore its client argument still fails both tests on 5.3.1 and 8.0.1. test_async_cluster_passes_async_only_kwargs now builds the real async cluster client and reads connection_kwargs off it, so it no longer needs a patched class factory; the constructor does no I/O. The retry-attempts test keeps its patch, since redis-py >= 6 stores no cluster_error_retry_attempts attribute on the built client and the constructor call is the only place the forwarded value shows up. The _get_redis_cluster_kwargs docstring cited the same two parameters as its examples of async-only kwargs, which is what made the test look reasonable; cluster_error_retry_attempts is on both classes and connection_error_retry_attempts is gone from 6.0 on, so it now names decode_responses instead. * test(redis): drop internal patches from the kwarg coercion tests The test-quality gate flagged the new patch() calls on litellm internals these tests added. Three of them faked litellm._redis.inspect.signature with a MagicMock to hand _coerce_redis_kwargs_types a synthetic parameter; that function already takes a client argument, so they pass stub functions instead, matching the _redis_signature_8x idiom the file uses elsewhere. The fourth patched _redis_kwargs_from_environment to {} to prove _get_redis_client_logic raises without a host or url, which clearing the real env keys through _get_redis_env_kwarg_mapping does without pinning the test to that call. Both files now sit one TQ008 below the merge base rather than six above it. * fix(redis): keep the sync client construction inside the basedpyright budget _get_redis_client_logic now returns dict[str, object] rather than an untyped dict, which is the honest type for operator-supplied config, but it turns the 33 reportUnknownArgumentType errors at redis.Redis(**redis_kwargs) into 33 reportArgumentType errors plus one reportCallIssue, both over their budget. No static type fits: redis-py's constructor declares 40-odd differently typed parameters and the values arrive from config and env, so the allow-list and coercion above are derived from that same signature and redis-py validates each value itself at runtime. The two suppressions name their exact rule and carry that reason. The file ends up 42 basedpyright errors below the merge base, with reportArgumentType and reportCallIssue back at the base counts of 3 and 0. * fix(redis): coerce cluster-only and None-default bool kwargs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mangabits <1457532+mangabits@users.noreply.github.com> Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-redis-compat.yml | 77 +++++++ litellm/_redis.py | 125 ++++++++++- .../caching/test_redis_connection_pool.py | 201 ++++++++++++++++-- tests/test_litellm/test_redis.py | 67 ++++++ 4 files changed, 443 insertions(+), 27 deletions(-) create mode 100644 .github/workflows/test-redis-compat.yml diff --git a/.github/workflows/test-redis-compat.yml b/.github/workflows/test-redis-compat.yml new file mode 100644 index 00000000000..f29755a74b1 --- /dev/null +++ b/.github/workflows/test-redis-compat.yml @@ -0,0 +1,77 @@ +name: "Unit Tests: Redis Client Version Compatibility" + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + paths: + - "litellm/_redis.py" + - "litellm/_redis_credential_provider.py" + - "tests/test_litellm/test_redis.py" + - "tests/test_litellm/caching/test_redis_connection_pool.py" + - ".github/workflows/test-redis-compat.yml" + - "pyproject.toml" + - "uv.lock" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + redis-compat: + name: "redis-py ${{ matrix.redis-version }}" + runs-on: ubuntu-latest + timeout-minutes: 15 + + strategy: + fail-fast: false + matrix: + # 5.3.1 is the version pinned in uv.lock (redisvl caps it below 6); the + # newer legs prove the inspect.signature introspection in litellm/_redis.py + # keeps extracting kwargs on the redis-py releases people actually run now. + # Only the exact release 6.0.0 is skipped: rq (pulled by the proxy extra) + # specifies `redis != 6`, which excludes 6.0.0 alone, so 6.4.0 stands in + # for the 6.x line. + redis-version: ["5.3.1", "6.4.0", "7.4.1", "8.0.1"] + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + + - name: Pin redis-py to the matrix version + env: + REDIS_VERSION: ${{ matrix.redis-version }} + run: | + uv pip install "redis==${REDIS_VERSION:?}" + uv run --no-sync python -c "import redis; assert redis.__version__ == '${REDIS_VERSION:?}', redis.__version__; print('redis-py', redis.__version__)" + + - name: Run redis unit tests + run: | + uv run --no-sync pytest \ + tests/test_litellm/test_redis.py \ + tests/test_litellm/caching/test_redis_connection_pool.py \ + --tb=short -vv \ + --reruns 2 \ + --reruns-delay 1 \ + --durations=20 diff --git a/litellm/_redis.py b/litellm/_redis.py index 9381357931e..3e68d50cf16 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -13,6 +13,7 @@ import json # s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation import os from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import Final from urllib.parse import urlsplit, urlunsplit @@ -38,9 +39,25 @@ from ._logging import verbose_logger AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" -def _get_redis_kwargs(): - arg_spec: Final = inspect.getfullargspec(redis.Redis) +def _unwrapped_init_args(cls: type) -> frozenset[str]: + """Every parameter on a single class's own ``__init__``, decorator-unwrapped. + Unlike ``_init_arg_names`` below, this does not walk the MRO: ``redis.Redis`` + and ``redis.RedisCluster`` (sync and async) each declare every real + constructor parameter directly on their own ``__init__``, so MRO-walking is + unnecessary — and it actively breaks the several tests here that mock the + class with ``patch(..., autospec=True)``, since ``inspect.getmro`` needs a + real ``__mro__`` that an autospec'd stand-in for a class does not provide. + + Still unwraps first: redis-py >= 7.4 decorates these ``__init__``s with + ``@deprecated_args`` too, which the same class of bug as ``_init_arg_names`` + would otherwise silently empty this allowlist through (see its docstring). + """ + spec: Final = inspect.getfullargspec(inspect.unwrap(cls.__init__)) + return frozenset(spec.args + spec.kwonlyargs) + + +def _get_redis_kwargs(): # Only allow primitive arguments exclude_args: Final = { "self", @@ -60,7 +77,7 @@ def _get_redis_kwargs(): "azure_client_secret", } - available_args: Final = {x for x in arg_spec.args if x not in exclude_args} | include_args + available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args return available_args @@ -120,15 +137,23 @@ def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]: return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args -def _get_redis_cluster_kwargs(client=None): +def _get_redis_cluster_kwargs(client: type | None = None): + """Config kwargs the target cluster client's constructor actually accepts. + + Defaults to the sync ``redis.RedisCluster``, but the async cluster client + (``redis.asyncio.cluster.RedisCluster``) declares connection settings such as + ``decode_responses`` on its own constructor, where the sync class takes them + through ``**kwargs`` and so never names them in its signature. Introspecting + only the sync class regardless of which client is actually built silently + drops those for every async cluster caller. + """ if client is None: - client = redis.Redis.from_url - arg_spec: Final = inspect.getfullargspec(redis.RedisCluster) + client = redis.RedisCluster # Only allow primitive arguments exclude_args: Final = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"} - available_args = {x for x in arg_spec.args if x not in exclude_args} + available_args = {x for x in _unwrapped_init_args(client) if x not in exclude_args} available_args |= { "password", "username", @@ -161,6 +186,79 @@ def _get_redis_env_kwarg_mapping(): return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs() if x not in exclude_from_environment} +def _str_to_bool(value: str) -> bool: + return value.lower() in ("true", "1", "yes") + + +def _coerce_redis_kwargs_types( + redis_kwargs: Mapping[str, object], + client: type | tuple[type, ...] = redis.Redis, +) -> dict[str, object]: # mutable-ok: a caller mutates the returned kwargs before constructing its client + """Coerces string values to the numeric/boolean type ``client``'s constructor + declares for that parameter. ``client`` may be a tuple of client classes; a + parameter's type is taken from the first signature that declares it, which + lets cluster callers coerce cluster-only kwargs such as + ``cluster_error_retry_attempts`` alongside the shared connection kwargs. + + Environment variables are always strings, and Helm ``--set`` stringifies values + too, so a config value like ``health_check_interval`` or ``socket_timeout`` + can arrive as ``"30"``/``"5.5"`` rather than a real number. redis-py's own + connection-health-check arithmetic (``loop.time() + self.health_check_interval``) + then raises ``TypeError`` on every Redis operation instead of connecting. + + ``max_connections``, ``socket_timeout``, and ``socket_connect_timeout`` use an + explicit target type rather than the parameter's own signature default: redis-py + 8.x changed the timeout defaults from ``None`` to int ``5``, so inferring the + type from the default would make a fractional ``"5.5"`` fail ``int()`` and get + silently dropped on 8.x while working on older versions. ``socket_keepalive`` + is explicit too: its signature default is ``None``, which carries no type to + infer from, and leaving it a string makes ``"false"`` truthy. + """ + signatures: Final = tuple(inspect.signature(c) for c in (client if isinstance(client, tuple) else (client,))) + explicit_param_types: Final = MappingProxyType( + { + "max_connections": int, + "socket_timeout": float, + "socket_connect_timeout": float, + "socket_keepalive": bool, + } + ) + result: Final = dict(redis_kwargs) # mutable-ok: per-key try/except coercion below needs to drop individual keys + for key, value in redis_kwargs.items(): + if not isinstance(value, str): + continue + param = next((sig.parameters[key] for sig in signatures if key in sig.parameters), None) + if param is None: + continue + explicit_type = explicit_param_types.get(key) + if explicit_type is bool: + result[key] = _str_to_bool(value) + continue + if explicit_type is not None: + try: + result[key] = explicit_type(value) + except (ValueError, TypeError): + del result[key] + continue + default: object = param.default # pyright: ignore[reportAny] # inspect.Parameter.default is stubbed as Any + if default is inspect.Parameter.empty: + continue + # bool must be checked before int, since bool subclasses int + if isinstance(default, bool): + result[key] = _str_to_bool(value) + elif isinstance(default, int): + try: + result[key] = int(value) + except (ValueError, TypeError): + del result[key] + elif isinstance(default, float): + try: + result[key] = float(value) + except (ValueError, TypeError): + del result[key] + return result + + def _redis_kwargs_from_environment(): mapping: Final = _get_redis_env_kwarg_mapping() @@ -505,7 +603,12 @@ def _get_redis_client_logic(**env_overrides): raise ValueError("Either 'host' or 'url' must be specified for redis.") # litellm.print_verbose(f"redis_kwargs: {redis_kwargs}") - return redis_kwargs + coercion_client: Final = ( + (redis.Redis, redis.RedisCluster, async_redis.RedisCluster) + if redis_kwargs.get("startup_nodes") + else redis.Redis + ) + return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client) def init_redis_cluster(redis_kwargs) -> redis.RedisCluster: @@ -657,7 +760,9 @@ def get_redis_client(**env_overrides): if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs: return _init_redis_sentinel(redis_kwargs) - return redis.Redis(**redis_kwargs) + return redis.Redis( # pyright: ignore[reportCallIssue] # object-valued kwargs match no overload statically + **redis_kwargs, # pyright: ignore[reportArgumentType] # allow-listed and coerced against this signature + ) def get_redis_async_client( @@ -669,7 +774,7 @@ def get_redis_async_client( if "startup_nodes" in redis_kwargs: from redis.cluster import ClusterNode - args = _get_redis_cluster_kwargs() + args = _get_redis_cluster_kwargs(async_redis.RedisCluster) cluster_kwargs: Final = {} for arg in redis_kwargs: if arg in args: diff --git a/tests/test_litellm/caching/test_redis_connection_pool.py b/tests/test_litellm/caching/test_redis_connection_pool.py index c824d3e7a0e..54dbe5361d7 100644 --- a/tests/test_litellm/caching/test_redis_connection_pool.py +++ b/tests/test_litellm/caching/test_redis_connection_pool.py @@ -1,15 +1,14 @@ -""" -Regression tests for Redis connection pool leak fixes (RC1-RC5). - -Tests are pure unit tests — no Redis server required. -""" - from unittest.mock import AsyncMock, MagicMock, patch import pytest -import redis.asyncio as async_redis -from litellm._redis import get_redis_async_client, get_redis_connection_pool +from litellm._redis import ( + _coerce_redis_kwargs_types, + _get_redis_client_logic, + _get_redis_env_kwarg_mapping, + get_redis_async_client, + get_redis_connection_pool, +) def test_url_config_uses_passed_pool(): @@ -60,16 +59,14 @@ def test_max_connections_url_config_string_value(monkeypatch): assert pool.max_connections == 25 -def test_max_connections_url_config_invalid_value(): - """Invalid max_connections should be silently ignored, falling back - to the pool default (50 for BlockingConnectionPool).""" - with patch("litellm._redis._get_redis_client_logic") as mock_logic: - mock_logic.return_value = { - "url": "redis://localhost:6379/0", - "max_connections": "not_a_number", - } +def test_max_connections_url_config_invalid_value(monkeypatch): + """Invalid max_connections from an env var should be silently dropped, + falling back to the pool default (50 for BlockingConnectionPool).""" + monkeypatch.setenv("REDIS_URL", "redis://localhost:6379/0") + monkeypatch.delenv("REDIS_HOST", raising=False) + monkeypatch.setenv("REDIS_MAX_CONNECTIONS", "not_a_number") - pool = get_redis_connection_pool() + pool = get_redis_connection_pool() # BlockingConnectionPool default is 50 assert pool.max_connections == 50 @@ -128,3 +125,173 @@ async def test_disconnect_idempotent(): await cache.disconnect() await cache.disconnect() # should not raise + + +def test_coerce_redis_kwargs_types_int(): + """String values for int-typed Redis params are coerced to int.""" + result = _coerce_redis_kwargs_types({"health_check_interval": "30", "port": "6380", "db": "1"}) + assert result["health_check_interval"] == 30 + assert isinstance(result["health_check_interval"], int) + assert result["port"] == 6380 + assert result["db"] == 1 + + +def test_coerce_redis_kwargs_types_bool(): + """String values for bool-typed Redis params are coerced to bool.""" + result = _coerce_redis_kwargs_types({"ssl": "true", "decode_responses": "false"}) + assert result["ssl"] is True + assert result["decode_responses"] is False + + +def test_coerce_redis_kwargs_types_none_default_numeric(): + """String values for known None-default numeric params are coerced.""" + result = _coerce_redis_kwargs_types({"max_connections": "20", "socket_timeout": "5.5"}) + assert result["max_connections"] == 20 + assert isinstance(result["max_connections"], int) + assert result["socket_timeout"] == 5.5 + assert isinstance(result["socket_timeout"], float) + + +def _redis_signature_pre_8x( + socket_timeout=None, + socket_connect_timeout=None, + max_connections=None, + health_check_interval=0, +): + """Stand-in for the redis-py <= 7.x Redis signature, where the timeout defaults are None.""" + + +def _redis_signature_8x( + socket_timeout=5, + socket_connect_timeout=5, + max_connections=None, + health_check_interval=0, +): + """Stand-in for the redis-py 8.x Redis signature, where the timeout defaults became int 5.""" + + +@pytest.mark.parametrize( + "client", + [_redis_signature_pre_8x, _redis_signature_8x], + ids=["redis-py<=7.x", "redis-py-8.x"], +) +def test_coerce_fractional_socket_timeout_survives_signature_default_change(client): + """redis-py 8.x changed socket_timeout's default from None to int 5. Deriving the + target type from the signature default made int("5.5") raise, so the key was dropped + and REDIS_SOCKET_TIMEOUT=5.5 silently disappeared on 8.x.""" + result = _coerce_redis_kwargs_types( + {"socket_timeout": "5.5", "socket_connect_timeout": "2.5", "max_connections": "20"}, + client=client, + ) + + assert result["socket_timeout"] == pytest.approx(5.5) + assert isinstance(result["socket_timeout"], float) + assert result["socket_connect_timeout"] == pytest.approx(2.5) + assert isinstance(result["socket_connect_timeout"], float) + assert result["max_connections"] == 20 + assert isinstance(result["max_connections"], int) + + +def test_coerce_invalid_socket_timeout_is_still_dropped(): + """Garbage must not survive the explicit-type path; Redis falls back to its own default.""" + result = _coerce_redis_kwargs_types({"socket_timeout": "not_a_number"}, client=_redis_signature_8x) + + assert "socket_timeout" not in result + + +def test_coerce_redis_kwargs_types_invalid_drops_key(): + """A string that cannot be coerced to the expected numeric type is dropped.""" + result = _coerce_redis_kwargs_types({"health_check_interval": "not_a_number"}) + assert "health_check_interval" not in result + + +def test_coerce_redis_kwargs_types_non_string_unchanged(): + """Non-string values pass through without modification.""" + result = _coerce_redis_kwargs_types({"health_check_interval": 30, "ssl": True}) + assert result["health_check_interval"] == 30 + assert result["ssl"] is True + + +def test_health_check_interval_from_env_is_int(monkeypatch): + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_HEALTH_CHECK_INTERVAL", "30") + + pool = get_redis_connection_pool() + + assert pool is not None + interval = pool.connection_kwargs.get("health_check_interval") + assert interval == 30 + assert isinstance(interval, int), f"Expected int, got {type(interval)}: {interval!r}" + + +def _signature_without_defaults(testkey): + """Stand-in for a client whose parameter declares no default at all.""" + + +def _signature_with_float_default(myparam=1.0): + """Stand-in for a client whose parameter declares a float default.""" + + +def test_coerce_redis_kwargs_types_empty_default_param_unchanged(): + """String params whose signature entry has no default (inspect.Parameter.empty) are left as-is.""" + result = _coerce_redis_kwargs_types({"testkey": "some_value"}, client=_signature_without_defaults) + + assert result["testkey"] == "some_value" + assert isinstance(result["testkey"], str) + + +def test_coerce_redis_kwargs_types_float_valid(): + """String values for params whose signature default is a float are coerced to float.""" + result = _coerce_redis_kwargs_types({"myparam": "3.14"}, client=_signature_with_float_default) + + assert result["myparam"] == pytest.approx(3.14) + assert isinstance(result["myparam"], float) + + +def test_coerce_redis_kwargs_types_float_invalid_drops_key(): + """An unconvertible string for a float-default param is dropped from the result.""" + result = _coerce_redis_kwargs_types({"myparam": "not_a_float"}, client=_signature_with_float_default) + + assert "myparam" not in result + + +@pytest.mark.parametrize( + ("raw", "expected"), + [("false", False), ("true", True), ("0", False), ("1", True)], +) +def test_coerce_socket_keepalive_string(raw, expected): + """socket_keepalive's signature default is None, so it needs an explicit bool + coercion: a leftover "false" string is truthy and enables keepalive.""" + result = _coerce_redis_kwargs_types({"socket_keepalive": raw}) + + assert result["socket_keepalive"] is expected + + +def test_get_redis_client_logic_coerces_cluster_only_kwargs(monkeypatch): + """Cluster-only kwargs (absent from redis.Redis's signature) must still be + coerced when routing to a cluster, or Helm-stringified values reach + RedisCluster as strings.""" + for envvar in (*_get_redis_env_kwarg_mapping(), "REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES"): + monkeypatch.delenv(envvar, raising=False) + + result = _get_redis_client_logic( + startup_nodes='[{"host": "localhost", "port": 7000}]', + cluster_error_retry_attempts="5", + require_full_coverage="false", + health_check_interval="30", + ) + + assert result["cluster_error_retry_attempts"] == 5 + assert isinstance(result["cluster_error_retry_attempts"], int) + assert result["require_full_coverage"] is False + assert result["health_check_interval"] == 30 + assert isinstance(result["health_check_interval"], int) + + +def test_get_redis_client_logic_raises_without_host_or_url(monkeypatch): + """_get_redis_client_logic raises ValueError when neither host nor url is provided.""" + for envvar in (*_get_redis_env_kwarg_mapping(), "REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES"): + monkeypatch.delenv(envvar, raising=False) + + with pytest.raises(ValueError, match="Either 'host' or 'url' must be specified for redis"): + _get_redis_client_logic() diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 826beb74a27..a96e8541e06 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -1,3 +1,4 @@ +import inspect import json from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -600,6 +601,72 @@ def test_reconnect_kwargs_in_cluster_kwargs(): assert "socket_keepalive" in kwargs +def test_retry_attempts_in_cluster_kwargs(): + """cluster_error_retry_attempts must survive the cluster kwarg allow-list so + operators can bound worst-case retry latency on a Redis Cluster: it was being + silently dropped because the allow-list was built from redis.RedisCluster's + decorated __init__ without unwrapping it, so getfullargspec saw an empty + (self, *args, **kwargs) wrapper signature.""" + kwargs = _get_redis_cluster_kwargs() + assert "cluster_error_retry_attempts" in kwargs + + +def test_async_only_kwargs_in_cluster_kwargs_when_async_client_requested(): + """decode_responses is on the async cluster client's constructor and not the sync + one, on every redis-py the matrix covers. Introspecting the sync class regardless + of which client is actually built silently drops it for every async cluster caller.""" + sync_kwargs = _get_redis_cluster_kwargs() + async_kwargs = _get_redis_cluster_kwargs(async_redis.RedisCluster) + + assert "decode_responses" not in sync_kwargs + assert "decode_responses" in async_kwargs + + +@patch( # test-quality-ok: redis-py >= 6 keeps no cluster_error_retry_attempts attribute on the built client, so the constructor call is the only place the value is observable + "litellm.caching.redis_cluster_node_isolation.get_litellm_async_redis_cluster_class" +) +def test_async_cluster_forwards_retry_attempts(mock_get_cluster_class): + """Regression: cluster_error_retry_attempts must reach the constructed async + cluster client. Silently dropping it removes an operator's only lever for + bounding a stuck node's worst-case retry latency, and the client falls back + to redis-py's own default (3 retries) instead.""" + mock_cluster_cls = mock_get_cluster_class.return_value + get_redis_async_client( + startup_nodes=[{"host": "cluster-node", "port": 6379}], + cluster_error_retry_attempts=2, + ) + + call_kwargs = mock_cluster_cls.call_args[1] + assert call_kwargs["cluster_error_retry_attempts"] == 2 + + +def test_async_cluster_passes_async_only_kwargs(): + """Regression: decode_responses is an async-cluster-only constructor arg. When + the allow-list came from the sync class it was filtered out and values came + back as bytes instead of str.""" + client = get_redis_async_client( + startup_nodes=[{"host": "cluster-node", "port": 6379}], + decode_responses=True, + ) + + assert client.connection_kwargs["decode_responses"] is True + + +@pytest.mark.parametrize("cluster_client", [redis.RedisCluster, async_redis.RedisCluster], ids=["sync", "async"]) +def test_cluster_kwargs_exclude_variadic_parameters(cluster_client): + """*args / **kwargs are signature placeholders, not connection settings, and + must never land in the allow-list regardless of which cluster client is + introspected.""" + variadic = { + name + for name, param in inspect.signature(cluster_client).parameters.items() + if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD) + } + + leaked = variadic & set(_get_redis_cluster_kwargs(cluster_client)) + assert not leaked, f"variadic params leaked into the allow-list: {leaked}" + + @patch("litellm.caching.redis_cluster_node_isolation.get_litellm_async_redis_cluster_class") def test_async_cluster_sets_reconnect_defaults(mock_get_cluster_class): """ From 760b864e43045032bd76348650cf15f6e207b2a9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 21:11:15 -0700 Subject: [PATCH 071/126] refactor(key): trim budget_limits_usage entries to current_spend max_budget and reset_at already live on the matching budget_limits entry, so repeating them (as budget_limit and reset_at) only invited confusion about which copy is authoritative. --- .../key_management_endpoints.py | 17 ++++-------- .../test_key_management_endpoints.py | 26 +++++-------------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 ++--- 3 files changed, 14 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c43b6ddf06a..99e201930c1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3608,23 +3608,16 @@ async def _budget_window_usage( duration: Final = window.get("budget_duration") if not isinstance(duration, str) or not duration: return None - max_budget: Final = _window_max_budget(window) spend: Final = await get_current_spend( counter_key=f"spend:key:{api_key_hash}:window:{duration}", fallback_spend=0.0, - max_budget=max_budget, + max_budget=_window_max_budget(window), window_entity_type="Key", window_entity_id=api_key_hash, window_duration=duration, window_start=get_budget_window_start(window), ) - return duration, MappingProxyType( - { - "current_spend": round(spend, 4), - "budget_limit": max_budget, - "reset_at": window.get("reset_at"), - } - ) + return duration, MappingProxyType({"current_spend": round(spend, 4)}) async def _build_budget_limits_usage( @@ -3771,9 +3764,9 @@ async def info_key_fn( - model_max_budget_usage: dict | None - Current-window spend per model, present only when the key has per-model budgets - budget_limits: list | None - Concurrent budget windows, exactly as stored - - budget_limits_usage: dict | None - Current-window spend per budget window, keyed by - budget_duration, present only when the key has budget windows (read from the same - cross-pod spend counter the budget enforcement uses) + - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. + {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows + (read from the same cross-pod spend counter the budget enforcement uses) - models: list - Model_name's the key is allowed to call - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index f6e57607717..b5a4204f0af 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -14205,13 +14205,7 @@ async def test_info_key_fn_reports_budget_limits_usage(monkeypatch): ) assert result["info"]["budget_limits"] == budget_limits - assert result["info"]["budget_limits_usage"] == { - "1h": { - "current_spend": 0.73, - "budget_limit": 2.0, - "reset_at": "2026-08-15T18:00:00+00:00", - } - } + assert result["info"]["budget_limits_usage"] == {"1h": {"current_spend": 0.73}} mock_get_current_spend.assert_awaited_once() call_kwargs = mock_get_current_spend.await_args.kwargs @@ -14342,16 +14336,8 @@ async def test_info_key_fn_v2_reports_budget_limits_usage(monkeypatch): assert len(result["info"]) == 1 assert result["info"][0]["budget_limits"] == budget_limits assert result["info"][0]["budget_limits_usage"] == { - "1h": { - "current_spend": 1.25, - "budget_limit": 2.0, - "reset_at": "2026-08-15T18:00:00+00:00", - }, - "1d": { - "current_spend": 1.25, - "budget_limit": 20.0, - "reset_at": "2026-08-16T00:00:00+00:00", - }, + "1h": {"current_spend": 1.25}, + "1d": {"current_spend": 1.25}, } assert mock_get_current_spend.await_count == 2 counter_keys = { @@ -14386,7 +14372,7 @@ async def test_build_budget_limits_usage_json_string_input(monkeypatch): ) result = await _build_budget_limits_usage(budget_limits=raw, api_key_hash="hash-1") - assert result == {"1h": {"current_spend": 0.5, "budget_limit": 2.0, "reset_at": None}} + assert result == {"1h": {"current_spend": 0.5}} mock_get_current_spend.assert_awaited_once() @@ -14427,7 +14413,7 @@ async def test_build_budget_limits_usage_window_without_max_budget(monkeypatch): budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" ) - assert result == {"2d": {"current_spend": 0.75, "budget_limit": None, "reset_at": None}} + assert result == {"2d": {"current_spend": 0.75}} call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-no-max:window:2d" assert call_kwargs["window_duration"] == "2d" @@ -14454,7 +14440,7 @@ async def test_build_budget_limits_usage_pydantic_windows(monkeypatch): api_key_hash="hash-2", ) - assert result == {"7d": {"current_spend": 1.0, "budget_limit": 10.0, "reset_at": None}} + assert result == {"7d": {"current_spend": 1.0}} call_kwargs = mock_get_current_spend.await_args.kwargs assert call_kwargs["counter_key"] == "spend:key:hash-2:window:7d" assert call_kwargs["window_duration"] == "7d" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7b6f428c5a4..77e42c09525 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7711,9 +7711,9 @@ export interface paths { * - model_max_budget_usage: dict | None - Current-window spend per model, present only when * the key has per-model budgets * - budget_limits: list | None - Concurrent budget windows, exactly as stored - * - budget_limits_usage: dict | None - Current-window spend per budget window, keyed by - * budget_duration, present only when the key has budget windows (read from the same - * cross-pod spend counter the budget enforcement uses) + * - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. + * {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows + * (read from the same cross-pod spend counter the budget enforcement uses) * - models: list - Model_name's the key is allowed to call * - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits * - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} From c1b5cacf1fe5eee475a01f87870558b915b5e2ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:12:53 -0700 Subject: [PATCH 072/126] refactor(speech): freeze httpx response header dicts (LIT002) --- .../speech/speech_to_completion_bridge/transformation.py | 4 +++- litellm/llms/vertex_ai/text_to_speech/transformation.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index e2d3fadf852..2ed140c0208 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -174,7 +174,9 @@ class SpeechToCompletionBridgeTransformationHandler: if self._is_gemini_tts_model(model) else (decoded_audio, "audio/mpeg") ) - response: Final = httpx.Response(status_code=200, content=content, headers={"Content-Type": content_type}) + response: Final = httpx.Response( + status_code=200, content=content, headers=MappingProxyType({"Content-Type": content_type}) + ) binary_response: Final = HttpxBinaryResponseContent(response) binary_response.set_response_cost(_completion_response_cost(model_response)) return binary_response diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index a5ff7eca021..332f892ae6b 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -7,6 +7,7 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s import base64 from collections.abc import Coroutine +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union import httpx @@ -464,7 +465,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): media_type: Final = speech_media_type_from_audio_bytes(binary_data) response: Final = httpx.Response( status_code=200, - headers={} if media_type is None else {"content-type": media_type}, + headers=None if media_type is None else MappingProxyType({"content-type": media_type}), content=binary_data, ) From a03378f6d18a09f68322f2b1b123195ec8965a1a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:14:54 -0700 Subject: [PATCH 073/126] fix(openai): forward reasoning_effort for unknown model aliases instead of failing closed --- .../llms/openai/chat/gpt_transformation.py | 4 ++ .../chat/test_openai_gpt_transformation.py | 41 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 5894658e5d2..255ee3159c3 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -178,6 +178,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): model_specific_params.append( "user" ) # user is not a param supported by all openai-compatible endpoints - e.g. azure ai + else: + model_specific_params.append( + "reasoning_effort" + ) # unknown model: likely a proxy alias for a reasoning-capable model, so forward and let the server decide return base_params + model_specific_params def _map_openai_params( diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 3f346b5e8e7..1ee53a90d7d 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -145,6 +145,47 @@ class TestGetOptionalParamsIntegration: assert regular_params.get("user") == "my-end-user" assert responses_params.get("user") == "my-end-user" + def test_reasoning_effort_supported_for_unknown_model_alias(self): + """An openai/-routed model litellm doesn't recognize is likely a proxy alias: + reasoning_effort must be forwarded so the server decides support.""" + supported_params = OpenAIGPTConfig().get_supported_openai_params( + "my-claude-alias" + ) + assert "reasoning_effort" in supported_params + + def test_reasoning_effort_not_supported_for_known_non_reasoning_models(self): + """Known OpenAI models keep failing closed client-side.""" + config = OpenAIGPTConfig() + assert "reasoning_effort" not in config.get_supported_openai_params("gpt-4o") + assert "reasoning_effort" not in config.get_supported_openai_params( + "responses/gpt-4.1-mini" + ) + + def test_reasoning_effort_forwarded_in_optional_params_for_unknown_model_alias( + self, + ): + """Regression test for reasoning_effort raising UnsupportedParamsError + client-side for openai/-prefixed proxy aliases before any HTTP request.""" + from litellm.utils import get_optional_params + + optional_params = get_optional_params( + model="my-claude-alias", + custom_llm_provider="openai", + reasoning_effort="low", + ) + assert optional_params.get("reasoning_effort") == "low" + + def test_reasoning_effort_still_rejected_for_known_non_reasoning_model(self): + """A real OpenAI model that doesn't reason still rejects the param client-side.""" + from litellm.utils import get_optional_params + + with pytest.raises(litellm.utils.UnsupportedParamsError): + get_optional_params( + model="gpt-4o", + custom_llm_provider="openai", + reasoning_effort="low", + ) + class TestOpenAIChatCompletionStreamingHandler: """Tests for OpenAIChatCompletionStreamingHandler.chunk_parser()""" From 76839ca9d8edbe356702f1811c3c19dc10bbf10a Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Mon, 31 Aug 2026 21:18:36 -0700 Subject: [PATCH 074/126] fix(bedrock): forward aws_external_id in files and batches credential loading --- litellm/llms/bedrock/common_utils.py | 1 + litellm/llms/bedrock/files/handler.py | 1 + litellm/llms/bedrock/files/transformation.py | 3 + .../files/test_bedrock_files_handler.py | 68 +++++++++++ .../test_bedrock_files_transformation.py | 108 ++++++++++++++++++ .../llms/bedrock/test_bedrock_common_utils.py | 53 +++++++++ 6 files changed, 234 insertions(+) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 72e3cc1b326..9cbceb4880c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1487,6 +1487,7 @@ class CommonBatchFilesUtils: aws_role_name=optional_params.get("aws_role_name"), aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + aws_external_id=optional_params.get("aws_external_id"), ) # Prepare the request data diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 13718d41cc1..e74c3802d20 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -113,6 +113,7 @@ class BedrockFilesHandler(BaseAWSLLM): aws_role_name=optional_params.get("aws_role_name"), aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + aws_external_id=optional_params.get("aws_external_id"), ) # Create S3 client diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index f442608a288..33b27943ad8 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -146,6 +146,7 @@ class _BedrockS3RequestParams(BaseModel): aws_role_name: str | None = None aws_web_identity_token: str | None = None aws_sts_endpoint: str | None = None + aws_external_id: str | None = None s3_region_name: str | None = None s3_endpoint_url: str | None = None @@ -1029,6 +1030,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): aws_role_name=optional_params.get("aws_role_name"), aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + aws_external_id=optional_params.get("aws_external_id"), ) # Calculate SHA256 hash of the content (REQUIRED for S3) @@ -1290,6 +1292,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): aws_role_name=request_params.aws_role_name, aws_web_identity_token=request_params.aws_web_identity_token, aws_sts_endpoint=request_params.aws_sts_endpoint, + aws_external_id=request_params.aws_external_id, ) empty_body_hash: Final = hashlib.sha256(b"").hexdigest() diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py index 7f91b49a6f5..a80e5dcc13b 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py @@ -204,3 +204,71 @@ def test_should_forward_trusted_model_credentials_to_retrieve_provider_config(): assert response is mock_response litellm_params = mock_retrieve_file.call_args.kwargs["litellm_params"] assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials + + +@pytest.mark.asyncio +async def test_afile_content_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-files-download": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAFILESDOWNLOADROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + class FakeS3Body: + def read(self): + return b'{"custom_id": "req-1"}' + + class FakeS3Client: + def get_object(self, Bucket, Key): + return {"Body": FakeS3Body()} + + s3_client_kwargs = {} + + def fake_boto3_client(service_name, **kwargs): + if service_name == "sts": + return FakeSTSClient() + s3_client_kwargs.update(kwargs) + return FakeS3Client() + + optional_params = { + "_litellm_internal_model_credentials": MappingProxyType({"s3_bucket_name": "safe-bucket"}), + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAFILESDOWNLOADCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-download-role", + "aws_session_name": "litellm-files-download-session", + "aws_external_id": "external-id-files-download", + } + + with patch.object(boto3, "client", side_effect=fake_boto3_client): + response = await BedrockFilesHandler().afile_content( + file_content_request={"file_id": "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl"}, + optional_params=optional_params, + timeout=10.0, + max_retries=None, + ) + + assert s3_client_kwargs["aws_access_key_id"] == "ASIAFILESDOWNLOADROLE" + assert s3_client_kwargs["aws_session_token"] == "assumed-session-token" + assert response.content == b'{"custom_id": "req-1"}' diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index da13f265ee4..541c0db15d8 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2404,3 +2404,111 @@ class TestBedrockFilesS3SignatureEncoding: body=None, headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM], ) + + +def test_sign_s3_request_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 upload request.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-files-put": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAFILESPUTROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAFILESPUTCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-put-role", + "aws_session_name": "litellm-files-put-session", + "aws_external_id": "external-id-files-put", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers, _signed_body = BedrockFilesConfig()._sign_s3_request( + content='{"custom_id": "req-1"}', + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + optional_params=optional_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "ASIAFILESPUTROLE" in authorization + + +def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.files.transformation import ( + BedrockFilesConfig, + _BedrockS3RequestParams, + ) + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-files-get": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAFILESGETROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + request_params = _BedrockS3RequestParams.model_validate( + { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAFILESGETCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-get-role", + "aws_session_name": "litellm-files-get-session", + "aws_external_id": "external-id-files-get", + } + ) + assert request_params.aws_external_id == "external-id-files-get" + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers = BedrockFilesConfig()._sign_s3_get_request( + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + aws_region_name="us-east-1", + request_params=request_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "ASIAFILESGETROLE" in authorization diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 389bf4a8e40..afd5e83ca52 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -520,3 +520,56 @@ def test_merge_bedrock_aws_request_params_keeps_caller_credentials_without_stati assert merged["aws_secret_access_key"] == "caller-secret" assert merged["aws_session_token"] == "caller-token" assert merged["aws_region_name"] == "us-west-2" + + +def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied when signing batch API requests.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-batch-sign": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIABATCHSIGNROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIABATCHSIGNCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-batch-sign-role", + "aws_session_name": "litellm-batch-sign-session", + "aws_external_id": "external-id-batch-sign", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers, signed_data = CommonBatchFilesUtils().sign_aws_request( + service_name="bedrock", + data={"jobName": "litellm-batch-job"}, + endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", + optional_params=optional_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "ASIABATCHSIGNROLE" in authorization + assert signed_data == b'{"jobName": "litellm-batch-job"}' From 60b24abd3e44cdc6800964f7808b463addb5a074 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:26:10 -0700 Subject: [PATCH 075/126] test(bedrock): capture s3 client kwargs from the boto3 mock instead of a mutable dict --- .../llms/bedrock/files/test_bedrock_files_handler.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py index a80e5dcc13b..639be272351 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py @@ -243,12 +243,9 @@ async def test_afile_content_assumes_role_with_external_id(monkeypatch): def get_object(self, Bucket, Key): return {"Body": FakeS3Body()} - s3_client_kwargs = {} - def fake_boto3_client(service_name, **kwargs): if service_name == "sts": return FakeSTSClient() - s3_client_kwargs.update(kwargs) return FakeS3Client() optional_params = { @@ -261,7 +258,7 @@ async def test_afile_content_assumes_role_with_external_id(monkeypatch): "aws_external_id": "external-id-files-download", } - with patch.object(boto3, "client", side_effect=fake_boto3_client): + with patch.object(boto3, "client", side_effect=fake_boto3_client) as mock_boto3_client: response = await BedrockFilesHandler().afile_content( file_content_request={"file_id": "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl"}, optional_params=optional_params, @@ -269,6 +266,7 @@ async def test_afile_content_assumes_role_with_external_id(monkeypatch): max_retries=None, ) + s3_client_kwargs = next(call.kwargs for call in mock_boto3_client.call_args_list if call.args[0] == "s3") assert s3_client_kwargs["aws_access_key_id"] == "ASIAFILESDOWNLOADROLE" assert s3_client_kwargs["aws_session_token"] == "assumed-session-token" assert response.content == b'{"custom_id": "req-1"}' From 65a46a5f32a824e5d42f6d92d4183ad7febf8fe4 Mon Sep 17 00:00:00 2001 From: George Pickett Date: Mon, 31 Aug 2026 21:28:45 -0700 Subject: [PATCH 076/126] fix(websearch): reject invalid explicit search tool selections (#38113) * fix(websearch): reject invalid explicit search tool selections * refactor(websearch): simplify explicit search tool validation --- .../websearch_interception/handler.py | 50 +++-- .../test_websearch_interception_handler.py | 186 +++++++++++++++++- 2 files changed, 217 insertions(+), 19 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 81310b9ddc3..aefd4fa3b47 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -416,15 +416,25 @@ class WebSearchInterceptionLogger(CustomLogger): if not tools: return None - if call_type in (CallTypes.responses, CallTypes.aresponses): - return self._convert_responses_tools(kwargs=kwargs, tools=tools) - - # Check if any tool is a web search tool (native or already LiteLLM standard) - has_websearch: Final = any(is_web_search_tool(t) for t in tools) - + is_responses_call: Final = call_type in (CallTypes.responses, CallTypes.aresponses) + has_websearch: Final = ( + any(is_web_search_tool_responses(tool) for tool in tools) + if is_responses_call + else any(is_web_search_tool(tool) for tool in tools) + ) if not has_websearch: return None + if self.search_tool_name: + try: + from litellm.proxy.proxy_server import llm_router + except ImportError: + llm_router = None + self._select_search_tool_from_router(llm_router=llm_router) + + if is_responses_call: + return self._convert_responses_tools(kwargs=kwargs, tools=tools) + verbose_logger.debug("WebSearchInterception: Converting native web_search tools to LiteLLM standard") # If the client sent an Anthropic-native web_search_* tool, mark the @@ -1631,9 +1641,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": - if llm_router is None or not hasattr(llm_router, "search_tools"): - return None - search_tools: Final = list(getattr(llm_router, "search_tools") or []) + search_tools: Final = list(getattr(llm_router, "search_tools", []) or []) return self._select_search_tool_from_list(search_tools=search_tools, source="router") def _select_search_tool_from_list( @@ -1643,20 +1651,26 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> "_SearchToolConfig | None": if self.search_tool_name: matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name] - if matching_tools: - search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider") - verbose_logger.debug( - "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", - self.search_tool_name, - source, - search_provider, + if not matching_tools: + raise ValueError(f"Configured search tool '{self.search_tool_name}' was not found") + + selected_tool: Final = matching_tools[0] + litellm_params: Final = selected_tool.get("litellm_params") + selected_search_provider: Final = ( + litellm_params.get("search_provider") if isinstance(litellm_params, Mapping) else None + ) + if not isinstance(selected_search_provider, str) or not selected_search_provider.strip(): + raise ValueError( + f"Configured search tool '{self.search_tool_name}' does not define a valid search provider" ) - return matching_tools[0] + verbose_logger.debug( - "WebSearchInterception: Search tool '%s' not found in %s, falling back to first available or perplexity", + "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", self.search_tool_name, source, + selected_search_provider, ) + return selected_tool if search_tools: first_tool: Final = search_tools[0] diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index f39f41a6d12..ec4bc1f49eb 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -14,7 +14,7 @@ from litellm.integrations.websearch_interception.handler import ( ) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, ProxyException, UserAPIKeyAuth -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders def test_initialize_from_proxy_config(): @@ -230,6 +230,124 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp assert forwarded_kwargs["max_retries"] == 2 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_tools", "error"), + [ + pytest.param(None, "was not found", id="router-not-configured"), + pytest.param( + [{"search_tool_name": "other-search", "litellm_params": {"search_provider": "tavily"}}], + "was not found", + id="requested-tool-not-configured", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": "not-a-mapping"}], + "does not define a valid search provider", + id="invalid-parameters", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": {}}], + "does not define a valid search provider", + id="missing-provider", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": {"search_provider": " "}}], + "does not define a valid search provider", + id="whitespace-provider", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": {"search_provider": 123}}], + "does not define a valid search provider", + id="invalid-provider", + ), + ], +) +async def test_execute_search_rejects_invalid_explicit_search_tool(monkeypatch, search_tools, error): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger(search_tool_name="parallel-search") + router = None if search_tools is None else MagicMock(search_tools=search_tools) + mock_asearch = AsyncMock() + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + with pytest.raises(ValueError, match=f"Configured search tool 'parallel-search' {error}"): + await logger._execute_search("what is litellm") + + mock_asearch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_execute_search_honors_explicit_parallel_search_tool(monkeypatch): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger(search_tool_name="parallel-search") + router = MagicMock( + search_tools=[ + { + "search_tool_name": "other-search", + "litellm_params": {"search_provider": "tavily", "api_key": "other-key"}, + }, + { + "search_tool_name": "parallel-search", + "litellm_params": {"search_provider": "parallel_ai", "api_key": "parallel-key"}, + }, + ], + ) + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm") + + mock_asearch.assert_awaited_once_with( + query="what is litellm", + search_provider="parallel_ai", + api_key="parallel-key", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_tools", "expected_search_kwargs"), + [ + pytest.param(None, {"search_provider": "perplexity"}, id="router-not-configured"), + pytest.param( + [ + { + "search_tool_name": "first-search", + "litellm_params": {"search_provider": "tavily", "api_key": "first-key"}, + }, + { + "search_tool_name": "parallel-search", + "litellm_params": {"search_provider": "parallel_ai", "api_key": "parallel-key"}, + }, + ], + {"search_provider": "tavily", "api_key": "first-key"}, + id="first-configured-tool", + ), + ], +) +async def test_execute_search_preserves_implicit_provider_selection(monkeypatch, search_tools, expected_search_kwargs): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + router = None if search_tools is None else MagicMock(search_tools=search_tools) + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm") + + mock_asearch.assert_awaited_once_with(query="what is litellm", **expected_search_kwargs) + + @pytest.mark.asyncio async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch): """An intercepted search is billed and logged against the key that made the LLM request. @@ -397,6 +515,72 @@ async def test_execute_search_enforces_team_search_tool_permission(monkeypatch): mock_asearch.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "web_search_tool"), + [ + pytest.param( + CallTypes.acompletion, + {"type": "web_search_20250305", "name": "web_search"}, + id="chat-completion", + ), + pytest.param(CallTypes.responses, {"type": "web_search"}, id="responses"), + pytest.param(CallTypes.aresponses, {"type": "web_search"}, id="async-responses"), + pytest.param( + CallTypes.anthropic_messages, + {"type": "web_search_20250305", "name": "web_search"}, + id="anthropic-messages", + ), + ], +) +async def test_deployment_hook_dispatcher_propagates_missing_explicit_search_tool( + monkeypatch, call_type, web_search_tool +): + import litellm + from litellm.proxy import proxy_server + from litellm.utils import async_pre_call_deployment_hook + + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"], search_tool_name="parallel-search") + mock_asearch = AsyncMock() + kwargs = { + "model": "bedrock/claude-sonnet-4", + "tools": [web_search_tool], + "custom_llm_provider": "bedrock", + } + + monkeypatch.setattr( + proxy_server, + "llm_router", + MagicMock(search_tools=[{"search_tool_name": "other-search", "litellm_params": {"search_provider": "tavily"}}]), + ) + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + with pytest.raises(ValueError, match="Configured search tool 'parallel-search' was not found"): + await async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type.value) + + assert kwargs["tools"] == [web_search_tool] + mock_asearch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_deployment_hook_skips_explicit_tool_validation_for_non_search_responses(monkeypatch): + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"], search_tool_name="parallel-search") + monkeypatch.setattr(proxy_server, "llm_router", MagicMock(search_tools=[])) + + result = await logger.async_pre_call_deployment_hook( + kwargs={ + "tools": [{"type": "function", "name": "calculator"}], + "custom_llm_provider": "bedrock", + }, + call_type=CallTypes.aresponses, + ) + + assert result is None + + @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_provider_from_top_level_kwargs(): """Test that async_pre_call_deployment_hook finds custom_llm_provider at top-level kwargs. From bfea8a8c19ac83e2f7457f46a4c494ead35642e9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 21:31:08 -0700 Subject: [PATCH 077/126] feat(shadow_eval): compare several auto-routers on one job's sampled traffic (#39028) --- .../migration.sql | 3 + .../litellm_proxy_extras/schema.prisma | 4 +- litellm/integrations/shadow_eval_logger.py | 151 ++++-- .../auto_router_endpoints.py | 53 ++- litellm/proxy/schema.prisma | 4 +- .../auto_router_endpoints.py | 75 ++- schema.prisma | 4 +- .../integrations/test_shadow_eval_logger.py | 180 +++++++- .../test_auto_router_endpoints.py | 129 +++++- .../_components/ShadowEvalSection.test.tsx | 126 ++++- .../_components/ShadowEvalSection.tsx | 348 +------------- .../_components/ShadowEvalStartForm.tsx | 432 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 35 +- 13 files changed, 1124 insertions(+), 420 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql new file mode 100644 index 00000000000..90b21205310 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql @@ -0,0 +1,3 @@ +ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "router_names" TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[]; + +ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "router_name" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 01a607b68a9..7604ceadf7a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1531,7 +1531,8 @@ model LiteLLM_ShadowEvalJob { group_id String // legs of one job share this; the API's job id target_type String @default("key") // key | team | user target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows - router_name String // the auto-router under evaluation, in either direction + router_name String // first (often only) auto-router under evaluation; router_names is the full set + router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name) direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String @@ -1555,6 +1556,7 @@ model LiteLLM_ShadowEvalAttempt { job_id String request_id String // the judged real request outcome String // real | shadow | tie | error + router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router tier String? // router's tier for the prompt, when classified real_model String? shadow_model String? diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 18bda0a9d55..27da785331a 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -1,8 +1,11 @@ """Shadow Eval Logger: samples a shadowed key's successful LLM requests (chat completions, Anthropic Messages, and Responses API surfaces, each normalized to chat shape), duplicates -each against the job's other arm in a detached task (the auto-router for a forward job, the -fixed baseline model for a reverse one), blind-judges real vs shadow, and appends one -``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write. +each through every shadow arm in one detached task (each candidate auto-router for a +forward job, the fixed baseline model for a reverse one), blind-judges real vs each arm, +and appends one ``LiteLLM_ShadowEvalAttempt`` row per arm (verdict or error) as the +feature's only hot-path write. A multi-router job's arms therefore score the identical +sampled requests against the identical real responses, which is what makes their win +rates comparable head-to-head. Counts, status, and spend derive from those rows at read time, so nothing can disagree across pods or stop races; the hook reads active jobs through a short-TTL cache.""" @@ -498,12 +501,16 @@ def _decision_classifier_cost(metadata: Mapping[str, object]) -> float: return float(raw) if isinstance(raw, (int, float)) else 0.0 -def _request_was_routed_by(request_metadata: Mapping[str, object], router_name: str) -> bool: - """Whether the router under evaluation served this request, which is what decides - the direction it belongs to. A forward job skips its own router's traffic, since - duplicating it would compare the router to itself: guaranteed ties, judge spend for - zero information. A reverse job samples exactly that traffic and nothing else.""" - return _routing_decision(request_metadata).get("router_model_name") == router_name +def _direction_admits(request_metadata: Mapping[str, object], job: "ActiveShadowEvalJob") -> bool: + """Whether this request belongs to the job's direction. A forward job skips traffic + any of its candidate routers served: duplicating a router's own request compares it + to itself (guaranteed ties), and judging a sibling against another candidate's live + response would score candidates against each other instead of against the incumbent. + A reverse job samples exactly its one router's traffic and nothing else.""" + routed_by: Final = _routing_decision(request_metadata).get("router_model_name") + if job.direction == "reverse": + return routed_by == job.router_name + return routed_by not in job.arm_router_names @dataclass(frozen=True, slots=True) @@ -546,6 +553,7 @@ class ActiveShadowEvalJob(BaseModel): id: str router_name: str + router_names: tuple[str, ...] = () direction: ShadowEvalDirection = "forward" baseline_model: str | None = None shadow_percentage: float @@ -567,12 +575,25 @@ class ActiveShadowEvalJob(BaseModel): raise ValueError("baseline_model is set for exactly the reverse jobs") return self + @model_validator(mode="after") + def _reverse_evaluates_one_router(self) -> "ActiveShadowEvalJob": + """A reverse row naming several routers is unsamplable (there is no one traffic + slice they share) and fails closed.""" + if self.direction == "reverse" and len(self.arm_router_names) > 1: + raise ValueError("a reverse job evaluates exactly one router") + return self + @property - def shadow_target(self) -> str: - """The model the duplicated arm calls: the router itself for a forward job, the - fixed baseline for a reverse one. Total because the validator above pins + def arm_router_names(self) -> tuple[str, ...]: + """The job's full router set; rows from before router_names existed hold it in + router_name alone. The one place that reading lives on the sampling side.""" + return self.router_names or (self.router_name,) + + def arm_target(self, arm_router: str) -> str: + """The model one duplicated arm calls: the candidate router itself for a forward + job, the fixed baseline for a reverse one. Total because the validator above pins baseline_model to reverse jobs and only those.""" - return self.baseline_model or self.router_name + return self.baseline_model or arm_router def _as_active_job(record: object, attempts: int, spend: float) -> ActiveShadowEvalJob | None: @@ -696,7 +717,7 @@ class ShadowEvalLogger(CustomLogger): now >= job.ends_at or job.attempts + self._job_starts.get(job.id, 0) >= job.max_turns or (job.max_budget is not None and job.spend >= job.max_budget) - or _request_was_routed_by(request_metadata, job.router_name) != (job.direction == "reverse") + or not _direction_admits(request_metadata, job) ): continue if not _sample_hits(request_id, job.id, job.shadow_percentage): @@ -773,7 +794,10 @@ class ShadowEvalLogger(CustomLogger): if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS: self._record_funnel(job.id, "shed") continue - self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1 + # One start writes one attempt row per arm, and max_turns is a row + # ceiling, so admission must pre-count every arm or a multi-router + # job overshoots the valve N-fold within a cache generation. + self._job_starts[job.id] = self._job_starts.get(job.id, 0) + len(job.arm_router_names) self._inflight_shadow_tasks += 1 asyncio.create_task( self._run_shadow_eval( @@ -812,32 +836,74 @@ class ShadowEvalLogger(CustomLogger): shadow_params: Mapping[str, object], parent_metadata: Mapping[str, object], ) -> None: - """Budget gate -> shadow call -> blind judge -> one attempt row, and every exit - in exactly one coverage bucket: the gates that decline to spend on an admitted - sample (no DB to record into, an over-budget key, an unverifiable or exhausted - eval budget) count it withheld, so eligible traffic still reconciles as - not_sampled + unjudgeable + shed + withheld + attempt rows. The prisma gate sits - above the dispatch so no provider spend happens without a place to record the - outcome, and the budget read lives here rather than in the success hook.""" + """Budget gates once per sampled request, then every router arm in turn: shadow + call -> blind judge -> one attempt row stamped with the arm. The gates that + decline to spend on an admitted sample (no DB to record into, an over-budget key, + an unverifiable or exhausted eval budget) count the REQUEST withheld before any + arm runs, so funnel counters stay per-request and a leg's eligible traffic still + reconciles as not_sampled + unjudgeable + shed + withheld + sampled requests, + where each sampled request writes one attempt row per arm. A budget crossed + mid-loop lets the remaining arms overshoot by one round, the same class of + overshoot as the samples already in flight when the cap is crossed. The prisma + gate sits above the dispatch so no provider spend happens without a place to + record the outcome, and the budget read lives here rather than in the success + hook.""" prisma: Final = self._prisma_provider() + if prisma is None: + self._record_funnel(job.id, "withheld") + return + if await _key_or_team_is_over_budget(parent_metadata): + self._record_funnel(job.id, "withheld") + return + if job.max_budget is not None: + try: + spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget) + except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it + verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e) + self._record_funnel(job.id, "withheld") + return + if spend >= job.max_budget: + self._record_funnel(job.id, "withheld") + return + for arm_router in job.arm_router_names: + await self._run_shadow_arm( + prisma=prisma, + job=job, + arm_router=arm_router, + request_id=request_id, + messages=messages, + real_text=real_text, + real_model=real_model, + real_cost=real_cost, + real_classifier_cost=real_classifier_cost, + real_cache_hit=real_cache_hit, + control_tier=control_tier, + shadow_params=shadow_params, + parent_metadata=parent_metadata, + ) + + async def _run_shadow_arm( + self, + prisma: "PrismaClient", + job: ActiveShadowEvalJob, + arm_router: str, + request_id: str, + messages: Sequence[Mapping[str, object]], + real_text: str, + real_model: str, + real_cost: float, + real_classifier_cost: float, + real_cache_hit: bool, + control_tier: str | None, + shadow_params: Mapping[str, object], + parent_metadata: Mapping[str, object], + ) -> None: + """One arm's pipeline: shadow call -> blind judge -> one attempt row, every exit + recording this arm's outcome, so one arm's fault never silences a sibling arm.""" try: - if prisma is None: - self._record_funnel(job.id, "withheld") - return - if await _key_or_team_is_over_budget(parent_metadata): - self._record_funnel(job.id, "withheld") - return - if job.max_budget is not None: - try: - spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget) - except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it - verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e) - self._record_funnel(job.id, "withheld") - return - if spend >= job.max_budget: - self._record_funnel(job.id, "withheld") - return - shadow: Final = await self._call_router_shadow(job.shadow_target, messages, shadow_params, parent_metadata) + shadow: Final = await self._call_router_shadow( + job.arm_target(arm_router), messages, shadow_params, parent_metadata + ) except Exception as e: # noqa: BLE001 # detached task: nothing billed yet, record and never raise verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e) await self._record_attempt( @@ -845,6 +911,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=f"pipeline error: {e}", real_cost=real_cost, @@ -858,6 +925,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=shadow.error, shadow_cost=shadow.cost, @@ -882,6 +950,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=verdict.error, shadow=shadow, @@ -898,6 +967,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome=verdict.preference, shadow=shadow, real_model=real_model, @@ -916,6 +986,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=f"pipeline error: {e}", shadow=shadow, @@ -933,6 +1004,7 @@ class ShadowEvalLogger(CustomLogger): request_id: str, control_tier: str | None, *, + router_name: str, outcome: str, real_cost: float, real_classifier_cost: float, @@ -955,6 +1027,7 @@ class ShadowEvalLogger(CustomLogger): data={ # mutable-ok: Prisma payload "job_id": job.id, "request_id": request_id, + "router_name": router_name, "outcome": outcome, "tier": control_tier if job.direction == "reverse" else (shadow.tier if shadow else None), "real_model": real_model or None, diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 44b0cdcca2e..21e652114bc 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -833,7 +833,7 @@ def _judge_collisions_for_team( return tuple( (role, model) for role, model in ( - *_router_arm_models(llm_router, data.router_name), + *(arm for name in data.router_names for arm in _router_arm_models(llm_router, name)), *((("baseline", data.baseline_model),) if data.baseline_model is not None else ()), ) if judge & judge_target(llm_router, model, team_id).models @@ -904,7 +904,7 @@ class _AttemptAggRow(BaseModel): _ATTEMPT_AGG_ROWS: Final = TypeAdapter(list[_AttemptAggRow]) -_ATTEMPT_AGG_SELECT: Final = """ +_ATTEMPT_AGG_COLUMNS: Final = """ COUNT(*)::int AS turn_count, COUNT(*) FILTER (WHERE outcome = 'real')::int AS real_wins, COUNT(*) FILTER (WHERE outcome = 'shadow')::int AS shadow_wins, @@ -913,15 +913,34 @@ _ATTEMPT_AGG_SELECT: Final = """ COALESCE(SUM(real_cost + real_classifier_cost) FILTER (WHERE real_cost IS NOT NULL AND NOT real_cache_hit), 0)::float AS real_spend, COALESCE(SUM(shadow_cost + shadow_classifier_cost) FILTER (WHERE real_cost IS NOT NULL AND NOT real_cache_hit), 0)::float AS shadow_spend, COUNT(*) FILTER (WHERE real_cache_hit)::int AS cache_hit_turns +""" + +_ATTEMPT_AGG_SELECT: Final = ( + _ATTEMPT_AGG_COLUMNS + + """ FROM "LiteLLM_ShadowEvalAttempt" WHERE job_id = ANY($1::text[]) AND outcome != 'error' GROUP BY 1 """ +) _ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT _ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT _ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT +# Attempt rows from before arm stamping carry no router_name; they belong to the job's +# own router, which the join reads off the leg. +_ATTEMPT_AGG_BY_ROUTER_SQL: Final = ( + "SELECT COALESCE(a.router_name, j.router_name) AS grp," + + _ATTEMPT_AGG_COLUMNS + + """ +FROM "LiteLLM_ShadowEvalAttempt" a +JOIN "LiteLLM_ShadowEvalJob" j ON j.id = a.job_id +WHERE a.job_id = ANY($1::text[]) AND a.outcome != 'error' +GROUP BY 1 +""" +) + # These guards derive spend from attempt rows, the cross-pod authority; the sampler also # reads the live counter, so admission can stop before a row-based guard would fire (safe # direction, and mid-deploy rows from old pods price as judge-only until the deploy ends). @@ -1060,6 +1079,7 @@ class _LegRow(BaseModel): target_type: ShadowEvalTargetType target_id: str router_name: str + router_names: tuple[str, ...] = () direction: ShadowEvalDirection baseline_model: str | None = None judge_model: str @@ -1071,6 +1091,12 @@ class _LegRow(BaseModel): stopped_at: datetime | None = None stopped_by: str | None = None + @property + def arm_router_names(self) -> tuple[str, ...]: + """The job's full router set; rows from before router_names existed hold it in + router_name alone. The one place that reading lives on the endpoint side.""" + return self.router_names or (self.router_name,) + @field_validator("created_at", "ends_at", "stopped_at") @classmethod def _as_aware_utc(cls, value: datetime | None) -> datetime | None: @@ -1123,7 +1149,7 @@ def _group_response( ) for leg in sorted(legs, key=lambda leg: (leg.target_type, leg.target_id)) ), - router_name=first.router_name, + router_names=first.arm_router_names, direction=first.direction, baseline_model=first.baseline_model, judge_model=first.judge_model, @@ -1252,6 +1278,9 @@ async def _shadow_eval_results( for slice in _slices(by_leg) } ) + by_router: Final = _ATTEMPT_AGG_ROWS.validate_python( + await _query_raw(prisma_client, _ATTEMPT_AGG_BY_ROUTER_SQL, leg_ids) or () + ) total_turns: Final = sum(r.turn_count for r in by_tier) funnel_rows: Final = await _query_raw(prisma_client, _FUNNEL_TOTALS_SQL, leg_ids) counted: Final = _FunnelTotalsRow.model_validate(funnel_rows[0]) if funnel_rows else None @@ -1261,6 +1290,7 @@ async def _shadow_eval_results( result: Final = ShadowEvalResult( by_tier=_slices(by_tier), by_current_model=_slices(by_model), + by_router=_slices(by_router), overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns), overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns), sampled_real_spend=sum(r.real_spend for r in by_tier), @@ -1314,8 +1344,15 @@ async def start_shadow_eval( _require_admin_writer(user_api_key_dict, "start a shadow eval") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, data.router_name): - raise HTTPException(status_code=400, detail=f"'{data.router_name}' is not a configured auto-router") + unconfigured: Final = tuple( + name + for name in data.router_names + if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, name) + ) + if unconfigured: + raise HTTPException( + status_code=400, detail=f"Not a configured auto-router: {', '.join(repr(n) for n in unconfigured)}" + ) token_rows: Final = ( await _verification_tokens(prisma_client).find_many( where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter @@ -1416,7 +1453,9 @@ async def start_shadow_eval( ends_at: Final = now + timedelta(days=data.duration_days) shared_config: Final = { # mutable-ok: Prisma payload "group_id": group_id, - "router_name": data.router_name, + # a pre-router_names pod samples router_name alone, so it must be a real arm + "router_name": data.router_names[0], + "router_names": list(data.router_names), # mutable-ok: Prisma payload "direction": data.direction, "baseline_model": data.baseline_model, "judge_model": data.judge_model, @@ -1477,7 +1516,7 @@ async def start_shadow_eval( ) for target_type, target_id in sorted(requested_targets) ), - router_name=data.router_name, + router_names=data.router_names, direction=data.direction, baseline_model=data.baseline_model, judge_model=data.judge_model, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 01a607b68a9..7604ceadf7a 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1531,7 +1531,8 @@ model LiteLLM_ShadowEvalJob { group_id String // legs of one job share this; the API's job id target_type String @default("key") // key | team | user target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows - router_name String // the auto-router under evaluation, in either direction + router_name String // first (often only) auto-router under evaluation; router_names is the full set + router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name) direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String @@ -1555,6 +1556,7 @@ model LiteLLM_ShadowEvalAttempt { job_id String request_id String // the judged real request outcome String // real | shadow | tie | error + router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router tier String? // router's tier for the prompt, when classified real_model String? shadow_model String? diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index dfc45ccf9bf..88869a1edfb 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -251,8 +251,12 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5" # Sample-count ceiling written on every new job: a zero-cost error loop (a shadow arm that # fails before billing) never consumes spend budget, so it must terminate on count instead. +# A multi-router job writes one attempt row per router arm, so the valve is reached +# proportionally sooner; it is a safety valve, not a sample budget. SHADOW_EVAL_TURN_VALVE: Final[int] = 10_000 +SHADOW_EVAL_MAX_ROUTERS: Final[int] = 4 + class StartShadowEvalRequest(BaseModel): """Start duplicating one or more targets' traffic for blind comparison against an auto-router. @@ -288,7 +292,24 @@ class StartShadowEvalRequest(BaseModel): "to across all their teams: JWT requests carrying their subject claim and virtual keys they own" ), ) - router_name: str = Field(description="The auto-router under evaluation, in either direction") + router_name: str | None = Field( + default=None, + description=( + "The auto-router under evaluation, in either direction: the single-router spelling of " + "router_names. Provide exactly one of the two fields" + ), + ) + router_names: tuple[str, ...] = Field( + default=(), + max_length=SHADOW_EVAL_MAX_ROUTERS, + description=( + "The auto-routers under evaluation, at most " + f"{SHADOW_EVAL_MAX_ROUTERS}. Every sampled request runs through every router listed and each " + "arm is judged independently against the same real response, so routers compare head-to-head " + "on identical traffic. More than one router requires direction 'forward'. After validation " + "this field always carries the full deduplicated set, whichever spelling the caller used" + ), + ) direction: ShadowEvalDirection = Field( default="forward", description=( @@ -332,7 +353,8 @@ class StartShadowEvalRequest(BaseModel): "Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with " "the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval " "spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight " - "samples can overshoot the cap by one sampling cache window" + "samples can overshoot the cap by one sampling cache window. Every router arm draws from the " + "same per-target budget, so a multi-router job reaches it proportionally sooner" ), ) @@ -373,6 +395,23 @@ class StartShadowEvalRequest(BaseModel): raise ValueError("baseline_model is only meaningful when direction is 'reverse'") return self + @model_validator(mode="after") + def _resolve_router_set(self) -> "StartShadowEvalRequest": + """Whichever spelling the caller used, router_names leaves validation as the full + deduplicated set, so every downstream reader consumes one field.""" + if (self.router_name is None) == (not self.router_names): + raise ValueError("provide exactly one of router_name or router_names") + single: Final = () if self.router_name is None else (self.router_name,) + routers: Final = tuple(dict.fromkeys(self.router_names or single)) + if not all(name.strip() for name in routers): + raise ValueError("router names must be non-empty strings") + if len(routers) > 1 and self.direction == "reverse": + raise ValueError("a reverse job evaluates one router against baseline_model; pass a single router") + # A returned model_copy is ignored on the __init__ construction path, so the + # normalization must land as a self attribute store to hold for every caller. + self.router_names = routers + return self + class ShadowEvalSlice(BaseModel): """Judge outcomes for one slice of a job's verdicts: a router tier, one of the @@ -428,15 +467,28 @@ class ShadowEvalResult(BaseModel): "and in reverse the models the router itself picked" ) ) + by_router: tuple[ShadowEvalSlice, ...] = Field( + default=(), + description=( + "One slice per router arm, grouped on the router name. Every arm of a multi-router job is " + "judged against the same real responses over the same sampled requests, so these slices " + "compare routers head-to-head: like-for-like win rates and spends on identical traffic. " + "Verdicts from before arm stamping existed count toward the job's own router" + ), + ) overall_shadow_win_rate_pct: float overall_tie_rate_pct: float sampled_real_spend: float = Field( default=0.0, - description="USD the real arm billed across all judged turns, cache-served turns excluded", + description=( + "USD the real arm billed across all judged turns, cache-served turns excluded. A judged turn " + "is one (request, router arm) verdict, so a multi-router job counts the real response once per " + "arm it was judged against; per-router comparisons read by_router" + ), ) sampled_shadow_spend: float = Field( default=0.0, - description="USD the shadow arm billed across the same turns, judge excluded, like for like", + description="USD the shadow arms billed across the same turns, judge excluded, like for like", ) not_sampled_count: int | None = Field( default=None, @@ -540,7 +592,13 @@ class ShadowEvalJobResponse(BaseModel): min_length=1, description="The targets whose traffic this job evaluates, and only theirs, each with its own budget", ) - router_name: str + router_names: tuple[str, ...] = Field( + min_length=1, + description=( + "Every auto-router this job runs as a shadow arm. Multi-router jobs sample one slice of " + "traffic and judge every arm against the same real responses" + ), + ) direction: ShadowEvalDirection = "forward" baseline_model: str | None = None judge_model: str @@ -562,6 +620,13 @@ class ShadowEvalJobResponse(BaseModel): last_error: str | None = Field(default=None, description="Most recent attempt error; detail endpoint only") results: ShadowEvalResult | None = Field(default=None, description="Stratified verdicts; detail endpoint only") + @computed_field + @property + def router_name(self) -> str: + """The first router, kept for callers that predate router_names; derived so the + two fields can never disagree.""" + return self.router_names[0] + @computed_field @property def status(self) -> ShadowEvalStatus: diff --git a/schema.prisma b/schema.prisma index 01a607b68a9..7604ceadf7a 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1531,7 +1531,8 @@ model LiteLLM_ShadowEvalJob { group_id String // legs of one job share this; the API's job id target_type String @default("key") // key | team | user target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows - router_name String // the auto-router under evaluation, in either direction + router_name String // first (often only) auto-router under evaluation; router_names is the full set + router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name) direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String @@ -1555,6 +1556,7 @@ model LiteLLM_ShadowEvalAttempt { job_id String request_id String // the judged real request outcome String // real | shadow | tie | error + router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router tier String? // router's tier for the prompt, when classified real_model String? shadow_model String? diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 1af3dd3f613..5628d69de26 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -65,6 +65,7 @@ def _job_record(job: ActiveShadowEvalJob, target_type="key", target_id="key-hash target_type=target_type, target_id=target_id, router_name=job.router_name, + router_names=job.router_names, direction=job.direction, baseline_model=job.baseline_model, shadow_percentage=job.shadow_percentage, @@ -81,6 +82,7 @@ def _router( shadow_text="shadow answer", judge_json='{"preference": "A", "confidence": 0.9, "reasoning": "x"}', classifier_cost=None, + sibling_router_texts=None, ): """One mock router serving the shadow call first, the judge call second, told apart by the internal-origin stamp rather than the model, since a reverse job's shadow arm names @@ -100,6 +102,15 @@ def _router( decision["classifier_cost"] = classifier_cost kwargs["metadata"]["routing_decision"] = decision return {"choices": [{"message": {"content": shadow_text}}], "usage": {"completion_tokens": 5}} + if sibling_router_texts and kwargs["model"] in sibling_router_texts: + kwargs["metadata"]["routing_decision"] = { + "tier_label": "MEDIUM", + "routed_model": f"{kwargs['model']}-pick", + } + return { + "choices": [{"message": {"content": sibling_router_texts[kwargs["model"]]}}], + "usage": {"completion_tokens": 5}, + } return ModelResponse( model=kwargs["model"], choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": shadow_text}}], @@ -1210,29 +1221,36 @@ class TestJobValidation: {"direction": "reverse"}, {"baseline_model": "baseline-model"}, {"direction": "sideways", "baseline_model": "baseline-model"}, + {"direction": "reverse", "baseline_model": "baseline-model", "router_names": ("a", "b")}, ], - ids=["reverse-without-baseline", "forward-with-baseline", "unknown-direction"], + ids=["reverse-without-baseline", "forward-with-baseline", "unknown-direction", "reverse-with-router-set"], ) def test_unsamplable_shapes_are_rejected(self, overrides): with pytest.raises(ValidationError): _job(**overrides) - def test_shadow_target_follows_direction(self): - assert _job().shadow_target == "my-router" - assert _reverse_job().shadow_target == "baseline-model" + def test_arm_target_follows_direction(self): + assert _job().arm_target("my-router") == "my-router" + assert _reverse_job().arm_target("my-router") == "baseline-model" + + def test_rows_from_before_router_names_carry_their_set_in_router_name(self): + assert _job().arm_router_names == ("my-router",) + assert _job(router_names=("my-router", "alt-router")).arm_router_names == ("my-router", "alt-router") @pytest.mark.asyncio class TestDirection: @pytest.mark.parametrize( - "job,routed_by,sampled", + "job,routed_by,attempt_rows", [ - (_job(), None, True), - (_job(), "my-router", False), - (_job(), "other-router", True), - (_reverse_job(), "my-router", True), - (_reverse_job(), None, False), - (_reverse_job(), "other-router", False), + (_job(), None, 1), + (_job(), "my-router", 0), + (_job(), "other-router", 1), + (_reverse_job(), "my-router", 1), + (_reverse_job(), None, 0), + (_reverse_job(), "other-router", 0), + (_job(router_names=("my-router", "alt-router")), "alt-router", 0), + (_job(router_names=("my-router", "alt-router")), "other-router", 2), ], ids=[ "forward-samples-unrouted", @@ -1241,20 +1259,24 @@ class TestDirection: "reverse-samples-its-own-router", "reverse-skips-unrouted", "reverse-skips-another-router", + "forward-skips-any-candidates-own-traffic", + "forward-multi-samples-once-per-arm", ], ) - async def test_direction_decides_which_traffic_is_sampled(self, job, routed_by, sampled): + async def test_direction_decides_which_traffic_is_sampled(self, job, routed_by, attempt_rows): """The two directions partition the key's traffic: whatever one samples, the other - skips, so a key running both never judges the same turn twice for the same reason.""" + skips, so a key running both never judges the same turn twice for the same reason. + A multi-router job extends the forward skip to every candidate: a request one + candidate served must not be judged as the incumbent against another candidate.""" prisma = _prisma() - logger = _logger(router=_router(), prisma=prisma, jobs=(job,)) + logger = _logger(router=_router(sibling_router_texts={"alt-router": "alt answer"}), prisma=prisma, jobs=(job,)) await logger.async_log_success_event( _success_kwargs(request_metadata=_routed_by(routed_by) if routed_by else {}), RESPONSE, None, None ) await _drain(logger) - assert prisma.db.litellm_shadowevalattempt.create.await_count == int(sampled) + assert prisma.db.litellm_shadowevalattempt.create.await_count == attempt_rows async def test_reverse_duplicates_against_the_baseline_model(self): prisma = _prisma() @@ -1316,6 +1338,134 @@ class TestDirection: assert logger._job_starts == {"forward-job": 1, "reverse-job": 1} +@pytest.mark.asyncio +class TestMultiRouterArms: + async def test_every_arm_judges_the_same_request_and_stamps_its_own_row(self): + """One sampled request, one row per candidate router, both judged against the same + real response: the paired comparison that makes multi-router win rates comparable.""" + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + logger = _logger(router=router, prisma=prisma) + + await logger._run_shadow_eval( + job=_job(router_names=("my-router", "alt-router")), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.001, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.await_args_list] + assert [row["router_name"] for row in rows] == ["my-router", "alt-router"] + assert {row["request_id"] for row in rows} == {"req-1"} + assert [row["shadow_model"] for row in rows] == ["cheap-model", "alt-router-pick"] + assert all(row["outcome"] in ("real", "shadow", "tie") for row in rows) + assert all(row["real_cost"] == 0.001 for row in rows) + + async def test_a_single_router_job_stamps_its_router_on_the_row(self): + prisma = _prisma() + logger = _logger(router=_router(), prisma=prisma) + + await logger._run_shadow_eval( + job=_job(), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + row = prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"] + assert row["router_name"] == "my-router" + + async def test_one_arms_failure_never_silences_the_sibling(self): + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + healthy = router.acompletion.side_effect + + async def first_arm_explodes(**kwargs): + if kwargs["model"] == "my-router": + raise RuntimeError("provider exploded") + return await healthy(**kwargs) + + router.acompletion.side_effect = first_arm_explodes + logger = _logger(router=router, prisma=prisma) + + await logger._run_shadow_eval( + job=_job(router_names=("my-router", "alt-router")), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.await_args_list] + assert [row["router_name"] for row in rows] == ["my-router", "alt-router"] + assert rows[0]["outcome"] == "error" + assert "provider exploded" in rows[0]["error"] + assert rows[1]["outcome"] in ("real", "shadow", "tie") + + async def test_the_turn_valve_counts_every_arm_a_start_will_write(self): + """max_turns is a row ceiling and one sampled request writes one row per arm, so + admission pre-counts the arms: a two-arm job with two turns of budget admits one + request, not two.""" + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + logger = _logger( + router=router, prisma=prisma, jobs=(_job(router_names=("my-router", "alt-router"), max_turns=2),) + ) + + await logger.async_log_success_event(_success_kwargs(request_id="req-1"), RESPONSE, None, None) + await logger.async_log_success_event(_success_kwargs(request_id="req-2"), RESPONSE, None, None) + await _drain(logger) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.await_args_list] + assert {row["request_id"] for row in rows} == {"req-1"} + assert len(rows) == 2 + + async def test_a_withheld_request_runs_no_arm_and_counts_once(self): + """The budget gates run once per sampled request, before any arm: funnel counters + stay per-request, so coverage math is arm-count independent.""" + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + logger = _logger(router=router, prisma=prisma) + + await logger._run_shadow_eval( + job=_job(router_names=("my-router", "alt-router"), max_budget=1.0, spend=2.0), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + router.acompletion.assert_not_called() + prisma.db.litellm_shadowevalattempt.create.assert_not_called() + assert logger._test_funnel == [("job-1", "withheld")] + + @pytest.mark.asyncio class TestActiveJobsFailClosed: async def test_a_row_the_sampler_cannot_read_is_dropped_not_guessed(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index a74aa553449..c525af84511 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -881,6 +881,7 @@ def _leg_record(**overrides: object) -> MagicMock: "target_type": "key", "target_id": "key-hash", "router_name": "my-router", + "router_names": (), "direction": "forward", "baseline_model": None, "judge_model": "anthropic/claude-sonnet-5", @@ -931,6 +932,7 @@ def _shadow_prisma( legs=(), agg_rows=None, by_leg_rows=None, + by_router_rows=None, known_keys=("key-hash", "key-hash-2"), key_teams=None, known_teams=None, @@ -1030,6 +1032,7 @@ def _shadow_prisma( "target_type", "target_id", "router_name", + "router_names", "direction", "baseline_model", "judge_model", @@ -1065,6 +1068,8 @@ def _shadow_prisma( return [{"judged_count": 10, "error_count": 2, "judge_spend": 0.031}] if "SELECT job_id AS grp" in sql: return by_leg_rows if by_leg_rows is not None else [] + if "COALESCE(a.router_name" in sql: + return by_router_rows if by_router_rows is not None else [] if 'FROM "LiteLLM_ShadowEvalFunnel"' in sql: return prisma.funnel_rows return agg_rows if agg_rows is not None else [] @@ -1121,7 +1126,15 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once() rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] assert [(row["target_type"], row["target_id"]) for row in rows] == [("key", "key-hash"), ("key", "key-hash-2")] - assert len({frozenset((k, v) for k, v in row.items() if k not in ("target_id", "id")) for row in rows}) == 1 + assert ( + len( + { + frozenset((k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id")) + for row in rows + } + ) + == 1 + ) assert len({row["id"] for row in rows}) == len(rows) assert len({row["group_id"] for row in rows}) == 1 assert all(row["max_turns"] == SHADOW_EVAL_TURN_VALVE and row["created_by"] == "admin" for row in rows) @@ -1138,6 +1151,63 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp assert all(target.max_turns == SHADOW_EVAL_TURN_VALVE for target in response.targets) +@pytest.mark.asyncio +async def test_start_shadow_eval_multi_router_writes_the_set_on_every_leg(monkeypatch: pytest.MonkeyPatch): + """A multi-router job stores the full set in router_names and the first router in + router_name, so a rolling-deploy pod that predates router_names still runs a valid + single-arm eval and its unstamped attempt rows attribute to that first router.""" + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + response = await start_shadow_eval( + _start_request(router_name=None, router_names=("my-router", "classifier-router")), ADMIN + ) + + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert all(row["router_name"] == "my-router" for row in rows) + assert all(row["router_names"] == ["my-router", "classifier-router"] for row in rows) + assert response.router_names == ("my-router", "classifier-router") + assert response.router_name == "my-router" + + +@pytest.mark.asyncio +async def test_start_shadow_eval_rejects_an_unconfigured_router_in_the_set(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException, match="not-a-router") as exc: + await start_shadow_eval(_start_request(router_name=None, router_names=("my-router", "not-a-router")), ADMIN) + + assert exc.value.status_code == 400 + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_judge_collision_is_found_on_every_router_of_the_set(monkeypatch: pytest.MonkeyPatch): + """The judge-as-candidate guard walks every candidate router: a judge that serves an + arm of the SECOND router still poisons the whole job's win rates.""" + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException, match="also an arm") as exc: + await start_shadow_eval(_start_request(router_name=None, router_names=("my-router", "sonnet-router")), ADMIN) + + assert exc.value.status_code == 400 + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + @pytest.mark.asyncio async def test_start_shadow_eval_rejects_an_uncredentialed_sdk_judge(monkeypatch: pytest.MonkeyPatch) -> None: import litellm @@ -1798,6 +1868,63 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke assert error_where == {"job_id": {"in": ["leg-1", "leg-2"]}, "outcome": "error"} +@pytest.mark.asyncio +async def test_get_shadow_eval_job_slices_results_per_router(monkeypatch: pytest.MonkeyPatch): + """A multi-router job's detail carries one slice per arm, aggregated by the arm + stamped on each attempt row, with unstamped legacy rows attributed to the job's own + router by the read (the COALESCE against the leg's router_name).""" + import litellm.proxy.proxy_server as proxy_server + + def agg(grp: str, wins: int) -> dict[str, object]: + return { + "grp": grp, + "turn_count": 4, + "real_wins": 4 - wins, + "shadow_wins": wins, + "ties": 0, + "avg_confidence": 0.8, + "real_spend": 0.08, + "shadow_spend": 0.02, + "cache_hit_turns": 0, + } + + prisma = _shadow_prisma( + legs=[_leg_record(router_names=("my-router", "alt-router"))], + agg_rows=[agg("SIMPLE", 3)], + by_router_rows=[agg("my-router", 1), agg("alt-router", 3)], + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response = await get_shadow_eval_job("job-1", VIEWER) + + assert response.router_names == ("my-router", "alt-router") + assert response.router_name == "my-router" + assert [(s.group, s.shadow_win_rate_pct) for s in response.results.by_router] == [ + ("my-router", 25.0), + ("alt-router", 75.0), + ] + router_sql = next( + call.args[0] for call in prisma.db.query_raw.await_args_list if "COALESCE(a.router_name" in call.args[0] + ) + assert "COALESCE(a.router_name, j.router_name)" in router_sql + assert 'JOIN "LiteLLM_ShadowEvalJob" j ON j.id = a.job_id' in router_sql + assert "a.job_id = ANY($1::text[])" in router_sql + + +@pytest.mark.asyncio +async def test_job_responses_resolve_router_names_with_legacy_fallback(monkeypatch: pytest.MonkeyPatch): + """Rows from before router_names existed carry their whole set in router_name.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record(router_names=())]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response = await get_shadow_eval_job("job-1", VIEWER) + + assert response.router_names == ("my-router",) + assert response.router_name == "my-router" + + @pytest.mark.asyncio async def test_get_shadow_eval_job_404s_and_gates_on_role(monkeypatch: pytest.MonkeyPatch): import litellm.proxy.proxy_server as proxy_server diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index 8b397f20552..64e03985f57 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -102,6 +102,7 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ job_id: "job-1", status: "running", router_name: "claude-auto", + router_names: ["claude-auto"], direction: "forward", baseline_model: null, judge_model: "anthropic/claude-sonnet-5", @@ -436,7 +437,7 @@ describe("ShadowEvalSection", () => { await user.click(within(keyList).getByText("prod-alpha")); await user.click(keyInput); await user.click(within(keyList).getByText("staging-beta")); - await user.click(screen.getByPlaceholderText("Select an auto-router")); + await user.click(screen.getByPlaceholderText("Select up to 4 auto-routers")); await user.click(await screen.findByText("gpt-auto")); expect(screen.getByText("Start shadow eval")).toBeDisabled(); @@ -449,7 +450,7 @@ describe("ShadowEvalSection", () => { api_key_ids: ["hash-alpha", "hash-beta"], team_ids: [], user_ids: [], - router_name: "gpt-auto", + router_names: ["gpt-auto"], direction: "forward", shadow_percentage: 10, duration_days: 7, @@ -469,7 +470,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByPlaceholderText("Search teams by alias")); const teamList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(teamList).getByText("engineering")); - await user.click(screen.getByPlaceholderText("Select an auto-router")); + await user.click(screen.getByPlaceholderText("Select up to 4 auto-routers")); await user.click(await screen.findByText("gpt-auto")); await user.click(screen.getByPlaceholderText("Select a judge model")); await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); @@ -479,7 +480,7 @@ describe("ShadowEvalSection", () => { api_key_ids: [], team_ids: ["team-eng"], user_ids: [], - router_name: "gpt-auto", + router_names: ["gpt-auto"], direction: "forward", shadow_percentage: 10, duration_days: 7, @@ -501,7 +502,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByPlaceholderText("Search keys by alias")); const keyList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(keyList).getByText("prod-alpha")); - await user.click(screen.getByPlaceholderText("Select an auto-router")); + await user.click(screen.getByPlaceholderText("Select up to 4 auto-routers")); await user.click(await screen.findByText("gpt-auto")); await user.click(screen.getByPlaceholderText("Select a judge model")); await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); @@ -517,7 +518,7 @@ describe("ShadowEvalSection", () => { api_key_ids: ["hash-alpha"], team_ids: [], user_ids: [], - router_name: "gpt-auto", + router_names: ["gpt-auto"], direction: "reverse", baseline_model: "prod-claude", shadow_percentage: 10, @@ -528,6 +529,119 @@ describe("ShadowEvalSection", () => { expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); + it("submits every picked auto-router so one job compares them on the same traffic", async () => { + const user = userEvent.setup(); + const { start } = mockHooks({}); + render(); + + await user.click(screen.getByPlaceholderText("Search keys by alias")); + const keyList = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(keyList).getByText("prod-alpha")); + const routerInput = screen.getByPlaceholderText("Select up to 4 auto-routers"); + await user.click(routerInput); + await user.click(await screen.findByText("gpt-auto")); + await user.click(routerInput); + await user.click(await screen.findByText("claude-auto")); + expect( + screen.getByText("Every router sees the same sampled requests, judged against the same live responses"), + ).toBeInTheDocument(); + await user.click(screen.getByPlaceholderText("Select a judge model")); + await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(screen.getByText("Start shadow eval")); + + const expectedBody = { + api_key_ids: ["hash-alpha"], + team_ids: [], + user_ids: [], + router_names: ["gpt-auto", "claude-auto"], + direction: "forward", + shadow_percentage: 10, + duration_days: 7, + max_budget: 10, + judge_model: "anthropic/claude-sonnet-5", + }; + expect(start.mutate).toHaveBeenCalledWith(expectedBody); + }); + + it("blocks starting a reverse job with more than one router and says why", async () => { + const user = userEvent.setup(); + mockHooks({}); + render(); + + await user.click(screen.getByPlaceholderText("Search keys by alias")); + const keyList = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(keyList).getByText("prod-alpha")); + const routerInput = screen.getByPlaceholderText("Select up to 4 auto-routers"); + await user.click(routerInput); + await user.click(await screen.findByText("gpt-auto")); + await user.click(routerInput); + await user.click(await screen.findByText("claude-auto")); + await user.click(screen.getByText("Adoption check: key's traffic vs the router")); + await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await user.click(screen.getByPlaceholderText("Select a judge model")); + await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(screen.getByPlaceholderText("Select a baseline model")); + await user.click(screen.getByRole("option", { name: /prod-claude/ })); + + expect(screen.getByText("A regression check compares one router to its baseline")).toBeInTheDocument(); + expect(screen.getByText("Start shadow eval")).toBeDisabled(); + }); + + it("renders a per-router comparison table only when the job ran several routers", () => { + const routerSlice = (group: string, wins: number) => ({ + group, + turn_count: 20, + real_win_rate_pct: 100 - wins - 10, + shadow_win_rate_pct: wins, + tie_rate_pct: 10, + avg_judge_confidence: 0.8, + real_spend: 0.4, + shadow_spend: 0.2, + cache_hit_turns: 0, + }); + const base = job(); + const multi = job({ + router_names: ["claude-auto", "gpt-auto"], + results: { ...base.results!, by_router: [routerSlice("claude-auto", 40), routerSlice("gpt-auto", 70)] }, + }); + mockHooks({ jobs: [multi], detailsById: { "job-1": multi } }); + render(); + + expect(screen.getByText("Router")).toBeInTheDocument(); + const rows = screen.getAllByRole("row").map((row) => row.textContent ?? ""); + expect(rows.some((text) => text.includes("claude-auto") && text.includes("40.0%"))).toBe(true); + expect(rows.some((text) => text.includes("gpt-auto") && text.includes("70.0%"))).toBe(true); + expect( + screen.getByText( + (_, element) => + element?.textContent === "Shadowing 10% of prod-alpha traffic via claude-auto, gpt-auto" && + element.tagName === "P", + ), + ).toBeInTheDocument(); + }); + + it("renders a job from an older proxy that predates router_names", () => { + const legacy = { ...job(), router_names: undefined } as unknown as ShadowEvalJob; + mockHooks({ jobs: [legacy], detailsById: { "job-1": legacy } }); + render(); + + expect( + screen.getByText( + (_, element) => + element?.textContent === "Shadowing 10% of prod-alpha traffic via claude-auto" && element.tagName === "P", + ), + ).toBeInTheDocument(); + }); + + it("keeps the per-router table hidden for a single-router job", () => { + const base = job(); + const single = job({ results: { ...base.results!, by_router: [] } }); + mockHooks({ jobs: [single], detailsById: { "job-1": single } }); + render(); + + expect(screen.queryByText("Router")).not.toBeInTheDocument(); + }); + it("flips the arm labels and headline for a reverse job's results", () => { const j = job({ direction: "reverse", baseline_model: "openai/gpt-4o" }); mockHooks({ jobs: [j], detailsById: { "job-1": j } }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx index ec145bc617a..c66d74074c2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx @@ -2,32 +2,21 @@ import React, { useMemo, useState } from "react"; -import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; -import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; -import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; -import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; -import TeamMultiSelect from "@/components/common_components/team_multi_select"; -import { userOptionLabel } from "@/components/common_components/UserDropdown"; -import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { CircleHelp } from "lucide-react"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Card } from "@/components/ui/card"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; -import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { ApiError } from "@/lib/http/client"; import { usd } from "./costOptimizationUtils"; +import { StartForm } from "./ShadowEvalStartForm"; import { useShadowEvalJob, useShadowEvalJobs, - useStartShadowEval, useStopShadowEval, type ShadowEvalJob, type ShadowEvalJobTarget, @@ -96,17 +85,19 @@ const targetStatus = (job: ShadowEvalJob, target: ShadowEvalJobTarget): string = return target.stopped_at != null ? "stopped" : "running"; }; +const jobRouters = (job: ShadowEvalJob): string => (job.router_names ?? [job.router_name]).join(", "); + const jobHeadline = (job: ShadowEvalJob): React.ReactNode => job.direction === "reverse" ? ( <> - Comparing {job.router_name} to{" "} + Comparing {jobRouters(job)} to{" "} {job.baseline_model} on {job.shadow_percentage}% of{" "} {shadowedTargetsLabel(job)} traffic ) : ( <> Shadowing {job.shadow_percentage}% of {shadowedTargetsLabel(job)}{" "} - traffic via {job.router_name} + traffic via {jobRouters(job)} ); @@ -352,6 +343,11 @@ const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({
+ {(results.by_router ?? []).length > 1 && ( +
+ +
+ )} {results.by_current_model.length > 0 && ( { - const { data: costMap } = useModelCostMap(); - return useMemo(() => { - if (!costMap) return []; - const chatModels = Object.entries(costMap as Record) - .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) - .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); - return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); - }, [costMap]); -}; - -const useJudgeModelOptions = (): SearchSelectOption[] => { - const chatModels = useChatModelNames(); - return useMemo(() => { - const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ - label: model, - value: model, - sublabel: "Recommended", - })); - const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); - const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); - return [...pinned, ...rest]; - }, [chatModels]); -}; - -const useBaselineModelOptions = (): SearchSelectOption[] => { - const configuredGroups = usePlainModelGroups(); - const chatModels = useChatModelNames(); - return useMemo(() => { - const configured = [...configuredGroups] - .toSorted((a, b) => a.localeCompare(b)) - .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); - const rest = chatModels - .filter((model) => !configuredGroups.has(model)) - .map((model) => ({ label: model, value: model })); - return [...configured, ...rest]; - }, [configuredGroups, chatModels]); -}; - -const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ - { value: "forward", label: "Adoption check: key's traffic vs the router" }, - { value: "reverse", label: "Regression check: router's picks vs a baseline" }, -] as const; - -const START_FORM_DESCRIPTION: Record = { - forward: - "Duplicates a sampled slice of the selected targets' traffic (keys, teams, or users) through the auto-router and has an LLM judge compare both answers blind. Each target gets its own spend budget. The router's answers are never served to users; judge calls bill to the sampled traffic's own identity.", - reverse: - "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. Each target gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the sampled traffic's own identity.", -}; - -const DURATION_OPTIONS = [ - { value: "1", label: "1 day" }, - { value: "3", label: "3 days" }, - { value: "7", label: "7 days" }, - { value: "14", label: "14 days" }, - { value: "30", label: "30 days" }, -] as const; - -const Field: React.FC<{ label: string; htmlFor?: string; className?: string; children: React.ReactNode }> = ({ - label, - htmlFor, - className, - children, -}) => ( -
- - {children} -
-); - -const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => void }> = ({ value, onChange }) => { - const [search, setSearch] = useState(""); - const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteKeys(50, { - selectedKeyAlias: search || null, - }); - const options = useMemo( - () => - (data?.pages ?? []) - .flatMap((page) => page.keys) - .map((key) => ({ - label: key.key_alias || key.key_name || key.token, - value: key.token, - sublabel: key.token, - })), - [data], - ); - return ( - void fetchNextPage()} - hasNextPage={hasNextPage} - isFetchingNextPage={isFetchingNextPage} - isLoading={isPending} - placeholder="Search keys by alias" - emptyText="No matching keys" - errorText={isError ? "Keys could not be loaded. Refresh the page to retry." : undefined} - /> - ); -}; - -const UserSelect: React.FC<{ value: string[]; onChange: (ids: string[]) => void }> = ({ value, onChange }) => { - const [search, setSearch] = useState(""); - const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteUsers( - 50, - search || undefined, - ); - const options = useMemo( - () => - Array.from( - new Map( - (data?.pages ?? []) - .flatMap((page) => page.users) - .map((user) => [user.user_id, { label: userOptionLabel(user), value: user.user_id }] as const), - ).values(), - ), - [data], - ); - return ( - void fetchNextPage()} - hasNextPage={hasNextPage} - isFetchingNextPage={isFetchingNextPage} - isLoading={isPending} - placeholder="Search users by email" - emptyText="No matching users" - errorText={isError ? "Users could not be loaded. Refresh the page to retry." : undefined} - /> - ); -}; - -const StartForm: React.FC = () => { - const { accessToken } = useAuthorized(); - const [apiKeyIds, setApiKeyIds] = useState([]); - const [teamIds, setTeamIds] = useState([]); - const [userIds, setUserIds] = useState([]); - const [routerName, setRouterName] = useState(""); - const [direction, setDirection] = useState("forward"); - const [baselineModel, setBaselineModel] = useState(""); - const [percentage, setPercentage] = useState("10"); - const [durationDays, setDurationDays] = useState("7"); - const [judgeModel, setJudgeModel] = useState(""); - const [maxBudget, setMaxBudget] = useState("10"); - const { data: autoRouters } = useAutoRouters(); - const judgeModelOptions = useJudgeModelOptions(); - const baselineModelOptions = useBaselineModelOptions(); - const start = useStartShadowEval(); - - const routerOptions = useMemo(() => { - const names = new Set( - (autoRouters ?? []).map((deployment) => deployment.model_name).filter((name): name is string => Boolean(name)), - ); - return [...names].toSorted().map((name) => ({ label: name, value: name })); - }, [autoRouters]); - - const parsedPct = Number.parseFloat(percentage); - const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; - const parsedMaxBudget = Number.parseFloat(maxBudget); - const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; - const baselinePicked = direction === "forward" || baselineModel !== ""; - const targetsPicked = apiKeyIds.length + teamIds.length + userIds.length > 0; - const filled = targetsPicked && [routerName, judgeModel].every((field) => field !== "") && baselinePicked; - const boundsValid = percentageValid && maxBudgetValid; - const valid = Boolean(accessToken) && filled && boundsValid; - const handleStart = () => { - const startBody = { - api_key_ids: apiKeyIds, - team_ids: teamIds, - user_ids: userIds, - router_name: routerName, - direction, - ...(direction === "reverse" ? { baseline_model: baselineModel } : {}), - shadow_percentage: parsedPct, - duration_days: Number.parseInt(durationDays, 10), - max_budget: parsedMaxBudget, - judge_model: judgeModel, - }; - start.mutate(startBody); - }; - - return ( - - - Start a shadow eval -

{START_FORM_DESCRIPTION[direction]}

- - -
- - - - - - - - - - - - - - - - -
- setPercentage(e.target.value)} - /> - % of traffic -
-
- {percentage.trim() !== "" && !percentageValid && ( -

Enter a value from 0.1 to 100

- )} -
- - - - - -
- $ - setMaxBudget(e.target.value)} - /> - max shadow + judge spend, per target -
- {maxBudget.trim() !== "" && !maxBudgetValid && ( -

Enter a value from 0.01 to 10000

- )} - - {direction === "reverse" && ( - - - - )} - - - -
- - - - ); -}; - const previousSummary = (job: ShadowEvalJob): string => { const results = job.results; if (results) return pct(routerMatchedOrBeatPct(job.direction, results)); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx new file mode 100644 index 00000000000..f96910a4ad6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -0,0 +1,432 @@ +"use client"; + +import React, { useMemo, useState } from "react"; + +import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; +import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; +import { MultiSelect } from "@/components/shared/MultiSelect"; +import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; +import TeamMultiSelect from "@/components/common_components/team_multi_select"; +import { userOptionLabel } from "@/components/common_components/UserDropdown"; +import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; + +import { useStartShadowEval, type ShadowEvalJob } from "./useShadowEval"; + +type ShadowEvalDirection = ShadowEvalJob["direction"]; + +const MAX_ROUTERS = 4; + +const RECOMMENDED_JUDGE_MODELS = ["anthropic/claude-sonnet-5", "openai/gpt-4o", "gemini/gemini-2.5-pro"] as const; + +interface CostMapEntry { + litellm_provider?: string; + mode?: string; +} + +const useChatModelNames = (): string[] => { + const { data: costMap } = useModelCostMap(); + return useMemo(() => { + if (!costMap) return []; + const chatModels = Object.entries(costMap as Record) + .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) + .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); + return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); + }, [costMap]); +}; + +const useJudgeModelOptions = (): SearchSelectOption[] => { + const chatModels = useChatModelNames(); + return useMemo(() => { + const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ + label: model, + value: model, + sublabel: "Recommended", + })); + const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); + const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); + return [...pinned, ...rest]; + }, [chatModels]); +}; + +const useBaselineModelOptions = (): SearchSelectOption[] => { + const configuredGroups = usePlainModelGroups(); + const chatModels = useChatModelNames(); + return useMemo(() => { + const configured = [...configuredGroups] + .toSorted((a, b) => a.localeCompare(b)) + .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); + const rest = chatModels + .filter((model) => !configuredGroups.has(model)) + .map((model) => ({ label: model, value: model })); + return [...configured, ...rest]; + }, [configuredGroups, chatModels]); +}; + +const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ + { value: "forward", label: "Adoption check: key's traffic vs the router" }, + { value: "reverse", label: "Regression check: router's picks vs a baseline" }, +] as const; + +const START_FORM_DESCRIPTION: Record = { + forward: + "Duplicates a sampled slice of the selected targets' traffic (keys, teams, or users) through the auto-router and has an LLM judge compare both answers blind. Each target gets its own spend budget. The router's answers are never served to users; judge calls bill to the sampled traffic's own identity.", + reverse: + "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. Each target gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the sampled traffic's own identity.", +}; + +const DURATION_OPTIONS = [ + { value: "1", label: "1 day" }, + { value: "3", label: "3 days" }, + { value: "7", label: "7 days" }, + { value: "14", label: "14 days" }, + { value: "30", label: "30 days" }, +] as const; + +const Field: React.FC<{ label: string; htmlFor?: string; className?: string; children: React.ReactNode }> = ({ + label, + htmlFor, + className, + children, +}) => ( +
+ + {children} +
+); + +const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => void }> = ({ value, onChange }) => { + const [search, setSearch] = useState(""); + const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteKeys(50, { + selectedKeyAlias: search || null, + }); + const options = useMemo( + () => + (data?.pages ?? []) + .flatMap((page) => page.keys) + .map((key) => ({ + label: key.key_alias || key.key_name || key.token, + value: key.token, + sublabel: key.token, + })), + [data], + ); + return ( + void fetchNextPage()} + hasNextPage={hasNextPage} + isFetchingNextPage={isFetchingNextPage} + isLoading={isPending} + placeholder="Search keys by alias" + emptyText="No matching keys" + errorText={isError ? "Keys could not be loaded. Refresh the page to retry." : undefined} + /> + ); +}; + +const UserSelect: React.FC<{ value: string[]; onChange: (ids: string[]) => void }> = ({ value, onChange }) => { + const [search, setSearch] = useState(""); + const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteUsers( + 50, + search || undefined, + ); + const options = useMemo( + () => + Array.from( + new Map( + (data?.pages ?? []) + .flatMap((page) => page.users) + .map((user) => [user.user_id, { label: userOptionLabel(user), value: user.user_id }] as const), + ).values(), + ), + [data], + ); + return ( + void fetchNextPage()} + hasNextPage={hasNextPage} + isFetchingNextPage={isFetchingNextPage} + isLoading={isPending} + placeholder="Search users by email" + emptyText="No matching users" + errorText={isError ? "Users could not be loaded. Refresh the page to retry." : undefined} + /> + ); +}; + +const RouterField: React.FC<{ + options: SearchSelectOption[]; + routerNames: string[]; + onChange: (names: string[]) => void; + direction: ShadowEvalDirection; +}> = ({ options, routerNames, onChange, direction }) => ( + + + {routerNames.length > MAX_ROUTERS && ( +

Pick at most {MAX_ROUTERS} auto-routers

+ )} + {direction === "reverse" && routerNames.length > 1 && ( +

A regression check compares one router to its baseline

+ )} + {direction === "forward" && routerNames.length > 1 && ( +

+ Every router sees the same sampled requests, judged against the same live responses +

+ )} + +); + +interface StartFormValidityInputs { + accessToken: string | null | undefined; + apiKeyIds: string[]; + teamIds: string[]; + userIds: string[]; + routerNames: string[]; + direction: ShadowEvalDirection; + baselineModel: string; + judgeModel: string; + percentage: string; + maxBudget: string; +} + +const startFormValidity = (inputs: StartFormValidityInputs) => { + const parsedPct = Number.parseFloat(inputs.percentage); + const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; + const parsedMaxBudget = Number.parseFloat(inputs.maxBudget); + const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; + const baselinePicked = inputs.direction === "forward" || inputs.baselineModel !== ""; + const targetsPicked = inputs.apiKeyIds.length + inputs.teamIds.length + inputs.userIds.length > 0; + const routerCountValid = inputs.routerNames.length >= 1 && inputs.routerNames.length <= MAX_ROUTERS; + const routersMatchDirection = inputs.direction === "forward" || inputs.routerNames.length === 1; + const routersValid = routerCountValid && routersMatchDirection; + const modelsPicked = routersValid && inputs.judgeModel !== "" && baselinePicked; + const filled = targetsPicked && modelsPicked; + const boundsValid = percentageValid && maxBudgetValid; + const valid = Boolean(inputs.accessToken) && filled && boundsValid; + return { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid }; +}; + +interface StartBodyInputs { + apiKeyIds: string[]; + teamIds: string[]; + userIds: string[]; + routerNames: string[]; + direction: ShadowEvalDirection; + baselineModel: string; + shadowPercentage: number; + durationDays: number; + maxBudget: number; + judgeModel: string; +} + +const buildStartBody = (inputs: StartBodyInputs) => ({ + api_key_ids: inputs.apiKeyIds, + team_ids: inputs.teamIds, + user_ids: inputs.userIds, + router_names: inputs.routerNames, + direction: inputs.direction, + ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel } : {}), + shadow_percentage: inputs.shadowPercentage, + duration_days: inputs.durationDays, + max_budget: inputs.maxBudget, + judge_model: inputs.judgeModel, +}); + +export const StartForm: React.FC = () => { + const { accessToken } = useAuthorized(); + const [apiKeyIds, setApiKeyIds] = useState([]); + const [teamIds, setTeamIds] = useState([]); + const [userIds, setUserIds] = useState([]); + const [routerNames, setRouterNames] = useState([]); + const [direction, setDirection] = useState("forward"); + const [baselineModel, setBaselineModel] = useState(""); + const [percentage, setPercentage] = useState("10"); + const [durationDays, setDurationDays] = useState("7"); + const [judgeModel, setJudgeModel] = useState(""); + const [maxBudget, setMaxBudget] = useState("10"); + const { data: autoRouters } = useAutoRouters(); + const judgeModelOptions = useJudgeModelOptions(); + const baselineModelOptions = useBaselineModelOptions(); + const start = useStartShadowEval(); + + const routerOptions = useMemo(() => { + const names = new Set( + (autoRouters ?? []).map((deployment) => deployment.model_name).filter((name): name is string => Boolean(name)), + ); + return [...names].toSorted().map((name) => ({ label: name, value: name })); + }, [autoRouters]); + + const validityInputs: StartFormValidityInputs = { + accessToken, + apiKeyIds, + teamIds, + userIds, + routerNames, + direction, + baselineModel, + judgeModel, + percentage, + maxBudget, + }; + const { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid } = startFormValidity(validityInputs); + const handleStart = () => { + const bodyInputs: StartBodyInputs = { + apiKeyIds, + teamIds, + userIds, + routerNames, + direction, + baselineModel, + shadowPercentage: parsedPct, + durationDays: Number.parseInt(durationDays, 10), + maxBudget: parsedMaxBudget, + judgeModel, + }; + start.mutate(buildStartBody(bodyInputs)); + }; + + return ( + + + Start a shadow eval +

{START_FORM_DESCRIPTION[direction]}

+ + +
+ + + + + + + + + + + + + + +
+ setPercentage(e.target.value)} + /> + % of traffic +
+
+ {percentage.trim() !== "" && !percentageValid && ( +

Enter a value from 0.1 to 100

+ )} +
+ + + + + +
+ $ + setMaxBudget(e.target.value)} + /> + max shadow + judge spend, per target +
+ {maxBudget.trim() !== "" && !maxBudgetValid && ( +

Enter a value from 0.01 to 10000

+ )} + + {direction === "reverse" && ( + + + + )} + + + +
+ + + + ); +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4ad3ca06501..e944062e15e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35331,8 +35331,17 @@ export interface components { last_error?: string | null; /** @description Stratified verdicts; detail endpoint only */ results?: components["schemas"]["ShadowEvalResult"] | null; - /** Router Name */ - router_name: string; + /** + * Router Name + * @description The first router, kept for callers that predate router_names; derived so the + * two fields can never disagree. + */ + readonly router_name: string; + /** + * Router Names + * @description Every auto-router this job runs as a shadow arm. Multi-router jobs sample one slice of traffic and judge every arm against the same real responses + */ + router_names: string[]; /** Shadow Percentage */ shadow_percentage: number; /** @@ -35420,6 +35429,12 @@ export interface components { * @description Sliced by the model that served the real arm: the keys' incumbent models in forward mode, and in reverse the models the router itself picked */ by_current_model: components["schemas"]["ShadowEvalSlice"][]; + /** + * By Router + * @description One slice per router arm, grouped on the router name. Every arm of a multi-router job is judged against the same real responses over the same sampled requests, so these slices compare routers head-to-head: like-for-like win rates and spends on identical traffic. Verdicts from before arm stamping existed count toward the job's own router + * @default [] + */ + by_router: components["schemas"]["ShadowEvalSlice"][]; /** By Tier */ by_tier: components["schemas"]["ShadowEvalSlice"][]; /** @@ -35433,13 +35448,13 @@ export interface components { overall_tie_rate_pct: number; /** * Sampled Real Spend - * @description USD the real arm billed across all judged turns, cache-served turns excluded + * @description USD the real arm billed across all judged turns, cache-served turns excluded. A judged turn is one (request, router arm) verdict, so a multi-router job counts the real response once per arm it was judged against; per-router comparisons read by_router * @default 0 */ sampled_real_spend: number; /** * Sampled Shadow Spend - * @description USD the shadow arm billed across the same turns, judge excluded, like for like + * @description USD the shadow arms billed across the same turns, judge excluded, like for like * @default 0 */ sampled_shadow_spend: number; @@ -35733,15 +35748,21 @@ export interface components { judge_model: string; /** * Max Budget - * @description Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window + * @description Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window. Every router arm draws from the same per-target budget, so a multi-router job reaches it proportionally sooner * @default 10 */ max_budget: number; /** * Router Name - * @description The auto-router under evaluation, in either direction + * @description The auto-router under evaluation, in either direction: the single-router spelling of router_names. Provide exactly one of the two fields */ - router_name: string; + router_name?: string | null; + /** + * Router Names + * @description The auto-routers under evaluation, at most 4. Every sampled request runs through every router listed and each arm is judged independently against the same real response, so routers compare head-to-head on identical traffic. More than one router requires direction 'forward'. After validation this field always carries the full deduplicated set, whichever spelling the caller used + * @default [] + */ + router_names: string[]; /** * Shadow Percentage * @description Percentage of each target's requests to duplicate through the router From 0a9676bd4f7af7880a49aafe665b3951343e1cd9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:42:31 -0700 Subject: [PATCH 078/126] fix(openai): scope unknown-model reasoning_effort forwarding to the plain openai provider --- .../llms/openai/chat/gpt_transformation.py | 26 ++++++++++++------- litellm/llms/openai/openai.py | 10 +++++-- litellm/utils.py | 7 +++++ .../chat/test_openai_gpt_transformation.py | 26 +++++++++++++++++-- 4 files changed, 56 insertions(+), 13 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 255ee3159c3..d4747b2fb06 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -170,20 +170,20 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): if model != "gpt-3.5-turbo-16k" and model != "gpt-4": # gpt-4 does not support 'response_format' model_specific_params.append("response_format") - # Normalize model name for responses API (e.g., "responses/gpt-4.1" -> "gpt-4.1") - model_for_check: Final = model.split("responses/", 1)[1] if "responses/" in model else model - if ( - model_for_check in litellm.open_ai_chat_completion_models - ) or model_for_check in litellm.open_ai_text_completion_models: + if OpenAIGPTConfig.is_openai_catalog_model(model): model_specific_params.append( "user" ) # user is not a param supported by all openai-compatible endpoints - e.g. azure ai - else: - model_specific_params.append( - "reasoning_effort" - ) # unknown model: likely a proxy alias for a reasoning-capable model, so forward and let the server decide return base_params + model_specific_params + @staticmethod + def is_openai_catalog_model(model: str) -> bool: + model_for_check: Final = model.split("responses/", 1)[1] if "responses/" in model else model + return ( + model_for_check in litellm.open_ai_chat_completion_models + or model_for_check in litellm.open_ai_text_completion_models + ) + def _map_openai_params( self, non_default_params: dict, @@ -759,6 +759,14 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) +class OpenAIUnknownModelConfig(OpenAIGPTConfig): + """A model the openai provider does not recognize is typically a LiteLLM proxy alias, so + forward reasoning_effort and let the server decide whether it is supported.""" + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: inherited contract + return super().get_supported_openai_params(model) + ["reasoning_effort"] # mutable-ok: inherited contract + + class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): def _map_reasoning_to_reasoning_content(self, choices: list) -> list: """ diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 16fa0017b23..56495f0097d 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -43,6 +43,7 @@ from litellm.utils import ( from ...types.llms.openai import * from ..base import BaseLLM from .chat.gpt_5_transformation import OpenAIGPT5Config +from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( BaseOpenAILLM, @@ -189,7 +190,12 @@ class OpenAIConfig(BaseConfig): elif litellm.openAIGPTAudioConfig.is_model_gpt_audio_model(model=model): return litellm.openAIGPTAudioConfig.get_supported_openai_params(model=model) else: - return litellm.openAIGPTConfig.get_supported_openai_params(model=model) + return self._gpt_config_for_model(model).get_supported_openai_params(model=model) + + def _gpt_config_for_model(self, model: str) -> OpenAIGPTConfig: + if type(self) is OpenAIConfig and not OpenAIGPTConfig.is_openai_catalog_model(model): + return OpenAIUnknownModelConfig() + return litellm.openAIGPTConfig def _map_openai_params(self, non_default_params: dict, optional_params: dict, model: str) -> dict: supported_openai_params: Final = self.get_supported_openai_params(model) @@ -231,7 +237,7 @@ class OpenAIConfig(BaseConfig): drop_params=drop_params, ) - return litellm.openAIGPTConfig.map_openai_params( + return self._gpt_config_for_model(model).map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, diff --git a/litellm/utils.py b/litellm/utils.py index f5adca8f272..15c1b0e9c0f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8278,10 +8278,17 @@ class ProviderConfigManager: """ # Handle OpenAI special cases (O-series and GPT-5 models) if provider == LlmProviders.OPENAI: + from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIGPTConfig, + OpenAIUnknownModelConfig, + ) + if litellm.openaiOSeriesConfig.is_model_o_series_model(model=model): return litellm.openaiOSeriesConfig if litellm.OpenAIGPT5Config.is_model_gpt_5_model(model=model): return litellm.OpenAIGPT5Config() + if not OpenAIGPTConfig.is_openai_catalog_model(model): + return OpenAIUnknownModelConfig() # Handle Azure before the generic map so base_model can be threaded through if provider == LlmProviders.AZURE: diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 1ee53a90d7d..3ef5e39fc5f 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -148,19 +148,41 @@ class TestGetOptionalParamsIntegration: def test_reasoning_effort_supported_for_unknown_model_alias(self): """An openai/-routed model litellm doesn't recognize is likely a proxy alias: reasoning_effort must be forwarded so the server decides support.""" - supported_params = OpenAIGPTConfig().get_supported_openai_params( + from litellm.llms.openai.openai import OpenAIConfig + + supported_params = OpenAIConfig().get_supported_openai_params( "my-claude-alias" ) assert "reasoning_effort" in supported_params def test_reasoning_effort_not_supported_for_known_non_reasoning_models(self): """Known OpenAI models keep failing closed client-side.""" - config = OpenAIGPTConfig() + from litellm.llms.openai.openai import OpenAIConfig + + config = OpenAIConfig() assert "reasoning_effort" not in config.get_supported_openai_params("gpt-4o") assert "reasoning_effort" not in config.get_supported_openai_params( "responses/gpt-4.1-mini" ) + def test_reasoning_effort_not_inherited_by_openai_compatible_subclasses(self): + """Providers subclassing either openai config keep their own reasoning_effort gating + for their models, which are all unknown to the openai catalog.""" + from litellm.llms.openai.openai import OpenAIConfig + + class InheritingDispatcherConfig(OpenAIConfig): + pass + + class InheritingGPTConfig(OpenAIGPTConfig): + pass + + assert "reasoning_effort" not in InheritingDispatcherConfig().get_supported_openai_params( + "some-unknown-model" + ) + assert "reasoning_effort" not in InheritingGPTConfig().get_supported_openai_params( + "some-unknown-model" + ) + def test_reasoning_effort_forwarded_in_optional_params_for_unknown_model_alias( self, ): From 8a83f9e3cc1f96652a3d23d8d88365e2d71a5aaa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 21:50:22 -0700 Subject: [PATCH 079/126] refactor(key_management): extract allowed_routes update gate to keep complexity budget --- .../key_management_endpoints.py | 31 +++++++++++++------ 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8b3d2ad6e6e..b18704d4ee9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -752,6 +752,23 @@ def _is_safe_preset_route_transition( ) +def _enforce_allowed_routes_update_permission( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + if _is_safe_preset_route_transition( + incoming_allowed_routes=data.allowed_routes, + existing_allowed_routes=existing_key_row.allowed_routes, + ): + return + _check_allowed_routes_caller_permission( + allowed_routes=data.allowed_routes, + user_api_key_dict=user_api_key_dict, + allowed_routes_was_provided="allowed_routes" in data.model_fields_set, + ) + + def _check_permissions_caller_permission( data: GenerateRequestBase, user_api_key_dict: UserAPIKeyAuth, @@ -2552,15 +2569,11 @@ async def _validate_update_key_data( _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - if not _is_safe_preset_route_transition( - incoming_allowed_routes=data.allowed_routes, - existing_allowed_routes=existing_key_row.allowed_routes, - ): - _check_allowed_routes_caller_permission( - allowed_routes=data.allowed_routes, - user_api_key_dict=user_api_key_dict, - allowed_routes_was_provided="allowed_routes" in data.model_fields_set, - ) + _enforce_allowed_routes_update_permission( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, + ) _check_passthrough_routes_caller_permission( data=data, user_api_key_dict=user_api_key_dict, From 4bfc6766e7b4bb9770a690251f9746630fede523 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 21:56:44 -0700 Subject: [PATCH 080/126] refactor(key_management): immutable types in preset transition helper for lint budgets --- .../proxy/management_endpoints/key_management_endpoints.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b18704d4ee9..1fc2850c4ae 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -737,8 +737,8 @@ def _check_allowed_routes_caller_permission( def _is_safe_preset_route_transition( - incoming_allowed_routes: list | None, - existing_allowed_routes: list | None, + incoming_allowed_routes: Sequence[str] | None, + existing_allowed_routes: Sequence[str] | None, ) -> bool: """ True when every route on BOTH sides is a safe `key_type` preset bucket @@ -748,7 +748,7 @@ def _is_safe_preset_route_transition( """ return all( route in _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS - for route in (*(incoming_allowed_routes or []), *(existing_allowed_routes or [])) + for route in (*(incoming_allowed_routes or ()), *(existing_allowed_routes or ())) ) From 81c48f810edf9412165539f278e51131dffb348d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 22:09:21 -0700 Subject: [PATCH 081/126] fix(key_management): keep read-only keys read-only in non-admin preset transitions A non-admin could widen a read-only (info_routes) key to llm_api or full access through the preset carve-out. Read-only keys now stay read-only unless a proxy admin widens them; the other preset transitions, including the LIT-4891 llm_api to full access switch, still work. Also converts the transition tests to assert on a returned outcome so the no-403 cases carry real assertions. --- .../key_management_endpoints.py | 18 ++-- .../test_key_management_endpoints.py | 94 ++++++++++++++----- 2 files changed, 81 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 1fc2850c4ae..5b237cccecf 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -736,6 +736,9 @@ def _check_allowed_routes_caller_permission( ) +_READ_ONLY_ALLOWED_ROUTES_PRESET: Final = frozenset(("info_routes",)) + + def _is_safe_preset_route_transition( incoming_allowed_routes: Sequence[str] | None, existing_allowed_routes: Sequence[str] | None, @@ -743,13 +746,16 @@ def _is_safe_preset_route_transition( """ True when every route on BOTH sides is a safe `key_type` preset bucket (empty = full access, which non-admins already get from a default - `/key/generate`). Requiring the existing side too keeps an owner from - clearing an admin-set custom route restriction (LIT-4139). + `/key/generate`), with one carve-out: a read-only (`info_routes`) key + stays read-only, so widening it needs an admin. Requiring the existing + side to be a safe preset keeps an owner from clearing an admin-set + custom route restriction (LIT-4139). """ - return all( - route in _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS - for route in (*(incoming_allowed_routes or ()), *(existing_allowed_routes or ())) - ) + incoming: Final = frozenset(incoming_allowed_routes or ()) + existing: Final = frozenset(existing_allowed_routes or ()) + if not (incoming | existing) <= _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS: + return False + return existing != _READ_ONLY_ALLOWED_ROUTES_PRESET or incoming == existing def _enforce_allowed_routes_update_permission( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 0ddc8c5a822..ff555d893f5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -11054,56 +11054,100 @@ class TestLIT4891SafePresetKeyTypeTransition: ) async def _run_update(self, data, existing_key_row): - await _validate_update_key_data( - data=data, - existing_key_row=existing_key_row, - user_api_key_dict=self._make_auth(), - llm_router=None, - premium_user=False, - prisma_client=AsyncMock(), - user_api_key_cache=MagicMock(), - ) + try: + await _validate_update_key_data( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=self._make_auth(), + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + except HTTPException as exc: + return exc + return None + + def _assert_routes_403(self, exc): + assert exc is not None + assert exc.status_code == 403 + assert "Only proxy admins can set" in str(exc.detail) @pytest.mark.asyncio async def test_non_admin_owner_can_clear_safe_preset_to_full_access(self): - await self._run_update( - data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), - existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + is None ) @pytest.mark.asyncio async def test_non_admin_owner_can_switch_full_access_to_safe_preset(self): - await self._run_update( - data=UpdateKeyRequest(key="sk-test", allowed_routes=["llm_api_routes"]), - existing_key_row=self._make_existing_key(allowed_routes=[]), + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["llm_api_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=[]), + ) + is None ) @pytest.mark.asyncio - async def test_non_admin_owner_can_switch_between_safe_presets(self): - await self._run_update( - data=UpdateKeyRequest(key="sk-test", allowed_routes=["info_routes"]), - existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + async def test_non_admin_owner_can_narrow_to_read_only_preset(self): + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["info_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + is None + ) + + @pytest.mark.asyncio + async def test_non_admin_can_resend_read_only_preset_unchanged(self): + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["info_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["info_routes"]), + ) + is None + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_widen_read_only_key_to_full_access(self): + self._assert_routes_403( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["info_routes"]), + ) + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_widen_read_only_key_to_llm_api(self): + self._assert_routes_403( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["llm_api_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["info_routes"]), + ) ) @pytest.mark.asyncio async def test_non_admin_cannot_clear_custom_route_restriction(self): - with pytest.raises(HTTPException) as exc_info: + self._assert_routes_403( await self._run_update( data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), existing_key_row=self._make_existing_key(allowed_routes=["/chat/completions"]), ) - assert exc_info.value.status_code == 403 - assert "Only proxy admins can set" in str(exc_info.value.detail) + ) @pytest.mark.asyncio async def test_non_admin_cannot_set_non_preset_routes(self): - with pytest.raises(HTTPException) as exc_info: + self._assert_routes_403( await self._run_update( data=UpdateKeyRequest(key="sk-test", allowed_routes=["management_routes"]), existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), ) - assert exc_info.value.status_code == 403 - assert "Only proxy admins can set" in str(exc_info.value.detail) + ) class TestKeyOwnerPrivilegeEscalation: From 4a24be886d4d76d06f355cb5164af36fb3f36b4d Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 22:12:59 -0700 Subject: [PATCH 082/126] feat(ui): one classification frequency picker for complexity auto-routers (#39042) Classification timing and session affinity are the same operator question, so Advanced: Classification Method now carries a single "How often to classify" radio: every request, every new user message, or once per session. The session choice writes session_affinity and stays disabled on custom tier sets, where the backend rejects it. Advanced: Affinity keeps the deployment switch alone. The serializer always writes classification_mode, matching session_affinity on the line below it, so an explicitly stored every_request survives an untouched save instead of being dropped back to the backend default. --- .../src/autorouter_presets.json | 4 + .../add_model/ClassificationMethodConfig.tsx | 51 ++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 96 +++++++++++++++++-- .../add_model/ComplexityRouterConfig.tsx | 41 +++++--- .../add_model/add_auto_router_tab.test.tsx | 46 ++++++++- .../add_model/add_auto_router_tab.tsx | 1 + .../build_complexity_router_config.test.ts | 19 ++++ .../build_complexity_router_config.ts | 6 ++ ...d_updated_complexity_router_config.test.ts | 28 ++++++ .../edit_auto_router_modal.test.ts | 2 + .../edit_auto_router_modal.test.tsx | 85 ++++++++++++---- .../edit_auto_router_modal.tsx | 7 ++ .../src/lib/autorouter_presets.test.ts | 32 +++++++ .../src/lib/autorouter_presets.ts | 2 + 14 files changed, 380 insertions(+), 40 deletions(-) diff --git a/ui/litellm-dashboard/src/autorouter_presets.json b/ui/litellm-dashboard/src/autorouter_presets.json index 4cbb548a855..da35d6171dc 100644 --- a/ui/litellm-dashboard/src/autorouter_presets.json +++ b/ui/litellm-dashboard/src/autorouter_presets.json @@ -14,6 +14,7 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, "deployment_affinity": true } @@ -30,6 +31,7 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, "deployment_affinity": true } @@ -56,6 +58,7 @@ }, "classifier_context_window_size": 0, "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, "deployment_affinity": true } @@ -72,6 +75,7 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, "deployment_affinity": true } diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 84f61c95eca..96c93306611 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -15,9 +15,12 @@ import { RestrictedSection, restrictedBy } from "./TierRestrictions"; import HeuristicScoringConfig from "./HeuristicScoringConfig"; import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults"; import { + ClassificationFrequency, ClassifierFallback, ClassifierType, ComplexityRouterConfigValue, + classificationFrequency, + withClassificationFrequency, DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, MIN_QUOTED_CONTEXT_TURN_CHARS, DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, @@ -211,6 +214,7 @@ const ClassificationMethodConfig: React.FC = ({ const [draft, setDraft] = React.useState<{ id: string; raw: string } | null>(null); const hasDefaultModel = Boolean(defaultModel); const classifierType = effectiveClassifierType(value); + const sessionFrequencyRestriction = restrictedBy(value, "sessionAffinity"); const classifierModelMissing = showValidationErrors && usesLlmClassifier(classifierType) && !value.classifier_llm_config?.model; const usesCustomPrompt = Boolean(value.classifier_llm_config?.system_prompt?.trim()); @@ -305,6 +309,10 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, classifier_fallback: fallback }); }; + const handleClassificationFrequencyChange = (frequency: ClassificationFrequency) => { + onChange(withClassificationFrequency(value, frequency)); + }; + const handleClassifierContextWindowSizeChange = (windowSize: number) => { onChange({ ...value, @@ -367,6 +375,49 @@ const ClassificationMethodConfig: React.FC = ({ )} +
+ How often to classify + + handleClassificationFrequencyChange(frequency as ClassificationFrequency) + } + > +
+ + + +
+ +

+ Holding the tier keeps an agent on one model for a whole tool loop and cuts scoring cost. A turn the router + cannot match to a held decision, such as one with no session id or an expired one, is scored again +

+
+ {usesLlmClassifier(classifierType) && (
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 218c99e32c5..bbee03135c7 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -598,6 +598,82 @@ describe("ComplexityRouterConfig classifier fallback", () => { }); }); +describe("ComplexityRouterConfig classification frequency", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + + it("defaults to every request, matching both backend field defaults", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Every request/ })).toBeChecked(); + expect(screen.getByRole("radio", { name: /Every new user message/ })).not.toBeChecked(); + expect(screen.getByRole("radio", { name: /Once per session/ })).not.toBeChecked(); + }); + + it("writes both wire fields when the frequency moves to every new user message", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + fireEvent.click(screen.getByRole("radio", { name: /Every new user message/ })); + expect(onChange).toHaveBeenCalledWith({ + ...llmValue, + classification_mode: "user_turn", + session_affinity: false, + }); + }); + + it("writes session affinity, not a classification mode, when the frequency moves to once per session", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + fireEvent.click(screen.getByRole("radio", { name: /Once per session/ })); + expect(onChange).toHaveBeenCalledWith({ + ...llmValue, + classification_mode: "every_request", + session_affinity: true, + }); + }); + + it("shows a hand-authored config that sets both fields as once per session, matching the backend", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Once per session/ })).toBeChecked(); + expect(screen.getByRole("radio", { name: /Every new user message/ })).not.toBeChecked(); + }); + + it("records a switch back to every request", () => { + const onChange = vi.fn(); + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Every new user message/ })).toBeChecked(); + fireEvent.click(screen.getByRole("radio", { name: /Every request/ })); + expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ classification_mode: "every_request" })); + }); + + it("offers the frequency on a heuristic router, where holding the tier still pins the model", () => { + // The backend pin is gated on the two fields alone, so a heuristic router that switches models + // mid tool loop is fixed by this control too. + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Every new user message/ })).toBeInTheDocument(); + }); +}); + describe("ComplexityRouterConfig classifier rubric", () => { const llmValue: ComplexityRouterConfigValue = { ...defaultValue, @@ -761,12 +837,12 @@ describe("ComplexityRouterConfig tier labels", () => { }); describe("ComplexityRouterConfig affinity panel", () => { - it("holds both affinity switches with their backend defaults", () => { + it("holds the deployment switch at its backend default, session pinning having moved to the frequency choice", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Affinity")); expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).toBeChecked(); - expect(screen.getByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + expect(screen.queryByRole("switch", { name: "Pin a session to its first model" })).not.toBeInTheDocument(); }); it("writes deployment_affinity through onChange without touching other keys", () => { @@ -1291,10 +1367,18 @@ describe("ComplexityRouterConfig tier editing", () => { expect(screen.getByLabelText("Fallback tier")).toBeInTheDocument(); }); - it("disables session pinning and says why, rather than letting a stripped value look saved", () => { - renderWithProviders(); - fireEvent.click(screen.getByText("Advanced: Affinity")); - expect(screen.getByLabelText("Pin a session to its first model")).toHaveAttribute("data-disabled"); + it("disables the once-per-session frequency and says why, rather than letting a stripped value look saved", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + const sessionOption = screen.getByRole("radio", { name: /Once per session/ }); + expect(sessionOption).toHaveAttribute("aria-disabled", "true"); + expect(sessionOption).not.toBeChecked(); expect( screen.getByText("Session pinning escalates along the built-in tier ladder", { exact: false }), ).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 111a7c9f10a..5fde87fd638 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -56,6 +56,16 @@ export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120; export const DEFAULT_SESSION_AFFINITY = false; export const DEFAULT_DEPLOYMENT_AFFINITY = true; +export type ClassificationMode = "every_request" | "user_turn"; + +export const DEFAULT_CLASSIFICATION_MODE: ClassificationMode = "every_request"; + +/** + * One operator-facing choice over the two wire fields that share the router's tier-pin machinery: + * session affinity pins every turn, user_turn pins every turn except a new human ask. + */ +export type ClassificationFrequency = ClassificationMode | "session"; + export type ComplexityTiers = { SIMPLE: string[]; MEDIUM: string[]; @@ -384,6 +394,7 @@ export interface ComplexityRouterConfigValue { classification_prompt?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + classification_mode?: ClassificationMode; session_affinity?: boolean; deployment_affinity?: boolean; /** Plan-mode floor as a tier ROW ID, unset meaning off. The wire carries the row's name. */ @@ -420,6 +431,21 @@ export interface ComplexityRouterConfigValue { tier_model_params?: TierModelParamsByTier; } +/** Session affinity wins where a hand-authored config sets both, matching the backend's own `or`. */ +export const classificationFrequency = (value: ComplexityRouterConfigValue): ClassificationFrequency => { + if (!value.custom_tier_set && (value.session_affinity ?? DEFAULT_SESSION_AFFINITY)) return "session"; + return value.classification_mode === "user_turn" ? "user_turn" : "every_request"; +}; + +export const withClassificationFrequency = ( + value: ComplexityRouterConfigValue, + frequency: ClassificationFrequency, +): ComplexityRouterConfigValue => ({ + ...value, + classification_mode: frequency === "user_turn" ? "user_turn" : "every_request", + session_affinity: frequency === "session", +}); + interface ComplexityRouterConfigProps { modelInfo: ModelGroup[]; value: ComplexityRouterConfigValue; @@ -498,23 +524,10 @@ const AffinityControls: React.FC<{ /> Pin a session to one deployment per model group
- + Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to load-balance every turn. -
- onChange({ ...value, session_affinity: sessionAffinity })} - aria-label="Pin a session to its first model" - /> - Pin a session to its first model -
- - {restrictedBy(value, "sessionAffinity")?.reason ?? - "Keeps a session on its first turn's model instead of re-classifying each turn. Also pins the deployment."} - ); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 71b454dbb06..d8a955b719c 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -362,8 +362,8 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); expandDetailedConfiguration(); - await user.click(screen.getByText("Advanced: Affinity")); - expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + await user.click(screen.getByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Once per session/ })).not.toBeChecked(); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -468,8 +468,8 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); expandDetailedConfiguration(); - await user.click(screen.getByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Once per session/ })); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -479,6 +479,44 @@ describe("AddAutoRouterTab", () => { }); }); + it("carries every new user message through to the create payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "user-turn-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every new user message/ })); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + classification_mode: "user_turn", + }); + }); + + it("writes every_request into the create payload when the default frequency stays selected", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "default-timing-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Every request/ })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect( + vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config.classification_mode, + ).toBe("every_request"); + }); + it("defaults a new router to deployment affinity on, matching the backend field default", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index ed584a4882b..c60ce5e4959 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC = ({ planModeMinTier: complexityRouterConfig.plan_mode_min_tier, classificationPrompt: complexityRouterConfig.classification_prompt, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, classifierLlmConfig: complexityRouterConfig.classifier_llm_config, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index feddaaa0eac..40addf086b6 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -52,6 +52,7 @@ describe("buildComplexityRouterConfig", () => { const expected = { tiers, classifier_type: "heuristic", + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, escalation_keywords: ["LITELLM ESCALATE"], @@ -790,6 +791,20 @@ describe("heuristic_first", () => { }); }); +describe("classification_mode", () => { + it("emits user_turn", () => { + const config = buildComplexityRouterConfig({ ...baseParams, classificationMode: "user_turn" }); + expect(config.classification_mode).toBe("user_turn"); + }); + + it("writes every_request explicitly, so a saved router never depends on the backend default", () => { + expect( + buildComplexityRouterConfig({ ...baseParams, classificationMode: "every_request" }).classification_mode, + ).toBe("every_request"); + expect(buildComplexityRouterConfig(baseParams).classification_mode).toBe("every_request"); + }); +}); + describe("buildComplexityRouterConfig with an edited tier set", () => { const customTierSet = { tiers: [ @@ -902,6 +917,10 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { expect(payload.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 }); }); + it("keeps classification_mode, which the backend accepts beside tier_definitions", () => { + expect(build({ classificationMode: "user_turn" }).classification_mode).toBe("user_turn"); + }); + it("carries the plan-mode floor as the row's name, not the row id the form holds", () => { expect(build({ planModeMinTier: "sec" }).plan_mode_min_tier).toBe("SECURITY_REVIEW"); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 9a14e956207..6a087649c00 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -19,10 +19,12 @@ import { import { AdaptiveEligible, AdaptiveRouterWeights, + ClassificationMode, ClassifierFallback, ClassifierLLMConfig, ClassifierType, ComplexityTierLabels, + DEFAULT_CLASSIFICATION_MODE, ComplexityRouterConfigValue, ComplexityTiers, DimensionWeights, @@ -105,6 +107,7 @@ export interface BuildComplexityRouterConfigParams { classifierFallback: ClassifierFallback | undefined; classificationPrompt: string | undefined; heuristicFirstMaxTier: string | undefined; + classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; deploymentAffinity: boolean; customTechnicalKeywords: string[]; @@ -159,6 +162,7 @@ export interface ComplexityRouterConfigPayload { classifier_fallback?: ClassifierFallback; classification_prompt?: string; heuristic_first_max_tier?: string; + classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; custom_technical_keywords?: string[]; @@ -393,6 +397,7 @@ export const buildComplexityRouterConfig = ({ classifierFallback, classificationPrompt, heuristicFirstMaxTier, + classificationMode, sessionAffinity, deploymentAffinity, customTechnicalKeywords, @@ -452,6 +457,7 @@ export const buildComplexityRouterConfig = ({ ...(cleanedTierLabels && { tier_labels: cleanedTierLabels }), classifier_type: classifierType, ...classifierWireFields(effectiveType, classifierInputs), + classification_mode: classificationMode ?? DEFAULT_CLASSIFICATION_MODE, session_affinity: sessionAffinity, deployment_affinity: deploymentAffinity, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index d8c2987ead5..3e63cbd3b32 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -257,6 +257,33 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => { }); }); +describe("buildUpdatedComplexityRouterConfig classification mode", () => { + it("round-trips a stored user_turn through hydrate then save", () => { + const stored = { ...STORED, classification_mode: "user_turn" }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + + expect(hydrated.classification_mode).toBe("user_turn"); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).classification_mode).toBe("user_turn"); + }); + + it("round-trips an explicitly stored every_request, so an untouched save leaves it as written", () => { + const stored = { ...STORED, classification_mode: "every_request" }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + + expect(hydrated.classification_mode).toBe("every_request"); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).classification_mode).toBe("every_request"); + }); + + it("rewrites a stored user_turn to every_request once the operator picks the default back", () => { + const stored = { ...STORED, classification_mode: "user_turn" }; + const result = buildUpdatedComplexityRouterConfig(stored, { + ...FORM_VALUE, + classification_mode: "every_request", + }); + expect(result.classification_mode).toBe("every_request"); + }); +}); + describe("buildUpdatedComplexityRouterConfig deployment affinity", () => { it("writes deployment_affinity=false when the toggle is off", () => { const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, deployment_affinity: false }); @@ -476,6 +503,7 @@ describe("managed keys survive an untouched open-and-save", () => { classifier_context_budget_chars: 4000, classifier_context_include_assistant_turns: true, classifier_fallback: "default_model", + classification_mode: "user_turn", session_affinity: true, deployment_affinity: false, adaptive: true, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index 027e01a9351..d199af51b41 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -47,6 +47,7 @@ const expectedClassifiedTierConfig = { semantic_keyword_matching: true, embedding_model: "voyage-4-large", match_threshold: 0.65, + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, adaptive: true, @@ -68,6 +69,7 @@ const expectedAdaptiveDisabledConfig = { semantic_keyword_matching: true, embedding_model: "voyage-4-large", match_threshold: 0.65, + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index a4921fcfcb5..96e8549eac8 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -364,7 +364,7 @@ describe("EditAutoRouterModal assistant turns", () => { }); }); -describe("EditAutoRouterModal session affinity", () => { +describe("EditAutoRouterModal classification frequency", () => { beforeEach(() => { modelPatchUpdateCall.mockClear(); }); @@ -381,15 +381,15 @@ describe("EditAutoRouterModal session affinity", () => { />, ); - // A stored config with no session_affinity key now runs with affinity OFF, because the backend - // field defaults to False. The toggle has to render what the router actually does, and an - // untouched save must not flip it. - it("shows a stored config with no session_affinity key as off", async () => { + // A stored config with neither key now runs with affinity OFF, because both backend fields + // default that way. The picker has to render what the router actually does, and an untouched + // save must not flip it. + it("shows a stored config with neither key as every request", async () => { const user = userEvent.setup(); renderWithStoredConfig(STORED_CONFIG); - await user.click(await screen.findByText("Advanced: Affinity")); - expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Every request/ })).toBeChecked(); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -397,12 +397,12 @@ describe("EditAutoRouterModal session affinity", () => { expect(savedConfig().session_affinity).toBe(false); }); - it("shows a stored session_affinity=true as on and preserves it through an untouched save", async () => { + it("shows a stored session_affinity=true as once per session and preserves it through an untouched save", async () => { const user = userEvent.setup(); renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); - await user.click(await screen.findByText("Advanced: Affinity")); - expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).toBeChecked(); + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Once per session/ })).toBeChecked(); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -410,12 +410,12 @@ describe("EditAutoRouterModal session affinity", () => { expect(savedConfig().session_affinity).toBe(true); }); - it("persists turning session affinity on", async () => { + it("persists picking once per session", async () => { const user = userEvent.setup(); renderWithStoredConfig(STORED_CONFIG); - await user.click(await screen.findByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Once per session/ })); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -423,18 +423,71 @@ describe("EditAutoRouterModal session affinity", () => { expect(savedConfig().session_affinity).toBe(true); }); - it("persists turning session affinity back off", async () => { + it("persists picking every request back over a stored session pin", async () => { const user = userEvent.setup(); renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); - await user.click(await screen.findByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every request/ })); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); expect(savedConfig().session_affinity).toBe(false); }); + + it("clears a stored session pin when the operator moves to every new user message", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every new user message/ })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity).toBe(false); + expect(savedConfig().classification_mode).toBe("user_turn"); + }); + + it("shows a stored user_turn as selected and preserves it through an untouched save", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, classification_mode: "user_turn" }); + + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Every new user message/ })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classification_mode).toBe("user_turn"); + }); + + it("persists switching a stored config to every new user message", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every new user message/ })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classification_mode).toBe("user_turn"); + }); + + it("rewrites the stored mode to every_request when the operator picks it back", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, classification_mode: "user_turn" }); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every request/ })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classification_mode).toBe("every_request"); + }); }); describe("EditAutoRouterModal deployment affinity", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 4751c2e64b7..2aecf0b3483 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -96,6 +96,7 @@ export interface StoredComplexityRouterConfig { classifier_context_budget_chars?: unknown; classifier_context_include_assistant_turns?: unknown; classifier_fallback?: unknown; + classification_mode?: unknown; tier_boundaries?: unknown; token_thresholds?: unknown; dimension_weights?: unknown; @@ -165,6 +166,10 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + classification_mode: + parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" + ? parsedConfig.classification_mode + : undefined, tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), @@ -207,6 +212,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_fallback", "classification_prompt", "heuristic_first_max_tier", + "classification_mode", "session_affinity", "deployment_affinity", "adaptive", @@ -294,6 +300,7 @@ export const buildUpdatedComplexityRouterConfig = ( planModeMinTier: value.plan_mode_min_tier, classificationPrompt: value.classification_prompt, heuristicFirstMaxTier: value.heuristic_first_max_tier, + classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, classifierLlmConfig: value.classifier_llm_config, diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index f14d6279e32..5632d6d947f 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -230,6 +230,7 @@ describe("autorouter_presets", () => { const config = { tiers: { SIMPLE: [presetModel], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -273,6 +274,7 @@ describe("autorouter_presets", () => { const config = { tiers: { SIMPLE: ["claude-opus-5"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -287,6 +289,7 @@ describe("autorouter_presets", () => { const config = { tiers: { SIMPLE: ["claude-opus-5"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -320,6 +323,7 @@ describe("autorouter_presets", () => { const simpleTierConfig = (presetModel: string) => ({ tiers: { SIMPLE: [presetModel], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }); @@ -563,6 +567,7 @@ describe("autorouter_presets", () => { const config = { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, match_threshold: 0, @@ -578,6 +583,7 @@ describe("autorouter_presets", () => { { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, enable_context_window_escalation: false, @@ -589,11 +595,29 @@ describe("autorouter_presets", () => { expect(prefill.complexityRouterConfig.context_window_escalation_buffer).toBe(0.9); }); + it("carries a preset's classification_mode and defaults it when the preset omits one", () => { + const tiers = { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }; + const base = { + tiers, + classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, + session_affinity: false, + deployment_affinity: true, + }; + const availability = groupsOnly(["gpt-5-nano"]); + expect( + buildPresetPrefill({ ...base, classification_mode: "user_turn" }, availability).complexityRouterConfig + .classification_mode, + ).toBe("user_turn"); + expect(buildPresetPrefill(base, availability).complexityRouterConfig.classification_mode).toBe("every_request"); + }); + it("falls back to the defaults when a preset omits match_threshold and escalation_keywords", () => { const prefill = buildPresetPrefill( { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }, @@ -610,6 +634,7 @@ describe("autorouter_presets", () => { const base = { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -625,6 +650,7 @@ describe("autorouter_presets", () => { const config = { tiers: { SIMPLE: ["claude-sonnet-4-5"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -639,6 +665,7 @@ describe("autorouter_presets", () => { REASONING: [{ model_name: "o3", litellm_params: { reasoning_effort: "high" } }], }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -658,6 +685,7 @@ describe("autorouter_presets", () => { REASONING: [{ model_name: "claude-sonnet-4-5", litellm_params: { reasoning_effort: "high" } }], }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; @@ -680,6 +708,9 @@ describe("autorouter_presets", () => { ], }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, + session_affinity: false, + deployment_affinity: true, }; const prefill = buildPresetPrefill(config, groupsOnly(["claude-sonnet-4.5"])); // temperature survives from the spelling that would otherwise have been overwritten; @@ -693,6 +724,7 @@ describe("autorouter_presets", () => { const config = { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic" as const, + classification_mode: "every_request" as const, session_affinity: false, deployment_affinity: true, }; diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index f96dd5ddb4c..4c1f08b62e9 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -6,6 +6,7 @@ import { ComplexityRouterConfigValue, ClassifierType, ClassifierLLMConfig, + DEFAULT_CLASSIFICATION_MODE, DEFAULT_SESSION_AFFINITY, DEFAULT_DEPLOYMENT_AFFINITY, usesLlmClassifier, @@ -288,6 +289,7 @@ export const buildPresetPrefill = ( classifier_context_budget_chars: config.classifier_context_budget_chars, classifier_context_per_turn_chars: config.classifier_context_per_turn_chars, classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns, + classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE, session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY, deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, adaptive: config.adaptive, From 88501a074d1aa5323815b3108ff22c8f76f419ec Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 22:20:36 -0700 Subject: [PATCH 083/126] test(e2e-ui): poll credential availability before Test Connect to deflake multi-instance runs --- .../e2e/ui/tests/modelsPage/addModel.spec.ts | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index a16040678f3..566d78549b2 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -226,6 +226,32 @@ 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 instance on a periodic sync + await expect + .poll( + async () => { + const probe = await page.request.post("/health/test_connection", { + headers: auth, + data: { + litellm_params: { + model: "openai/fake-gpt-4", + custom_llm_provider: "openai", + litellm_credential_name: credentialName, + }, + model_info: {}, + mode: "chat", + }, + }); + if (!probe.ok()) return false; + return (await probe.json()).status === "success"; + }, + { + message: `stored credential ${credentialName} never became usable for a connection test`, + timeout: 60_000, + }, + ) + .toBe(true); + try { await navigateToPage(page, Page.Models); await page.getByRole("tab", { name: "Add Model" }).click(); From 191313e756fe606db5dceac80bed72618ff6679c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:22:05 -0700 Subject: [PATCH 084/126] test(websearch): register configured search tool in pre-request hook test PR #38113 made a configured search_tool_name fail fast when the router does not carry a matching search tool, which broke test_pre_request_hook_modifies_request_body: it names test-search-tool but never registers it. Stub the proxy router with that tool so the test exercises the conversion path again. --- .../test_websearch_interception_e2e.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py index 091ea106b91..cc7901b1710 100644 --- a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py +++ b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py @@ -937,6 +937,14 @@ async def test_pre_request_hook_modifies_request_body(): print("✅ WebSearchInterceptionLogger initialized") + mock_router = MagicMock() + mock_router.search_tools = [ + { + "search_tool_name": "test-search-tool", + "litellm_params": {"search_provider": "tavily"}, + } + ] + # Track what actually gets sent to the API captured_request = {} @@ -987,7 +995,7 @@ async def test_pre_request_hook_modifies_request_body(): with patch( "litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages_handler", side_effect=mock_anthropic_messages_handler, - ): + ), patch("litellm.proxy.proxy_server.llm_router", mock_router): print( "\n📝 Making request with native web_search_20250305 tool (stream=True)..." From f65bee6d74322804dea454ef0e6c904a26f89198 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:32:05 -0700 Subject: [PATCH 085/126] test(websearch): carry a reasoned test-quality suppression on the router patch --- .../test_websearch_interception_e2e.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py index cc7901b1710..fd95b7fa8f2 100644 --- a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py +++ b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py @@ -995,7 +995,10 @@ async def test_pre_request_hook_modifies_request_body(): with patch( "litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages_handler", side_effect=mock_anthropic_messages_handler, - ), patch("litellm.proxy.proxy_server.llm_router", mock_router): + ), patch( # test-quality-ok: the hook imports this process-global router at call time; no injection seam exists to register search_tools + "litellm.proxy.proxy_server.llm_router", + mock_router, + ): print( "\n📝 Making request with native web_search_20250305 tool (stream=True)..." From a48953a0a86b1db88e607af90540c9263780229c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 22:37:35 -0700 Subject: [PATCH 086/126] test(e2e-ui): require consecutive credential probe successes to cover multi-replica routing --- tests/e2e/ui/tests/modelsPage/addModel.spec.ts | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index 566d78549b2..073d3c0b79c 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -226,7 +226,9 @@ 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 instance on a periodic sync + // 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 + let consecutiveProbeSuccesses = 0; await expect .poll( async () => { @@ -242,15 +244,16 @@ test.describe("Add Model", () => { mode: "chat", }, }); - if (!probe.ok()) return false; - return (await probe.json()).status === "success"; + const healthy = probe.ok() && (await probe.json()).status === "success"; + consecutiveProbeSuccesses = healthy ? consecutiveProbeSuccesses + 1 : 0; + return consecutiveProbeSuccesses; }, { message: `stored credential ${credentialName} never became usable for a connection test`, timeout: 60_000, }, ) - .toBe(true); + .toBeGreaterThanOrEqual(3); try { await navigateToPage(page, Page.Models); From db46973ec4b0009d42d0530f50b5d72e5692291b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 23:08:40 -0700 Subject: [PATCH 087/126] feat(ui): modality routing toggle on the auto-router create and edit forms (#39059) --- .../src/autorouter_presets.json | 4 +++ .../components/add_model/AffinityControls.tsx | 26 +++++++++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 21 ++++++++++++++ .../add_model/ComplexityRouterConfig.tsx | 29 ++++++------------- .../add_model/ModalityRoutingControls.tsx | 26 +++++++++++++++++ .../add_model/add_auto_router_tab.tsx | 1 + .../build_complexity_router_config.test.ts | 7 +++++ .../build_complexity_router_config.ts | 4 +++ .../edit_auto_router_modal.test.ts | 24 ++++++++++++++- .../edit_auto_router_modal.tsx | 4 +++ .../src/lib/autorouter_presets.test.ts | 12 ++++++++ .../src/lib/autorouter_presets.ts | 1 + 12 files changed, 138 insertions(+), 21 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx diff --git a/ui/litellm-dashboard/src/autorouter_presets.json b/ui/litellm-dashboard/src/autorouter_presets.json index da35d6171dc..b977bd484ac 100644 --- a/ui/litellm-dashboard/src/autorouter_presets.json +++ b/ui/litellm-dashboard/src/autorouter_presets.json @@ -16,6 +16,7 @@ "escalation_keywords": ["LITELLM ESCALATE"], "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } }, @@ -33,6 +34,7 @@ "escalation_keywords": ["LITELLM ESCALATE"], "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } }, @@ -60,6 +62,7 @@ "escalation_keywords": ["LITELLM ESCALATE"], "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } }, @@ -77,6 +80,7 @@ "escalation_keywords": ["LITELLM ESCALATE"], "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } } diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx new file mode 100644 index 00000000000..4d9e2122739 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx @@ -0,0 +1,26 @@ +import React from "react"; + +import { Switch } from "@/components/ui/switch"; + +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { DEFAULT_DEPLOYMENT_AFFINITY } from "./ComplexityRouterConfig"; + +export const AffinityControls: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => ( + <> +
+ onChange({ ...value, deployment_affinity: deploymentAffinity })} + aria-label="Pin a session to one deployment per model group" + /> + Pin a session to one deployment per model group +
+ + Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to + load-balance every turn. + + +); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index bbee03135c7..751f8870561 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -836,6 +836,27 @@ describe("ComplexityRouterConfig tier labels", () => { }); }); +describe("ComplexityRouterConfig modality panel", () => { + it("defaults the image-routing switch off and writes modality_routing through onChange", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + const toggle = screen.getByRole("switch", { name: "Route image requests to vision-capable models" }); + expect(toggle).not.toBeChecked(); + fireEvent.click(toggle); + + expect(onChange).toHaveBeenCalledWith({ ...defaultValue, modality_routing: true }); + }); + + it("renders a stored modality_routing=true as on", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + expect(screen.getByRole("switch", { name: "Route image requests to vision-capable models" })).toBeChecked(); + }); +}); + describe("ComplexityRouterConfig affinity panel", () => { it("holds the deployment switch at its backend default, session pinning having moved to the frequency choice", () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 5fde87fd638..6062705aa82 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -4,6 +4,9 @@ import { SearchSelect } from "@/components/shared/SearchSelect"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react"; import { Switch } from "@/components/ui/switch"; + +import { AffinityControls } from "./AffinityControls"; +import { ModalityRoutingControls } from "./ModalityRoutingControls"; import { Card, CardContent } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; @@ -396,6 +399,7 @@ export interface ComplexityRouterConfigValue { heuristic_first_max_tier?: string; classification_mode?: ClassificationMode; session_affinity?: boolean; + modality_routing?: boolean; deployment_affinity?: boolean; /** Plan-mode floor as a tier ROW ID, unset meaning off. The wire carries the row's name. */ plan_mode_min_tier?: string; @@ -511,26 +515,6 @@ export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE"; */ export const HEURISTIC_FIRST_MAX_TIER_KEYS = TIER_KEYS.slice(0, -1); -const AffinityControls: React.FC<{ - value: ComplexityRouterConfigValue; - onChange: (value: ComplexityRouterConfigValue) => void; -}> = ({ value, onChange }) => ( - <> -
- onChange({ ...value, deployment_affinity: deploymentAffinity })} - aria-label="Pin a session to one deployment per model group" - /> - Pin a session to one deployment per model group -
- - Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to - load-balance every turn. - - -); - const PlanModeOverrideControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; @@ -841,6 +825,11 @@ const ComplexityRouterConfig: React.FC = ({ label: Advanced: Affinity, children: , }, + { + key: "modality", + label: Advanced: Modality Routing, + children: , + }, { key: "plan-mode", label: Advanced: Plan-Mode Override, diff --git a/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx new file mode 100644 index 00000000000..dd697b35239 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx @@ -0,0 +1,26 @@ +import React from "react"; + +import { Switch } from "@/components/ui/switch"; + +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; + +export const ModalityRoutingControls: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => ( + <> +
+ onChange({ ...value, modality_routing: modalityRouting })} + aria-label="Route image requests to vision-capable models" + /> + Route image requests to vision-capable models +
+ + Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default + model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are replaced, + and a kept session pin still wins. + + +); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index c60ce5e4959..318adcce369 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -351,6 +351,7 @@ const AddAutoRouterTab: React.FC = ({ classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns, classifierFallback: complexityRouterConfig.classifier_fallback, sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY, + modalityRouting: complexityRouterConfig.modality_routing ?? false, deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords, keywordTierRules, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 40addf086b6..af87cc6c8fb 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -55,6 +55,7 @@ describe("buildComplexityRouterConfig", () => { classification_mode: "every_request", session_affinity: false, deployment_affinity: true, + modality_routing: false, escalation_keywords: ["LITELLM ESCALATE"], }; expect(config).toEqual(expected); @@ -248,6 +249,12 @@ describe("buildComplexityRouterConfig", () => { expect(config.return_raw_model_name).toBeUndefined(); }); + it("writes modality_routing explicitly both ways, so the stored config never relies on the backend default", () => { + expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: true }).modality_routing).toBe(true); + expect(buildComplexityRouterConfig(baseParams).modality_routing).toBe(false); + expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: false }).modality_routing).toBe(false); + }); + it("writes session_affinity=true so turning the toggle on overrides the backend's off-by-default", () => { const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: true }); expect(config.session_affinity).toBe(true); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 6a087649c00..af34dc92c0f 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -109,6 +109,7 @@ export interface BuildComplexityRouterConfigParams { heuristicFirstMaxTier: string | undefined; classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; + modalityRouting?: boolean; deploymentAffinity: boolean; customTechnicalKeywords: string[]; keywordTierRules: KeywordTierRule[]; @@ -165,6 +166,7 @@ export interface ComplexityRouterConfigPayload { classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; + modality_routing: boolean; custom_technical_keywords?: string[]; keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[]; semantic_keyword_matching?: boolean; @@ -399,6 +401,7 @@ export const buildComplexityRouterConfig = ({ heuristicFirstMaxTier, classificationMode, sessionAffinity, + modalityRouting, deploymentAffinity, customTechnicalKeywords, keywordTierRules, @@ -460,6 +463,7 @@ export const buildComplexityRouterConfig = ({ classification_mode: classificationMode ?? DEFAULT_CLASSIFICATION_MODE, session_affinity: sessionAffinity, deployment_affinity: deploymentAffinity, + modality_routing: modalityRouting ?? false, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), escalation_keywords: cleanedEscalationKeywords, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index d199af51b41..481ba2b6b00 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -1,4 +1,4 @@ -import { buildUpdatedComplexityRouterConfig } from "./edit_auto_router_modal"; +import { buildUpdatedComplexityRouterConfig, hydrateComplexityRouterConfig } from "./edit_auto_router_modal"; const storedConfigValue = { tiers: { @@ -50,6 +50,7 @@ const expectedClassifiedTierConfig = { classification_mode: "every_request", session_affinity: false, deployment_affinity: true, + modality_routing: false, adaptive: true, adaptive_weights: { quality: 0.4, cost: 0.6 }, adaptive_eligible: "classified_tier", @@ -72,6 +73,7 @@ const expectedAdaptiveDisabledConfig = { classification_mode: "every_request", session_affinity: false, deployment_affinity: true, + modality_routing: false, }; describe("buildUpdatedComplexityRouterConfig", () => { @@ -87,6 +89,26 @@ describe("buildUpdatedComplexityRouterConfig", () => { expect(updatedConfig).toEqual(expectedAdaptiveDisabledConfig); }); + it("hydrates a stored modality_routing into form state and defaults absent to off", () => { + expect(hydrateComplexityRouterConfig({ ...storedConfig, modality_routing: true }, null).modality_routing).toBe( + true, + ); + expect(hydrateComplexityRouterConfig(storedConfig, null).modality_routing).toBe(false); + }); + + it("round-trips modality_routing explicitly in both directions", () => { + const enabled = buildUpdatedComplexityRouterConfig(storedConfig, { + ...classifiedTierValue, + modality_routing: true, + }); + expect(enabled.modality_routing).toBe(true); + const disabled = buildUpdatedComplexityRouterConfig( + { ...storedConfig, modality_routing: true }, + { ...classifiedTierValue, modality_routing: false }, + ); + expect(disabled.modality_routing).toBe(false); + }); + it("includes return_raw_model_name only when enabled", () => { const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, { ...classifiedTierValue, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 2aecf0b3483..e18582f77a0 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -102,6 +102,7 @@ export interface StoredComplexityRouterConfig { dimension_weights?: unknown; reasoning_override_min_score?: unknown; session_affinity?: unknown; + modality_routing?: unknown; deployment_affinity?: unknown; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; @@ -176,6 +177,7 @@ export const hydrateComplexityRouterConfig = ( reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), session_affinity: typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, + modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, deployment_affinity: typeof parsedConfig.deployment_affinity === "boolean" ? parsedConfig.deployment_affinity @@ -214,6 +216,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "heuristic_first_max_tier", "classification_mode", "session_affinity", + "modality_routing", "deployment_affinity", "adaptive", "adaptive_weights", @@ -309,6 +312,7 @@ export const buildUpdatedComplexityRouterConfig = ( classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, classifierFallback: value.classifier_fallback, sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, + modalityRouting: value.modality_routing ?? false, deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords: customTechnicalKeywords ?? [], keywordTierRules: keywordMatching?.keywordTierRules ?? [], diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 5632d6d947f..2a8306473b8 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -129,6 +129,18 @@ describe("autorouter_presets", () => { } }); + it("carries a preset's modality_routing into the prefilled form state", () => { + const preset = getPresetByKey("anthropic_family")!; + const withFlag = { ...preset.complexity_router_config, modality_routing: true }; + const prefill = buildPresetPrefill(withFlag, groupsOnly(getRequiredModelsInPreset(preset))); + expect(prefill.complexityRouterConfig.modality_routing).toBe(true); + const withoutFlag = buildPresetPrefill( + preset.complexity_router_config, + groupsOnly(getRequiredModelsInPreset(preset)), + ); + expect(withoutFlag.complexityRouterConfig.modality_routing).toBe(false); + }); + it("prefills the anthropic preset's effort through to tier_model_params", () => { const preset = getPresetByKey("anthropic_family")!; const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset))); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 4c1f08b62e9..721bd6f2b2a 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -292,6 +292,7 @@ export const buildPresetPrefill = ( classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE, session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY, deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + modality_routing: config.modality_routing ?? false, adaptive: config.adaptive, adaptive_weights: config.adaptive_weights, tier_distance_penalty: config.tier_distance_penalty, From a27e12367e2c3574586128a55e970fa5d17d5379 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 31 Aug 2026 21:50:40 -0700 Subject: [PATCH 088/126] fix(bedrock): forward native structured outputs on Invoke instead of silently inlining the schema --- litellm/llms/anthropic/chat/transformation.py | 24 +- .../anthropic_claude3_transformation.py | 44 +--- litellm/llms/bedrock/common_utils.py | 89 ++++++++ .../anthropic_claude3_transformation.py | 60 ++--- ...odel_prices_and_context_window_backup.json | 24 +- model_prices_and_context_window.json | 24 +- .../test_anthropic_chat_transformation.py | 42 ++++ ...ations_anthropic_claude3_transformation.py | 131 +++++++++-- .../test_anthropic_claude3_transformation.py | 206 +++++++++++++++--- .../llms/bedrock/test_bedrock_common_utils.py | 41 ++++ 10 files changed, 535 insertions(+), 150 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e1387a9068c..a3c76d6a29b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1992,19 +1992,35 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return data def _apply_output_config(self, data: dict, model: str, optional_params: dict) -> None: - """Validate and apply output_config to the request data.""" + """Validate and apply output_config to the request data. + + The ``drop_params`` gate here is an effort gate: ``format`` is a + structured-output field, not an effort field, so it survives the drop + and is vetted where it is consumed (the map's + ``supports_native_structured_output`` flag on emission paths). + """ if "output_config" not in optional_params: return output_config: Final = optional_params.get("output_config") if not output_config or not isinstance(output_config, dict): return - if litellm.drop_params is True and not self._model_supports_effort_param(model, self._resolved_provider): + if ( + litellm.drop_params is True + and any(key != "format" for key in output_config) + and not self._model_supports_effort_param(model, self._resolved_provider) + ): litellm.verbose_logger.warning( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, model, ) - optional_params.pop("output_config", None) - data.pop("output_config", None) + preserved_format: Final = output_config.get("format") + if preserved_format is None: + optional_params.pop("output_config", None) + data.pop("output_config", None) + return + format_only: Final = {"format": preserved_format} # mutable-ok: json body + optional_params["output_config"] = format_only # rebind-ok: out-param store + data["output_config"] = format_only # rebind-ok: out-param store return effort: Final = output_config.get("effort") valid_efforts: Final = ["high", "medium", "low", "xhigh", "max"] diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 40b90014f3b..8e709349400 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers -from litellm.litellm_core_utils.litellm_logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -16,17 +15,16 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( - convert_bedrock_invoke_output_format_to_inline_schema, + apply_bedrock_invoke_structured_output, get_anthropic_beta_from_headers, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - pop_bedrock_invoke_output_config_format, + strip_unsupported_bedrock_invoke_output_config_keys, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse -from litellm.utils import _supports_factory if TYPE_CHECKING: import tiktoken @@ -212,36 +210,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) anthropic_request.pop("stream_chunk_size", None) - output_format: Final = anthropic_request.pop("output_format", None) - output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_request) - if output_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_format, - request_body=anthropic_request, - ) - elif output_config_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_config_format, - request_body=anthropic_request, - ) - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model, "bedrock") - ): - if anthropic_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) + apply_bedrock_invoke_structured_output( + model=model, + request_body=anthropic_request, + ) + strip_unsupported_bedrock_invoke_output_config_keys( + model=model, + request_body=anthropic_request, + ) if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 72e3cc1b326..df65df642a2 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -177,6 +177,95 @@ def convert_bedrock_invoke_output_format_to_inline_schema( request_body["messages"] = new_messages +def _bedrock_model_supports(model: str, key: str) -> bool: + from litellm.utils import _supports_factory + + return _supports_factory(model=model, custom_llm_provider="bedrock", key=key) + + +def apply_bedrock_invoke_structured_output( + model: str, + request_body: dict[str, object], # mutable-ok: edited in place like siblings +) -> None: + """ + Route Anthropic structured-output params to what the Bedrock model supports. + + Consumes the legacy top-level ``output_format`` and the newer + ``output_config.format``, keeping the pre-existing precedence of the legacy + field when a request carries both. Models flagged + ``supports_native_structured_output`` in the model map get the schema + forwarded as ``output_config.format``, which Bedrock relays to the model for + enforced structured output. For every other model the schema is inlined into + the last user message as best-effort text, with a warning because nothing + enforces it. + """ + legacy_output_format: Final = request_body.pop("output_format", None) + output_config_format: Final = pop_bedrock_invoke_output_config_format(request_body) + schema_format: Final = legacy_output_format if isinstance(legacy_output_format, dict) else output_config_format + if schema_format is None: + return + + if _bedrock_model_supports(model, "supports_native_structured_output"): + existing_output_config: Final = request_body.get("output_config") + if isinstance(existing_output_config, dict): + existing_output_config["format"] = schema_format + else: + request_body["output_config"] = {"format": schema_format} # rebind-ok: out-param # mutable-ok: json + return + + verbose_logger.warning( + "Bedrock Invoke: model=%s does not advertise `supports_native_structured_output` " + "in model_prices_and_context_window.json, so the JSON schema was inlined into " + "the last user message and is NOT enforced by the model.", + model, + ) + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=schema_format, + request_body=request_body, + ) + + +def strip_unsupported_bedrock_invoke_output_config_keys( + model: str, + request_body: dict[str, object], # mutable-ok: edited in place like siblings +) -> None: + """ + Drop ``output_config`` keys the Bedrock model does not accept. + + ``format`` survives unconditionally: it is only attached for models whose map + entry advertises ``supports_native_structured_output``. Effort-bearing keys + survive only when the map flags ``supports_output_config`` or a + ``supports_*_reasoning_effort`` tier; otherwise they are dropped with a + warning so Bedrock does not reject the request. + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + output_config: Final = request_body.get("output_config") + if not isinstance(output_config, dict): + return + if all(key == "format" for key in output_config): + return + if _bedrock_model_supports(model, "supports_output_config") or AnthropicConfig._model_supports_effort_param( + model, "bedrock" + ): + return + + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` keys for " + "model=%s: neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) + preserved_format: Final = output_config.get("format") + if preserved_format is None: + request_body.pop("output_config", None) + else: + request_body["output_config"] = {"format": preserved_format} # rebind-ok: out-param # mutable-ok: json + + def normalize_custom_field_on_tools(request_body: dict) -> None: """ Drop the ``custom`` field from each tool, first hoisting a boolean diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index f74a290d773..6ff9f0155f9 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -29,14 +29,14 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( - convert_bedrock_invoke_output_format_to_inline_schema, + apply_bedrock_invoke_structured_output, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - pop_bedrock_invoke_output_config_format, + strip_unsupported_bedrock_invoke_output_config_keys, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -51,7 +51,6 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk, ModelResponseStream from litellm.types.utils import GenericStreamingChunk as GChunk -from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -708,52 +707,25 @@ class AmazonAnthropicClaudeMessagesConfig( # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models) self._remove_ttl_from_cache_control(anthropic_messages_request=anthropic_messages_request, model=model) - # 5. Convert structured-output params to inline schema. - # Bedrock Invoke doesn't support top-level `output_format`; its - # accepted `output_config` subset is also narrower than Anthropic's, so - # consume the newer `output_config.format` shape here instead of - # forwarding it as an unknown nested key. + # 5. Route structured-output params (`output_format` / + # `output_config.format`) to native enforcement or the inline-schema + # fallback, then strip `output_config` keys the model does not accept. + # Ref: https://github.com/BerriAI/litellm/issues/22797 existing_output_config: Final = anthropic_messages_request.get("output_config") if isinstance(existing_output_config, dict): anthropic_messages_request["output_config"] = dict(existing_output_config) - output_format: Final = anthropic_messages_request.pop("output_format", None) - output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_messages_request) - if output_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_format, - request_body=anthropic_messages_request, - ) - elif output_config_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_config_format, - request_body=anthropic_messages_request, - ) + apply_bedrock_invoke_structured_output( + model=model, + request_body=anthropic_messages_request, + ) normalize_bedrock_opus_output_config_effort( model=model, output_config=anthropic_messages_request.get("output_config"), ) - - # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, - # but older models do not — strip it to avoid request rejection. - # Ref: https://github.com/BerriAI/litellm/issues/22797 - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model, "bedrock") - ): - if anthropic_messages_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) + strip_unsupported_bedrock_invoke_output_config_keys( + model=model, + request_body=anthropic_messages_request, + ) # 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) # Ref: https://github.com/BerriAI/litellm/issues/22847 @@ -774,9 +746,11 @@ class AmazonAnthropicClaudeMessagesConfig( if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + remaining_output_config: Final = anthropic_messages_request.get("output_config") if ( litellm.drop_params is True - and "output_config" in anthropic_messages_request + and isinstance(remaining_output_config, dict) + and any(key != "format" for key in remaining_output_config) and not AnthropicConfig._model_supports_effort_param(model, "bedrock") ): verbose_logger.warning( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 718e6c489fd..80dc49a770b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1591,7 +1591,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1627,7 +1627,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1663,7 +1663,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1699,7 +1699,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1735,7 +1735,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1771,7 +1771,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2064,7 +2064,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2101,7 +2101,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2138,7 +2138,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2175,7 +2175,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2212,7 +2212,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2249,7 +2249,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 718e6c489fd..80dc49a770b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1591,7 +1591,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1627,7 +1627,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1663,7 +1663,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1699,7 +1699,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1735,7 +1735,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1771,7 +1771,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2064,7 +2064,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2101,7 +2101,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2138,7 +2138,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2175,7 +2175,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2212,7 +2212,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2249,7 +2249,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 25e2c3cda80..c4df46dea83 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6207,3 +6207,45 @@ def test_disabled_thinking_omitted_only_for_always_on_models( assert "thinking" not in request else: assert request["thinking"] == {"type": "disabled"} + + +def test_anthropic_drop_params_keeps_format_only_output_config(monkeypatch): + """``drop_params=True`` must not consume ``output_config.format``: the drop + gate is an effort gate and ``format`` is a structured-output field.""" + monkeypatch.setattr(litellm, "drop_params", True) + config = AnthropicConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"z": {"type": "integer"}}}, + } + + result = config.transform_request( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_config": {"format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_anthropic_drop_params_reduces_mixed_output_config_to_format(monkeypatch): + """``drop_params=True`` drops the effort key on unsupported models but keeps + ``format`` so structured outputs still reach the provider.""" + monkeypatch.setattr(litellm, "drop_params", True) + config = AnthropicConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"z": {"type": "integer"}}}, + } + + result = config.transform_request( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_config": {"effort": "low", "format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index cea299280f8..a122d97a0f0 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -428,30 +428,58 @@ def test_output_config_forwarded_for_bedrock_chat_invoke_request(): def test_output_config_format_converted_for_bedrock_chat_invoke_request(): - """Bedrock Invoke chat path consumes ``output_config.format`` before forwarding.""" + """Bedrock Invoke chat path inlines ``output_config.format`` for models + without native structured-output support and keeps the effort key.""" config = AmazonAnthropicClaudeConfig() schema = { "type": "object", "properties": {"answer": {"type": "string"}}, } - result = config.transform_request( + with patch( # test-quality-ok: pin non-native path + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", + ): + result = config.transform_request( + model="anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "test"}], + optional_params={ + "max_tokens": 100, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"effort": "xhigh"} + last_content = result["messages"][0]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_output_config_format_forwarded_for_bedrock_chat_invoke_request(): + """Bedrock Invoke chat path forwards ``output_config.format`` alongside effort + for models with native structured-output support (Claude Opus 4.7).""" + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + } + + result = AmazonAnthropicClaudeConfig().transform_request( model="anthropic.claude-opus-4-7", messages=[{"role": "user", "content": "test"}], optional_params={ "max_tokens": 100, - "output_config": { - "effort": "xhigh", - "format": {"type": "json_schema", "schema": schema}, - }, + "output_config": {"effort": "xhigh", "format": schema_format}, }, litellm_params={}, headers={}, ) - assert result.get("output_config") == {"effort": "xhigh"} - last_content = result["messages"][0]["content"] - assert json.loads(last_content[-1]["text"]) == schema + assert result.get("output_config") == {"effort": "xhigh", "format": schema_format} + assert "answer" not in json.dumps(result["messages"]) @pytest.mark.parametrize( @@ -488,7 +516,7 @@ def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider( optional_params = {"max_tokens": 100, "output_config": {"effort": "high"}} with patch( - "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ) as mock_supports_factory: result = config.transform_request( @@ -499,11 +527,7 @@ def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider( headers={}, ) - mock_supports_factory.assert_called_once_with( - model="us.anthropic.claude-opus-4-7", - custom_llm_provider="bedrock", - key="supports_output_config", - ) + mock_supports_factory.assert_called_once_with("us.anthropic.claude-opus-4-7", "supports_output_config") assert result["output_config"] == {"effort": "high"} @@ -542,3 +566,80 @@ def test_output_format_removed_from_bedrock_invoke_request(): assert ( "output_format" not in result ), f"output_format should be removed for Bedrock Invoke, got keys: {result.keys()}" + + +def test_bedrock_chat_invoke_forwards_output_config_format_natively(local_model_cost_map): + """Regression: ``output_config.format`` is forwarded verbatim on models Bedrock + enforces structured outputs for, instead of being inlined as prompt text.""" + import json + + config = AmazonAnthropicClaudeConfig() + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"zebra_count": {"type": "integer"}}, + "required": ["zebra_count"], + "additionalProperties": False, + }, + } + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + assert "zebra_count" not in json.dumps(result["messages"]) + + +def test_bedrock_chat_invoke_drop_params_keeps_native_output_config_format(local_model_cost_map, monkeypatch): + """``drop_params=True`` must not eat ``output_config.format`` before the + native-forwarding router runs (Sonnet 4.5 has no effort flags).""" + import litellm + + monkeypatch.setattr(litellm, "drop_params", True) + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = AmazonAnthropicClaudeConfig().transform_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={"max_tokens": 100, "output_config": {"format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_bedrock_chat_invoke_drop_params_still_inlines_for_non_native(local_model_cost_map, monkeypatch): + """``drop_params=True`` on a model without native structured-output support + still reaches the inline-schema fallback instead of losing the schema.""" + import litellm + + monkeypatch.setattr(litellm, "drop_params", True) + schema = {"type": "object", "properties": {"zebra_count": {"type": "integer"}}} + + result = AmazonAnthropicClaudeConfig().transform_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params={}, + headers={}, + ) + + assert "output_config" not in result + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 8d07d38b1b6..09ebc1a3c95 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -935,7 +935,7 @@ def test_bedrock_messages_strips_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=False, ): result = cfg.transform_anthropic_messages_request( @@ -970,7 +970,7 @@ def test_bedrock_messages_preserves_output_config_for_claude_4_6(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1003,7 +1003,7 @@ def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ) as mock_supports_factory: result = cfg.transform_anthropic_messages_request( @@ -1014,11 +1014,7 @@ def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): headers={}, ) - mock_supports_factory.assert_called_with( - model="us.anthropic.claude-opus-4-7", - custom_llm_provider="bedrock", - key="supports_output_config", - ) + mock_supports_factory.assert_called_with("us.anthropic.claude-opus-4-7", "supports_output_config") assert result["output_config"] == {"effort": "high"} @@ -1038,7 +1034,7 @@ def test_bedrock_messages_forwards_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1054,27 +1050,29 @@ def test_bedrock_messages_forwards_output_config(): def test_bedrock_messages_forwards_output_config_with_output_format(): - """``output_config`` is forwarded; ``output_format`` is converted to inline schema.""" + """Legacy ``output_format`` is forwarded as ``output_config.format`` on models + that support native structured outputs, alongside the effort key.""" from unittest.mock import patch from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } optional_params = { "max_tokens": 4096, "output_config": {"effort": "low"}, - "output_format": { - "type": "json_schema", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - }, + "output_format": schema_format, } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1085,12 +1083,14 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): headers={}, ) - assert result.get("output_config") == {"effort": "low"} + assert result.get("output_config") == {"effort": "low", "format": schema_format} assert "output_format" not in result + assert "answer" not in json.dumps(result["messages"]) def test_bedrock_messages_converts_output_config_format_to_inline_schema(): - """``output_config.format`` is consumed so Bedrock does not see an unknown nested key.""" + """Without native structured-output support, ``output_config.format`` falls back + to the inline schema so Bedrock does not see an unknown nested key.""" from unittest.mock import patch from litellm.types.router import GenericLiteLLMParams @@ -1110,8 +1110,8 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", - return_value=True, + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1146,7 +1146,7 @@ def test_bedrock_messages_normalizes_output_config_effort_for_opus( cfg = AmazonAnthropicClaudeMessagesConfig() with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1184,8 +1184,8 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", - return_value=True, + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1229,7 +1229,7 @@ def test_bedrock_messages_does_not_mutate_callers_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): cfg.transform_anthropic_messages_request( @@ -1271,7 +1271,7 @@ def test_bedrock_messages_strips_output_config_with_output_format(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=False, ): result = cfg.transform_anthropic_messages_request( @@ -1332,7 +1332,7 @@ def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): litellm.drop_params = True try: with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1375,7 +1375,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1482,7 +1482,7 @@ def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -3104,3 +3104,149 @@ async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): break await asyncio.sleep(0.01) assert logging_obj.completion_start_time is not None + + +def test_bedrock_messages_forwards_output_config_format_natively(local_model_cost_map): + """Regression: on a model Bedrock enforces structured outputs for (Claude + Sonnet 4.5), ``output_config.format`` must be forwarded verbatim, not + silently rewritten into inline prompt text.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "zebra_count": {"type": "integer"}, + "is_tuesday": {"type": "boolean"}, + }, + "required": ["zebra_count", "is_tuesday"], + "additionalProperties": False, + }, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + assert "zebra_count" not in json.dumps(result["messages"]) + + +def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): + """Bedrock rejects ``output_config.format`` for the Claude 5 family, so the + schema falls back to the inline-text path instead of a deterministic 400.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema = { + "type": "object", + "properties": {"zebra_count": {"type": "integer"}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "output_config" not in result + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(local_model_cost_map): + """When a request carries both schema forms, the legacy top-level + ``output_format`` keeps winning, matching the pre-existing precedence.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + legacy_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"legacy_field": {"type": "string"}}}, + } + newer_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"newer_field": {"type": "string"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_format": legacy_format, + "output_config": {"format": newer_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": legacy_format} + assert "output_format" not in result + assert "newer_field" not in json.dumps(result) + + +def test_bedrock_messages_drop_params_keeps_native_output_config_format(local_model_cost_map, monkeypatch): + """``drop_params=True`` must not strip a natively forwarded + ``output_config.format`` on models without effort support (Sonnet 4.5).""" + import litellm + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setattr(litellm, "drop_params", True) + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_model_cost_map): + """Sonnet 4.5 has native structured-output support but no effort support, so + a mixed ``output_config`` keeps ``format`` and drops ``effort``.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 4096, + "output_config": {"format": schema_format, "effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 389bf4a8e40..609b5c75801 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -520,3 +520,44 @@ def test_merge_bedrock_aws_request_params_keeps_caller_credentials_without_stati assert merged["aws_secret_access_key"] == "caller-secret" assert merged["aws_session_token"] == "caller-token" assert merged["aws_region_name"] == "us-west-2" + + +def test_strip_unsupported_output_config_keeps_format_drops_effort(local_model_cost_map): + """On a model with neither effort flag, only the ``format`` key survives.""" + from litellm.llms.bedrock.common_utils import ( + strip_unsupported_bedrock_invoke_output_config_keys, + ) + + schema_format = {"type": "json_schema", "schema": {"type": "object"}} + body = {"output_config": {"effort": "high", "format": schema_format}} + + strip_unsupported_bedrock_invoke_output_config_keys( + model="anthropic.claude-3-haiku-20240307-v1:0", + request_body=body, + ) + + assert body["output_config"] == {"format": schema_format} + + +def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_map): + """The legacy ``output_format`` wins over ``output_config.format`` when a + request carries both, matching the pre-existing precedence.""" + from litellm.llms.bedrock.common_utils import ( + apply_bedrock_invoke_structured_output, + ) + + legacy = {"type": "json_schema", "schema": {"type": "object", "properties": {"a": {"type": "string"}}}} + newer = {"type": "json_schema", "schema": {"type": "object", "properties": {"b": {"type": "string"}}}} + body = { + "messages": [{"role": "user", "content": "hi"}], + "output_format": legacy, + "output_config": {"format": newer}, + } + + apply_bedrock_invoke_structured_output( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + request_body=body, + ) + + assert body["output_config"] == {"format": legacy} + assert "output_format" not in body From b11f0bcb9211d23fd05fbbc842b516b362156fce Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 31 Aug 2026 23:19:45 -0700 Subject: [PATCH 089/126] fix(proxy): include litellm_model_table in GET /v2/team/list (#39045) * fix(proxy): include litellm_model_table in GET /v2/team/list GET /v2/team/list built its find_many queries without joining the LiteLLM_ModelTable relation, so litellm_model_table (and the model_aliases it carries) always read back as null there, same bug class as GH #26312 which PR #33047 fixed on /team/info and /team/list but never touched this endpoint. * fix(test): assert observable output, not mock calls, in v2 team list test The test-quality gate flagged the regression test for asserting on find_many's call args instead of what the caller gets back. Rewritten so the fake find_many only attaches litellm_model_table when its own include kwarg asks for it, so the assertions are on the response. * fix(proxy): drop invalid litellm_model_table include on deleted-team query Greptile caught that LiteLLM_DeletedTeamTable has no litellm_model_table relation in the Prisma schema, so passing that include on the deleted-team find_many raised UnknownRelationalFieldError against a real database on every GET /v2/team/list?status=deleted call. Confirmed live against Postgres. Scope the fix to the active-team branch only, where the relation exists; update the test to reflect that and assert the deleted branch no longer requests it. --- .../management_endpoints/team_endpoints.py | 2 + .../test_team_endpoints.py | 87 +++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c6d7975b75e..714cf252e69 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5148,6 +5148,7 @@ async def list_team_v2( # Get teams with pagination if use_deleted_table: + # LiteLLM_DeletedTeamTable has no litellm_model_table relation, unlike below teams = await _deleted_team_db(prisma_client).find_many( where=where_conditions, skip=skip, @@ -5162,6 +5163,7 @@ async def list_team_v2( skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort + include=_INCLUDE_MODEL_TABLE, ) # Get total count for pagination total_count = await _team_db(prisma_client).count(where=where_conditions) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ffa6bc601e9..30b2ab86b9a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3741,6 +3741,93 @@ async def test_list_team_v2_with_status_deleted(): assert len(result["teams"]) == 2 +@pytest.mark.asyncio +async def test_list_team_v2_includes_litellm_model_table(): + """ + Regression test for GH #26312: GET /v2/team/list must eagerly load the + litellm_model_table relation for active teams, same as /team/info and + /team/list, or a team's model_aliases always read back as null from this + endpoint. Deleted teams are excluded: LiteLLM_DeletedTeamTable has no such + relation in the Prisma schema, so requesting it there raises + UnknownRelationalFieldError against a real database. + + The fake find_many below only attaches litellm_model_table when its own + `include` kwarg actually asks for the relation, so the assertions below + are on what the caller gets back, not on how find_many was called. + """ + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user_123", + ) + + def _team_row(team_id: str, include) -> Mock: + model_table = ( + { + "id": 1, + "model_aliases": {"my-fast-model": "fake-model"}, + "created_by": "u", + "updated_by": "u", + "team": None, + } + if (include or {}).get("litellm_model_table") + else None + ) + return Mock( + team_id=team_id, + model_dump=lambda: { + "team_id": team_id, + "team_alias": "t", + "litellm_model_table": model_table, + }, + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: # test-quality-ok: this file's DB-mock convention + mock_db = Mock() + mock_prisma_client.db = mock_db + + mock_db.litellm_teamtable.find_many = AsyncMock( + side_effect=lambda **kw: [_team_row("team_1", kw.get("include"))] + ) + mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + + result = await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status=None, + ) + + assert result["teams"][0].litellm_model_table is not None + assert result["teams"][0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} + + mock_db.litellm_deletedteamtable.find_many = AsyncMock( + side_effect=lambda **kw: [_team_row("team_2", kw.get("include"))] + ) + mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=1) + + await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status="deleted", + ) + + assert "include" not in mock_db.litellm_deletedteamtable.find_many.call_args.kwargs + + @pytest.mark.asyncio async def test_list_team_v2_org_admin_sees_org_teams(): """ From 847d737b8e93dd3abcff819e0537e84550b8c436 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 00:13:05 -0700 Subject: [PATCH 090/126] test(ui): budget DOM-structure assertions in dashboard tests Turn on testing-library/no-node-access, no-container and prefer-screen-queries as warnings and baseline them in eslint-budgets.json so the counts can only go down. These three rules catch tests that assert on DOM structure rather than on what a user can observe: reaching through parentElement chains, querying the container by CSS selector, and destructuring queries off render instead of going through screen. Those assertions break on refactors that change nothing a user sees, and stay green when the behaviour underneath is broken. Baselines are the current counts, so nothing fails today. --- ui/litellm-dashboard/eslint-budgets.json | 5 ++++- ui/litellm-dashboard/eslint.config.mjs | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index bbf69c4a77a..d986722c3b3 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -4,5 +4,8 @@ "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, - "local/no-long-condition-chain": { "max": 265, "target": 120 } + "local/no-long-condition-chain": { "max": 265, "target": 120 }, + "testing-library/no-container": { "max": 150, "target": 50 }, + "testing-library/no-node-access": { "max": 760, "target": 500 }, + "testing-library/prefer-screen-queries": { "max": 221, "target": 0 } } diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index 23cc5096bb1..f5e3b23b3ec 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -104,10 +104,13 @@ const eslintConfig = [ plugins: { "testing-library": testingLibrary, "jest-dom": jestDom }, rules: { "testing-library/await-async-queries": "error", + "testing-library/no-container": "warn", + "testing-library/no-node-access": "warn", "testing-library/no-wait-for-multiple-assertions": "error", "testing-library/no-wait-for-side-effects": "error", "testing-library/prefer-find-by": "error", "testing-library/prefer-presence-queries": "error", + "testing-library/prefer-screen-queries": "warn", "jest-dom/prefer-checked": "error", "jest-dom/prefer-empty": "error", "jest-dom/prefer-enabled-disabled": "error", From 2fbea77afa31f488c442d347ff96da06edf30c9d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 00:19:27 -0700 Subject: [PATCH 091/126] test(ui): assert DataTable behavior instead of DOM structure The shared DataTable test reached for elements by CSS selector and by walking parentElement chains, then asserted on Tailwind class strings. It had no role queries at all, so a wrapper div anywhere in the render tree broke it while changing nothing a user sees. Columns, rows and headers are now found the way a user finds them: by role and by the text on screen. The compact skeleton row is compared against the loaded row's height rather than a hard-coded h-8, so renaming the class no longer breaks the test but shrinking the row still does. The fillHeight and maxBodyHeight cases stay class assertions. jsdom has no layout engine, so there is nothing behavioural to assert there. What they no longer do is derive their elements from incidental nesting: the three layout wrappers and the header now publish a stable test id, which is also why the resizer's write-only data-resizer attribute became one. Budgets drop with the counts: no-container 150 to 133, no-node-access 760 to 723. --- ui/litellm-dashboard/eslint-budgets.json | 4 +- .../shared/DataTable/DataTable.test.tsx | 114 +++++++++--------- .../components/shared/DataTable/DataTable.tsx | 13 +- 3 files changed, 66 insertions(+), 65 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index d986722c3b3..3bfe56724eb 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -5,7 +5,7 @@ "max-depth": { "max": 70, "target": 30 }, "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, "local/no-long-condition-chain": { "max": 265, "target": 120 }, - "testing-library/no-container": { "max": 150, "target": 50 }, - "testing-library/no-node-access": { "max": 760, "target": 500 }, + "testing-library/no-container": { "max": 133, "target": 50 }, + "testing-library/no-node-access": { "max": 723, "target": 500 }, "testing-library/prefer-screen-queries": { "max": 221, "target": 0 } } diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 58a0cd94997..7ead9bb64d4 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -1,5 +1,5 @@ import type { ColumnDef, ExpandedState } from "@tanstack/react-table"; -import { render, screen, waitFor } from "@testing-library/react"; +import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -21,6 +21,12 @@ function person(id: string, name: string, flagged = false): Person { const names = (): (string | null)[] => screen.getAllByTestId("name-cell").map((el) => el.textContent); +const heightClassesOf = (el: HTMLElement | undefined): string[] => + (el?.className ?? "") + .split(/\s+/) + .filter((cls) => cls.startsWith("h-")) + .sort(); + const nameCellColumns: ColumnDef[] = [ { accessorKey: "name", @@ -229,12 +235,10 @@ describe("DataTable sorting", () => { describe("DataTable layout", () => { it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => { - const { container } = render(); + render(); - const table = container.querySelector("table"); - expect(table).not.toBeNull(); // width pins the natural column total (horizontal scroll on overflow); minWidth:100% fills the gap on underflow. - expect(table?.style.minWidth).toBe("100%"); + expect(screen.getByRole("table")).toHaveStyle({ minWidth: "100%" }); }); }); @@ -354,17 +358,21 @@ describe("DataTable loading", () => { const { rerender } = render( , ); - const skeletonRow = screen.getAllByTestId("skeleton-row").at(0); - const loadedRowHeight = "h-8"; - expect(skeletonRow?.className).toContain(loadedRowHeight); + const skeletonHeight = heightClassesOf(screen.getAllByRole("row").at(-1)); rerender(); - expect(document.querySelector("[data-row-id]")?.className).toContain(loadedRowHeight); + const loadedHeight = heightClassesOf(screen.getByRole("row", { name: /Charlie/ })); + + expect(loadedHeight).not.toEqual([]); + expect(skeletonHeight).toEqual(loadedHeight); }); it("does not force the compact height on default-size skeleton rows", () => { - render(); - expect(screen.getAllByTestId("skeleton-row").at(0)?.className).not.toContain("h-8"); + const { rerender } = render(); + const skeletonHeight = heightClassesOf(screen.getAllByRole("row").at(-1)); + + rerender(); + expect(heightClassesOf(screen.getAllByRole("row").at(-1))).not.toEqual(skeletonHeight); }); it("varies skeleton shape and width per column instead of one fixed bar", () => { @@ -420,7 +428,7 @@ describe("DataTable loading", () => { describe("DataTable column visibility", () => { it("hides a column when toggled off in the view-options menu", async () => { const user = userEvent.setup(); - const { container } = render( + render( { />, ); - expect(container.querySelector('th[data-header-id="email"]')).not.toBeNull(); + expect(screen.getByRole("columnheader", { name: "Email" })).toBeInTheDocument(); await user.click(screen.getByTestId("view-options-trigger")); await user.click(await screen.findByTestId("view-option-email")); - await waitFor(() => expect(container.querySelector('th[data-header-id="email"]')).toBeNull()); + await waitFor(() => expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument()); await user.click(screen.getByTestId("view-option-email")); - await waitFor(() => expect(container.querySelector('th[data-header-id="email"]')).not.toBeNull()); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); }); it("omits columns that opt out of hiding from the menu", async () => { @@ -468,14 +476,10 @@ describe("DataTable column visibility", () => { describe("DataTable pinned columns", () => { it("applies sticky positioning to a pinned column only", () => { - const { container } = render(); + render(); - const pinnedHead = container.querySelector('th[data-header-id="name"]'); - const normalHead = container.querySelector('th[data-header-id="email"]'); - - expect(pinnedHead?.style.position).toBe("sticky"); - expect(pinnedHead?.style.left).toBe("0px"); - expect(normalHead?.style.position).toBe(""); + expect(screen.getByRole("columnheader", { name: "Name" })).toHaveStyle({ position: "sticky", left: "0px" }); + expect(screen.getByRole("columnheader", { name: "Email" })).not.toHaveStyle({ position: "sticky" }); }); }); @@ -570,7 +574,7 @@ describe("DataTable expansion", () => { describe("DataTable row styling and footer", () => { it("applies rowClassName to the matching row only", () => { const data = [person("a", "Alice", true), person("b", "Bob", false)]; - const { container } = render( + render( { />, ); - expect(container.querySelector('tr[data-row-id="a"]')?.className).toContain("flagged-row"); - expect(container.querySelector('tr[data-row-id="b"]')?.className).not.toContain("flagged-row"); + expect(screen.getByRole("row", { name: /Alice/ })).toHaveClass("flagged-row"); + expect(screen.getByRole("row", { name: /Bob/ })).not.toHaveClass("flagged-row"); }); it("renders the footer slot inside a tfoot element", () => { @@ -596,63 +600,57 @@ describe("DataTable row styling and footer", () => { />, ); - expect(screen.getByTestId("footer-row").closest("tfoot")).not.toBeNull(); + const rowGroups = screen.getAllByRole("rowgroup"); + expect(within(rowGroups.at(-1) as HTMLElement).getByText("Total: 3")).toBeInTheDocument(); }); }); describe("DataTable layout", () => { it("exposes resize handles with stable selectors only when resizing is enabled", () => { - const { container, rerender } = render( - , - ); - expect(container.querySelectorAll("[data-resizer][data-header-id]").length).toBe(2); + const { rerender } = render(); + expect(screen.getByTestId("column-resizer-name")).toBeInTheDocument(); + expect(screen.getByTestId("column-resizer-email")).toBeInTheDocument(); rerender(); - expect(container.querySelectorAll("[data-resizer]").length).toBe(0); + expect(screen.queryByTestId("column-resizer-name")).not.toBeInTheDocument(); }); it("makes the header sticky and constrains body height when maxBodyHeight is set", () => { - const { container } = render(); - expect(container.querySelector("thead")?.className).toContain("sticky"); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; - expect(scroller).toHaveStyle({ maxHeight: "240px" }); + render(); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky"); + expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" }); }); it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => { - const { container } = render(); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; - const frame = scroller.parentElement as HTMLElement; - const outer = frame.parentElement as HTMLElement; + render(); + const outer = screen.getByTestId("data-table-root"); + const frame = screen.getByTestId("data-table-frame"); + const scroller = screen.getByTestId("data-table-scroller"); // A ceiling, not a stretch: flex-1 here would hold the footer at the bottom on a two-row table. - expect(outer.className).toContain("max-h-full"); - expect(outer.className).not.toContain("flex-1"); - expect(frame.className).not.toContain("flex-1"); - expect(scroller.className).not.toContain("flex-1"); + expect(outer).toHaveClass("max-h-full", "flex-col"); + expect(outer).not.toHaveClass("flex-1"); + expect(frame).toHaveClass("flex-col"); + expect(frame).not.toHaveClass("flex-1"); + expect(scroller).not.toHaveClass("flex-1"); - expect(outer.className).toContain("flex-col"); - expect(frame.className).toContain("flex-col"); - expect(scroller.className).toContain("min-h-0"); - expect(scroller.className).toContain("overflow-auto"); + expect(scroller).toHaveClass("min-h-0", "overflow-auto"); expect(scroller).toHaveStyle({ maxHeight: "" }); // Without this the Table primitive's own overflow container captures the sticky header. - expect(scroller.className).toContain("[&_[data-slot=table-container]]:overflow-visible"); + expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible"); - const thead = container.querySelector("thead") as HTMLElement; - expect(thead.className).toContain("sticky"); // Rows pass under the header, so the semi-transparent row tint alone would let them show through. - expect(thead.className).toContain("bg-background"); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background"); }); it("leaves the default layout untouched when neither height mode is set", () => { - const { container } = render(); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; + render(); + const scroller = screen.getByTestId("data-table-scroller"); - expect(scroller.className).toContain("overflow-x-auto"); - expect(scroller.className).not.toContain("min-h-0"); + expect(scroller).toHaveClass("overflow-x-auto"); + expect(scroller).not.toHaveClass("min-h-0"); expect(scroller).toHaveStyle({ maxHeight: "" }); - expect((scroller.parentElement as HTMLElement).className).not.toContain("flex-col"); - expect(container.querySelector("thead")?.className).not.toContain("sticky"); - expect(container.querySelector("thead")?.className).not.toContain("bg-background"); + expect(screen.getByTestId("data-table-frame")).not.toHaveClass("flex-col"); + expect(screen.getByTestId("data-table-head")).not.toHaveClass("sticky", "bg-background"); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index e8357b37715..60267606951 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -195,8 +195,7 @@ function DataTableHeadCell({ header, size, stickyHeader, enableColumnResi )} {canResize && (
column.resetSize()} @@ -589,15 +588,19 @@ export function DataTable(props: DataTableProps -
+
+
{toolbar !== undefined &&
{toolbar(table)}
}
- + {table.getHeaderGroups().map((headerGroup) => ( {headerGroup.headers.map((header) => ( From c5ba2b5fcf6f698b0e0f06a190ebd9d3b553c51b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 00:26:49 -0700 Subject: [PATCH 092/126] test(ui): query the screen instead of the render result Two changes, both about finding elements the way a user finds them. Twenty-six test files destructured queries off render and called them bare. Those queries are scoped to the render container, so they quietly miss anything portalled into the body, and they read as if they were free functions. They now go through screen. ChatMessageBubble and the key info panel derived elements by walking closest/parentElement/firstElementChild and then asserted on the classes they found. A wrapper element anywhere in between broke them. The bubble surface, the avatar and the budget reset value now publish a test id, so the assertions survive markup changes and still fail when the styling they check actually regresses. Budgets drop with the counts: prefer-screen-queries 221 to 21, no-node-access 723 to 716. The 21 remaining prefer-screen-queries are not all fixable: 18 of them are within(dialog) results in MCPToolsetsTab, which the rule cannot tell apart from a render result. Target is 18, not 0. --- ui/litellm-dashboard/eslint-budgets.json | 4 +- .../_components/APIReferenceView.test.tsx | 12 +- .../cache_settings/RedisTypeSelector.test.tsx | 8 +- .../_components/CacheLeakageCard.test.tsx | 50 ++++---- .../CostOptimizationView.activity.test.tsx | 16 +-- .../_components/CostOptimizationView.test.tsx | 50 ++++---- .../_components/PromptCachingTab.test.tsx | 8 +- .../_components/UsageTab.test.tsx | 118 +++++++++--------- .../_components/guardrail_info.test.tsx | 78 +++++------- .../_components/pii_components.test.tsx | 18 ++- .../_components/pii_configuration.test.tsx | 6 +- .../_components/mcp_servers.test.tsx | 28 ++--- .../PriceDataManagementTab.test.tsx | 6 +- .../models-and-endpoints/page.test.tsx | 50 ++++---- .../chat_ui/ChatMessageBubble.test.tsx | 8 +- .../components/chat_ui/ChatMessageBubble.tsx | 2 + .../components/compareUI/CompareUI.test.tsx | 30 ++--- .../components/MessageDisplay.test.tsx | 30 ++--- .../EntityUsageExportModal.test.tsx | 13 +- .../add_model/advanced_settings.test.tsx | 34 ++--- .../add_model/litellm_model_name.test.tsx | 12 +- .../bulk_create_users_button.test.tsx | 4 +- ...cost_optimization_feedback_banner.test.tsx | 18 +-- .../organization/organization_view.test.tsx | 4 +- .../src/components/settings.test.tsx | 20 +-- .../shared/PaginationStatusAlerts.test.tsx | 16 +-- .../shared/charts/area_chart.test.tsx | 6 +- .../shared/charts/bar_chart.test.tsx | 4 +- .../KeyInfoView.handleKeyUpdate.test.tsx | 4 +- .../key_info_view.budget_display.test.tsx | 2 +- .../components/templates/key_info_view.tsx | 2 +- 31 files changed, 323 insertions(+), 338 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index 3bfe56724eb..e7545eb2383 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -6,6 +6,6 @@ "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, "local/no-long-condition-chain": { "max": 265, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, - "testing-library/no-node-access": { "max": 723, "target": 500 }, - "testing-library/prefer-screen-queries": { "max": 221, "target": 0 } + "testing-library/no-node-access": { "max": 716, "target": 500 }, + "testing-library/prefer-screen-queries": { "max": 21, "target": 18 } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx index b2e0a42eecb..dae19032e20 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx @@ -13,17 +13,17 @@ describe("APIReferenceView", () => { it("uses the API doc base url when provided", () => { const apiDocUrl = "https://docs.litellm.test"; - const { getAllByTestId } = render(); + render(); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); expect(codeBlocks[0]).toHaveTextContent(new RegExp(apiDocUrl)); }); it("falls back to the proxy base url when the docs url is missing", () => { const proxyUrl = "https://proxy.litellm.test"; - const { getAllByTestId } = render(); + render(); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); expect(codeBlocks[0]).toHaveTextContent(new RegExp(proxyUrl)); }); @@ -31,7 +31,7 @@ describe("APIReferenceView", () => { const apiDocUrl = "https://docs-preferred.litellm.test"; const proxyUrl = "https://proxy-backup.litellm.test"; - const { getAllByTestId } = render( + render( { />, ); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); const renderedCode = codeBlocks[0].textContent ?? ""; expect(renderedCode).toContain(apiDocUrl); expect(renderedCode).not.toContain(proxyUrl); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx index 9d4d5a6d425..372f27e2be1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx @@ -1,12 +1,10 @@ import { describe, expect, it } from "vitest"; import RedisTypeSelector from "./RedisTypeSelector"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; describe("RedisTypeSelector", () => { it("should render the component", () => { - const { getAllByText } = render( - {}} />, - ); - expect(getAllByText(/Redis/i).length).toBeGreaterThan(0); + render( {}} />); + expect(screen.getAllByText(/Redis/i).length).toBeGreaterThan(0); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx index 8c36b934789..f320d8e0f97 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import type { DailyData, KeyMetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; @@ -79,25 +79,25 @@ const renderWith = (results: DailyData[], overrides: Partial describe("CacheLeakageCard", () => { it("ranks leaking keys by uncached prompt tokens and shows cache hit ratio", () => { - const { getByText, getByLabelText } = renderWith([ + renderWith([ dayWithKeys("2026-07-12", { "hash-caching": key("caching-key", { prompt_tokens: 1000, cache_read_input_tokens: 900 }), "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }), ]); - expect(getByText("leaky-key")).toBeInTheDocument(); - expect(getByText("0.0%")).toBeInTheDocument(); - expect(getByText("90.0%")).toBeInTheDocument(); + expect(screen.getByText("leaky-key")).toBeInTheDocument(); + expect(screen.getByText("0.0%")).toBeInTheDocument(); + expect(screen.getByText("90.0%")).toBeInTheDocument(); [ "Input tokens you sent in this range that weren't served from or written to the cache", "Share of your input tokens that were served from the cache", "About how much you'd save if this uncached input used prompt caching. Estimated as uncached input tokens times what your cached traffic already nets per cached token (realized cache savings, after write premiums, ÷ cache read and write tokens). Blank when caching is not currently saving anything overall.", - ].forEach((info) => expect(getByLabelText(info)).toBeInTheDocument()); + ].forEach((info) => expect(screen.getByLabelText(info)).toBeInTheDocument()); }); it("sorts by the clicked column, worst cache hit rate first", () => { - const { getAllByRole, getByText } = renderWith([ + renderWith([ dayWithKeys("2026-07-12", { "hash-a": key("alpha", { prompt_tokens: 10000, @@ -111,48 +111,48 @@ describe("CacheLeakageCard", () => { }), }), ]); - const firstDataRow = () => getAllByRole("row")[1]; + const firstDataRow = () => screen.getAllByRole("row")[1]; expect(firstDataRow()).toHaveTextContent("alpha"); - fireEvent.click(getByText("Cache hit rate")); + fireEvent.click(screen.getByText("Cache hit rate")); expect(firstDataRow()).toHaveTextContent("bravo"); - fireEvent.click(getByText("Cache hit rate")); + fireEvent.click(screen.getByText("Cache hit rate")); expect(firstDataRow()).toHaveTextContent("alpha"); }); it("switches to the model view and lists only Anthropic models", () => { - const { getByText, queryByText } = renderWith([ + renderWith([ dayWithModels("2026-07-12", { "claude-sonnet-5": { prompt_tokens: 5000, cache_read_input_tokens: 0 }, "gpt-4o": { prompt_tokens: 8000, cache_read_input_tokens: 0 }, }), ]); - fireEvent.click(getByText("By model")); + fireEvent.click(screen.getByText("By model")); - expect(getByText("Cache leakage by model")).toBeInTheDocument(); - expect(getByText("claude-sonnet-5")).toBeInTheDocument(); - expect(queryByText("gpt-4o")).not.toBeInTheDocument(); + expect(screen.getByText("Cache leakage by model")).toBeInTheDocument(); + expect(screen.getByText("claude-sonnet-5")).toBeInTheDocument(); + expect(screen.queryByText("gpt-4o")).not.toBeInTheDocument(); }); it("shows an empty state when no key used tokens in the range", () => { - const { getByText, queryByRole } = renderWith([dayWithKeys("2026-07-12", {})]); + renderWith([dayWithKeys("2026-07-12", {})]); - expect(getByText("No key usage in this range.")).toBeInTheDocument(); - expect(queryByRole("table")).not.toBeInTheDocument(); + expect(screen.getByText("No key usage in this range.")).toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); }); it("tells the user the table is still filling in while fallback pages stream", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { getByText, getByRole } = renderWith([day], { isFetchingMore: true }); + renderWith([day], { isFetchingMore: true }); - expect(getByRole("table")).toBeInTheDocument(); + expect(screen.getByRole("table")).toBeInTheDocument(); expect( - getByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.getByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).toBeInTheDocument(); }); @@ -160,10 +160,10 @@ describe("CacheLeakageCard", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { queryByText } = renderWith([day], { loading: true }); + renderWith([day], { loading: true }); expect( - queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); @@ -171,10 +171,10 @@ describe("CacheLeakageCard", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { queryByText } = renderWith([day]); + renderWith([day]); expect( - queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 1cc7bec13d1..03250e3e53b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render, waitFor } from "@testing-library/react"; +import { fireEvent, render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -56,7 +56,7 @@ describe("CostOptimizationView daily activity", () => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const { getByRole, getByTestId, findByTestId, queryByText } = render( + render( , @@ -64,12 +64,12 @@ describe("CostOptimizationView daily activity", () => { await waitFor(() => expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1)); - fireEvent.click(getByRole("tab", { name: "Prompt Caching" })); - await findByTestId("caching-settings"); + fireEvent.click(screen.getByRole("tab", { name: "Prompt Caching" })); + await screen.findByTestId("caching-settings"); expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1); expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); - expect(queryByText(/Currently fetching spend data/)).not.toBeInTheDocument(); + expect(screen.queryByText(/Currently fetching spend data/)).not.toBeInTheDocument(); }); it("shows the fetch-progress banner while the paginated fallback streams pages in", async () => { @@ -84,13 +84,13 @@ describe("CostOptimizationView daily activity", () => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const { findByText, getByRole } = render( + render( , ); - expect(await findByText(/Currently fetching spend data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); - expect(getByRole("button", { name: "Stop" })).toBeInTheDocument(); + expect(await screen.findByText(/Currently fetching spend data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Stop" })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index d5df5aa75da..028367555a1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -45,32 +45,32 @@ describe("CostOptimizationView", () => { }); it("renders the standard page header with the sidebar's Cost Optimization icon", () => { - const { container, getByRole, getByText } = renderView(); + const { container } = renderView(); - expect(getByRole("heading", { level: 1, name: "Cost Optimization" })).toBeInTheDocument(); - expect(getByText(/Track and configure the mechanisms that save you money/)).toBeInTheDocument(); + expect(screen.getByRole("heading", { level: 1, name: "Cost Optimization" })).toBeInTheDocument(); + expect(screen.getByText(/Track and configure the mechanisms that save you money/)).toBeInTheDocument(); expect(container.querySelector(".lucide-piggy-bank")).not.toBeNull(); }); it("renders the four cost-optimization tabs", () => { - const { getByText } = renderView(); + renderView(); - expect(getByText("Overall")).toBeInTheDocument(); - expect(getByText("Prompt Compression")).toBeInTheDocument(); - expect(getByText("Prompt Caching")).toBeInTheDocument(); - expect(getByText("Auto-Router")).toBeInTheDocument(); + expect(screen.getByText("Overall")).toBeInTheDocument(); + expect(screen.getByText("Prompt Compression")).toBeInTheDocument(); + expect(screen.getByText("Prompt Caching")).toBeInTheDocument(); + expect(screen.getByText("Auto-Router")).toBeInTheDocument(); }); it("defaults to the Overall tab and switches the active tab on click", () => { - const { getByRole } = renderView(); + renderView(); - expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); - expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); - fireEvent.click(getByRole("tab", { name: "Prompt Compression" })); + fireEvent.click(screen.getByRole("tab", { name: "Prompt Compression" })); - expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); - expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); // Unlike the other three pages in this cleanup, Cost Optimization keeps its @@ -80,21 +80,21 @@ describe("CostOptimizationView", () => { // are proxy-admin-only, so those are what disappear. describe("proxy-admin-only tabs", () => { it.each(["Internal User", "Internal Viewer", "Org Admin"])("shows %s the Overall tab only", (userRole) => { - const { getByRole, queryByRole } = renderView(userRole); + renderView(userRole); - expect(getByRole("tab", { name: "Overall" })).toBeInTheDocument(); - expect(queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Overall" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); }); it("never mounts the panels behind the admin-only endpoints for an internal user", () => { - const { getByTestId, queryByTestId } = renderView("Internal User"); + renderView("Internal User"); - expect(getByTestId("usage-tab")).toBeInTheDocument(); - expect(queryByTestId("compression-tab")).not.toBeInTheDocument(); - expect(queryByTestId("caching-tab")).not.toBeInTheDocument(); - expect(queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); + expect(screen.getByTestId("usage-tab")).toBeInTheDocument(); + expect(screen.queryByTestId("compression-tab")).not.toBeInTheDocument(); + expect(screen.queryByTestId("caching-tab")).not.toBeInTheDocument(); + expect(screen.queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx index 38517dab0ab..2c602033171 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx @@ -1,4 +1,4 @@ -import { render, waitFor } from "@testing-library/react"; +import { render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; const mockGetGeneralSettingsCall = vi.fn(); @@ -37,10 +37,10 @@ describe("PromptCachingTab", () => { cancelled: false, cancel: vi.fn(), }; - const { getByTestId } = render(); + render(); - expect(getByTestId("caching-settings")).toBeInTheDocument(); - expect(getByTestId("cache-leakage-card")).toBeInTheDocument(); + expect(screen.getByTestId("caching-settings")).toBeInTheDocument(); + expect(screen.getByTestId("cache-leakage-card")).toBeInTheDocument(); await waitFor(() => expect(mockCacheLeakageCard).toHaveBeenCalledWith(expect.objectContaining({ activity }))); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index df23e5509bf..f85a667a074 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ToolSpendResponse } from "@/components/networking"; @@ -152,13 +152,13 @@ describe("UsageTab", () => { gateway_injected_caching_savings_spend: 0.006, compression_saved_tokens: 100000, }; - const { getByText } = renderWith([day("2026-07-12", firstDay), day("2026-07-13", secondDay)]); + renderWith([day("2026-07-12", firstDay), day("2026-07-13", secondDay)]); - expect(getByText("$0.1500")).toBeInTheDocument(); - expect(getByText("$0.1400")).toBeInTheDocument(); - expect(getByText("$0.0100")).toBeInTheDocument(); - expect(getByText("$0.0160")).toBeInTheDocument(); - expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + expect(screen.getByText("$0.1500")).toBeInTheDocument(); + expect(screen.getByText("$0.1400")).toBeInTheDocument(); + expect(screen.getByText("$0.0100")).toBeInTheDocument(); + expect(screen.getByText("$0.0160")).toBeInTheDocument(); + expect(screen.getByText("140,000 tokens compressed")).toBeInTheDocument(); }); const twoDays = () => [ @@ -167,11 +167,11 @@ describe("UsageTab", () => { ]; it("opens on a running total anchored at $0 at the start of the range", () => { - const { getByTestId } = renderWith(twoDays()); + renderWith(twoDays()); // Cumulative prepends a synthetic $0 point at the range start (Jul 1) so the // line rises from zero rather than floating; the daily running totals follow. - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series).toHaveLength(3); expect(series[0]).toMatchObject({ date: "Jul 1", Compression: 0, "Prompt caching": 0 }); expect(series[1]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); @@ -183,12 +183,12 @@ describe("UsageTab", () => { // The original complaint: a one-day range plotted a single floating dot. The // synthetic start anchor gives the line a zero origin to climb from. const oneDay = new Date(2026, 6, 24); - const { getByTestId } = renderWith( - [day("2026-07-24", { compression_savings_spend: 0.2, gateway_injected_caching_savings_spend: 0.05 })], - { from: oneDay, to: oneDay }, - ); + renderWith([day("2026-07-24", { compression_savings_spend: 0.2, gateway_injected_caching_savings_spend: 0.05 })], { + from: oneDay, + to: oneDay, + }); - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series).toHaveLength(2); expect(series[0]).toMatchObject({ date: "Jul 24", Compression: 0, "Prompt caching": 0 }); expect(series[1]).toMatchObject({ date: "Jul 24", Compression: 0.2, "Prompt caching": 0.05 }); @@ -202,49 +202,49 @@ describe("UsageTab", () => { day("2026-07-13", { gateway_injected_caching_savings_spend: 0.1 }), day("2026-07-12", { gateway_injected_caching_savings_spend: 0.04 }), ]; - const { getByTestId, getByRole } = renderWith(newestFirst); + renderWith(newestFirst); // The $0 anchor leads, then the days climb oldest to newest. - const cumulative = readSeries(getByTestId("area-chart")); + const cumulative = readSeries(screen.getByTestId("area-chart")); expect(cumulative.map((p: { date: string }) => p.date)).toEqual(["Jul 1", "Jul 12", "Jul 13"]); expect(cumulative[1]["Prompt caching"]).toBeCloseTo(0.04, 5); expect(cumulative[2]["Prompt caching"]).toBeCloseTo(0.14, 5); expect(cumulative[2]["Prompt caching"]).toBeGreaterThan(cumulative[1]["Prompt caching"]); - await userEvent.click(getByRole("tab", { name: "Per day" })); - const perDay = readSeries(getByTestId("bar-chart")); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + const perDay = readSeries(screen.getByTestId("bar-chart")); expect(perDay.map((p: { date: string }) => p.date)).toEqual(["Jul 12", "Jul 13"]); }); it("draws bars of the raw per-interval readings on the other tab", async () => { - const { getByRole, getByTestId, queryByTestId } = renderWith(twoDays()); + renderWith(twoDays()); // Cumulative opens on the area line. - expect(getByTestId("area-chart")).toBeInTheDocument(); + expect(screen.getByTestId("area-chart")).toBeInTheDocument(); - await userEvent.click(getByRole("tab", { name: "Per day" })); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); // Per day switches to a bar chart of the unaccumulated daily savings, with no // synthetic anchor prepended. - expect(queryByTestId("area-chart")).not.toBeInTheDocument(); - const series = readSeries(getByTestId("bar-chart")); + expect(screen.queryByTestId("area-chart")).not.toBeInTheDocument(); + const series = readSeries(screen.getByTestId("bar-chart")); expect(series).toHaveLength(2); expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); }); it("says what the line means and over what range", async () => { - const { getByText, getByRole } = renderWith(twoDays()); + renderWith(twoDays()); - expect(getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); - await userEvent.click(getByRole("tab", { name: "Per day" })); - expect(getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); + expect(screen.getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + expect(screen.getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); }); it("builds the per-driver donut from the range totals, not the running total", () => { - const { getByTestId } = renderWith(twoDays()); + renderWith(twoDays()); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([ { driver: "Compression", color: "emerald", usd: expect.closeTo(0.14, 5) }, { driver: "Prompt caching", color: "blue", usd: expect.closeTo(0.016, 5) }, @@ -252,9 +252,9 @@ describe("UsageTab", () => { }); it("omits a driver slice when that driver has no savings", () => { - const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([{ driver: "Compression", color: "emerald", usd: expect.closeTo(0.04, 5) }]); }); @@ -262,7 +262,7 @@ describe("UsageTab", () => { // Stacking sums the series into one bar. Auto-router savings go negative when a // model switch pays for a cold cache, and that segment would be drawn below the // axis while the rest of the bar still read as the day's total. - const { getByRole, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.1, gateway_injected_caching_savings_spend: 0.02, @@ -270,8 +270,8 @@ describe("UsageTab", () => { }), ]); - await userEvent.click(getByRole("tab", { name: "Per day" })); - const bars = getByTestId("bar-chart"); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + const bars = screen.getByTestId("bar-chart"); expect(bars).toHaveAttribute("data-stack", "false"); expect(readSeries(bars)[0]).toMatchObject({ "Auto-router": -0.05 }); }); @@ -281,10 +281,10 @@ describe("UsageTab", () => { // per day"). Hand-rolled rows made it compete with the legend and the toggle for // width, so the header grew a line on one tab and the chart moved with it. CardHeader // sizes the action column to its content and gives the rest to the title column. - const { getByRole, getByTestId, container } = renderWith(twoDays()); + const { container } = renderWith(twoDays()); const header = () => { - const legend = getByTestId("chart-legend"); + const legend = screen.getByTestId("chart-legend"); const action = legend.closest('[data-slot="card-action"]') as HTMLElement; const cardHeader = action.parentElement as HTMLElement; const description = cardHeader.querySelector('[data-slot="card-description"]') as HTMLElement; @@ -295,12 +295,12 @@ describe("UsageTab", () => { expect(before.action).toBeTruthy(); expect(before.description).toBeTruthy(); // the toggle rides in the same action slot as the legend, so neither moves alone - expect(before.action.contains(getByRole("tablist"))).toBe(true); + expect(before.action.contains(screen.getByRole("tablist"))).toBe(true); // the subtitle lives outside that slot, so its length cannot reposition the controls expect(before.action.contains(before.description)).toBe(false); expect(before.description).toHaveTextContent(/Running total saved/); - await userEvent.click(getByRole("tab", { name: "Per day" })); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); const after = header(); expect(after.action).toBe(before.action); @@ -314,7 +314,7 @@ describe("UsageTab", () => { // Switching models leaves the new one with a cold cache, so a route can cost more // than the baseline would have. A negative slice is meaningless in a donut, but the // total has to keep the loss or the page can only ever report good news. - const { getByText, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.1, gateway_injected_caching_savings_spend: 0.02, @@ -322,16 +322,16 @@ describe("UsageTab", () => { }), ]); - expect(getByText("$0.0700")).toBeInTheDocument(); - expect(getByText("-$0.0500")).toBeInTheDocument(); + expect(screen.getByText("$0.0700")).toBeInTheDocument(); + expect(screen.getByText("-$0.0500")).toBeInTheDocument(); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices.map((d: { driver: string }) => d.driver)).toEqual(["Compression", "Prompt caching"]); - expect(getByTestId("donut-chart")).toHaveAttribute("data-label", "$0.1200"); + expect(screen.getByTestId("donut-chart")).toHaveAttribute("data-label", "$0.1200"); }); it("carries auto-router savings into the summary card, donut slice, and cumulative series", () => { - const { getByText, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.04, gateway_injected_caching_savings_spend: 0.006, @@ -345,11 +345,11 @@ describe("UsageTab", () => { ]); // Total saved now sums three drivers, and the auto-router card carries its own total. - expect(getByText("$0.2260")).toBeInTheDocument(); - expect(getByText("$0.0700")).toBeInTheDocument(); + expect(screen.getByText("$0.2260")).toBeInTheDocument(); + expect(screen.getByText("$0.0700")).toBeInTheDocument(); // The driver donut gains a third slice priced from the range totals. - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([ { driver: "Compression", color: "emerald", usd: expect.closeTo(0.14, 5) }, { driver: "Prompt caching", color: "blue", usd: expect.closeTo(0.016, 5) }, @@ -357,7 +357,7 @@ describe("UsageTab", () => { ]); // And the cumulative line accumulates the auto-router series alongside the others. - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series[2]["Auto-router"]).toBeCloseTo(0.07, 5); }); @@ -371,9 +371,9 @@ describe("UsageTab", () => { start_date: "2026-07-12", end_date: "2026-07-12", }; - const { findAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + renderWith([day("2026-07-12", {})], { toolSpend }); - const bars = await findAllByTestId("bar-chart"); + const bars = await screen.findAllByTestId("bar-chart"); const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); // The 64px bar cap is this card's opt-in; the shared BarChart must not cap @@ -391,14 +391,16 @@ describe("UsageTab", () => { start_date: "2026-07-12", end_date: "2026-07-12", }; - const { findAllByTestId, getAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + renderWith([day("2026-07-12", {})], { toolSpend }); - const bars = await findAllByTestId("bar-chart"); + const bars = await screen.findAllByTestId("bar-chart"); const [totalByTool, dailyByTool] = bars.slice(-2); expect(dailyByTool).toHaveAttribute("data-show-legend", "false"); expect(totalByTool).toHaveAttribute("data-colors", dailyByTool.getAttribute("data-colors")); - const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); + const toolLegends = screen + .getAllByTestId("chart-legend") + .filter((legend) => legend.textContent === "search,read_file"); expect(toolLegends).toHaveLength(1); }); @@ -415,23 +417,23 @@ describe("UsageTab", () => { it.each(["Internal User", "Internal Viewer", "Org Admin"])( "hides the card and never calls the endpoint for %s", async (userRole) => { - const { queryByText, getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend, userRole, }); // Liveness gate: the daily-activity charts still render for this role, // so the absence below is the gate, not an empty tab. - expect(getByTestId("donut-chart")).toBeInTheDocument(); - expect(queryByText("Spend by tool")).not.toBeInTheDocument(); + expect(screen.getByTestId("donut-chart")).toBeInTheDocument(); + expect(screen.queryByText("Spend by tool")).not.toBeInTheDocument(); await vi.waitFor(() => expect(mockGetToolSpend).not.toHaveBeenCalled()); }, ); it("keeps the card and the endpoint call for an admin", async () => { - const { findByText } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); - expect(await findByText("Spend by tool")).toBeInTheDocument(); + expect(await screen.findByText("Spend by tool")).toBeInTheDocument(); expect(mockGetToolSpend).toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 2b90a1d8cbc..fcffc2122e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -1,5 +1,5 @@ import * as networking from "@/components/networking"; -import { fireEvent, render, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, waitFor, within, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import GuardrailInfoView from "./guardrail_info"; @@ -65,21 +65,19 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getAllByText, getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); // Wait for the loading to complete and data to be rendered await waitFor(() => { // The guardrail name appears in multiple places (title and settings tab) - const elements = getAllByText("Test Guardrail"); + const elements = screen.getAllByText("Test Guardrail"); expect(elements.length).toBeGreaterThan(0); }); // Verify other key elements are present - expect(getByText("Back to Guardrails")).toBeInTheDocument(); - expect(getByText("Overview")).toBeInTheDocument(); - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Back to Guardrails")).toBeInTheDocument(); + expect(screen.getByText("Overview")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); it("should render a tag-based mode object rather than crashing the detail view", async () => { @@ -105,11 +103,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findAllByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + expect(await screen.findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); }); it("should render the provider logo from the bundled guardrail logo map", async () => { @@ -135,11 +131,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findByAltText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - const logo = await findByAltText("Presidio PII logo"); + const logo = await screen.findByAltText("Presidio PII logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("microsoft_azure.svg")); }); @@ -167,25 +161,27 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText, findByText, container } = render( + const { container } = render( {}} accessToken="123" isAdmin={true} />, ); await waitFor(() => { - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); // Click the Settings tab - fireEvent.click(getByText("Settings")); + fireEvent.click(screen.getByText("Settings")); // Wait for the Settings panel to render await waitFor(() => { - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); await userEvent.hover(within(container).getByRole("img", { name: "Config guardrail details" })); - expect(await findByText("Guardrail is defined in the config file and cannot be edited.")).toBeInTheDocument(); + expect( + await screen.findByText("Guardrail is defined in the config file and cannot be edited."), + ).toBeInTheDocument(); }); it("should render the guardrail info", async () => { @@ -216,12 +212,10 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); await waitFor(() => { - expect(getByText("PII Entity Configuration")).toBeInTheDocument(); + expect(screen.getByText("PII Entity Configuration")).toBeInTheDocument(); }); }); it("should handle content filter updates correctly", async () => { @@ -251,30 +245,28 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); vi.mocked(networking.updateGuardrailCall).mockResolvedValue({ status: "success" }); - const { getByText, getByLabelText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); await waitFor(() => { - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); // Go to Settings tab - fireEvent.click(getByText("Settings")); + fireEvent.click(screen.getByText("Settings")); await waitFor(() => { - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); // Enter Edit Mode - fireEvent.click(getByText("Edit Settings")); + fireEvent.click(screen.getByText("Edit Settings")); // Modify Guardrail Name to force an update - const nameInput = getByLabelText("Guardrail Name"); + const nameInput = screen.getByLabelText("Guardrail Name"); fireEvent.change(nameInput, { target: { value: "Updated Name" } }); // Save with only name change - const saveButton = getByText("Save Changes"); + const saveButton = screen.getByText("Save Changes"); fireEvent.click(saveButton); await waitFor(() => { @@ -300,16 +292,16 @@ describe("Guardrail Info", () => { // Enter Edit Mode again to make changes await waitFor(() => { - expect(getByText("Edit Settings")).toBeInTheDocument(); + expect(screen.getByText("Edit Settings")).toBeInTheDocument(); }); - fireEvent.click(getByText("Edit Settings")); + fireEvent.click(screen.getByText("Edit Settings")); // Now modify the values using the mock button - const simulateChangeButton = getByText("Simulate Change"); + const simulateChangeButton = screen.getByText("Simulate Change"); fireEvent.click(simulateChangeButton); // Save again - fireEvent.click(getByText("Save Changes")); + fireEvent.click(screen.getByText("Save Changes")); await waitFor(() => { expect(networking.updateGuardrailCall).toHaveBeenCalled(); @@ -339,12 +331,10 @@ describe("Guardrail Info", () => { }); vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findByRole, getByRole, getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - expect(await findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); - expect(getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx index 2f839487111..30ff2cf7c5f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import { CategoryFilter, QuickActions, PiiEntityList } from "./pii_components"; import type { PiiEntityCategory } from "@/components/guardrails/types"; @@ -6,25 +6,21 @@ import type { PiiEntityCategory } from "@/components/guardrails/types"; describe("CategoryFilter", () => { it("should render", () => { const emptyCategories: PiiEntityCategory[] = []; - const { getByText } = render( - {}} />, - ); - expect(getByText("Filter by category")).toBeInTheDocument(); + render( {}} />); + expect(screen.getByText("Filter by category")).toBeInTheDocument(); }); }); describe("QuickActions", () => { it("should render", () => { - const { getByText } = render( - {}} onUnselectAll={() => {}} hasSelectedEntities={false} />, - ); - expect(getByText("Quick Actions")).toBeInTheDocument(); + render( {}} onUnselectAll={() => {}} hasSelectedEntities={false} />); + expect(screen.getByText("Quick Actions")).toBeInTheDocument(); }); }); describe("PiiEntityList", () => { it("should render", () => { - const { getByText } = render( + render( { entityToCategoryMap={new Map()} />, ); - expect(getByText("No PII types match your filter criteria")).toBeInTheDocument(); + expect(screen.getByText("No PII types match your filter criteria")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx index 00c568ef35b..4f822578fe8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx @@ -1,10 +1,10 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import PiiConfiguration from "./pii_configuration"; describe("PiiConfiguration", () => { it("should render", () => { - const { getByText } = render( + render( { entityCategories={[]} />, ); - expect(getByText("Configure PII Protection")).toBeInTheDocument(); + expect(screen.getByText("Configure PII Protection")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 8dca87e24ef..8abb8855e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -45,7 +45,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -53,11 +53,11 @@ describe("MCPServers", () => { // Wait for the component to load and check if title renders await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify the title is rendered - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); it("should render mocked MCP servers data in the table", async () => { @@ -96,7 +96,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); const queryClient = createQueryClient(); - const { getByText, getAllByText } = render( + render( , @@ -104,19 +104,19 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Wait for the mocked data to render in the table await waitFor(() => { - expect(getByText("Test Server 1")).toBeInTheDocument(); + expect(screen.getByText("Test Server 1")).toBeInTheDocument(); }); // Verify the mocked server data is rendered in the table - expect(getByText("Test Server 1")).toBeInTheDocument(); - expect(getByText("Test Server 2")).toBeInTheDocument(); - expect(getAllByText("test-server-1").length).toBeGreaterThan(0); - expect(getAllByText("test-server-2").length).toBeGreaterThan(0); + expect(screen.getByText("Test Server 1")).toBeInTheDocument(); + expect(screen.getByText("Test Server 2")).toBeInTheDocument(); + expect(screen.getAllByText("test-server-1").length).toBeGreaterThan(0); + expect(screen.getAllByText("test-server-2").length).toBeGreaterThan(0); // Verify the API was called // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock @@ -168,7 +168,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -176,7 +176,7 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify the health check API was called (without a server ID filter — the hook always @@ -211,7 +211,7 @@ describe("MCPServers", () => { ); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -219,7 +219,7 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify that health check was initiated diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx index 8b34d61ebad..7cd418bc176 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx @@ -1,5 +1,5 @@ /* @vitest-environment jsdom */ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PriceDataManagementTab from "./PriceDataManagementTab"; @@ -11,7 +11,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ describe("PriceDataManagementTab", () => { it("renders its content standalone, without a tab-panel ancestor", () => { - const { getByText } = render(); - expect(getByText("Price Data Management")).toBeInTheDocument(); + render(); + expect(screen.getByText("Price Data Management")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 521f89a39f2..a504d75bb63 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -1,6 +1,6 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsPage from "./page"; @@ -64,48 +64,48 @@ describe("ModelsAndEndpointsPage", () => { }); it("renders the admin tab bar and the All Models panel by default", () => { - const { getByRole, getByTestId } = renderPage(); - expect(getByRole("tab", { name: "All Models" })).toBeInTheDocument(); - expect(getByRole("tab", { name: "LLM Credentials" })).toBeInTheDocument(); - expect(getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); - expect(getByTestId("panel-all-models")).toBeInTheDocument(); + renderPage(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "LLM Credentials" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); + expect(screen.getByTestId("panel-all-models")).toBeInTheDocument(); }); it("switches tabs in-memory, mounting only the active panel", async () => { const user = userEvent.setup(); - const { getByRole, getByTestId, queryByTestId } = renderPage(); - await user.click(getByRole("tab", { name: "Health Status" })); - expect(getByTestId("panel-health")).toBeInTheDocument(); - expect(queryByTestId("panel-all-models")).not.toBeInTheDocument(); + renderPage(); + await user.click(screen.getByRole("tab", { name: "Health Status" })); + expect(screen.getByTestId("panel-health")).toBeInTheDocument(); + expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); }); it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; - const { getByTestId, queryByRole } = renderPage(); - expect(getByTestId("model-info")).toHaveTextContent("model:abc-123"); - expect(queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); + renderPage(); + expect(screen.getByTestId("model-info")).toHaveTextContent("model:abc-123"); + expect(screen.queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); }); it("renders the team detail overlay from the ?team drill-in", () => { detailState.teamId = "team-9"; - const { getByTestId } = renderPage(); - expect(getByTestId("team-info")).toHaveTextContent("team:team-9"); + renderPage(); + expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9"); }); it("hides admin-only tabs for a non-admin user", () => { mockUseAuthorized.mockReturnValue(NON_ADMIN); - const { queryByRole } = renderPage(); - expect(queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); + renderPage(); + expect(screen.queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); }); // Auto-routers are excluded from the All Models table, so this tab is their home: the only // place in the product to list, create, edit or delete one. describe("Auto-Routers tab", () => { it("sits third, after All Models and Add Model", () => { - const { getAllByRole } = renderPage(); + renderPage(); - const tabs = getAllByRole("tab").map((tab) => tab.textContent); + const tabs = screen.getAllByRole("tab").map((tab) => tab.textContent); expect(tabs[0]).toContain("All Models"); expect(tabs[1]).toBe("Add Model"); expect(tabs[2]).toContain("Auto-Routers"); @@ -115,17 +115,17 @@ describe("ModelsAndEndpointsPage", () => { it("renders its panel when selected", async () => { const user = userEvent.setup(); - const { getByRole, getByTestId } = renderPage(); + renderPage(); - await user.click(getByRole("tab", { name: /Auto-Routers/ })); - expect(getByTestId("panel-auto-routers")).toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: /Auto-Routers/ })); + expect(screen.getByTestId("panel-auto-routers")).toBeInTheDocument(); }); it("is hidden from non-admins, who cannot write models", () => { mockUseAuthorized.mockReturnValue(NON_ADMIN); - const { queryByRole } = renderPage(); + renderPage(); - expect(queryByRole("tab", { name: /Auto-Routers/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: /Auto-Routers/ })).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx index a83c11d1444..ed02e7e16cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx @@ -93,13 +93,13 @@ describe("ChatMessageBubble", () => { ])("should paint the $role surface from theme tokens, not fixed colours", ({ role, bubble, avatar }) => { render(); - const header = screen.getByText(role).closest("div") as HTMLElement; - const surface = header.parentElement as HTMLElement; + const surface = screen.getByTestId("message-surface"); + const avatarEl = screen.getByTestId("message-avatar"); expect(surface).toHaveClass(...bubble); expect(surface).not.toHaveAttribute("style"); - expect(header.firstElementChild).toHaveClass(avatar); - expect(header.firstElementChild).not.toHaveAttribute("style"); + expect(avatarEl).toHaveClass(avatar); + expect(avatarEl).not.toHaveAttribute("style"); }); it("should show model badge for assistant messages when model is provided", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index 8c54d9e89fa..bd6a4bc49a4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -46,6 +46,7 @@ function ChatMessageBubble({ return (
{ describe("CompareUI", () => { it("should render", () => { - const { getByTestId } = render(); - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); - expect(getByTestId("message-input")).toBeInTheDocument(); + render(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("message-input")).toBeInTheDocument(); }); it("adds a comparison when Add Comparison button is clicked", async () => { const user = userEvent.setup(); - const { container, getByTestId } = render( - , - ); + const { container } = render(); // Verify initial state: 2 comparison panels - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); let comparisonPanels = container.querySelectorAll('[data-testid^="comparison-panel-"]'); expect(comparisonPanels).toHaveLength(2); @@ -117,15 +115,13 @@ describe("CompareUI", () => { }); // Verify the original 2 panels are still there - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); }); it("should handle image upload and send message with attachment", async () => { const user = userEvent.setup(); - const { getByTestId, queryByTestId } = render( - , - ); + render(); const file = new File(["test content"], "test-image.png", { type: "image/png" }); @@ -138,13 +134,13 @@ describe("CompareUI", () => { } await waitFor(() => { - expect(getByTestId("has-attachment")).toBeInTheDocument(); + expect(screen.getByTestId("has-attachment")).toBeInTheDocument(); }); - const textarea = getByTestId("message-textarea"); + const textarea = screen.getByTestId("message-textarea"); fireEvent.change(textarea, { target: { value: "Describe this image" } }); - const sendButton = getByTestId("send-button"); + const sendButton = screen.getByTestId("send-button"); expect(sendButton).toBeEnabled(); await user.click(sendButton); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx index 72a1d41f9fe..69b50d4088f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import type { MessageType } from "@/components/chat_ui/types"; import { MessageDisplay } from "./MessageDisplay"; @@ -39,9 +39,9 @@ describe("MessageDisplay", () => { model: "gpt-4", }, ]; - const { getByText } = render(); - expect(getByText("Hello")).toBeInTheDocument(); - expect(getByText("Hi there!")).toBeInTheDocument(); + render(); + expect(screen.getByText("Hello")).toBeInTheDocument(); + expect(screen.getByText("Hi there!")).toBeInTheDocument(); }); it("displays user and assistant messages with proper grouping and shows loading state", () => { @@ -64,13 +64,13 @@ describe("MessageDisplay", () => { }, }, ]; - const { getByText, getByTestId } = render(); - expect(getByText("You")).toBeInTheDocument(); - expect(getByText("What is 2+2?")).toBeInTheDocument(); - expect(getByText("gpt-4")).toBeInTheDocument(); - expect(getByText("calculator")).toBeInTheDocument(); - expect(getByText("2+2 equals 4")).toBeInTheDocument(); - expect(getByTestId("response-metrics")).toBeInTheDocument(); + render(); + expect(screen.getByText("You")).toBeInTheDocument(); + expect(screen.getByText("What is 2+2?")).toBeInTheDocument(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByText("calculator")).toBeInTheDocument(); + expect(screen.getByText("2+2 equals 4")).toBeInTheDocument(); + expect(screen.getByTestId("response-metrics")).toBeInTheDocument(); }); it("should display image attachment in user message", () => { @@ -86,10 +86,10 @@ describe("MessageDisplay", () => { model: "gpt-4", }, ]; - const { getByTestId, getByText } = render(); - expect(getByText("What is in this image? [Image attached]")).toBeInTheDocument(); - expect(getByTestId("chat-image-renderer")).toBeInTheDocument(); - const image = getByTestId("chat-image-renderer").querySelector("img"); + render(); + expect(screen.getByText("What is in this image? [Image attached]")).toBeInTheDocument(); + expect(screen.getByTestId("chat-image-renderer")).toBeInTheDocument(); + const image = screen.getByTestId("chat-image-renderer").querySelector("img"); expect(image).toHaveAttribute("src", "blob:test-image-url"); }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx index 8edf2174eee..cb04323e4c9 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx @@ -10,6 +10,7 @@ */ import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import EntityUsageExportModal from "./EntityUsageExportModal"; @@ -73,13 +74,13 @@ describe("EntityUsageExportModal", () => { const user = userEvent.setup(); const { handleExportCSV } = await import("./utils"); - const { getByRole } = renderWithProviders(); + renderWithProviders(); // Default primary action reflects CSV export - expect(getByRole("button", { name: /Export CSV/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Export CSV/i })).toBeInTheDocument(); // Click export - await user.click(getByRole("button", { name: /Export CSV/i })); + await user.click(screen.getByRole("button", { name: /Export CSV/i })); // Verifies export function was invoked with correct parameters expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily", "Tag", "tag", {}); @@ -97,14 +98,14 @@ describe("EntityUsageExportModal", () => { const user = userEvent.setup(); const { handleExportCSV } = await import("./utils"); - const { getByText, getByRole } = renderWithProviders(); + renderWithProviders(); // Choose the alternate export type - click the label to trigger radio - const dailyModelLabel = getByText(/Day-by-day by tag and model/i); + const dailyModelLabel = screen.getByText(/Day-by-day by tag and model/i); await user.click(dailyModelLabel); // Export with default CSV format - const exportBtn = getByRole("button", { name: /Export CSV/i }); + const exportBtn = screen.getByRole("button", { name: /Export CSV/i }); await user.click(exportBtn); // Ensure the selected scope flowed through diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index 01e00d903aa..ceb9d1ca713 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -1,4 +1,4 @@ -import { act, fireEvent, render, waitFor } from "@testing-library/react"; +import { act, fireEvent, render, waitFor, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { MountedFormHost } from "../../../tests/mounted-form-host"; import AdvancedSettings from "./advanced_settings"; @@ -35,51 +35,51 @@ describe("AdvancedSettings", () => { }); it("should render tags list", async () => { - const { getByText } = renderAdvancedSettings(); - fireEvent.click(getByText("Advanced Settings")); + renderAdvancedSettings(); + fireEvent.click(screen.getByText("Advanced Settings")); await waitFor(() => { - expect(getByText("Tags")).toBeInTheDocument(); + expect(screen.getByText("Tags")).toBeInTheDocument(); }); }); it("should render the litellm params", async () => { - const { getByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("LiteLLM Params")).toBeInTheDocument(); + expect(screen.getByText("LiteLLM Params")).toBeInTheDocument(); }); }); it("hides every PTU field when PTU cost attribution is disabled", async () => { - const { getByText, queryByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("Tags")).toBeInTheDocument(); + expect(screen.getByText("Tags")).toBeInTheDocument(); }); for (const label of PTU_LABELS) { - expect(queryByText(label)).not.toBeInTheDocument(); + expect(screen.queryByText(label)).not.toBeInTheDocument(); } - expect(queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); }); it("shows every PTU field when PTU cost attribution is enabled", async () => { mockUsePtuCostAttributionEnabled.mockReturnValue(true); - const { getByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("PTU Count")).toBeInTheDocument(); + expect(screen.getByText("PTU Count")).toBeInTheDocument(); }); for (const label of PTU_LABELS) { - expect(getByText(label)).toBeInTheDocument(); + expect(screen.getByText(label)).toBeInTheDocument(); } - expect(getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective To (UTC)")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx index 64074481603..d6476089cd7 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; import { getPlaceholder, Providers } from "../provider_info_helpers"; import { MountedFormHost } from "../../../tests/mounted-form-host"; @@ -6,7 +6,7 @@ import LiteLLMModelNameField from "./litellm_model_name"; describe("LitellmModelNameField", () => { it("should render", () => { - const { getByText } = render( + render( { /> , ); - expect(getByText("LiteLLM Model Name(s)")).toBeInTheDocument(); + expect(screen.getByText("LiteLLM Model Name(s)")).toBeInTheDocument(); }); it("should show Azure placeholder as 'my-deployment'", () => { - const { getByPlaceholderText, queryByPlaceholderText } = render( + render( , ); - expect(getByPlaceholderText("my-deployment")).toBeInTheDocument(); - expect(queryByPlaceholderText("gpt-3.5-turbo")).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("my-deployment")).toBeInTheDocument(); + expect(screen.queryByPlaceholderText("gpt-3.5-turbo")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx index 24685f9c129..e1faf8be1df 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx @@ -26,8 +26,8 @@ const openUploadStep = async () => { describe("BulkCreateUsersButton", () => { it("should render", () => { - const { getByText } = render(); - expect(getByText("+ Bulk Invite Users")).toBeInTheDocument(); + render(); + expect(screen.getByText("+ Bulk Invite Users")).toBeInTheDocument(); }); it("parses a CSV chosen through the file input", async () => { diff --git a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx index bad92555bd5..3fc1d41dfdd 100644 --- a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx +++ b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it } from "vitest"; import CostOptimizationFeedbackBanner from "./cost_optimization_feedback_banner"; @@ -10,24 +10,24 @@ describe("CostOptimizationFeedbackBanner", () => { }); it("renders with a link to the feedback discussion", () => { - const { getByText } = render(); - const link = getByText("Share Feedback").closest("a"); + render(); + const link = screen.getByText("Share Feedback").closest("a"); expect(link).toHaveAttribute("href", "https://github.com/BerriAI/litellm/discussions/32172"); }); it("hides itself and persists the dismissal when the dismiss button is clicked", () => { - const { getByText, queryByText, getByLabelText } = render(); - expect(getByText("Help shape cost optimization")).toBeInTheDocument(); + render(); + expect(screen.getByText("Help shape cost optimization")).toBeInTheDocument(); - fireEvent.click(getByLabelText("Dismiss banner")); + fireEvent.click(screen.getByLabelText("Dismiss banner")); - expect(queryByText("Help shape cost optimization")).not.toBeInTheDocument(); + expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); expect(localStorage.getItem(STORAGE_KEY)).toBe("true"); }); it("stays dismissed on remount once persisted", () => { localStorage.setItem(STORAGE_KEY, "true"); - const { queryByText } = render(); - expect(queryByText("Help shape cost optimization")).not.toBeInTheDocument(); + render(); + expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx index b8d8e3ba9c8..700d19eb13c 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx @@ -108,7 +108,7 @@ beforeEach(() => { test("renders organization view after loading data", async () => { mockUseOrganization.mockReturnValue({ data: mockOrg, isLoading: false } as any); - const { findAllByText } = renderWithProviders( + renderWithProviders( {}} @@ -120,7 +120,7 @@ test("renders organization view after loading data", async () => { />, ); - const [orgName] = await findAllByText("Acme Corp"); + const [orgName] = await screen.findAllByText("Acme Corp"); expect(orgName).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index b97fef32402..762b23f413e 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -77,21 +77,21 @@ describe("Settings", () => { }); it("should render the logging callbacks tab when access token is provided", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); }); it("should display additional settings tabs", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("CloudZero Cost Tracking")).toBeInTheDocument(); - expect(getByText("Alerting Types")).toBeInTheDocument(); - expect(getByText("Alerting Settings")).toBeInTheDocument(); - expect(getByText("Email Alerts")).toBeInTheDocument(); + expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); + expect(screen.getByText("Alerting Types")).toBeInTheDocument(); + expect(screen.getByText("Alerting Settings")).toBeInTheDocument(); + expect(screen.getByText("Email Alerts")).toBeInTheDocument(); }); }); @@ -279,13 +279,13 @@ describe("Settings", () => { }); it("should display CloudZero Cost Tracking tab", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); - expect(getByText("CloudZero Cost Tracking")).toBeInTheDocument(); + expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx index 3698b68155e..5bd48b1aa4c 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PaginationStatusAlerts from "./PaginationStatusAlerts"; @@ -6,7 +6,7 @@ import PaginationStatusAlerts from "./PaginationStatusAlerts"; describe("PaginationStatusAlerts", () => { it("shows page progress and wires the Stop button while fetching", () => { const cancel = vi.fn(); - const { getByRole, getByText } = render( + render( { />, ); - expect(getByText(/Currently fetching spend data: fetched 7 \/ 42 pages/)).toBeInTheDocument(); - fireEvent.click(getByRole("button", { name: "Stop" })); + expect(screen.getByText(/Currently fetching spend data: fetched 7 \/ 42 pages/)).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Stop" })); expect(cancel).toHaveBeenCalledTimes(1); }); it("shows the partial-data notice after a cancel, frozen at the last fetched page", () => { - const { getByText } = render( + render( { />, ); - expect(getByText("Showing partial spend data (7/42 pages loaded)")).toBeInTheDocument(); + expect(screen.getByText("Showing partial spend data (7/42 pages loaded)")).toBeInTheDocument(); }); it("names the subject it is fetching", () => { - const { getByText } = render( + render( { />, ); - expect(getByText(/Currently fetching agent data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); + expect(screen.getByText(/Currently fetching agent data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); }); it("renders nothing when idle", () => { diff --git a/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx index 6e4ab14be33..c81ade526dd 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import React from "react"; import { describe, expect, it } from "vitest"; import { AreaChart } from "./area_chart"; @@ -21,9 +21,9 @@ describe("AreaChart", () => { }); it("renders the No data placeholder instead of a chart when data is empty", () => { - const { container, getByText } = render(); + const { container } = render(); - expect(getByText("No data")).toBeInTheDocument(); + expect(screen.getByText("No data")).toBeInTheDocument(); expect(container.querySelector('[data-slot="chart"]')).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx index a322eeb3ad0..3cd7bd2fd08 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx @@ -21,9 +21,9 @@ describe("BarChart", () => { }); it("renders the No data placeholder instead of a chart when data is empty", () => { - const { container, getByText } = render(); + const { container } = render(); - expect(getByText("No data")).toBeInTheDocument(); + expect(screen.getByText("No data")).toBeInTheDocument(); expect(container.querySelector('[data-slot="chart"]')).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 6592ff8c357..705679fe2e6 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -445,7 +445,7 @@ describe("KeyInfoView handleKeyUpdate budget_duration", () => { ); fireEvent.click(screen.getByText("Settings")); - expect(screen.getByText("Budget Reset").parentElement?.textContent).toContain("Every 30d"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Every 30d"); fireEvent.click(screen.getByText("Edit Settings")); (globalThis as any).__TEST_FORM_VALUES = { @@ -456,7 +456,7 @@ describe("KeyInfoView handleKeyUpdate budget_duration", () => { fireEvent.click(screen.getByText("Mock Submit")); await waitFor(() => { - expect(screen.getByText("Budget Reset").parentElement?.textContent).toBe("Budget ResetNever"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Never"); }); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx index f506c0e51d7..a02f80f7800 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx @@ -381,6 +381,6 @@ describe("KeyInfoView budget reset visibility", () => { await waitFor(() => { expect(screen.getByText("Budget Reset")).toBeInTheDocument(); }); - expect(screen.getByText("Budget Reset").parentElement).toHaveTextContent("Never"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Never"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 0be80c3e173..0dd0dd6d6af 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -895,7 +895,7 @@ export default function KeyInfoView({

Budget Reset

-

+

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

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

{skill.name}

+
+

{skill.name}

{skill.description && ( -

{skill.description}

+

{skill.description}

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

Skill Details

-

Metadata registered with this skill

- +
+

Skill Details

+

Metadata registered with this skill

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

Using this skill

-

+

+

Using this skill

+

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

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

+

+

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

- 
+            
               /plugin marketplace update litellm
-

+

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

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

- One-time marketplace setup -

+
+

One-time marketplace setup

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

+

Run this command in Claude Code to register the marketplace:

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

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

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

-
-
- ~/.claude/settings.json +
+
+ ~/.claude/settings.json
- 
-              {settingsSnippet}
-
+ 
{settingsSnippet}
)} From a5f941068144ff7f0acc3838e2ace35202f74ce2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 09:46:43 -0700 Subject: [PATCH 106/126] test(ui): wait for the select popup before clicking its option key_edit_view opened a select and then clicked the option it found by title text or by raw text. Both queries match the moment the option enters the DOM, which is one render before the popup finishes entering. Until then the positioner still carries an inline pointer-events: none, and user-event refuses to click through it. That is a race, and a fast machine loses it. Five of the file's 84 tests failed on every local run while CI stayed green, which is the worst shape for a test to have: it is only ever red on the machine of whoever is trying to change the code. tests/test-utils.tsx already ships chooseSelectOption for exactly this. It finds the option by role and waits for the positioner to release pointer events before clicking. The five call sites now use it, and the helper takes the direct user-event API as well as a setup() instance so callers do not have to restructure to use it. Five consecutive full-file runs pass where every previous run failed. Also finishes this file's screen queries, which brings prefer-screen-queries to its target of 18. --- ui/litellm-dashboard/eslint-budgets.json | 2 +- .../templates/key_edit_view.test.tsx | 31 +++++++------------ ui/litellm-dashboard/tests/test-utils.tsx | 2 +- 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index e7545eb2383..e8207d179bd 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -7,5 +7,5 @@ "local/no-long-condition-chain": { "max": 265, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, "testing-library/no-node-access": { "max": 716, "target": 500 }, - "testing-library/prefer-screen-queries": { "max": 21, "target": 18 } + "testing-library/prefer-screen-queries": { "max": 18, "target": 18 } } diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index bf8f43b8b5f..97b09e00808 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1,7 +1,7 @@ import { fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { KeyResponse } from "../key_team_helpers/key_list"; import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "../key_team_helpers/ModelMaxBudgetEditor"; import { @@ -300,7 +300,7 @@ describe("KeyEditView", () => { }); it("should render", async () => { - const { getByText } = renderWithProviders( + renderWithProviders( {}} @@ -313,12 +313,12 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("Save Changes")).toBeInTheDocument(); + expect(screen.getByText("Save Changes")).toBeInTheDocument(); }); }); it("should render tags", async () => { - const { getByText } = renderWithProviders( + renderWithProviders( {}} @@ -331,12 +331,12 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("test-tag")).toBeInTheDocument(); + expect(screen.getByText("test-tag")).toBeInTheDocument(); }); }); it("should not render tags in metadata textarea", async () => { - const { getByLabelText } = renderWithProviders( + renderWithProviders( {}} @@ -348,7 +348,7 @@ describe("KeyEditView", () => { />, ); - const metadataTextarea = getByLabelText("Metadata") as HTMLTextAreaElement; + const metadataTextarea = screen.getByLabelText("Metadata") as HTMLTextAreaElement; await waitFor(() => { expect(metadataTextarea).toHaveValue("{}"); }); @@ -963,10 +963,7 @@ describe("KeyEditView", () => { />, ); - await userEvent.click(await screen.findByLabelText("Reset Budget")); - - const weeklyOption = await screen.findByText("weekly"); - await userEvent.click(weeklyOption); + await chooseSelectOption(userEvent, await screen.findByLabelText("Reset Budget"), "weekly"); const submitButton = screen.getByRole("button", { name: /save changes/i }); await userEvent.click(submitButton); @@ -1042,8 +1039,7 @@ describe("KeyEditView", () => { ); const resetBudget = await screen.findByLabelText("Reset Budget"); - await userEvent.click(resetBudget); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, resetBudget, "Never resets"); await waitFor(() => { expect(resetBudget).toHaveTextContent("Never resets"); @@ -1074,8 +1070,7 @@ describe("KeyEditView", () => { />, ); - await userEvent.click(await screen.findByLabelText("Reset Budget")); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, await screen.findByLabelText("Reset Budget"), "Never resets"); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -1946,8 +1941,7 @@ describe("KeyEditView", () => { await userEvent.clear(duration); await userEvent.type(duration, "45d"); - await userEvent.click(screen.getByLabelText(/TPM Rate Limit Type/)); - await userEvent.click(await screen.findByTitle("Guaranteed throughput")); + await chooseSelectOption(userEvent, screen.getByLabelText(/TPM Rate Limit Type/), /^Guaranteed throughput/); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -2103,8 +2097,7 @@ describe("KeyEditView", () => { renderForPayload(onSubmitMock); await screen.findByRole("button", { name: /save changes/i }); - await userEvent.click(screen.getByLabelText(/RPM Rate Limit Type/)); - await userEvent.click(await screen.findByTitle("Guaranteed throughput")); + await chooseSelectOption(userEvent, screen.getByLabelText(/RPM Rate Limit Type/), /^Guaranteed throughput/); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index 66966201a9c..162b8a3df7b 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -52,7 +52,7 @@ const pointerBlocked = (element: HTMLElement): boolean => { * the option text alone is a race that React 19's flush timing loses. */ export const chooseSelectOption = async ( - user: ReturnType, + user: Pick, "click">, trigger: HTMLElement, optionName: string | RegExp, ) => { From 06d4521fc0b7a69aa0db0d6ba1aa54dc4ab1af28 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:15:18 -0700 Subject: [PATCH 107/126] fix(registry): add vertex veo 3.1 resolution tier pricing per vertex pricing page --- ...odel_prices_and_context_window_backup.json | 10 ++++-- model_prices_and_context_window.json | 10 ++++-- tests/test_litellm/test_video_generation.py | 35 +++++++++++++++++++ 3 files changed, 51 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dd45495884d..2cfa1f93f79 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -43509,7 +43509,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43523,6 +43524,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" @@ -43538,7 +43541,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43553,6 +43557,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dd45495884d..2cfa1f93f79 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -43509,7 +43509,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43523,6 +43524,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" @@ -43538,7 +43541,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43553,6 +43557,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index b166e902d6e..2a60ff9c4b5 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -532,6 +532,41 @@ class TestVideoGeneration: assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch): + """The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, provider: str, resolution: str | None, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = { + "duration_seconds": duration, + **({"video_resolution": resolution} if resolution else {}), + } + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider=provider, + ) + + for provider in ("gemini", "vertex_ai"): + for suffix in ("generate-preview", "generate-001"): + standard = f"{provider}/veo-3.1-{suffix}" + fast = f"{provider}/veo-3.1-fast-{suffix}" + assert abs(cost_for(standard, provider, None, 8.0) - 3.2) < 1e-6 + assert abs(cost_for(standard, provider, "1080p", 8.0) - 3.2) < 1e-6 + assert abs(cost_for(standard, provider, "4k", 8.0) - 4.8) < 1e-6 + assert abs(cost_for(fast, provider, "720p", 8.0) - 0.8) < 1e-6 + assert abs(cost_for(fast, provider, "1080p", 8.0) - 0.96) < 1e-6 + assert abs(cost_for(fast, provider, "4k", 8.0) - 2.4) < 1e-6 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig() From 8acdb9208756fb306680878d06f53bbd75c2f04a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 10:27:43 -0700 Subject: [PATCH 108/126] fix(ui): keep the skill detail copy buttons transparent bg-none only clears background-image, so the buttons fell back to the browser's default button background instead of the transparent one the inline style had. --- .../src/components/claude_code_plugins/skill_detail.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 809272c2ac0..7cbbb08b623 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -182,7 +182,7 @@ const SkillDetail: React.FC = ({ skill, onBack }) => {
- Property - - {skill.name} -
Property{skill.name}
{row.property}{row.value}
{row.property}{row.value}