From 48bde68781c20df4d98915cc970eb65b23e343fe Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:47:14 +0000 Subject: [PATCH 001/146] fix(key_generate): use user's budget for UI session personal keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 13 ++-- .../test_key_management_endpoints.py | 77 +++++++++++++++++++ 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 802a7c3e469..4f4e56d7418 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1233,11 +1233,10 @@ async def _common_key_generation_helper( # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. - is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None - # Session tokens (lite login) carry max_budget=None to avoid a per-session - # LLM spend cap, but that None must not be read as "unlimited delegation - # authority". A personal key (no team) has no team-budget enforcement at - # request time, so a session token cannot delegate any budget for one. + # Session tokens (lite login) use their session max_budget for team keys, but + # personal keys are capped by user_max_budget when it is available. + is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID + is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None if ( user_api_key_dict.is_session_token and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -1255,7 +1254,9 @@ async def _common_key_generation_helper( }, ) delegation_ceiling: Final = ( - user_api_key_dict.max_budget + user_api_key_dict.user_max_budget + if is_ui_session_token and user_api_key_dict.user_max_budget is not None + else user_api_key_dict.max_budget if user_api_key_dict.max_budget is not None else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None) ) 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 cc0a7631b59..809c4183e13 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 @@ -15509,6 +15509,83 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped(): assert "cannot exceed" in msg.lower() +@pytest.mark.asyncio +async def test_ui_session_token_personal_key_ceiling_is_user_budget(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=100) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-ui-session", + user_id="user-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + max_budget=1.0, + user_max_budget=500.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly + patch( # test-quality-ok: helper has no dependency injection seam for key persistence + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"} + try: + await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + team_table=None, + ) + except (HTTPException, ProxyException) as err: + msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) + assert "cannot exceed" not in msg.lower() + + +@pytest.mark.asyncio +async def test_ui_session_token_personal_key_above_user_budget_rejected(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=600) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-ui-session", + user_id="user-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + max_budget=1.0, + user_max_budget=500.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly + patch( # test-quality-ok: helper has no dependency injection seam for key persistence + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"} + with pytest.raises((HTTPException, ProxyException)) as exc_info: + await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + team_table=None, + ) + err = exc_info.value + code = getattr(err, "status_code", None) or getattr(err, "code", None) + msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) + assert str(code) == "400" + assert "cannot exceed" in msg.lower() + assert "500.0" in msg + + @pytest.mark.asyncio async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(): """ From 0ac0362b42890d5d40698dbd5c206065b0015e5f Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 21:01:16 +0000 Subject: [PATCH 002/146] fix(proxy): enforce virtual key budgets for JEV test routing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 2 +- .../test_auto_router_endpoints.py | 76 ++++++++++++++++++- 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 200ed6c3bf3..f79425d2e97 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -319,7 +319,7 @@ async def _authorize_models_this_test_can_call( its calls through the proxy. Team and member budgets are already enforced on every route. """ models: Final = _models_this_test_can_call(config) - if not models: + if not models and config.classifier_type != "jev": return from litellm.proxy.proxy_server import proxy_logging_obj diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 067f30c2fd7..36130137c64 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -3,23 +3,33 @@ Unit tests for auto router management endpoints """ from collections.abc import Mapping, Sequence +from functools import partial from pathlib import Path from typing import Final +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException, Request from pydantic import ValidationError +from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.management_endpoints import auto_router_endpoints from litellm.proxy.management_endpoints.auto_router_endpoints import ( preview_auto_router_routing, ) from litellm.router import Router +from litellm.router_strategy.complexity_router import ComplexityRouter +from litellm.router_strategy.complexity_router.jev_classifier import ( + JevChoiceAnswer, + JevClassifierClient, + JevSystemOneResponse, +) from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, @@ -422,8 +432,67 @@ async def test_a_key_over_its_budget_cannot_run_a_classifier_config(monkeypatch: assert calls == [] +@pytest.mark.parametrize( + "max_budget, spend, denied", + ( + pytest.param(0.0, 0.0, True, id="zero-budget"), + pytest.param(1.0, 1.0, True, id="budget-reached"), + pytest.param(1.0, 2.0, True, id="budget-exceeded"), + pytest.param(1.0, 0.5, False, id="budget-remaining"), + pytest.param(None, 2.0, False, id="unlimited"), + ), +) @pytest.mark.asyncio -async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.MonkeyPatch): +async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( + monkeypatch: pytest.MonkeyPatch, max_budget: float | None, spend: float, denied: bool +) -> None: + client: Final = AsyncMock(spec=JevClassifierClient) + client.evaluate.return_value = JevSystemOneResponse( + model="jev-test", + answers={ + "tier": JevChoiceAnswer(type="choice", choice="SIMPLE", probabilities={"SIMPLE": 1.0}, confidence=1.0) + }, + ) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + monkeypatch.setattr(auto_router_endpoints, "ComplexityRouter", partial(ComplexityRouter, jev_client=client)) + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-budget-test", + user_id="admin", + models=["cheap-model"], + max_budget=max_budget, + spend=spend, + ) + request: Final = _request( + "what is 2+2", + classifier_type="jev", + jev_classifier_config={"model": "jev-test"}, + ) + + if denied: + with pytest.raises(ProxyException) as exc_info: + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.code == "400" + assert exc_info.value.param is None + assert "Budget has been exceeded!" in exc_info.value.message + client.evaluate.assert_not_called() + return + + response: Final = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor + ) + assert response.routed_model == "cheap-model" + assert response.routing_decision["cause"] == "jev_classifier" + assert response.routing_decision["classifier_model"] == "typesafe/jev-test" + client.evaluate.assert_awaited_once() + + +@pytest.mark.parametrize("max_budget, spend", ((0.0, 0.0), (1.0, 2.0))) +@pytest.mark.asyncio +async def test_a_heuristic_config_does_not_need_a_budget( + monkeypatch: pytest.MonkeyPatch, max_budget: float, spend: float +): import litellm.proxy.proxy_server as proxy_server monkeypatch.setattr(proxy_server, "llm_router", _router()) @@ -435,8 +504,8 @@ async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.Mon user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-broke", user_id="admin", - max_budget=1.0, - spend=2.0, + max_budget=max_budget, + spend=spend, models=["cheap-model"], ), ) @@ -851,7 +920,6 @@ class TestAutoRouterBenchmarks: # --------------------------------------------------------------------------- from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock from litellm.proxy.management_endpoints.auto_router_endpoints import ( get_shadow_eval_job, From ee7d2b50946b4d84f249e7119c346b253abc8bef Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Fri, 18 Sep 2026 14:31:34 -0700 Subject: [PATCH 003/146] Update litellm/proxy/management_endpoints/key_management_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- 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 4f4e56d7418..f60e68bcfe7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1233,8 +1233,7 @@ async def _common_key_generation_helper( # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. - # Session tokens (lite login) use their session max_budget for team keys, but - # personal keys are capped by user_max_budget when it is available. + # UI session personal keys are capped by user_max_budget when it is available. is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None if ( From b9e5bb3abb0f2f0ddc06dcaf9edb63563cac2a2a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 22:03:37 +0000 Subject: [PATCH 004/146] test(proxy): allow JEV dependency in budget fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_auto_router_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 36130137c64..a5c93c41a84 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -459,7 +459,7 @@ async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-jev-budget-test", user_id="admin", - models=["cheap-model"], + models=["cheap-model", "typesafe/jev-test"], max_budget=max_budget, spend=spend, ) From b32d1112a6c9af25f0b5a66eae7ba33d8e201c74 Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:31:55 +0000 Subject: [PATCH 005/146] feat(team): show whether a member follows the team default budget and allow resetting to it Adds budget_source (team_default, custom, none) to each membership in /team/info and a POST /team/{team_id}/member/{user_id}/reset_budget route that relinks a member to the team's shared team_member_budget row without touching their spend. The Admin UI team members table shows a Team default or Custom badge next to each member's budget and offers a Use team default action on customized members Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 19 +- .../management_endpoints/team_endpoints.py | 112 ++++++++- .../test_team_endpoints.py | 214 ++++++++++++++++++ .../hooks/teams/useResetTeamMemberBudget.ts | 16 ++ .../src/components/team/TeamInfo.tsx | 7 +- .../components/team/TeamMemberTab.test.tsx | 150 ++++++++++++ .../src/components/team/TeamMemberTab.tsx | 103 ++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 71 ++++++ 8 files changed, 679 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 76a51627d0c..70a93676d4e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4,7 +4,7 @@ import os from collections.abc import Callable, Mapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias import httpx from pydantic import ( @@ -4588,11 +4588,26 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone) +TeamMemberBudgetSource: TypeAlias = Literal["team_default", "custom", "none"] + + +class TeamInfoMembership(LiteLLM_TeamMembership): + budget_source: TeamMemberBudgetSource + + class TeamInfoResponseObject(TypedDict): team_id: str team_info: TeamInfoResponseObjectTeamTable keys: list - team_memberships: list[LiteLLM_TeamMembership] + team_memberships: ReadOnly[tuple[TeamInfoMembership, ...]] + + +class TeamMemberResetBudgetResponse(BaseModel): + team_id: str + user_id: str + budget_id: str | None + previous_budget_id: str | None + budget_source: TeamMemberBudgetSource class TeamListResponseObject(LiteLLM_TeamTable): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 28c12173ea7..b8d95167045 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -80,11 +80,14 @@ from litellm.proxy._types import ( TeamEditNone, TeamEditUnrestricted, TeamInfoMember, + TeamInfoMembership, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, TeamListResponseObject, TeamMemberAddRequest, + TeamMemberBudgetSource, TeamMemberDeleteRequest, + TeamMemberResetBudgetResponse, TeamMemberUpdateRequest, TeamMemberUpdateResponse, TeamModelAddRequest, @@ -3954,6 +3957,99 @@ async def reset_team_member_spend_fn( } +class _TeamMetadataView(BaseModel): + metadata: Mapping[str, object] | None = None + + +def _team_default_budget_id(team: LiteLLM_TeamTable) -> str | None: + view: Final = _TeamMetadataView.model_validate(team, from_attributes=True) + raw: Final = view.metadata.get("team_member_budget_id") if view.metadata is not None else None + return raw if isinstance(raw, str) else None + + +async def _existing_team_default_budget_id(team: LiteLLM_TeamTable, prisma_client: PrismaClient) -> str | None: + budget_id: Final = _team_default_budget_id(team) + if budget_id is None: + return None + row: Final = await _budget_db(prisma_client).find_unique( + where={"budget_id": budget_id}, # mutable-ok: prisma client requires a plain dict where= argument + ) + return budget_id if row is not None else None + + +def _member_budget_source(budget_id: str | None, team_default_budget_id: str | None) -> TeamMemberBudgetSource: + if budget_id is not None and budget_id != team_default_budget_id: + return "custom" + return "team_default" if team_default_budget_id is not None else "none" + + +@router.post( + "/team/{team_id}/member/{user_id}/reset_budget", + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence + dependencies=(Depends(user_api_key_auth),), + response_model=TeamMemberResetBudgetResponse, +) +@management_endpoint_wrapper +async def reset_team_member_budget_fn( + team_id: str, + user_id: str, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> TeamMemberResetBudgetResponse: + """ + Put a team member back on the team's shared default member budget (`team_member_budget`). + + Drops the member's own budget row link so team-wide changes made through /team/update + reach them again. Leaves the member with no budget when the team has no default. Spend is untouched. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None") + + team_obj: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) + + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument + } + membership_row: Final = await _team_membership_db(prisma_client).find_unique(where=membership_where) + if membership_row is None: + _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.") + + team_default_budget_id: Final = await _existing_team_default_budget_id(team_obj, prisma_client) + budget_link: Final = ( + { + "connect": {"budget_id": team_default_budget_id} + } # mutable-ok: prisma client requires a plain dict data= argument + if team_default_budget_id is not None + else {"disconnect": True} # mutable-ok: same prisma data= argument + ) + await _team_membership_db(prisma_client).update( + where=membership_where, + data={"litellm_budget_table": budget_link}, # mutable-ok: prisma client requires a plain dict data= argument + ) + await invalidate_team_member_spend_state( + user_id=user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + ) + + return TeamMemberResetBudgetResponse( + team_id=team_id, + user_id=user_id, + budget_id=team_default_budget_id, + previous_budget_id=membership_row.budget_id, + budget_source=_member_budget_source(team_default_budget_id, team_default_budget_id), + ) + + def _create_results_from_response( members: list[Member], response: TeamAddMemberResponse, @@ -4722,9 +4818,7 @@ async def team_info( _team_info = TeamInfoResponseObjectTeamTable() ## GET TEAM BUDGET (if exists) ## - team_member_budget_id: Final = ( - _team_info.metadata.get("team_member_budget_id") if _team_info.metadata is not None else None - ) + team_member_budget_id: Final = _team_default_budget_id(_team_info) if team_member_budget_id is not None: _team_info = await _add_team_member_budget_table( team_member_budget_id=team_member_budget_id, @@ -4757,7 +4851,17 @@ async def team_info( team_id=team_id, team_info=hydrated_team_info, keys=keys, - team_memberships=returned_tm, + team_memberships=tuple( + TeamInfoMembership.model_validate( + MappingProxyType( + { + **tm.model_dump(), + "budget_source": _member_budget_source(tm.budget_id, team_member_budget_id), + } + ) + ) + for tm in returned_tm + ), ) return response_object 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 690b5ae80b6..d55b9f79b5f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _verify_team_access, delete_team, list_available_teams, + reset_team_member_budget_fn, reset_team_member_spend_fn, router, team_member_add_duplication_check, @@ -14432,6 +14433,219 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey assert response["spend"] == 0.0 +def _reset_budget_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user") + + +def _team_with_default_budget(team_id: str, budget_id: str) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": budget_id}) + + +@pytest.mark.asyncio +async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default(monkeypatch): + """An admin undoing a per-member budget must put the membership back on the team's shared + default row (a connect, not a copy) so later /team/update changes reach the member again, + and must drop the cached membership so the old cap stops being enforced. The shared row and + the member's tracked spend are never written.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + + membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0, budget_id="custom-b1") + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0) + ) + mock_prisma_client.db.litellm_budgettable.update = AsyncMock() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")), + ): + response = await reset_team_member_budget_fn( + team_id="team-1", user_id="member-1", user_api_key_dict=_reset_budget_admin() + ) + + assert response.budget_id == "team-default-b" + assert response.previous_budget_id == "custom-b1" + assert response.budget_source == "team_default" + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"litellm_budget_table": {"connect": {"budget_id": "team-default-b"}}}, + ) + mock_prisma_client.db.litellm_budgettable.update.assert_not_awaited() + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_obj, default_row", + [ + (LiteLLM_TeamTable(team_id="team-1"), None), + (_team_with_default_budget("team-1", "gone-b"), None), + ], + ids=["no_default_configured", "configured_default_row_missing"], +) +async def test_reset_team_member_budget_fn_detaches_member_when_team_has_no_usable_default( + monkeypatch, team_obj, default_row +): + """With no shared default to link to, reset leaves the member exactly where a freshly added + member would be: no budget row at all, reported as budget_source='none', rather than + connecting to a budget_id that does not exist or leaving the custom cap in place.""" + mock_prisma_client = MagicMock() + membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id="custom-b1") + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=team_obj), + ): + response = await reset_team_member_budget_fn( + team_id="team-1", user_id="member-1", user_api_key_dict=_reset_budget_admin() + ) + + assert response.budget_id is None + assert response.previous_budget_id == "custom-b1" + assert response.budget_source == "none" + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"litellm_budget_table": {"disconnect": True}}, + ) + + +@pytest.mark.asyncio +async def test_reset_team_member_budget_fn_membership_not_found(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_budget_fn( + team_id="team-1", user_id="ghost-user", user_api_key_dict=_reset_budget_admin() + ) + assert exc.value.status_code == 404 + mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_budget_fn( + team_id="team-1", + user_id="member-1", + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user" + ), + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_info_reports_whether_each_member_follows_the_team_default_budget(): + """/team/info must tell the caller which members still follow the team's shared member budget + and which carry their own row, since budget_id alone only means something to a reader who + also knows the team's team_member_budget_id.""" + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = _team_with_default_budget("team-1", "team-default-b") + memberships = [ + LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"), + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0) + ) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + "inherits": "team_default", + "customized": "custom", + "unlinked": "team_default", + } + + +@pytest.mark.asyncio +async def test_team_info_reports_no_budget_source_when_team_has_no_default(): + """A team that never set team_member_budget has nothing for members to inherit, so an + unlinked member is 'none' rather than 'team_default', while a member with their own row is + still 'custom'.""" + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + memberships = [ + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + "customized": "custom", + "unlinked": "none", + } + + @pytest.mark.asyncio async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch): """Raising a stuck member's max_budget_in_team via the documented /team/member_update diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts new file mode 100644 index 00000000000..e7cf95440a5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts @@ -0,0 +1,16 @@ +import { useMutation } from "@tanstack/react-query"; +import { fetchClient } from "@/lib/http/api"; + +export interface ResetTeamMemberBudgetParams { + teamId: string; + userId: string; +} + +export const resetTeamMemberBudget = async ({ teamId, userId }: ResetTeamMemberBudgetParams): Promise => { + await fetchClient.POST("/team/{team_id}/member/{user_id}/reset_budget", { + params: { path: { team_id: teamId, user_id: userId } }, + }); +}; + +export const useResetTeamMemberBudget = () => + useMutation({ mutationFn: resetTeamMemberBudget }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3b2c344c2f3..22cc99b32c8 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1,4 +1,5 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import type { components } from "@/lib/http/schema"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useQueryClient } from "@tanstack/react-query"; @@ -247,10 +248,13 @@ export const retainedMcpToolPermissions = ( export const mcpUnresolvableSaveError = (reason: string): string => `Cannot save MCP tool permissions because ${reason}. Retry once the page has finished loading`; +export type TeamMemberBudgetSource = components["schemas"]["TeamMemberResetBudgetResponse"]["budget_source"]; + export interface TeamMembership { user_id: string; team_id: string; - budget_id: string; + budget_id: string | null; + budget_source: TeamMemberBudgetSource; spend: number; total_spend: number | null; litellm_budget_table: { @@ -1361,6 +1365,7 @@ const TeamInfoView: React.FC = ({ canEditTeam={canEditTeam} handleMemberDelete={handleMemberDelete} onMemberSpendReset={refreshTeamData} + onMemberBudgetReset={refreshTeamData} setSelectedEditMember={setSelectedEditMember} setIsEditMemberModalVisible={setIsEditMemberModalVisible} setIsAddMemberModalVisible={setIsAddMemberModalVisible} diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 52cba1e6330..8652ffa7de2 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -30,6 +30,7 @@ const mockSetSelectedEditMember = vi.fn(); const mockSetIsEditMemberModalVisible = vi.fn(); const mockSetIsAddMemberModalVisible = vi.fn(); const mockOnMemberSpendReset = vi.fn(); +const mockOnMemberBudgetReset = vi.fn(); const budgetResetIso = new Date(2026, 6, 15, 12, 0, 0).toISOString(); @@ -74,6 +75,7 @@ const createMockTeamData = (overrides: Partial = {}): TeamData => ({ user_id: "user1@test.com", team_id: "team-123", budget_id: "budget1", + budget_source: "custom", spend: 100.5, total_spend: 1538.2608, litellm_budget_table: { @@ -126,6 +128,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -142,6 +145,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -161,6 +165,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -180,6 +185,7 @@ describe("TeamMembersComponent", () => { canEditTeam: false, handleMemberDelete: mockHandleMemberDelete, onMemberSpendReset: mockOnMemberSpendReset, + onMemberBudgetReset: mockOnMemberBudgetReset, setSelectedEditMember: mockSetSelectedEditMember, setIsEditMemberModalVisible: mockSetIsEditMemberModalVisible, setIsAddMemberModalVisible: mockSetIsAddMemberModalVisible, @@ -204,6 +210,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -231,6 +238,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -258,6 +266,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -274,6 +283,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -293,6 +303,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -309,6 +320,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -326,6 +338,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -346,6 +359,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -381,6 +395,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -435,6 +450,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -466,6 +482,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -486,6 +503,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -503,6 +521,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -521,6 +540,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -603,4 +623,134 @@ describe("TeamMembersComponent", () => { expect(screen.getByTestId("reset-member-spend")).toBeVisible(); }); }); + + describe("budget source", () => { + const teamDataWithDefault = () => { + const base = createMockTeamData(); + return createMockTeamData({ + team_info: { + ...base.team_info, + team_member_budget_table: { max_budget: 25, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }, + team_memberships: [ + base.team_memberships[0], + { + user_id: "user2@test.com", + team_id: "team-123", + budget_id: "team-default-budget", + budget_source: "team_default", + spend: 0, + total_spend: null, + litellm_budget_table: { + budget_id: "team-default-budget", + soft_budget: null, + max_budget: 25, + max_parallel_requests: null, + tpm_limit: null, + rpm_limit: null, + model_max_budget: null, + budget_duration: null, + budget_reset_at: null, + }, + }, + ], + }); + }; + + const renderTab = (teamData: TeamData, canEditTeam = true) => + renderWithProviders( + , + ); + + it("labels each member's budget as Custom or Team default and shows the team amount for inherited members", () => { + renderTab(teamDataWithDefault()); + + const customRow = screen.getByRole("row", { name: /user1@test\.com/ }); + const inheritedRow = screen.getByRole("row", { name: /user2@test\.com/ }); + expect(within(customRow).getByTestId("member-budget-source")).toHaveTextContent("Custom"); + expect(customRow).toHaveTextContent("$1,000.00"); + expect(within(inheritedRow).getByTestId("member-budget-source")).toHaveTextContent("Team default"); + expect(inheritedRow).toHaveTextContent("$25.00"); + }); + + it("shows no source label for a member with neither a custom nor a team budget", () => { + renderTab(createMockTeamData({ team_memberships: [] })); + + expect(screen.queryByTestId("member-budget-source")).not.toBeInTheDocument(); + expect(screen.queryByTestId("reset-member-budget")).not.toBeInTheDocument(); + }); + + it("only offers Use team default on customized members, and only to editors", () => { + const { unmount } = renderTab(teamDataWithDefault()); + + expect( + within(screen.getByRole("row", { name: /user1@test\.com/ })).getByTestId("reset-member-budget"), + ).toBeVisible(); + expect( + within(screen.getByRole("row", { name: /user2@test\.com/ })).queryByTestId("reset-member-budget"), + ).not.toBeInTheDocument(); + + unmount(); + renderTab(teamDataWithDefault(), false); + expect(screen.queryByTestId("reset-member-budget")).not.toBeInTheDocument(); + }); + + it("puts the member back on the team default after confirming, then refreshes the team", async () => { + const user = userEvent.setup(); + POST.mockResolvedValue({ data: {} }); + renderTab(teamDataWithDefault()); + + await user.click(screen.getByTestId("reset-member-budget")); + + const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" }); + expect(dialog).toHaveTextContent("user1@test.com"); + expect(dialog).toHaveTextContent("team default of $25.00"); + expect(dialog).toHaveTextContent("Custom budget: $1,000.00"); + expect(POST).not.toHaveBeenCalled(); + + await user.click(within(dialog).getByRole("button", { name: "Use team default" })); + + await waitFor(() => expect(mockOnMemberBudgetReset).toHaveBeenCalledTimes(1)); + expect(POST).toHaveBeenCalledExactlyOnceWith("/team/{team_id}/member/{user_id}/reset_budget", { + params: { path: { team_id: "team-123", user_id: "user1@test.com" } }, + }); + expect(mockOnMemberSpendReset).not.toHaveBeenCalled(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("keeps the dialog open and does not refresh the team when the reset fails", async () => { + const user = userEvent.setup(); + POST.mockRejectedValue(new Error("Team admin cannot reset budgets")); + renderTab(teamDataWithDefault()); + + await user.click(screen.getByTestId("reset-member-budget")); + const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" }); + await user.click(within(dialog).getByRole("button", { name: "Use team default" })); + + await waitFor(() => expect(POST).toHaveBeenCalledTimes(1)); + expect(mockOnMemberBudgetReset).not.toHaveBeenCalled(); + expect(screen.getByRole("dialog", { name: "Reset Team Member Budget" })).toBeInTheDocument(); + }); + + it("does not call the API when the dialog is cancelled", async () => { + const user = userEvent.setup(); + renderTab(teamDataWithDefault()); + + await user.click(screen.getByTestId("reset-member-budget")); + const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" }); + await user.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(POST).not.toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index a869c1ad624..660416504fe 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -1,6 +1,8 @@ +import { useResetTeamMemberBudget } from "@/app/(dashboard)/hooks/teams/useResetTeamMemberBudget"; import { useResetTeamMemberSpend } from "@/app/(dashboard)/hooks/teams/useResetTeamMemberSpend"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { SimpleTooltip } from "@/components/ui/tooltip"; @@ -13,7 +15,15 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles"; import { CircleHelp } from "lucide-react"; import { useState, type ComponentProps } from "react"; -import { TeamData, TeamMembership } from "./TeamInfo"; +import { TeamData, TeamMemberBudgetSource, TeamMembership } from "./TeamInfo"; + +const BUDGET_SOURCE_LABELS: Record, string> = { + team_default: "Team default", + custom: "Custom", +}; + +const formatBudget = (value: number | null): string => + value === null ? "Unlimited" : `$${formatNumberWithCommas(value, 2)}`; export const seedMemberBudgetFields = ( record: Member, @@ -37,6 +47,7 @@ interface TeamMemberTabProps { setIsEditMemberModalVisible: (visible: boolean) => void; setIsAddMemberModalVisible: (visible: boolean) => void; onMemberSpendReset: () => void; + onMemberBudgetReset: () => void; } export default function TeamMemberTab({ @@ -47,9 +58,13 @@ export default function TeamMemberTab({ setIsEditMemberModalVisible, setIsAddMemberModalVisible, onMemberSpendReset, + onMemberBudgetReset, }: TeamMemberTabProps) { const [memberToResetSpend, setMemberToResetSpend] = useState(null); + const [memberToResetBudget, setMemberToResetBudget] = useState(null); const { mutate: resetMemberSpend, isPending: isResettingSpend } = useResetTeamMemberSpend(); + const { mutate: resetMemberBudget, isPending: isResettingBudget } = useResetTeamMemberBudget(); + const teamDefaultBudget = teamData.team_info.team_member_budget_table?.max_budget ?? null; const formatNumber = (value: number | null): string => { if (value === null || value === undefined) return "0"; @@ -82,10 +97,19 @@ export default function TeamMemberTab({ return membership?.total_spend ?? 0; }; + const getUserBudgetSource = (userId: string | null): TeamMemberBudgetSource => { + if (!userId) return "none"; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + return membership?.budget_source ?? "none"; + }; + const getUserBudget = (userId: string | null): number | null => { if (!userId) return null; const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.litellm_budget_table?.max_budget ?? null; + return ( + membership?.litellm_budget_table?.max_budget ?? + (membership?.budget_source === "team_default" ? teamDefaultBudget : null) + ); }; // Helper function to get rate limits for a user @@ -182,12 +206,40 @@ export default function TeamMemberTab({ render: (record: Member) => , }, { - title: "Team Member Budget (USD)", + title: ( + + Team Member Budget (USD) + + + + + ), key: "budget", sortValue: (record: Member) => getUserBudget(record.user_id), - render: (record: Member) => ( - - ), + render: (record: Member) => { + const source = getUserBudgetSource(record.user_id); + return ( + + + {source !== "none" && ( + + {BUDGET_SOURCE_LABELS[source]} + + )} + {source === "custom" && canEditTeam && ( + + )} + + ); + }, }, { title: "Budget Reset", @@ -224,6 +276,21 @@ export default function TeamMemberTab({ ); }; + const handleResetBudget = () => { + if (!memberToResetBudget?.user_id) return; + resetMemberBudget( + { teamId: teamData.team_id, userId: memberToResetBudget.user_id }, + { + onSuccess: () => { + toast.success("Team member budget reset to the team default"); + setMemberToResetBudget(null); + onMemberBudgetReset(); + }, + onError: (error) => toast.fromError(parseErrorMessage(error)), + }, + ); + }; + return ( <> + !open && setMemberToResetBudget(null)}> + + + Reset Team Member Budget + +

+ Remove the custom budget for{" "} + {memberToResetBudget?.user_email || memberToResetBudget?.user_id} and put them back on the + team default of {formatBudget(teamDefaultBudget)}? +

+

+ Custom budget: {formatBudget(getUserBudget(memberToResetBudget?.user_id ?? null))}. Their + spend is kept. Future changes to the team's member budget will apply to them again. +

+ + + + + + ); } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 136b8f26784..aa6776d12bd 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16272,6 +16272,29 @@ export interface paths { patch?: never; trace?: never; }; + "/team/{team_id}/member/{user_id}/reset_budget": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Reset Team Member Budget Fn + * @description Put a team member back on the team's shared default member budget (`team_member_budget`). + * + * Drops the member's own budget row link so team-wide changes made through /team/update + * reach them again. Leaves the member with no budget when the team has no default. Spend is untouched. + */ + post: operations["reset_team_member_budget_fn_team__team_id__member__user_id__reset_budget_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/{team_id}/member/{user_id}/reset_spend": { parameters: { query?: never; @@ -38547,6 +38570,22 @@ export interface components { /** User Id */ user_id?: string | null; }; + /** TeamMemberResetBudgetResponse */ + TeamMemberResetBudgetResponse: { + /** Budget Id */ + budget_id: string | null; + /** + * Budget Source + * @enum {string} + */ + budget_source: "team_default" | "custom" | "none"; + /** Previous Budget Id */ + previous_budget_id: string | null; + /** Team Id */ + team_id: string; + /** User Id */ + user_id: string; + }; /** TeamMemberUpdateRequest */ TeamMemberUpdateRequest: { /** @@ -61730,6 +61769,38 @@ export interface operations { }; }; }; + reset_team_member_budget_fn_team__team_id__member__user_id__reset_budget_post: { + parameters: { + query?: never; + header?: never; + path: { + team_id: string; + user_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["TeamMemberResetBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post: { parameters: { query?: never; From d3a364d74f8bee7f6133d51e68c3036cde7f8136 Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:49:58 +0000 Subject: [PATCH 006/146] fix(team): report no budget source when the team default row was deleted Derive budget_source from the budget row /team/info actually loaded, so a metadata id whose row was removed via /budget/delete reads as none instead of team_default. Share the /team/info test scaffolding so the added patch calls stay within the TQ008 budget, and allowlist the imperative reset_budget route in the provider endpoint audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 5 +- .../endpointaudit/coverage_allowlist.txt | 1 + .../test_team_endpoints.py | 112 ++++++++++-------- 3 files changed, 70 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b8d95167045..a441b3834ed 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4825,6 +4825,9 @@ async def team_info( prisma_client=prisma_client, team_info_response_object=_team_info, ) + active_default_budget_id: Final = ( + team_member_budget_id if _team_info.team_member_budget_table is not None else None + ) # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) @@ -4856,7 +4859,7 @@ async def team_info( MappingProxyType( { **tm.model_dump(), - "budget_source": _member_budget_source(tm.budget_id, team_member_budget_id), + "budget_source": _member_budget_source(tm.budget_id, active_default_budget_id), } ) ) diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 6bc8947e89f..4ea64b152f1 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -81,6 +81,7 @@ POST /prompts/test POST /search_tools/test_connection POST /team/bulk_member_add POST /team/{team_id}/member/{user_id}/reset_spend +POST /team/{team_id}/member/{user_id}/reset_budget POST /team/key/bulk_update POST /team/permissions_bulk_update POST /team/{team_id}/disable_logging 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 d55b9f79b5f..484c054fa54 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14572,40 +14572,52 @@ async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch): mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() +async def _team_info_budget_sources( + team_row: LiteLLM_TeamTable, + memberships: list[LiteLLM_TeamMembership], + default_budget_row: LiteLLM_BudgetTable | None, +) -> dict[str, str]: + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(return_value=default_budget_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch.object( # test-quality-ok: membership lookup is a module-level DB query with no injection point + team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships) + ), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id=team_row.team_id, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + return {tm.user_id: tm.budget_source for tm in response["team_memberships"]} + + @pytest.mark.asyncio async def test_team_info_reports_whether_each_member_follows_the_team_default_budget(): """/team/info must tell the caller which members still follow the team's shared member budget and which carry their own row, since budget_id alone only means something to a reader who also knows the team's team_member_budget_id.""" - from fastapi import Request - - from litellm.proxy.management_endpoints import team_endpoints - - team_row = _team_with_default_budget("team-1", "team-default-b") - memberships = [ - LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"), - LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), - LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), - ] - - mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) - mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( - return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0) + sources = await _team_info_budget_sources( + team_row=_team_with_default_budget("team-1", "team-default-b"), + memberships=[ + LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"), + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ], + default_budget_row=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0), ) - mock_prisma.get_data = AsyncMock(return_value=[]) - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), - ): - response = await team_endpoints.team_info( - http_request=MagicMock(spec=Request), - team_id="team-1", - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - - assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + assert sources == { "inherits": "team_default", "customized": "custom", "unlinked": "team_default", @@ -14617,30 +14629,36 @@ async def test_team_info_reports_no_budget_source_when_team_has_no_default(): """A team that never set team_member_budget has nothing for members to inherit, so an unlinked member is 'none' rather than 'team_default', while a member with their own row is still 'custom'.""" - from fastapi import Request + sources = await _team_info_budget_sources( + team_row=LiteLLM_TeamTable(team_id="team-1"), + memberships=[ + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ], + default_budget_row=None, + ) - from litellm.proxy.management_endpoints import team_endpoints + assert sources == { + "customized": "custom", + "unlinked": "none", + } - memberships = [ - LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), - LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), - ] - mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")) - mock_prisma.get_data = AsyncMock(return_value=[]) +@pytest.mark.asyncio +async def test_team_info_reports_no_budget_source_when_team_default_row_was_deleted(): + """If the budget row named by team_member_budget_id was removed via /budget/delete, nothing is + enforced for unlinked members any more, so /team/info must not keep advertising a team default + that no longer exists.""" + sources = await _team_info_budget_sources( + team_row=_team_with_default_budget("team-1", "deleted-b"), + memberships=[ + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ], + default_budget_row=None, + ) - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), - ): - response = await team_endpoints.team_info( - http_request=MagicMock(spec=Request), - team_id="team-1", - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - - assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + assert sources == { "customized": "custom", "unlinked": "none", } From 52a71ff68188b0d6781204144472c613c1c8240b Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:58:39 +0000 Subject: [PATCH 007/146] fix(team): let team admins reach the member reset_budget route and cover it in the behavior suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + .../test_team_member_reset_budget.py | 201 ++++++++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 tests/proxy_behavior/management/test_team_member_reset_budget.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a624234cf5f..c21113f8c29 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -865,6 +865,7 @@ class LiteLLMRoutes(enum.Enum): "/management/v1/teams/{team_id}/members/bulk_update", "/team/member_update", "/team/{team_id}/member/{user_id}/reset_spend", + "/team/{team_id}/member/{user_id}/reset_budget", "/team/permissions_list", "/team/permissions_update", "/team/daily/activity", diff --git a/tests/proxy_behavior/management/test_team_member_reset_budget.py b/tests/proxy_behavior/management/test_team_member_reset_budget.py new file mode 100644 index 00000000000..42f327c33ef --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_reset_budget.py @@ -0,0 +1,201 @@ +import uuid + +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_TEAM_DEFAULT_MAX_BUDGET = 100.0 +_CUSTOM_MAX_BUDGET = 50.0 + +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_budget(prisma, budget_id: str, max_budget: float) -> str: + await prisma.db.litellm_budgettable.create( + data={ + "budget_id": budget_id, + "max_budget": max_budget, + "created_by": "phase4-scratch", + "updated_by": "phase4-scratch", + } + ) + return budget_id + + +async def _seed_team_with_default_budget(prisma, world, shape: str, team_id: str, scratch) -> str: + default_budget_id = await _seed_budget(prisma, scratch.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET) + metadata = {"team_member_budget_id": default_budget_id} + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + metadata=metadata, + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id, metadata=metadata) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + return default_budget_id + + +async def _seed_custom_member(prisma, team_id: str, member_id: str, scratch) -> str: + custom_budget_id = await _seed_budget(prisma, scratch.tag("custom-budget"), _CUSTOM_MAX_BUDGET) + await prisma.db.litellm_teammembership.create( + data={ + "user_id": member_id, + "team_id": team_id, + "spend": _SEED_SPEND, + "litellm_budget_table": {"connect": {"budget_id": custom_budget_id}}, + } + ) + return custom_budget_id + + +async def _membership(prisma, team_id: str, member_id: str): + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}} + ) + assert row is not None + return row + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_reset_budget_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + default_budget_id = await _seed_team_with_default_budget(prisma, world, shape, scratch.prefix, scratch) + custom_budget_id = await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + caller = world.keys[actor] + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.spend == _SEED_SPEND, "reset_budget must never touch spend" + if expected_status == 200: + assert row.budget_id == default_budget_id + body = resp.json() + assert body["budget_id"] == default_budget_id + assert body["previous_budget_id"] == custom_budget_id + assert body["budget_source"] == "team_default" + else: + assert row.budget_id == custom_budget_id, "denied but budget relinked" + + +async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world): + """Relinking must point the member at the shared row, not copy or edit it, so a later + /team/update to team_member_budget reaches this member again.""" + member_id = scratch.tag("member") + default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + + default_row = await prisma.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id}) + assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET + + info = await proxy_client.get( + f"/team/info?team_id={scratch.prefix}", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert info.status_code == 200, info.text + memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]} + assert memberships[member_id]["budget_source"] == "team_default" + assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET + + +async def test_team_member_reset_budget_without_team_default_detaches_member(proxy_client, prisma, scratch, world): + member_id = scratch.tag("member") + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["budget_id"] is None + assert resp.json()["budget_source"] == "none" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.budget_id is None + assert row.spend == _SEED_SPEND + + +async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world): + """metadata.team_member_budget_id can outlive its budget row; a stale id must not be + relinked to (the FK would fail) and must read as no budget, not as the team default.""" + member_id = scratch.tag("member") + await create_scratch_team( + prisma, + scratch.prefix, + organization_id=world.org_a_id, + metadata={"team_member_budget_id": scratch.tag("deleted-budget")}, + ) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["budget_id"] is None + assert resp.json()["budget_source"] == "none" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.budget_id is None + + +async def test_team_member_reset_budget_missing_team_is_404(proxy_client, world): + resp = await proxy_client.post( + f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_budget_missing_membership_is_404(proxy_client, prisma, scratch, world): + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 404, resp.text From 8bd9d356dcc10632a8efdc1a2229646e97cb301f Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 01:00:31 +0000 Subject: [PATCH 008/146] test(team): drop docstrings that restate the budget source and reset assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management/test_team_member_reset_budget.py | 4 ---- .../management_endpoints/test_team_endpoints.py | 16 ---------------- 2 files changed, 20 deletions(-) diff --git a/tests/proxy_behavior/management/test_team_member_reset_budget.py b/tests/proxy_behavior/management/test_team_member_reset_budget.py index 42f327c33ef..1e55b8b6b15 100644 --- a/tests/proxy_behavior/management/test_team_member_reset_budget.py +++ b/tests/proxy_behavior/management/test_team_member_reset_budget.py @@ -117,8 +117,6 @@ async def test_team_member_reset_budget_authz_matrix( async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world): - """Relinking must point the member at the shared row, not copy or edit it, so a later - /team/update to team_member_budget reaches this member again.""" member_id = scratch.tag("member") default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch) await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) @@ -161,8 +159,6 @@ async def test_team_member_reset_budget_without_team_default_detaches_member(pro async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world): - """metadata.team_member_budget_id can outlive its budget row; a stale id must not be - relinked to (the FK would fail) and must read as no budget, not as the team default.""" member_id = scratch.tag("member") await create_scratch_team( prisma, 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 484c054fa54..1f19163933a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14443,10 +14443,6 @@ def _team_with_default_budget(team_id: str, budget_id: str) -> LiteLLM_TeamTable @pytest.mark.asyncio async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default(monkeypatch): - """An admin undoing a per-member budget must put the membership back on the team's shared - default row (a connect, not a copy) so later /team/update changes reach the member again, - and must drop the cached membership so the old cap stops being enforced. The shared row and - the member's tracked spend are never written.""" from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache mock_prisma_client = MagicMock() @@ -14498,9 +14494,6 @@ async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default async def test_reset_team_member_budget_fn_detaches_member_when_team_has_no_usable_default( monkeypatch, team_obj, default_row ): - """With no shared default to link to, reset leaves the member exactly where a freshly added - member would be: no budget row at all, reported as budget_source='none', rather than - connecting to a budget_id that does not exist or leaving the custom cap in place.""" mock_prisma_client = MagicMock() membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id="custom-b1") mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) @@ -14604,9 +14597,6 @@ async def _team_info_budget_sources( @pytest.mark.asyncio async def test_team_info_reports_whether_each_member_follows_the_team_default_budget(): - """/team/info must tell the caller which members still follow the team's shared member budget - and which carry their own row, since budget_id alone only means something to a reader who - also knows the team's team_member_budget_id.""" sources = await _team_info_budget_sources( team_row=_team_with_default_budget("team-1", "team-default-b"), memberships=[ @@ -14626,9 +14616,6 @@ async def test_team_info_reports_whether_each_member_follows_the_team_default_bu @pytest.mark.asyncio async def test_team_info_reports_no_budget_source_when_team_has_no_default(): - """A team that never set team_member_budget has nothing for members to inherit, so an - unlinked member is 'none' rather than 'team_default', while a member with their own row is - still 'custom'.""" sources = await _team_info_budget_sources( team_row=LiteLLM_TeamTable(team_id="team-1"), memberships=[ @@ -14646,9 +14633,6 @@ async def test_team_info_reports_no_budget_source_when_team_has_no_default(): @pytest.mark.asyncio async def test_team_info_reports_no_budget_source_when_team_default_row_was_deleted(): - """If the budget row named by team_member_budget_id was removed via /budget/delete, nothing is - enforced for unlinked members any more, so /team/info must not keep advertising a team default - that no longer exists.""" sources = await _team_info_budget_sources( team_row=_team_with_default_budget("team-1", "deleted-b"), memberships=[ From 5f54f87d9887c13b75c42ffca5142ef2b7bcf2f1 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:40:10 +0000 Subject: [PATCH 009/146] feat(fal_ai): add Seedance video generation via fal queue API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/__init__.py | 3 + litellm/llms/fal_ai/videos/transformation.py | 512 ++++++++++++++++++ ...odel_prices_and_context_window_backup.json | 121 +++++ litellm/utils.py | 4 + model_prices_and_context_window.json | 121 +++++ .../test_fal_ai_video_transformation.py | 231 ++++++++ 6 files changed, 992 insertions(+) create mode 100644 litellm/llms/fal_ai/videos/__init__.py create mode 100644 litellm/llms/fal_ai/videos/transformation.py create mode 100644 tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py diff --git a/litellm/llms/fal_ai/videos/__init__.py b/litellm/llms/fal_ai/videos/__init__.py new file mode 100644 index 00000000000..c7e8f76c75b --- /dev/null +++ b/litellm/llms/fal_ai/videos/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig + +__all__ = ("FalAIVideoConfig",) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py new file mode 100644 index 00000000000..f8ebf828d68 --- /dev/null +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -0,0 +1,512 @@ +import math +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from httpx._types import FileContent, RequestFiles +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared HTTP factory is private + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.main import ( + CharacterObject, + VideoCreateOptionalRequestParams, + VideoObject, +) +from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_video_id_with_provider, +) + + +class FalAIVideoError(BaseLLMException): + pass + + +_ALLOWED_ASPECT_RATIOS: Final[frozenset[str]] = frozenset({"auto", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}) +_ALLOWED_RESOLUTIONS: Final[frozenset[str]] = frozenset({"480p", "720p", "1080p", "4k"}) +_RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( + (480, "480p"), + (720, "720p"), + (1080, "1080p"), +) +_FAL_AI_PROVIDER: Final[str] = LlmProviders.FAL_AI.value + + +def _queue_request_base_path(model: str) -> str: + segments: Final[tuple[str, ...]] = tuple(model.split("/")) + segment_count: Final[int] = 3 if segments and segments[0] in frozenset(("workflows", "comfy")) else 2 + return "/".join(segments[:segment_count]) + + +def _duration_value(value: object) -> str | None: + if isinstance(value, str) and value == "auto": + return value + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + return str(int(float(value))) + except (TypeError, ValueError): + return None + + +def _resolution_for_height(height: int) -> str: + return next((resolution for threshold, resolution in _RESOLUTION_TIERS if height <= threshold), "4k") + + +def _size_params(size: object) -> Mapping[str, str]: + if not isinstance(size, str): + return MappingProxyType({}) + if size in _ALLOWED_RESOLUTIONS: + return MappingProxyType({"resolution": size}) + if size.count("x") != 1: + return MappingProxyType({}) + width_text, height_text = size.split("x") + if not (width_text.isdigit() and height_text.isdigit()): + return MappingProxyType({}) + width: Final[int] = int(width_text) + height: Final[int] = int(height_text) + if width <= 0 or height <= 0: + return MappingProxyType({}) + reduced_gcd: Final[int] = math.gcd(width, height) + aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" + resolution: Final[str] = _resolution_for_height(height) + if aspect_ratio in _ALLOWED_ASPECT_RATIOS: + return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) + return MappingProxyType({"resolution": resolution}) + + +def _numeric_duration(value: object) -> float | None: + duration: Final[str | None] = _duration_value(value) + if duration is None or duration == "auto": + return None + return float(duration) + + +def _response_data(raw_response: httpx.Response) -> Mapping[str, object]: + return TypeAdapter(Mapping[str, object]).validate_python(raw_response.json()) + + +def _response_string(response_data: Mapping[str, object], key: str, default: str = "") -> str: + value: Final[object] = response_data.get(key) + return value if isinstance(value, str) else default + + +class FalAIVideoConfig(BaseVideoConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: API contract requires a list + return [ # mutable-ok: API contract requires a list + "model", + "prompt", + "input_reference", + "seconds", + "size", + "user", + "extra_headers", + ] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseVideoConfig requires a mutable mapping + supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) + input_reference: Final[object] = video_create_optional_params.get("input_reference") + input_reference_params: Final[Mapping[str, str]] = ( + MappingProxyType({}) + if "input_reference" not in video_create_optional_params + else ( + MappingProxyType({"image_url": input_reference}) + if isinstance(input_reference, str) + else self._invalid_input_reference() + ) + ) + duration_params: Final[Mapping[str, str]] = ( + MappingProxyType({}) + if "seconds" not in video_create_optional_params + else self._duration_params(video_create_optional_params["seconds"]) + ) + size_params: Final[Mapping[str, str]] = ( + self._size_params(video_create_optional_params["size"]) + if "size" in video_create_optional_params + else MappingProxyType({}) + ) + user_params: Final[Mapping[str, str]] = ( + MappingProxyType({"end_user_id": user}) + if isinstance(user := video_create_optional_params.get("user"), str) + else MappingProxyType({}) + ) + return dict( # mutable-ok: BaseVideoConfig requires a mutable mapping + MappingProxyType( + { + **input_reference_params, + **duration_params, + **size_params, + **user_params, + **{ # mutable-ok: dynamic passthrough fields require a mapping + key: value for key, value in video_create_optional_params.items() if key not in supported_params + }, + } + ) + ) # mutable-ok: BaseVideoConfig requires a mutable mapping + + @staticmethod + def _invalid_input_reference() -> Mapping[str, str]: + raise ValueError("fal.ai needs a public image URL for input_reference") + + @staticmethod + def _duration_params(seconds: object) -> Mapping[str, str]: + duration: Final[str | None] = _duration_value(seconds) + if duration is None: + raise ValueError("fal.ai seconds must be a numeric value") + return MappingProxyType({"duration": duration}) + + @staticmethod + def _size_params(size: object) -> Mapping[str, str]: + return _size_params(size) + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable headers + final_api_key: Final[str | None] = ( + api_key + or (litellm_params.api_key if litellm_params is not None else None) + or get_secret_str("FAL_AI_API_KEY") + or get_secret_str("FAL_KEY") + ) + if not final_api_key: + raise ValueError("fal.ai API key is required") + return dict( # mutable-ok: BaseVideoConfig requires mutable headers + MappingProxyType( + { + **headers, + "Authorization": f"Key {final_api_key}", + "Content-Type": "application/json", + } + ) + ) # mutable-ok: BaseVideoConfig requires mutable headers + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig requires mutable parameters + ) -> str: + return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict[ # mutable-ok: BaseVideoConfig requires mutable parameters + str, object + ], # mutable-ok: BaseVideoConfig requires mutable parameters + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_data: Final[dict[str, object]] = dict( # mutable-ok: HTTP JSON payload requires mutable data + MappingProxyType( + { + "prompt": prompt, + **{ # mutable-ok: dynamic request fields require a mapping + key: value for key, value in video_create_optional_request_params.items() if key != "model" + }, + } + ) + ) + return request_data, [], f"{api_base.rstrip('/')}/{model}" # mutable-ok: HTTP files payload requires a list + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + request_params: Final[Mapping[str, object]] = request_data or MappingProxyType({}) + request_id: Final[str] = _response_string(response_data, "request_id") + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + duration: Final[float | None] = _numeric_duration(request_params.get("duration")) + resolution: Final[object] = request_params.get("resolution") + seconds: Final[str | None] = _duration_value(request_params["duration"]) if duration is not None else None + size: Final[str | None] = resolution if isinstance(resolution, str) else None + usage: Final[dict[str, object]] = dict( # mutable-ok: VideoObject requires a mutable usage mapping + MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", duration), + ("video_resolution", resolution if isinstance(resolution, str) else "720p"), + ) + if value is not None + } + ) + ) # mutable-ok: VideoObject requires a mutable usage mapping + video_object: Final[VideoObject] = VideoObject( + id=encode_video_id_with_provider(request_id, provider, model), + object="video", + status="queued", + created_at=int(time.time()), + model=model, + seconds=seconds, + size=size, + ) + video_object.usage = usage + return video_object + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_id, model_id = self._decode_video_id(video_id) + encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") + return ( + f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status", + {}, # mutable-ok: BaseVideoConfig requires a mutable mapping + ) + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") + status: Final[str] = MappingProxyType( + { + "IN_QUEUE": "queued", + "IN_PROGRESS": "in_progress", + "COMPLETED": "completed", + } + ).get(raw_status, "queued") + error_value: Final[object] = response_data.get("error") + error: Final[str | None] = error_value if isinstance(error_value, str) else None + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + return VideoObject( + id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider), + object="video", + status="failed" if error else status, + created_at=0, + error=( + {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict + ), # mutable-ok: VideoObject requires a dict + ) + + @staticmethod + def _decode_video_id(video_id: str) -> tuple[str, str]: + decoded: Final = decode_video_id_with_provider(video_id) + request_id: Final[str] = decoded.get("video_id", video_id) + model_id: Final[str | None] = decoded.get("model_id") + if not model_id: + raise ValueError("fal.ai video ids must be created through litellm with a model") + return request_id, model_id + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + variant: str | None = None, + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_id, model_id = self._decode_video_id(video_id) + encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") + return ( + f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}", + {}, # mutable-ok: BaseVideoConfig requires a mutable mapping + ) + + @staticmethod + def _extract_video_url(response_data: Mapping[str, object]) -> str: + raw_video_data: Final[object] = response_data.get("video") + video_data: Final[Mapping[str, object] | None] = ( + TypeAdapter(Mapping[str, object]).validate_python(raw_video_data) + if isinstance(raw_video_data, Mapping) + else None + ) + if video_data is not None: + video_url: Final[object] = video_data.get("url") + if isinstance(video_url, str) and video_url: + return video_url + error_message: Final[str | None] = next( + (value for key in ("error", "detail") if isinstance(value := response_data.get(key), str)), + None, + ) + if error_message: + raise ValueError(f"fal.ai video result did not include a video URL: {error_message}") + raise ValueError("fal.ai video result did not include a video URL") + + def transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) + httpx_client: Final[HTTPHandler] = _get_httpx_client() + video_response: Final[httpx.Response] = httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + video_url + ) + video_response.raise_for_status() + return video_response.content + + async def async_transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) + async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI) + video_response: Final[httpx.Response] = await async_httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + video_url + ) + video_response.raise_for_status() + return video_response.content + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video remix is not supported for fal.ai") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video remix is not supported for fal.ai") + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video listing is not supported for fal.ai") + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video listing is not supported for fal.ai") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video delete is not supported for fal.ai") + + def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: + raise NotImplementedError("video delete is not supported for fal.ai") + + def transform_video_create_character_request( + self, + name: str, + video: object, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, list[object]]: # mutable-ok: BaseVideoConfig requires mutable lists + raise NotImplementedError("video character creation is not supported for fal.ai") + + def transform_video_create_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character creation is not supported for fal.ai") + + def transform_video_get_character_request( + self, + character_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video character retrieval is not supported for fal.ai") + + def transform_video_get_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character retrieval is not supported for fal.ai") + + def transform_video_edit_request( + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + video_file: FileContent | None = None, + extra_body: Mapping[str, object] | None = None, + prefetched_source_data: Mapping[str, object] | None = None, + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: + raise NotImplementedError("video edit is not supported for fal.ai") + + def transform_video_edit_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + raise NotImplementedError("video edit is not supported for fal.ai") + + def transform_video_extension_request( + self, + prompt: str, + video_id: str, + seconds: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video extension is not supported for fal.ai") + + def transform_video_extension_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video extension is not supported for fal.ai") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseLLMException requires mutable headers + ) -> BaseLLMException: + return FalAIVideoError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4dbf0337894..324eb6b2d66 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22807,6 +22807,127 @@ "/v1/images/generations" ] }, + "fal_ai/bytedance/seedance-2.5/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/image-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/text-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/image-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/reference-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/fal-ai/ideogram/v3": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/litellm/utils.py b/litellm/utils.py index 48d13bc16af..3991cecdac6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9403,6 +9403,10 @@ class ProviderConfigManager: from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig return RunwayMLVideoConfig() + elif LlmProviders.FAL_AI == provider: + from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig + + return FalAIVideoConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4dbf0337894..324eb6b2d66 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22807,6 +22807,127 @@ "/v1/images/generations" ] }, + "fal_ai/bytedance/seedance-2.5/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/image-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/text-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/image-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/reference-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/fal-ai/ideogram/v3": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py new file mode 100644 index 00000000000..d4b058ddcf6 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -0,0 +1,231 @@ +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +import litellm.llms.fal_ai.videos.transformation as fal_video_module +from litellm.cost_calculator import default_video_cost_calculator +from litellm.llms.fal_ai.videos.transformation import ( + FalAIVideoConfig, + FalAIVideoError, + _queue_request_base_path, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.utils import decode_video_id_with_provider +from litellm.utils import ProviderConfigManager + +MODEL = "bytedance/seedance-2.5/text-to-video" + + +class TestFalAIVideoTransformation: + def setup_method(self): + self.config = FalAIVideoConfig() + self.logging_obj = Mock() + + def test_map_openai_params(self): + mapped = self.config.map_openai_params( + { + "seconds": "5", + "size": "1280x720", + "input_reference": "https://example.com/image.png", + "user": "user-123", + "generate_audio": False, + }, + MODEL, + False, + ) + + assert mapped == { + "duration": "5", + "resolution": "720p", + "aspect_ratio": "16:9", + "image_url": "https://example.com/image.png", + "end_user_id": "user-123", + "generate_audio": False, + } + + assert self.config.map_openai_params({"size": "1080x1080"}, MODEL, False) == { + "resolution": "1080p", + "aspect_ratio": "1:1", + } + assert self.config.map_openai_params({"size": "720p"}, MODEL, False) == {"resolution": "720p"} + + def test_map_openai_params_rejects_non_url_input_reference(self): + with pytest.raises(ValueError, match="public image URL"): + self.config.map_openai_params({"input_reference": b"image"}, MODEL, False) + + def test_transform_video_create_request(self): + body, files, url = self.config.transform_video_create_request( + model=MODEL, + prompt="A quiet ocean at sunrise", + api_base="https://queue.fal.run", + video_create_optional_request_params={ + "duration": "5", + "resolution": "480p", + "aspect_ratio": "16:9", + "generate_audio": False, + "model": MODEL, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == f"https://queue.fal.run/{MODEL}" + assert files == [] + assert body == { + "prompt": "A quiet ocean at sunrise", + "duration": "5", + "resolution": "480p", + "aspect_ratio": "16:9", + "generate_audio": False, + } + assert "model" not in body + + def test_transform_video_create_response_encodes_model_and_usage(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + + video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": "5", "resolution": "480p"}, + ) + + decoded = decode_video_id_with_provider(video.id) + assert decoded["custom_llm_provider"] == "fal_ai" + assert decoded["model_id"] == MODEL + assert decoded["video_id"] == "abc" + assert video.status == "queued" + assert video.usage == {"duration_seconds": 5.0, "video_resolution": "480p"} + + auto_video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": "auto"}, + ) + assert auto_video.usage == {"video_resolution": "720p"} + assert auto_video.seconds is None + assert auto_video.size is None + + def test_status_request_uses_queue_base_path(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={}, + ) + + url, params = self.config.transform_video_status_retrieve_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + assert params == {} + assert _queue_request_base_path("workflows/owner/app/x") == "workflows/owner/app" + assert _queue_request_base_path("comfy/owner/app/x") == "comfy/owner/app" + + def test_status_request_rejects_unencoded_video_id(self): + with pytest.raises(ValueError, match="must be created through litellm"): + self.config.transform_video_status_retrieve_request( + video_id="abc", + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + @pytest.mark.parametrize( + ("response_data", "expected_status"), + [ + ({"request_id": "abc", "status": "IN_QUEUE"}, "queued"), + ({"request_id": "abc", "status": "IN_PROGRESS"}, "in_progress"), + ({"request_id": "abc", "status": "COMPLETED"}, "completed"), + ], + ) + def test_status_response_mapping(self, response_data, expected_status): + response = Mock(spec=httpx.Response) + response.json.return_value = response_data + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == expected_status + assert video.created_at == 0 + + def test_status_response_error(self): + response = Mock(spec=httpx.Response) + response.json.return_value = { + "request_id": "abc", + "status": "COMPLETED", + "error": "generation failed", + } + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "failed" + assert video.error == {"code": "fal_error", "message": "generation failed"} + + def test_content_response_downloads_video_url(self, monkeypatch): + content_response = httpx.Response( + 200, + content=b"video-bytes", + request=httpx.Request("GET", "https://cdn.example.com/video.mp4"), + ) + + class FakeHTTPClient: + def get(self, url): + assert url == "https://cdn.example.com/video.mp4" + return content_response + + monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: FakeHTTPClient()) + response = Mock(spec=httpx.Response) + response.json.return_value = {"video": {"url": "https://cdn.example.com/video.mp4"}} + + assert self.config.transform_video_content_response(response, self.logging_obj) == b"video-bytes" + + def test_content_response_rejects_missing_video(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"error": "generation failed"} + + with pytest.raises(ValueError, match="generation failed"): + self.config.transform_video_content_response(response, self.logging_obj) + + def test_provider_config_and_error_class(self): + provider_config = ProviderConfigManager.get_provider_video_config( + model=MODEL, + provider=LlmProviders.FAL_AI, + ) + assert isinstance(provider_config, FalAIVideoConfig) + assert isinstance(self.config.get_error_class("bad key", 401, {}), FalAIVideoError) + + def test_video_cost_uses_tiered_rows(self): + rows = { + model: row + for model, row in litellm.model_cost.items() + if row.get("litellm_provider") == "fal_ai" and row.get("mode") == "video_generation" + } + assert rows + for model, row in rows.items(): + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="480p") == ( + 5 * row["output_cost_per_second_480p"] + ) + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="720p") == ( + 5 * row["output_cost_per_second"] + ) From 0f4ce95492cde57b1f2eb29a1ea119dda0b13e24 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:45:53 +0000 Subject: [PATCH 010/146] refactor(fal_ai): simplify video config mappings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 178 +++++++++---------- 1 file changed, 80 insertions(+), 98 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index f8ebf828d68..f0142ddc3de 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -2,7 +2,7 @@ import math import time from collections.abc import Mapping from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import httpx from httpx._types import FileContent, RequestFiles @@ -42,12 +42,25 @@ _RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( (720, "720p"), (1080, "1080p"), ) +_QUEUE_NAMESPACES: Final[frozenset[str]] = frozenset(("workflows", "comfy")) +_STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "IN_QUEUE": "queued", + "IN_PROGRESS": "in_progress", + "COMPLETED": "completed", + } +) _FAL_AI_PROVIDER: Final[str] = LlmProviders.FAL_AI.value +_SupportedParams: TypeAlias = list[str] +_VideoParams: TypeAlias = dict[str, object] +_VideoHeaders: TypeAlias = dict[str, str] +_VideoStringParams: TypeAlias = dict[str, str] +_VideoFiles: TypeAlias = list[object] def _queue_request_base_path(model: str) -> str: segments: Final[tuple[str, ...]] = tuple(model.split("/")) - segment_count: Final[int] = 3 if segments and segments[0] in frozenset(("workflows", "comfy")) else 2 + segment_count: Final[int] = 3 if segments and segments[0] in _QUEUE_NAMESPACES else 2 return "/".join(segments[:segment_count]) @@ -105,8 +118,8 @@ def _response_string(response_data: Mapping[str, object], key: str, default: str class FalAIVideoConfig(BaseVideoConfig): - def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: API contract requires a list - return [ # mutable-ok: API contract requires a list + def get_supported_openai_params(self, model: str) -> _SupportedParams: + supported_params: Final[_SupportedParams] = [ # mutable-ok: BaseVideoConfig requires a list "model", "prompt", "input_reference", @@ -115,23 +128,22 @@ class FalAIVideoConfig(BaseVideoConfig): "user", "extra_headers", ] + return supported_params def map_openai_params( self, video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, object]: # mutable-ok: BaseVideoConfig requires a mutable mapping + ) -> _VideoParams: supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) input_reference: Final[object] = video_create_optional_params.get("input_reference") + if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): + raise ValueError("fal.ai needs a public image URL for input_reference") input_reference_params: Final[Mapping[str, str]] = ( MappingProxyType({}) - if "input_reference" not in video_create_optional_params - else ( - MappingProxyType({"image_url": input_reference}) - if isinstance(input_reference, str) - else self._invalid_input_reference() - ) + if not isinstance(input_reference, str) + else MappingProxyType({"image_url": input_reference}) ) duration_params: Final[Mapping[str, str]] = ( MappingProxyType({}) @@ -139,7 +151,7 @@ class FalAIVideoConfig(BaseVideoConfig): else self._duration_params(video_create_optional_params["seconds"]) ) size_params: Final[Mapping[str, str]] = ( - self._size_params(video_create_optional_params["size"]) + _size_params(video_create_optional_params["size"]) if "size" in video_create_optional_params else MappingProxyType({}) ) @@ -148,23 +160,16 @@ class FalAIVideoConfig(BaseVideoConfig): if isinstance(user := video_create_optional_params.get("user"), str) else MappingProxyType({}) ) - return dict( # mutable-ok: BaseVideoConfig requires a mutable mapping - MappingProxyType( - { - **input_reference_params, - **duration_params, - **size_params, - **user_params, - **{ # mutable-ok: dynamic passthrough fields require a mapping - key: value for key, value in video_create_optional_params.items() if key not in supported_params - }, - } - ) - ) # mutable-ok: BaseVideoConfig requires a mutable mapping - - @staticmethod - def _invalid_input_reference() -> Mapping[str, str]: - raise ValueError("fal.ai needs a public image URL for input_reference") + mapped_params: Final[_VideoParams] = { + **input_reference_params, + **duration_params, + **size_params, + **user_params, + **{ # mutable-ok: BaseVideoConfig requires a mutable parameter mapping + key: value for key, value in video_create_optional_params.items() if key not in supported_params + }, + } + return mapped_params @staticmethod def _duration_params(seconds: object) -> Mapping[str, str]: @@ -173,17 +178,13 @@ class FalAIVideoConfig(BaseVideoConfig): raise ValueError("fal.ai seconds must be a numeric value") return MappingProxyType({"duration": duration}) - @staticmethod - def _size_params(size: object) -> Mapping[str, str]: - return _size_params(size) - def validate_environment( self, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, model: str, api_key: str | None = None, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable headers + ) -> _VideoHeaders: final_api_key: Final[str | None] = ( api_key or (litellm_params.api_key if litellm_params is not None else None) @@ -192,21 +193,18 @@ class FalAIVideoConfig(BaseVideoConfig): ) if not final_api_key: raise ValueError("fal.ai API key is required") - return dict( # mutable-ok: BaseVideoConfig requires mutable headers - MappingProxyType( - { - **headers, - "Authorization": f"Key {final_api_key}", - "Content-Type": "application/json", - } - ) - ) # mutable-ok: BaseVideoConfig requires mutable headers + validated_headers: Final[_VideoHeaders] = { + **headers, + "Authorization": f"Key {final_api_key}", + "Content-Type": "application/json", + } + return validated_headers def get_complete_url( self, model: str, api_base: str | None, - litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig requires mutable parameters + litellm_params: _VideoParams, ) -> str: return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") @@ -215,22 +213,16 @@ class FalAIVideoConfig(BaseVideoConfig): model: str, prompt: str, api_base: str, - video_create_optional_request_params: dict[ # mutable-ok: BaseVideoConfig requires mutable parameters - str, object - ], # mutable-ok: BaseVideoConfig requires mutable parameters + video_create_optional_request_params: _VideoParams, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig requires mutable mappings - request_data: Final[dict[str, object]] = dict( # mutable-ok: HTTP JSON payload requires mutable data - MappingProxyType( - { - "prompt": prompt, - **{ # mutable-ok: dynamic request fields require a mapping - key: value for key, value in video_create_optional_request_params.items() if key != "model" - }, - } - ) - ) + headers: _VideoHeaders, + ) -> tuple[_VideoParams, RequestFiles, str]: + request_data: Final[_VideoParams] = { + "prompt": prompt, + **{ # mutable-ok: HTTP JSON payload requires a mutable mapping + key: value for key, value in video_create_optional_request_params.items() if key != "model" + }, + } return request_data, [], f"{api_base.rstrip('/')}/{model}" # mutable-ok: HTTP files payload requires a list def transform_video_create_response( @@ -249,18 +241,14 @@ class FalAIVideoConfig(BaseVideoConfig): resolution: Final[object] = request_params.get("resolution") seconds: Final[str | None] = _duration_value(request_params["duration"]) if duration is not None else None size: Final[str | None] = resolution if isinstance(resolution, str) else None - usage: Final[dict[str, object]] = dict( # mutable-ok: VideoObject requires a mutable usage mapping - MappingProxyType( - { - key: value - for key, value in ( - ("duration_seconds", duration), - ("video_resolution", resolution if isinstance(resolution, str) else "720p"), - ) - if value is not None - } + usage: Final[_VideoParams] = { # mutable-ok: VideoObject requires a mutable usage mapping + key: value + for key, value in ( + ("duration_seconds", duration), + ("video_resolution", resolution if isinstance(resolution, str) else "720p"), ) - ) # mutable-ok: VideoObject requires a mutable usage mapping + if value is not None + } video_object: Final[VideoObject] = VideoObject( id=encode_video_id_with_provider(request_id, provider, model), object="video", @@ -278,8 +266,8 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( @@ -295,13 +283,7 @@ class FalAIVideoConfig(BaseVideoConfig): ) -> VideoObject: response_data: Final[Mapping[str, object]] = _response_data(raw_response) raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") - status: Final[str] = MappingProxyType( - { - "IN_QUEUE": "queued", - "IN_PROGRESS": "in_progress", - "COMPLETED": "completed", - } - ).get(raw_status, "queued") + status: Final[str] = _STATUS_MAP.get(raw_status, "queued") error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER @@ -312,7 +294,7 @@ class FalAIVideoConfig(BaseVideoConfig): created_at=0, error=( {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict - ), # mutable-ok: VideoObject requires a dict + ), ) @staticmethod @@ -329,9 +311,9 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, variant: str | None = None, - ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoStringParams]: request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( @@ -383,9 +365,9 @@ class FalAIVideoConfig(BaseVideoConfig): prompt: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, extra_body: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video remix is not supported for fal.ai") def transform_video_remix_response( @@ -400,12 +382,12 @@ class FalAIVideoConfig(BaseVideoConfig): self, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, after: str | None = None, limit: int | None = None, order: str | None = None, extra_query: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video listing is not supported for fal.ai") def transform_video_list_response( @@ -413,7 +395,7 @@ class FalAIVideoConfig(BaseVideoConfig): raw_response: httpx.Response, logging_obj: object, custom_llm_provider: str | None = None, - ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> _VideoStringParams: raise NotImplementedError("video listing is not supported for fal.ai") def transform_video_delete_request( @@ -421,8 +403,8 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video delete is not supported for fal.ai") def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: @@ -434,8 +416,8 @@ class FalAIVideoConfig(BaseVideoConfig): video: object, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, list[object]]: # mutable-ok: BaseVideoConfig requires mutable lists + headers: _VideoHeaders, + ) -> tuple[str, _VideoFiles]: raise NotImplementedError("video character creation is not supported for fal.ai") def transform_video_create_character_response( @@ -450,8 +432,8 @@ class FalAIVideoConfig(BaseVideoConfig): character_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video character retrieval is not supported for fal.ai") def transform_video_get_character_response( @@ -467,7 +449,7 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, video_file: FileContent | None = None, extra_body: Mapping[str, object] | None = None, prefetched_source_data: Mapping[str, object] | None = None, @@ -490,9 +472,9 @@ class FalAIVideoConfig(BaseVideoConfig): seconds: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, extra_body: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video extension is not supported for fal.ai") def transform_video_extension_response( @@ -507,6 +489,6 @@ class FalAIVideoConfig(BaseVideoConfig): self, error_message: str, status_code: int, - headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseLLMException requires mutable headers + headers: _VideoHeaders | httpx.Headers, ) -> BaseLLMException: return FalAIVideoError(status_code=status_code, message=error_message, headers=headers) From 141548dcf3ec7588e86c2e304fca33be6c4cdc7e Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:54:11 +0000 Subject: [PATCH 011/146] fix(fal_ai): keep status ids pollable and size resolution by the short side Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 17 +++++++-- .../test_fal_ai_video_transformation.py | 37 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index f0142ddc3de..4ca1f918c90 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -75,8 +75,16 @@ def _duration_value(value: object) -> str | None: return None -def _resolution_for_height(height: int) -> str: - return next((resolution for threshold, resolution in _RESOLUTION_TIERS if height <= threshold), "4k") +def _resolution_for_short_side(short_side: int) -> str: + return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") + + +def _model_path_from_queue_url(url: object) -> str | None: + if not isinstance(url, str) or not url: + return None + path: Final[str] = httpx.URL(url).path.strip("/") + model_path, separator, _ = path.partition("/requests/") + return model_path if separator and model_path else None def _size_params(size: object) -> Mapping[str, str]: @@ -95,7 +103,7 @@ def _size_params(size: object) -> Mapping[str, str]: return MappingProxyType({}) reduced_gcd: Final[int] = math.gcd(width, height) aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" - resolution: Final[str] = _resolution_for_height(height) + resolution: Final[str] = _resolution_for_short_side(min(width, height)) if aspect_ratio in _ALLOWED_ASPECT_RATIOS: return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) return MappingProxyType({"resolution": resolution}) @@ -287,8 +295,9 @@ class FalAIVideoConfig(BaseVideoConfig): error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + model_path: Final[str | None] = _model_path_from_queue_url(response_data.get("response_url")) return VideoObject( - id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider), + id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider, model_path), object="video", status="failed" if error else status, created_at=0, diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index d4b058ddcf6..f367fa331d5 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -51,6 +51,14 @@ class TestFalAIVideoTransformation: "aspect_ratio": "1:1", } assert self.config.map_openai_params({"size": "720p"}, MODEL, False) == {"resolution": "720p"} + assert self.config.map_openai_params({"size": "720x1280"}, MODEL, False) == { + "resolution": "720p", + "aspect_ratio": "9:16", + } + assert self.config.map_openai_params({"size": "1080x1920"}, MODEL, False) == { + "resolution": "1080p", + "aspect_ratio": "9:16", + } def test_map_openai_params_rejects_non_url_input_reference(self): with pytest.raises(ValueError, match="public image URL"): @@ -165,6 +173,35 @@ class TestFalAIVideoTransformation: assert video.status == expected_status assert video.created_at == 0 + def test_status_response_id_stays_pollable(self): + response = Mock(spec=httpx.Response) + response.json.return_value = { + "request_id": "abc", + "status": "IN_PROGRESS", + "response_url": "https://queue.fal.run/bytedance/seedance-2.5/requests/abc", + } + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + status_url, _ = self.config.transform_video_status_retrieve_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + content_url, _ = self.config.transform_video_content_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert status_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + assert content_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc" + def test_status_response_error(self): response = Mock(spec=httpx.Response) response.json.return_value = { From c359ef763eae44523cace811bcb1999992250bde Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:01:09 +0000 Subject: [PATCH 012/146] fix(fal_ai): keep model in polled video ids and pick resolution from the short side Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 28 ++++++-- .../test_fal_ai_video_transformation.py | 64 +++++++++++-------- 2 files changed, 58 insertions(+), 34 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 4ca1f918c90..766316a5b18 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -79,12 +79,22 @@ def _resolution_for_short_side(short_side: int) -> str: return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") -def _model_path_from_queue_url(url: object) -> str | None: - if not isinstance(url, str) or not url: +def _model_path_from_request_url(raw_response: httpx.Response) -> str | None: + segments: Final[tuple[str, ...]] = tuple(segment for segment in raw_response.request.url.path.split("/") if segment) + if "requests" not in segments: return None - path: Final[str] = httpx.URL(url).path.strip("/") - model_path, separator, _ = path.partition("/requests/") - return model_path if separator and model_path else None + model_segments: Final[tuple[str, ...]] = segments[: segments.index("requests")] + segment_count: Final[int] = 3 if len(model_segments) >= 3 and model_segments[-3] in _QUEUE_NAMESPACES else 2 + return "/".join(model_segments[-segment_count:]) if len(model_segments) >= segment_count else None + + +def _request_id_from_request_url(raw_response: httpx.Response) -> str | None: + segments: Final[tuple[str, ...]] = tuple(segment for segment in raw_response.request.url.path.split("/") if segment) + if "requests" not in segments: + return None + request_index: Final[int] = segments.index("requests") + request_id_index: Final[int] = request_index + 1 + return segments[request_id_index] if len(segments) > request_id_index else None def _size_params(size: object) -> Mapping[str, str]: @@ -295,12 +305,16 @@ class FalAIVideoConfig(BaseVideoConfig): error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER - model_path: Final[str | None] = _model_path_from_queue_url(response_data.get("response_url")) + model_path: Final[str | None] = _model_path_from_request_url(raw_response) + request_id: Final[str] = _response_string(response_data, "request_id") or ( + _request_id_from_request_url(raw_response) or "" + ) return VideoObject( - id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider, model_path), + id=encode_video_id_with_provider(request_id, provider, model_path), object="video", status="failed" if error else status, created_at=0, + model=model_path, error=( {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict ), diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index f367fa331d5..8e0c68e30bb 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -161,8 +161,8 @@ class TestFalAIVideoTransformation: ], ) def test_status_response_mapping(self, response_data, expected_status): - response = Mock(spec=httpx.Response) - response.json.return_value = response_data + status_url = "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + response = httpx.Response(200, json=response_data, request=httpx.Request("GET", status_url)) video = self.config.transform_video_status_retrieve_response( raw_response=response, @@ -172,43 +172,32 @@ class TestFalAIVideoTransformation: assert video.status == expected_status assert video.created_at == 0 + decoded = decode_video_id_with_provider(video.id) + assert decoded["model_id"] == "bytedance/seedance-2.5" + assert decoded["video_id"] == "abc" - def test_status_response_id_stays_pollable(self): - response = Mock(spec=httpx.Response) - response.json.return_value = { - "request_id": "abc", - "status": "IN_PROGRESS", - "response_url": "https://queue.fal.run/bytedance/seedance-2.5/requests/abc", - } - - video = self.config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=self.logging_obj, - custom_llm_provider="fal_ai", - ) - - status_url, _ = self.config.transform_video_status_retrieve_request( + poll_url, _ = self.config.transform_video_status_retrieve_request( video_id=video.id, api_base="https://queue.fal.run", litellm_params=GenericLiteLLMParams(), headers={}, ) - content_url, _ = self.config.transform_video_content_request( - video_id=video.id, - api_base="https://queue.fal.run", - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - assert status_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" - assert content_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc" + assert poll_url == status_url def test_status_response_error(self): - response = Mock(spec=httpx.Response) - response.json.return_value = { + response_data = { "request_id": "abc", "status": "COMPLETED", "error": "generation failed", } + response = httpx.Response( + 200, + json=response_data, + request=httpx.Request( + "GET", + "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status", + ), + ) video = self.config.transform_video_status_retrieve_response( raw_response=response, @@ -219,6 +208,27 @@ class TestFalAIVideoTransformation: assert video.status == "failed" assert video.error == {"code": "fal_error", "message": "generation failed"} + def test_status_response_uses_namespaced_request_url(self): + response = httpx.Response( + 200, + json={"status": "IN_PROGRESS"}, + request=httpx.Request( + "GET", + "https://example.com/proxy/workflows/owner/app/requests/xyz/status", + ), + ) + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + decoded = decode_video_id_with_provider(video.id) + assert decoded["model_id"] == "workflows/owner/app" + assert decoded["video_id"] == "xyz" + assert video.model == "workflows/owner/app" + def test_content_response_downloads_video_url(self, monkeypatch): content_response = httpx.Response( 200, From aa5f0858f75b3e074264e0266f87b70a6cb70391 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:13:41 +0000 Subject: [PATCH 013/146] test(pricing): allow video endpoint and rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b40c10de428..d694c08510a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -940,6 +940,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "/v1/audio/transcriptions", "/v1/audio/speech", "/v1/ocr", + "/v1/videos", "/vertex_ai/live", "/v1/listen", "/v1beta/interactions", @@ -1069,6 +1070,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): # Add any model IDs that should be exempt from the cost validation # Example: "expensive-model-id", "runwayml/seedance2", # 4K output is 150 credits/second = $1.50/second + "fal_ai/bytedance/seedance-2.0/text-to-video", + "fal_ai/bytedance/seedance-2.0/image-to-video", + "fal_ai/bytedance/seedance-2.0/reference-to-video", ] is_valid, violations = validate_model_cost_values(actual_json, exceptions) From e0b455e94e83dfedad0364aedc6b4cdfcb59acb0 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:27:44 +0000 Subject: [PATCH 014/146] fix(fal_ai): read only the documented FAL_AI_API_KEY env var Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 5 ++--- .../test_fal_ai_video_transformation.py | 20 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 766316a5b18..98528c82f6b 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -207,10 +207,9 @@ class FalAIVideoConfig(BaseVideoConfig): api_key or (litellm_params.api_key if litellm_params is not None else None) or get_secret_str("FAL_AI_API_KEY") - or get_secret_str("FAL_KEY") ) if not final_api_key: - raise ValueError("fal.ai API key is required") + raise ValueError("FAL_AI_API_KEY is not set") validated_headers: Final[_VideoHeaders] = { **headers, "Authorization": f"Key {final_api_key}", @@ -224,7 +223,7 @@ class FalAIVideoConfig(BaseVideoConfig): api_base: str | None, litellm_params: _VideoParams, ) -> str: - return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") + return (api_base or "https://queue.fal.run").rstrip("/") def transform_video_create_request( self, diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index 8e0c68e30bb..5e2e4532265 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -91,6 +91,26 @@ class TestFalAIVideoTransformation: } assert "model" not in body + def test_get_complete_url_respects_api_base_override(self): + url = self.config.get_complete_url( + model=MODEL, + api_base="https://proxy.internal/", + litellm_params={}, + ) + + assert url == "https://proxy.internal" + + def test_validate_environment_requires_fal_ai_api_key(self, monkeypatch): + monkeypatch.setattr(fal_video_module, "get_secret_str", lambda _: None) + + with pytest.raises(ValueError, match="FAL_AI_API_KEY is not set"): + self.config.validate_environment( + headers={}, + model=MODEL, + api_key=None, + litellm_params=GenericLiteLLMParams(), + ) + def test_transform_video_create_response_encodes_model_and_usage(self): response = Mock(spec=httpx.Response) response.json.return_value = {"request_id": "abc"} From 3b0d32ec6f14979687d4bb76f51db5f9427a131a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 19 Sep 2026 18:11:50 +0000 Subject: [PATCH 015/146] fix(proxy): reject throttled exhausted budgets in JEV previews Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 8 +++ .../test_auto_router_endpoints.py | 49 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index f79425d2e97..c5a10cf5c80 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -345,6 +345,14 @@ async def _authorize_models_this_test_can_call( code=status.HTTP_400_BAD_REQUEST, ) from e + if config.classifier_type == "jev" and user_api_key_dict.budget_throttle_pct is not None: + raise ProxyException( + message="Budget has been exceeded! JEV Test Routing requires available budget.", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=status.HTTP_400_BAD_REQUEST, + ) + @router.post( "/auto_router/validate_complexity_router_config", diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index a5c93c41a84..eb9076a9d2d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -12,6 +12,7 @@ import pytest from fastapi import HTTPException, Request from pydantic import ValidationError +import litellm from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, @@ -488,6 +489,54 @@ async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( client.evaluate.assert_awaited_once() +@pytest.mark.parametrize( + "max_budget, spend, denied", + ((0.0, 0.0, True), (1.0, 2.0, True), (1.0, 0.5, False), (None, 2.0, False)), +) +@pytest.mark.asyncio +async def test_jev_test_routing_hard_blocks_exhausted_throttle_enabled_keys( + monkeypatch: pytest.MonkeyPatch, max_budget: float | None, spend: float, denied: bool +) -> None: + client: Final = AsyncMock(spec=JevClassifierClient) + client.evaluate.return_value = JevSystemOneResponse( + model="jev-test", + answers={ + "tier": JevChoiceAnswer(type="choice", choice="SIMPLE", probabilities={"SIMPLE": 1.0}, confidence=1.0) + }, + ) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + monkeypatch.setattr(auto_router_endpoints, "ComplexityRouter", partial(ComplexityRouter, jev_client=client)) + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-throttle-test", + user_id="admin", + models=["cheap-model", "typesafe/jev-test"], + max_budget=max_budget, + spend=spend, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + request: Final = _request( + "what is 2+2", + classifier_type="jev", + jev_classifier_config={"model": "jev-test"}, + ) + if denied: + with pytest.raises(ProxyException) as exc_info: + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.code == "400" + client.evaluate.assert_not_called() + return + + response: Final = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor + ) + assert response.routing_decision["cause"] == "jev_classifier" + client.evaluate.assert_awaited_once() + + @pytest.mark.parametrize("max_budget, spend", ((0.0, 0.0), (1.0, 2.0))) @pytest.mark.asyncio async def test_a_heuristic_config_does_not_need_a_budget( From 85a6a8e2063ef0932dc40b19c97eb0120a0c4235 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 18:22:58 +0000 Subject: [PATCH 016/146] test(e2e): cover fal Seedance video create, poll and download Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_nonconversational.yaml | 1 + tests/e2e/coverage_registry/schema.py | 2 + .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 2 + tests/e2e/llm_translation/endpoints_client.py | 36 +++++++++- .../test_video_generation_e2e.py | 69 +++++++++++++++++++ 5 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/llm_translation/test_video_generation_e2e.py diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 50f9b9808b2..6970567b6f0 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -80,6 +80,7 @@ - {id: llm.images_generations.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/image_generation/image_generation_handler.py", rationale: "Vertex Imagen"} - {id: llm.images_generations.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "bedrock/image_generation/image_handler.py", rationale: "Bedrock Titan Image"} - {id: llm.images_generations.black_forest_labs.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "black_forest_labs/image_generation/handler.py", rationale: "BFL Flux via OpenAI-compat"} +- {id: llm.videos.fal_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: videos, route: fal_ai, capability: basic, streaming: nonstream, assertions: [works], source: "test_video_generation_e2e.py", rationale: "fal queue video create, poll, content download"} - {id: llm.audio_speech.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_audio_speech_e2e.py:22", rationale: "OpenAI TTS binary audio"} - {id: llm.audio_speech.openai.basic.stream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:9043", rationale: "TTS streaming chunk generator"} - {id: llm.audio_speech.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.6 / LIT-4778", rationale: "TTS missing input/model, invalid voice, empty input rejected"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index fa6dad90126..3ae17432863 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -44,6 +44,7 @@ LlmEndpoint = Literal[ "vector_stores", "ocr", "bedrock_native", + "videos", ] LlmRoute = Literal[ @@ -53,6 +54,7 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", + "fal_ai", "gemini", "hosted_vllm", "openai", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 44d6e79122e..178af054f2b 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,6 +48,7 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | +| Video | live (fal.ai Seedance) | n/a | n/a | - | partial | | Responses / image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -61,6 +62,7 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | +| `test_fal_seedance_video_completes_and_downloads` | fal.ai Seedance video create, poll, and content download | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 4d2c73e7078..165a83e76c0 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Literal from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, Result, StreamingResponse +from e2e_http import BinaryStream, NoBody, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock from proxy_client import ProxyClient from pydantic import BaseModel @@ -26,6 +26,8 @@ __all__ = [ "TextBlock", "TranscriptionForm", "TranscriptionResult", + "VideoObject", + "VideoRequest", ] @@ -127,6 +129,13 @@ class ImageRequest(BaseModel): size: str = "1024x1024" +class VideoRequest(BaseModel): + model: str + prompt: str + seconds: str = "4" + size: str = "1280x720" + + class ImageEditForm(BaseModel): model: str prompt: str @@ -267,6 +276,12 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] +class VideoObject(BaseModel): + id: str + status: str + model: str | None = None + + class TranscriptionResult(BaseModel): text: str = "" @@ -440,6 +455,25 @@ class EndpointsClient: "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) ) + def videos(self, key: str, model: str, prompt: str) -> StreamingResponse: + return self._send( + "/v1/videos", key, VideoRequest(model=model, prompt=prompt) + ) + + def video_status(self, key: str, video_id: str) -> Result[VideoObject]: + return self.proxy.transport.get( + f"/v1/videos/{video_id}", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=VideoObject, + ) + + def video_content(self, key: str, video_id: str) -> StreamingResponse: + return self.proxy.transport.download( + f"/v1/videos/{video_id}/content", + headers=self.proxy.transport.bearer(key), + ) + def image_edit( self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" ) -> Result[ImagesResult]: diff --git a/tests/e2e/llm_translation/test_video_generation_e2e.py b/tests/e2e/llm_translation/test_video_generation_e2e.py new file mode 100644 index 00000000000..b65529aa260 --- /dev/null +++ b/tests/e2e/llm_translation/test_video_generation_e2e.py @@ -0,0 +1,69 @@ +"""Live e2e: POST /v1/videos creates a video and serves its content. + +Registers a fal.ai Seedance deployment at runtime, polls the queued video until it +completes, and asserts the generated content is returned as binary data. +""" + +from __future__ import annotations + +import time +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import require_successful_call, unwrap +from endpoints_client import EndpointsClient, VideoObject +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +_POLL_INTERVAL_SECONDS: Final[float] = 5.0 +_POLL_TIMEOUT_SECONDS: Final[float] = 600.0 + + +def _wait_for_completion( + endpoints_client: EndpointsClient, key: str, created: VideoObject +) -> VideoObject: + deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS + while time.monotonic() < deadline: + status = unwrap(endpoints_client.video_status(key, created.id)) + assert status.id == created.id + if status.status == "completed": + return status + if status.status == "failed": + pytest.fail(f"fal.ai video generation failed: {status}") + time.sleep(_POLL_INTERVAL_SECONDS) + pytest.fail(f"fal.ai video {created.id!r} did not complete within {_POLL_TIMEOUT_SECONDS}s") + + +class TestVideoGeneration: + @pytest.mark.covers("llm.videos.fal_ai.basic.nonstream.works") + def test_fal_seedance_video_completes_and_downloads( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-fal-video-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="fal_ai/bytedance/seedance-2.5/text-to-video", + api_key="os.environ/FAL_AI_API_KEY", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.videos( + key, model, "a red fox running through snow at dawn" + ) + require_successful_call(result) + created = VideoObject.model_validate_json(result.body) + assert created.id + assert created.model + + _wait_for_completion(endpoints_client, key, created) + + content = endpoints_client.video_content(key, created.id) + require_successful_call(content) + assert len(content.body) > 0 + assert not (content.content_type or "").startswith("application/json") From d268c8b58ae61c9fe5280a1915a0f17a85e5f1a8 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 18:34:09 +0000 Subject: [PATCH 017/146] feat(azure_ai): add MAI-Image-2.5-Pro image generation pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 13 ++++++++ model_prices_and_context_window.json | 13 ++++++++ .../test_mai_image_generation.py | 32 +++++++++++++++++++ 3 files changed, 58 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 48dded6a323..fff4c2b7e1b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11159,6 +11159,19 @@ ], "deprecation_date": "2026-10-01" }, + "azure_ai/MAI-Image-2.5-Pro": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.1085, + "output_cost_per_image_token": 0.000106, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-mai-image-2-5-pro-and-mai-voice-2-flash-in-microsoft-foundry/4539446", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "azure_ai/MAI-Image-2e": { "deprecation_date": "2026-08-15", "input_cost_per_token": 5e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 48dded6a323..fff4c2b7e1b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11159,6 +11159,19 @@ ], "deprecation_date": "2026-10-01" }, + "azure_ai/MAI-Image-2.5-Pro": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.1085, + "output_cost_per_image_token": 0.000106, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-mai-image-2-5-pro-and-mai-voice-2-flash-in-microsoft-foundry/4539446", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "azure_ai/MAI-Image-2e": { "deprecation_date": "2026-08-15", "input_cost_per_token": 5e-06, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 55656b97c57..27e78d35c69 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -453,6 +453,38 @@ class TestAzureMAIImageGeneration: ) assert round(cost, 10) == round(expected_cost, 10) + def test_mai_image_pro_edit_cost_splits_text_and_image_input(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "azure_ai/MAI-Image-2.5-Pro" + model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") + text_tokens = 37 + image_tokens = 1024 + output_image_tokens = 1024 + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=text_tokens + image_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=text_tokens, + image_tokens=image_tokens, + ), + output_tokens=output_image_tokens, + total_tokens=text_tokens + image_tokens + output_image_tokens, + ), + ) + + cost = azure_ai_image_cost_calculator(model=model, image_response=image_response) + + expected_cost = ( + text_tokens * model_info["input_cost_per_token"] + + image_tokens * model_info["input_cost_per_image_token"] + + output_image_tokens * model_info["output_cost_per_image_token"] + ) + assert round(cost, 10) == round(expected_cost, 10) + assert model_info["input_cost_per_image_token"] != model_info["input_cost_per_token"] + def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self, monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") From c55e9a492441e435c7ff3fe57561a355c96fdfb8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:40:27 +0000 Subject: [PATCH 018/146] registry audit: fireworks/together/openrouter fixes, absorb #28853 #27064 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 87 +++++++++++++++++-- model_prices_and_context_window.json | 87 +++++++++++++++++-- model_prices_and_context_window.schema.json | 4 + 3 files changed, 166 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index accabcf85d3..d3ff6e97c4a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3740,6 +3740,21 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true + }, "azure_ai/codex-mini": { "cache_read_input_token_cost": 3.75e-07, "deprecation_date": "2026-11-15", @@ -21891,6 +21906,7 @@ "supports_tool_choice": true }, "deepseek/deepseek-coder": { + "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 1.4e-07, "input_cost_per_token_cache_hit": 1.4e-08, "litellm_provider": "deepseek", @@ -21905,6 +21921,7 @@ "supports_tool_choice": true }, "deepseek/deepseek-r1": { + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 5.5e-07, "input_cost_per_token_cache_hit": 1.4e-07, "litellm_provider": "deepseek", @@ -21960,6 +21977,7 @@ "supports_tool_choice": true }, "deepseek/deepseek-v3.2": { + "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.8e-07, "input_cost_per_token_cache_hit": 2.8e-08, "litellm_provider": "deepseek", @@ -23678,6 +23696,25 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/deepseek-v4-pro-0813": { + "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost_priority": 5.5e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_priority": 1.65e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "output_cost_per_token_priority": 4.95e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/firefunction-v2": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", @@ -24064,7 +24101,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -24390,7 +24427,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": false }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, @@ -41353,6 +41390,7 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v3.2-exp": { + "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2.7e-07, "input_cost_per_token_cache_hit": 2e-08, "litellm_provider": "openrouter", @@ -41374,6 +41412,7 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-r1": { + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 7e-07, "input_cost_per_token_cache_hit": 1.4e-07, "litellm_provider": "openrouter", @@ -46169,8 +46208,8 @@ "together_ai/zai-org/GLM-4.6": { "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", - "max_input_tokens": 200000, - "max_tokens": 200000, + "max_input_tokens": 202752, + "max_tokens": 202752, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" }, @@ -46186,8 +46225,8 @@ "deprecation_date": "2026-04-02", "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", - "max_input_tokens": 200000, - "max_tokens": 200000, + "max_input_tokens": 202752, + "max_tokens": 202752, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" }, @@ -64093,6 +64132,25 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/glm-5p3": { + "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_priority": 3.25e-07, + "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { "cache_read_input_token_cost": 3.9e-07, "input_cost_per_token": 2.1e-06, @@ -64140,6 +64198,23 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/glm-5p3-flash": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_priority": 3.75e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.875e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_priority": 6.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/inkling": { "cache_read_input_token_cost": 1.7e-07, "input_cost_per_token": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index accabcf85d3..d3ff6e97c4a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3740,6 +3740,21 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true + }, "azure_ai/codex-mini": { "cache_read_input_token_cost": 3.75e-07, "deprecation_date": "2026-11-15", @@ -21891,6 +21906,7 @@ "supports_tool_choice": true }, "deepseek/deepseek-coder": { + "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 1.4e-07, "input_cost_per_token_cache_hit": 1.4e-08, "litellm_provider": "deepseek", @@ -21905,6 +21921,7 @@ "supports_tool_choice": true }, "deepseek/deepseek-r1": { + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 5.5e-07, "input_cost_per_token_cache_hit": 1.4e-07, "litellm_provider": "deepseek", @@ -21960,6 +21977,7 @@ "supports_tool_choice": true }, "deepseek/deepseek-v3.2": { + "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.8e-07, "input_cost_per_token_cache_hit": 2.8e-08, "litellm_provider": "deepseek", @@ -23678,6 +23696,25 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/deepseek-v4-pro-0813": { + "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost_priority": 5.5e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_priority": 1.65e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "output_cost_per_token_priority": 4.95e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/firefunction-v2": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", @@ -24064,7 +24101,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -24390,7 +24427,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": false }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, @@ -41353,6 +41390,7 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v3.2-exp": { + "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2.7e-07, "input_cost_per_token_cache_hit": 2e-08, "litellm_provider": "openrouter", @@ -41374,6 +41412,7 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-r1": { + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 7e-07, "input_cost_per_token_cache_hit": 1.4e-07, "litellm_provider": "openrouter", @@ -46169,8 +46208,8 @@ "together_ai/zai-org/GLM-4.6": { "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", - "max_input_tokens": 200000, - "max_tokens": 200000, + "max_input_tokens": 202752, + "max_tokens": 202752, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" }, @@ -46186,8 +46225,8 @@ "deprecation_date": "2026-04-02", "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", - "max_input_tokens": 200000, - "max_tokens": 200000, + "max_input_tokens": 202752, + "max_tokens": 202752, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" }, @@ -64093,6 +64132,25 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/glm-5p3": { + "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_priority": 3.25e-07, + "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { "cache_read_input_token_cost": 3.9e-07, "input_cost_per_token": 2.1e-06, @@ -64140,6 +64198,23 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/glm-5p3-flash": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_priority": 3.75e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.875e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_priority": 6.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/inkling": { "cache_read_input_token_cost": 1.7e-07, "input_cost_per_token": 1e-06, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 44b2569defd..aaf4d81bcc7 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -137,6 +137,10 @@ "type": "number", "minimum": 0 }, + "cache_read_input_image_token_cost": { + "type": "number", + "minimum": 0 + }, "cache_read_input_token_cost": { "type": "number", "minimum": 0, From 7b3e8afaece0b1aa4f1d9101c11daff3f6ab75c3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:52:56 +0000 Subject: [PATCH 019/146] registry: add cache_read_input_image_token_cost field for azure_ai/gpt-image-2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 1 + tests/test_litellm/test_utils.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d416e2af33a..de4126a0947 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -254,6 +254,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None cache_read_input_audio_token_cost: ReadOnly[float | None] + cache_read_input_image_token_cost: ReadOnly[float | None] cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b40c10de428..8986753af3e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -652,6 +652,7 @@ def validate_model_cost_values(model_data, exceptions=None): "cache_creation_input_audio_token_cost", "cache_read_input_token_cost", "cache_read_input_audio_token_cost", + "cache_read_input_image_token_cost", "input_dbu_cost_per_token", "output_db_cost_per_token", "output_dbu_cost_per_token", @@ -740,6 +741,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, "cache_read_input_audio_token_cost": {"type": "number"}, + "cache_read_input_image_token_cost": {"type": "number"}, "audio_transcription_config": {"type": "string"}, "deprecation_date": {"type": "string"}, "input_cost_per_audio_per_second": {"type": "number"}, From 4a4475fd7046ba0a5eda547aa5a1b14dc717a67d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:03:58 +0000 Subject: [PATCH 020/146] registry: add cache_read_input_image_token_cost to CustomPricingLiteLLMParams denylist Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index de4126a0947..9b78628c726 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3619,6 +3619,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None cache_read_input_audio_token_cost: float | None = None + cache_read_input_image_token_cost: float | None = None input_cost_per_character_above_128k_tokens: float | None = None input_cost_per_audio_token: float | None = None input_cost_per_token_cache_hit: float | None = None From a65b0c213660447f064a0bc0bc16653e25843894 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:08:25 +0000 Subject: [PATCH 021/146] registry: regen schema.d.ts for cache_read_input_image_token_cost Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7aa34c5752c..f300492e902 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30756,6 +30756,8 @@ export interface components { cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; + /** Cache Read Input Image Token Cost */ + cache_read_input_image_token_cost?: number | null; /** Cache Read Input Token Cost */ cache_read_input_token_cost?: number | null; /** Cache Read Input Token Cost Above 200K Tokens */ @@ -41410,6 +41412,8 @@ export interface components { cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; + /** Cache Read Input Image Token Cost */ + cache_read_input_image_token_cost?: number | null; /** Cache Read Input Token Cost */ cache_read_input_token_cost?: number | null; /** Cache Read Input Token Cost Above 200K Tokens */ From 9a63e06c63e60b3b30126aa7461115a2ccc041fe Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 23:28:51 +0000 Subject: [PATCH 022/146] test(integration): cover fal Seedance video queue wire contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_nonconversational.yaml | 1 - tests/e2e/coverage_registry/schema.py | 2 - .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 2 - tests/e2e/llm_translation/endpoints_client.py | 36 +--------- .../test_video_generation_e2e.py | 69 ------------------ tests/integration/contracts.json | 3 + .../providers/test_fal_ai_video_wire.py | 72 +++++++++++++++++++ 7 files changed, 76 insertions(+), 109 deletions(-) delete mode 100644 tests/e2e/llm_translation/test_video_generation_e2e.py create mode 100644 tests/integration/providers/test_fal_ai_video_wire.py diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 6970567b6f0..50f9b9808b2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -80,7 +80,6 @@ - {id: llm.images_generations.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/image_generation/image_generation_handler.py", rationale: "Vertex Imagen"} - {id: llm.images_generations.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "bedrock/image_generation/image_handler.py", rationale: "Bedrock Titan Image"} - {id: llm.images_generations.black_forest_labs.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "black_forest_labs/image_generation/handler.py", rationale: "BFL Flux via OpenAI-compat"} -- {id: llm.videos.fal_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: videos, route: fal_ai, capability: basic, streaming: nonstream, assertions: [works], source: "test_video_generation_e2e.py", rationale: "fal queue video create, poll, content download"} - {id: llm.audio_speech.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_audio_speech_e2e.py:22", rationale: "OpenAI TTS binary audio"} - {id: llm.audio_speech.openai.basic.stream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:9043", rationale: "TTS streaming chunk generator"} - {id: llm.audio_speech.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.6 / LIT-4778", rationale: "TTS missing input/model, invalid voice, empty input rejected"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 3ae17432863..fa6dad90126 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -44,7 +44,6 @@ LlmEndpoint = Literal[ "vector_stores", "ocr", "bedrock_native", - "videos", ] LlmRoute = Literal[ @@ -54,7 +53,6 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", - "fal_ai", "gemini", "hosted_vllm", "openai", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 178af054f2b..44d6e79122e 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,7 +48,6 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | -| Video | live (fal.ai Seedance) | n/a | n/a | - | partial | | Responses / image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -62,7 +61,6 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | -| `test_fal_seedance_video_completes_and_downloads` | fal.ai Seedance video create, poll, and content download | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 165a83e76c0..4d2c73e7078 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Literal from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, NoBody, Result, StreamingResponse +from e2e_http import BinaryStream, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock from proxy_client import ProxyClient from pydantic import BaseModel @@ -26,8 +26,6 @@ __all__ = [ "TextBlock", "TranscriptionForm", "TranscriptionResult", - "VideoObject", - "VideoRequest", ] @@ -129,13 +127,6 @@ class ImageRequest(BaseModel): size: str = "1024x1024" -class VideoRequest(BaseModel): - model: str - prompt: str - seconds: str = "4" - size: str = "1280x720" - - class ImageEditForm(BaseModel): model: str prompt: str @@ -276,12 +267,6 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] -class VideoObject(BaseModel): - id: str - status: str - model: str | None = None - - class TranscriptionResult(BaseModel): text: str = "" @@ -455,25 +440,6 @@ class EndpointsClient: "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) ) - def videos(self, key: str, model: str, prompt: str) -> StreamingResponse: - return self._send( - "/v1/videos", key, VideoRequest(model=model, prompt=prompt) - ) - - def video_status(self, key: str, video_id: str) -> Result[VideoObject]: - return self.proxy.transport.get( - f"/v1/videos/{video_id}", - headers=self.proxy.transport.bearer(key), - params=NoBody(), - response_type=VideoObject, - ) - - def video_content(self, key: str, video_id: str) -> StreamingResponse: - return self.proxy.transport.download( - f"/v1/videos/{video_id}/content", - headers=self.proxy.transport.bearer(key), - ) - def image_edit( self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" ) -> Result[ImagesResult]: diff --git a/tests/e2e/llm_translation/test_video_generation_e2e.py b/tests/e2e/llm_translation/test_video_generation_e2e.py deleted file mode 100644 index b65529aa260..00000000000 --- a/tests/e2e/llm_translation/test_video_generation_e2e.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Live e2e: POST /v1/videos creates a video and serves its content. - -Registers a fal.ai Seedance deployment at runtime, polls the queued video until it -completes, and asserts the generated content is returned as binary data. -""" - -from __future__ import annotations - -import time -from typing import Final - -import pytest -from e2e_config import unique_marker -from e2e_http import require_successful_call, unwrap -from endpoints_client import EndpointsClient, VideoObject -from lifecycle import ResourceManager -from models import LiteLLMParamsBody - -pytestmark = pytest.mark.e2e - -_POLL_INTERVAL_SECONDS: Final[float] = 5.0 -_POLL_TIMEOUT_SECONDS: Final[float] = 600.0 - - -def _wait_for_completion( - endpoints_client: EndpointsClient, key: str, created: VideoObject -) -> VideoObject: - deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS - while time.monotonic() < deadline: - status = unwrap(endpoints_client.video_status(key, created.id)) - assert status.id == created.id - if status.status == "completed": - return status - if status.status == "failed": - pytest.fail(f"fal.ai video generation failed: {status}") - time.sleep(_POLL_INTERVAL_SECONDS) - pytest.fail(f"fal.ai video {created.id!r} did not complete within {_POLL_TIMEOUT_SECONDS}s") - - -class TestVideoGeneration: - @pytest.mark.covers("llm.videos.fal_ai.basic.nonstream.works") - def test_fal_seedance_video_completes_and_downloads( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-fal-video-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="fal_ai/bytedance/seedance-2.5/text-to-video", - api_key="os.environ/FAL_AI_API_KEY", - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.videos( - key, model, "a red fox running through snow at dawn" - ) - require_successful_call(result) - created = VideoObject.model_validate_json(result.body) - assert created.id - assert created.model - - _wait_for_completion(endpoints_client, key, created) - - content = endpoints_client.video_content(key, created.id) - require_successful_call(content) - assert len(content.body) > 0 - assert not (content.content_type or "").startswith("application/json") diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 6958ade50f7..01f7af6e8fe 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -160,6 +160,9 @@ "other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates" ], + "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_video_create_status_and_content_follow_queue_wire_contract": [ + "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" + ], "tests/integration/mcp/test_mcp_lifecycle.py::test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit": [ "mcp.call_tool.saved_headers.reach_actual_transport" ], diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py new file mode 100644 index 00000000000..c1a2655f0aa --- /dev/null +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -0,0 +1,72 @@ +import json +import sys +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "bytedance/seedance-2.5/text-to-video" +_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + uuid.uuid4().bytes * 4 + + +@pytest.mark.covers("other.provider_wire.fal_ai.video_queue_create_status_and_content_download") +def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: Gateway) -> None: + request_id: Final = "fal-req-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.target == f"/files/{request_id}.mp4": + assert request.method == "GET" + return Reply(body=_MP4, content_type="video/mp4") + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_MODEL}" + assert json.loads(request.body) == { + "prompt": "a cat playing volleyball on a beach", + "duration": "4", + "resolution": "720p", + "aspect_ratio": "16:9", + } + return Reply( + body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode() + ) + assert request.method == "GET" + if request.target == f"/bytedance/seedance-2.5/requests/{request_id}/status": + return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode()) + assert request.target == f"/bytedance/seedance-2.5/requests/{request_id}" + return Reply(body=json.dumps({"video": {"url": f"{wire_url}/files/{request_id}.mp4"}}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model=f"fal_ai/{_MODEL}", + api_base=wire.url, + api_key="synthetic-fal-key", + ) + created: Final = gateway.post( + "/v1/videos", + { + "model": model, + "prompt": "a cat playing volleyball on a beach", + "seconds": "4", + "size": "1280x720", + }, + ) + assert created["status"] == "queued" + video_id: Final = created["id"] + assert isinstance(video_id, str) and video_id + status: Final = gateway.get(f"/v1/videos/{video_id}") + assert status["status"] == "completed" + status_id_matches_created_id: Final = status["id"] == video_id + sys.stdout.write(f"status_id_matches_created_id={status_id_matches_created_id}\n") + content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.headers["content-type"].startswith("video/mp4") + assert content.content == _MP4 + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"/{_MODEL}"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}/status"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}"), + ("GET", f"/files/{request_id}.mp4"), + ] From 21a2ed62448ebda3ab9de1245b2550aa0bf164e0 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 23:29:25 +0000 Subject: [PATCH 023/146] test(integration): drop id diagnostic from fal video wire test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/providers/test_fal_ai_video_wire.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py index c1a2655f0aa..8c72810ffb6 100644 --- a/tests/integration/providers/test_fal_ai_video_wire.py +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -1,5 +1,4 @@ import json -import sys import uuid from typing import Final @@ -58,8 +57,6 @@ def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: assert isinstance(video_id, str) and video_id status: Final = gateway.get(f"/v1/videos/{video_id}") assert status["status"] == "completed" - status_id_matches_created_id: Final = status["id"] == video_id - sys.stdout.write(f"status_id_matches_created_id={status_id_matches_created_id}\n") content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") assert content.status_code == 200, content.text assert content.headers["content-type"].startswith("video/mp4") From 2e23c2d6536b883e2a37d3aa4dcd5e5647f8c040 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 00:50:18 +0000 Subject: [PATCH 024/146] fix(user_update): evict cached user on max_budget change so the personal key ceiling refreshes on every worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../internal_user_endpoints.py | 5 ++- .../test_internal_user_endpoints.py | 43 +++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ba7a3309a90..0e5028d797b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -101,6 +101,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 +_USER_BUDGET_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget"}) def _user_table( @@ -1561,7 +1562,7 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) - if "model_max_budget" in non_default_values: + if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values): await evict_and_broadcast( cache_keys=(non_default_values["user_id"],), user_api_key_cache=user_api_key_cache, @@ -1892,7 +1893,7 @@ async def bulk_user_update( ), ) - if "model_max_budget" in non_default_values: + if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values): for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE): await asyncio.gather( *( diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0d8b19345f1..7655e6f80ff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2269,6 +2269,49 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) +@pytest.mark.asyncio +@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"]) +async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest + + saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", max_budget=500.0) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + admin: Final = UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN) + + if all_users: + await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"max_budget": 50.0}), + user_api_key_dict=admin, + litellm_changed_by=None, + ) + prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"max_budget": 50.0}) + else: + await _update_single_user_helper( + user_request=UpdateUserRequest(user_id=saved_user.user_id, max_budget=50.0), + user_api_key_dict=admin, + ) + assert prisma_client.update_data.call_args.kwargs["data"]["max_budget"] == 50.0 + + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) + + def test_generate_request_base_validator(): """ Test that GenerateRequestBase validator converts empty string to None for max_budget From a82f0a0bd2d31a435755640953ccd43476ac961d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:10:49 -0700 Subject: [PATCH 025/146] fix(auth): reject deactivated JWT users and invalidate cached status --- litellm/proxy/auth/user_api_key_auth.py | 10 +++ .../internal_user_endpoints.py | 2 +- .../proxy/auth/test_user_api_key_auth.py | 67 ++++++++++++++++--- .../test_internal_user_endpoints.py | 45 +++++++++++++ 4 files changed, 115 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index de0131772bc..49aab20e461 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1730,6 +1730,16 @@ async def _user_api_key_auth_builder( ) return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) + if ( + user_object is not None + and isinstance(user_object.metadata, dict) + and user_object.metadata.get("scim_active") is False + ): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", + ) + valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 4832c2f4c21..029a968e156 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1571,7 +1571,7 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) - if "model_max_budget" in non_default_values: + if "model_max_budget" in non_default_values or "metadata" in data_json: await evict_and_broadcast( cache_keys=(non_default_values["user_id"],), user_api_key_cache=user_api_key_cache, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 8593be751fa..dcbc0713404 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2091,7 +2091,8 @@ async def test_auto_register_binds_api_key_to_token_hash(): @pytest.mark.asyncio -async def test_auto_register_first_request_propagates_user_email(): +@pytest.mark.parametrize("active", [True, False]) +async def test_auto_register_first_request_propagates_user_email(active: bool) -> None: """ The first auto-registered JWT request must also carry user_email (resolved from the validated LiteLLM_UserTable), so attribution is consistent with the @@ -2120,6 +2121,7 @@ async def test_auto_register_first_request_propagates_user_email(): user_id="validated-user", user_email="validated@example.com", user_role="internal_user", + metadata={"scim_active": active}, ) mock_jwt_result = { "is_proxy_admin": False, @@ -2150,7 +2152,7 @@ async def test_auto_register_first_request_propagates_user_email(): patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))), patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), patch( "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", @@ -2170,8 +2172,22 @@ async def test_auto_register_first_request_propagates_user_email(): "litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping", new_callable=AsyncMock, return_value=auto_registered_key, - ), + ) as auto_register, ): + if not active: + with pytest.raises(ProxyException, match="deactivated via SCIM") as exc: + await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + assert int(exc.value.code) == 401 + auto_register.assert_not_awaited() + return result = await _user_api_key_auth_builder( request=mock_request, api_key=jwt_token, @@ -7315,15 +7331,15 @@ class TestJWTAuthUserEmail: the Prometheus `user_email` label and `user_api_key_user_email` in StandardLogging/SpendLogs metadata, which were always None for JWT traffic.""" - def _jwt_request(self, jwt_token): + def _jwt_request(self, jwt_token, route="/v1/chat/completions"): mock_request = MagicMock() - mock_request.url.path = "/v1/chat/completions" - mock_request.method = "POST" + mock_request.url.path = route + mock_request.method = "GET" if route.endswith("/list") else "POST" mock_request.headers = {"authorization": f"Bearer {jwt_token}"} mock_request.query_params = {} return mock_request - async def _run_jwt_auth(self, mock_jwt_result, jwt_token): + async def _run_jwt_auth(self, mock_jwt_result, jwt_token, route="/v1/chat/completions"): with ( patch( "litellm.proxy.proxy_server.general_settings", @@ -7344,7 +7360,7 @@ class TestJWTAuthUserEmail: litellm_jwtauth=LiteLLM_JWTAuth(), ) return await user_api_key_auth( - request=self._jwt_request(jwt_token), + request=self._jwt_request(jwt_token, route), api_key=f"Bearer {jwt_token}", ) @@ -7376,6 +7392,41 @@ class TestJWTAuthUserEmail: assert result.user_id == "jwt-human-user" assert result.user_email == "resolved@example.com" + @pytest.mark.asyncio + @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions"]) + @pytest.mark.parametrize("active", [False, True, None, "false", 0]) + async def test_jwt_auth_rejects_deactivated_user(self, route: str, active: bool | str | int | None) -> None: + from typing import Final + + jwt_token: Final = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + result: Final = { + "is_proxy_admin": False, + "team_object": None, + "user_object": LiteLLM_UserTable( + user_id="jwt-human-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + metadata={} if active is None else {"scim_active": active}, + ), + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": None, + "user_id": "jwt-human-user", + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "user1"}, + } + + if active is False: + with pytest.raises(ProxyException, match="deactivated via SCIM") as exc: + await self._run_jwt_auth(result, jwt_token, route) + assert int(exc.value.code) == 401 + else: + token: Final = await self._run_jwt_auth(result, jwt_token, route) + assert token.user_id == "jwt-human-user" + @pytest.mark.asyncio async def test_jwt_auth_populates_user_email_on_proxy_admin(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 3f2ba365a04..b64f7c8fb7c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2228,6 +2228,51 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) +@pytest.mark.asyncio +@pytest.mark.parametrize("by_email", [False, True]) +@pytest.mark.parametrize("active", [False, True, None]) +async def test_user_status_update_refreshes_cached_user( + mocker: MockerFixture, by_email: bool, active: bool | None +) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper + + saved_user: Final = LiteLLM_UserTable( + user_id="user-spruce", + user_email="spruce@example.test", + metadata={"scim_active": False if active is None else not active, "department": "engineering"}, + ) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id=None if by_email else saved_user.user_id, + user_email=saved_user.user_email if by_email else None, + metadata={"department": "engineering"} if active is None else {"scim_active": active}, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert prisma_client.update_data.call_args.kwargs["user_id"] == saved_user.user_id + assert prisma_client.update_data.call_args.kwargs["data"]["metadata"] == ( + {"department": "engineering"} if active is None else {"scim_active": active} + ) + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) + + @pytest.mark.asyncio async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None: from litellm.proxy._types import LiteLLM_UserTable From 6ea74d70f1b67ccf49311271228d1fd9bcd86b7c Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:33:47 -0700 Subject: [PATCH 026/146] fix(auth): enforce SCIM status for admin JWTs and refresh SCIM caches --- litellm/proxy/auth/handle_jwt.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 20 +++---- .../management_endpoints/scim/scim_v2.py | 7 +++ .../proxy/auth/test_handle_jwt.py | 43 +++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 11 ++-- .../scim/test_scim_key_deactivation.py | 60 +++++++++++++++++++ 6 files changed, 128 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 803093ff93a..d621bc634e6 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2463,7 +2463,7 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - if provisioning is None: + if provisioning is None or prisma_client is not None: identity: Final = await JWTAuthManager._resolve_claim_identity( jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 49aab20e461..4371ce4fda8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1714,6 +1714,16 @@ async def _user_api_key_auth_builder( jwt_claims = result.get("jwt_claims", None) agent_id: Final[str | None] = result.get("agent_id") + if ( + user_object is not None + and isinstance(user_object.metadata, dict) + and user_object.metadata.get("scim_active") is False + ): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", + ) + if is_proxy_admin: # Proxy admins authenticate via auth_builder (full # access), not via a mapped virtual key. If @@ -1730,16 +1740,6 @@ async def _user_api_key_auth_builder( ) return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) - if ( - user_object is not None - and isinstance(user_object.metadata, dict) - and user_object.metadata.get("scim_active") is False - ): - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", - ) - valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 2b74dc1e838..b676c0ddb82 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.scim.scim_transformations import ( @@ -1804,6 +1805,9 @@ async def update_user( where={"user_id": user_id}, data=update_data, ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) if client_set_active: new_active: Final = _scim_active_value(metadata) @@ -2375,6 +2379,9 @@ async def patch_user( where={"user_id": user_id}, data=update_data, ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) if new_active is not None and new_active != (True if prev_active is None else prev_active): await _set_user_keys_blocked(user_id=user_id, blocked=not new_active) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 15defb196af..232884d89fa 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -7144,3 +7144,46 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc else: create_team.assert_not_awaited() assert result["team_id"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("existing_user", [False, True]) +@pytest.mark.parametrize("warm_cache", [False, True]) +async def test_scope_admin_admission_resolves_existing_user_without_provisioning( + monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool +) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + private_key, jwk = _get_rsa_key_and_jwk("admin-status") + cache: Final = UserApiKeyCache() + cache.set_cache("litellm_jwt_auth_keys_https://admin.example/jwks", [jwk]) + user_id: Final = f"admin-status-{existing_user}-{warm_cache}" + user: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": False}, organization_memberships=[]) + if existing_user and warm_cache: + cache.set_cache(user_id, user) + database: Final = MagicMock() + users: Final = database.db.litellm_usertable + users.find_unique = AsyncMock(return_value=user if existing_user else None) + users.find_first = AsyncMock(return_value=None) + users.create = AsyncMock() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", user_id_upsert=True), + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://admin.example/jwks") + monkeypatch.setenv("JWT_ISSUER", "https://admin.example") + monkeypatch.setenv("JWT_AUDIENCE", "gateway") + token: Final = _encode_rsa_jwt( + private_key, "https://admin.example", "gateway", "admin-status", + {"sub": user_id, "scope": "litellm_proxy_admin"}, + ) + result: Final = await JWTAuthManager.auth_builder( + api_key=token, jwt_handler=handler, prisma_client=database, user_api_key_cache=cache, + parent_otel_span=None, proxy_logging_obj=MagicMock(), request_data={}, general_settings={}, route="/user/info", + ) + assert result["is_proxy_admin"] is True + assert result["user_id"] == user_id + assert result["user_object"] == (user if existing_user else None) + users.create.assert_not_awaited() diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index dcbc0713404..da36071a5b4 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7393,18 +7393,21 @@ class TestJWTAuthUserEmail: assert result.user_email == "resolved@example.com" @pytest.mark.asyncio - @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions"]) + @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"]) @pytest.mark.parametrize("active", [False, True, None, "false", 0]) - async def test_jwt_auth_rejects_deactivated_user(self, route: str, active: bool | str | int | None) -> None: + @pytest.mark.parametrize("is_admin", [False, True]) + async def test_jwt_auth_rejects_deactivated_user( + self, route: str, active: bool | str | int | None, is_admin: bool + ) -> None: from typing import Final jwt_token: Final = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" result: Final = { - "is_proxy_admin": False, + "is_proxy_admin": is_admin, "team_object": None, "user_object": LiteLLM_UserTable( user_id="jwt-human-user", - user_role=LitellmUserRoles.INTERNAL_USER.value, + user_role=LitellmUserRoles.PROXY_ADMIN.value if is_admin else LitellmUserRoles.INTERNAL_USER.value, metadata={} if active is None else {"scim_active": active}, ), "end_user_object": None, diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py index 0a9cf8b84cb..2cfbfbbb3cb 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py @@ -541,3 +541,63 @@ async def test_scim_put_user_explicit_active_false_blocks_keys(): assert update_kwargs["where"] == {"token": "hash-block-me"} assert update_kwargs["data"]["blocked"] is True assert '"scim_blocked": true' in update_kwargs["data"]["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["PUT", "PATCH"]) +@pytest.mark.parametrize("active", [False, True]) +@pytest.mark.parametrize("failure", [None, "write", "keys"]) +@pytest.mark.parametrize("status_change", [False, True]) +async def test_scim_status_write_refreshes_user_cache( + method: str, active: bool, failure: str | None, status_change: bool +) -> None: + import json + from typing import Final + + from litellm.proxy._types import ProxyException + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + user_id: Final = "scim-cache-user" + saved: Final = LiteLLM_UserTable( + user_id=user_id, user_email="x@example.com", teams=[], metadata={"scim_active": not active if status_change else active}, + ) + updated: Final = LiteLLM_UserTable( + user_id=user_id, user_email="x@example.com", teams=[], metadata={"scim_active": active}, + ) + client, db = _build_prisma_with_keys([], mock_user=saved.model_copy(deep=True), updated_user=updated) + if failure == "write": + db.litellm_usertable.update.side_effect = RuntimeError("status write failed") + if failure == "keys": + db.litellm_verificationtoken.find_many.side_effect = RuntimeError("key update failed") + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=user_id, value=saved, model_type=LiteLLM_UserTable) + with ( + patch("litellm.proxy.proxy_server.prisma_client", client), # test-quality-ok: substitute the database dependency + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: exercise a real isolated cache + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: isolate the logging dependency + patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=AsyncMock, + ) as broadcast, + ): + request: Final = ( + update_user(user_id=user_id, user=SCIMUser.model_validate(_build_put_user_payload(user_id, active=active))) + if method == "PUT" else + patch_user(user_id=user_id, patch_ops=SCIMPatchOp( + Operations=[SCIMPatchOperation(op="replace", path="active", value=active)] + )) + ) + if failure == "write" or (failure == "keys" and status_change): + with pytest.raises(ProxyException, match="status write failed" if failure == "write" else "key update failed"): + await request + else: + response: Final = await request + assert response.active is active + assert json.loads(db.litellm_usertable.update.await_args.kwargs["data"]["metadata"])["scim_active"] is active + cached: Final = await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) + if failure == "write": + assert cached == saved + broadcast.assert_not_awaited() + else: + assert cached is None + broadcast.assert_awaited_once_with(cache_key=user_id) From 124196cbaa68832894e1491e5d0f28e0b38a86dd Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:48:53 -0700 Subject: [PATCH 027/146] fix(auth): preserve scope-admin email policy during status lookup --- litellm/proxy/auth/handle_jwt.py | 19 +++++++++++++++++-- .../proxy/auth/test_handle_jwt.py | 16 +++++++++++----- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index d621bc634e6..996911cdfaa 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2463,12 +2463,27 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - if provisioning is None or prisma_client is not None: + if provisioning is None: identity: Final = await JWTAuthManager._resolve_claim_identity( jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj ) return {**admin_result, "user_object": identity.user_object} - return admin_result + if prisma_client is None: + return admin_result + try: + admin_user: Final = await get_user_object( + user_id=user_id, + user_email=user_email, + sso_user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return admin_result + return {**admin_result, "user_object": admin_user} # Get team with model access ## Check if team_id is specified via x-litellm-team-id header diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 232884d89fa..e768139f04a 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -7149,16 +7149,17 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc @pytest.mark.asyncio @pytest.mark.parametrize("existing_user", [False, True]) @pytest.mark.parametrize("warm_cache", [False, True]) +@pytest.mark.parametrize("email", [None, "admin@external.example", "admin@allowed.example"]) async def test_scope_admin_admission_resolves_existing_user_without_provisioning( - monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool + monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool, email: str | None ) -> None: from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache private_key, jwk = _get_rsa_key_and_jwk("admin-status") cache: Final = UserApiKeyCache() cache.set_cache("litellm_jwt_auth_keys_https://admin.example/jwks", [jwk]) - user_id: Final = f"admin-status-{existing_user}-{warm_cache}" - user: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": False}, organization_memberships=[]) + user_id: Final = f"admin-status-{existing_user}-{warm_cache}-{email}" + user: Final = LiteLLM_UserTable(user_id=user_id, user_email="admin@allowed.example", metadata={"scim_active": False}, organization_memberships=[]) if existing_user and warm_cache: cache.set_cache(user_id, user) database: Final = MagicMock() @@ -7170,14 +7171,17 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning handler.update_environment( prisma_client=database, user_api_key_cache=cache, - litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", user_id_upsert=True), + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", user_id_upsert=True, user_email_jwt_field="email", + user_allowed_email_domain="allowed.example", + ), ) monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://admin.example/jwks") monkeypatch.setenv("JWT_ISSUER", "https://admin.example") monkeypatch.setenv("JWT_AUDIENCE", "gateway") token: Final = _encode_rsa_jwt( private_key, "https://admin.example", "gateway", "admin-status", - {"sub": user_id, "scope": "litellm_proxy_admin"}, + {"sub": user_id, "scope": "litellm_proxy_admin", **({"email": email} if email else {})}, ) result: Final = await JWTAuthManager.auth_builder( api_key=token, jwt_handler=handler, prisma_client=database, user_api_key_cache=cache, @@ -7187,3 +7191,5 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning assert result["user_id"] == user_id assert result["user_object"] == (user if existing_user else None) users.create.assert_not_awaited() + if existing_user: + assert users.find_unique.await_count == (0 if warm_cache else 1) From 4fda0092d3aa21df1da30689c9c7660ac0bdcd80 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:52:13 -0700 Subject: [PATCH 028/146] fix(mcp): explain missing public client dependencies --- litellm/experimental_mcp_client/Readme.md | 17 +++++++++ litellm/experimental_mcp_client/__init__.py | 7 +++- .../base_sdk_tests/check_base_sdk_install.py | 12 ++++++ .../test_mcp_client.py | 37 +++++++++++++++++++ 4 files changed, 72 insertions(+), 1 deletion(-) diff --git a/litellm/experimental_mcp_client/Readme.md b/litellm/experimental_mcp_client/Readme.md index 14decce0256..3385a37cf69 100644 --- a/litellm/experimental_mcp_client/Readme.md +++ b/litellm/experimental_mcp_client/Readme.md @@ -2,6 +2,17 @@ LiteLLM MCP Client allows you to use MCP tools with LiteLLM +Install the optional dependencies with `pip install 'litellm[mcp]'`, then use the existing public imports: + +```python +from litellm.experimental_mcp_client import call_openai_tool, load_mcp_tools +from litellm.experimental_mcp_client.client import MCPClient + +client = MCPClient(server_url="https://mcp.example.com/mcp") +``` + +Core `import litellm` works without the MCP extra. Importing the experimental MCP client without its MCP or HTTPX2 dependency raises an error with this installation command + ## MCP Python SDK compatibility The `mcp` and `proxy` extras require MCP Python SDK 2.2 or newer within the 2.x release line. Installing core LiteLLM without these extras does not require MCP @@ -16,6 +27,12 @@ The shared unit-test workflow runs the MCP integration suite once, with SDK2 in See the official [SDK migration guide](https://py.sdk.modelcontextprotocol.io/migration/) for Python API changes +## Custom HTTP clients and authentication + +MCP HTTP and SSE transports now use `httpx2`. Custom authentication passed through `aws_auth` or `resolved_auth` must implement `httpx2.Auth`. Integrations that override the client's HTTP client factory or customize its event hooks must use `httpx2.AsyncClient`, request, response, timeout and transport types + +HTTPX1 clients, auth objects and hooks are not adapted by a compatibility shim. Migrate those integrations to HTTPX2 before upgrading. Ordinary `MCPClient` construction and LiteLLM's existing helper imports remain supported; this does not restore SDK1 Python imports or camelCase SDK model attributes in the shared Python environment + ## HTTP redirects For streamable HTTP POST requests, the MCP SDK follows method-preserving redirects such as HTTP 307/308 within the configured endpoint's origin. Redirects to another path on the same scheme, host and port work. The SDK also permits an HTTP-to-HTTPS upgrade on the same host using the default ports diff --git a/litellm/experimental_mcp_client/__init__.py b/litellm/experimental_mcp_client/__init__.py index 5399968ff74..7a3917a50ea 100644 --- a/litellm/experimental_mcp_client/__init__.py +++ b/litellm/experimental_mcp_client/__init__.py @@ -1,3 +1,8 @@ -from .tools import call_openai_tool, load_mcp_tools +try: + from .tools import call_openai_tool, load_mcp_tools +except ModuleNotFoundError as exc: + if exc.name not in ("mcp", "httpx2"): + raise + raise ImportError("MCP client dependencies are missing. Install them with: pip install 'litellm[mcp]'") from exc __all__ = ["call_openai_tool", "load_mcp_tools"] diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 190a900faf9..f680ba92645 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -50,6 +50,17 @@ def check_completion() -> str: return "mock completion round-trips" +def check_mcp_install_guidance() -> str: + try: + import litellm.experimental_mcp_client + except ImportError as error: + _require("pip install 'litellm[mcp]'" in str(error), f"missing MCP installation guidance: {error}") + _require(isinstance(error.__cause__, ModuleNotFoundError), "original missing-dependency cause was lost") + _require(error.__cause__.name == "mcp", f"unexpected missing dependency: {error.__cause__}") + return "optional MCP client explains how to install litellm[mcp]" + raise AssertionError("MCP client imported without the MCP extra") + + def check_embedding() -> str: import litellm @@ -109,6 +120,7 @@ def check_bedrock_credential_resolution() -> str: CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("environment is base-only", check_environment_is_base_only), ("import litellm", check_import), + ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), ("embedding", check_embedding), ("bundled model metadata", check_bundled_model_metadata), diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 7e4598c2e58..4b698f1258d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,10 +1,12 @@ import asyncio import base64 +import importlib import json import os import sys from collections.abc import AsyncIterator from pathlib import Path +from types import ModuleType from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2160,3 +2162,38 @@ async def test_404_before_session_initialization_preserves_method_not_found() -> ) assert caught.value.error.code == METHOD_NOT_FOUND assert caught.value.error.message == "Not Found" + + +@pytest.mark.parametrize("missing_module", ("mcp", "httpx2", "mcp.types", "openai.types.chat")) +def test_public_mcp_import_missing_dependency(missing_module: str) -> None: + with patch.dict(sys.modules): + for name in tuple(sys.modules): + if name.startswith(("litellm.experimental_mcp_client", "mcp.", "mcp_types.")) or name == "mcp": + del sys.modules[name] + with patch.dict(sys.modules, {missing_module: None}): + with pytest.raises(ImportError) as caught: + importlib.import_module("litellm.experimental_mcp_client.client") + + if missing_module in ("mcp", "httpx2"): + assert "pip install 'litellm[mcp]'" in str(caught.value) + assert isinstance(caught.value.__cause__, ModuleNotFoundError) + assert caught.value.__cause__.name == missing_module + else: + assert isinstance(caught.value, ModuleNotFoundError) + assert caught.value.name == missing_module + assert caught.value.__cause__ is None + assert "litellm[mcp]" not in str(caught.value) + + +def test_public_mcp_import_preserves_incompatible_sdk_error() -> None: + with patch.dict(sys.modules): + for name in tuple(sys.modules): + if name.startswith("litellm.experimental_mcp_client"): + del sys.modules[name] + with patch.dict(sys.modules, {"mcp": ModuleType("mcp")}): + with pytest.raises(ImportError, match="cannot import name 'ClientSession'") as caught: + importlib.import_module("litellm.experimental_mcp_client.client") + + assert not isinstance(caught.value, ModuleNotFoundError) + assert caught.value.__cause__ is None + assert "litellm[mcp]" not in str(caught.value) From f9244749e089a1fb0fcd157f5d5d8b2895d65e19 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 06:00:06 +0000 Subject: [PATCH 029/146] fix(proxy): return 422 instead of 429 for BudgetExceededError Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/exceptions.py | 2 +- .../_experimental/mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../test_litellm/litellm_core_utils/test_litellm_logging.py | 6 +++--- .../mcp_server/auth/test_user_api_key_auth_mcp.py | 6 +++--- .../test_litellm/proxy/auth/test_auth_exception_handler.py | 6 +++--- tests/test_litellm/proxy/auth/test_multi_budget_windows.py | 4 ++-- .../management_endpoints/test_key_management_endpoints.py | 4 ++-- tests/test_litellm/proxy/test_common_request_processing.py | 4 ++-- 8 files changed, 17 insertions(+), 17 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 14cc16452f0..4b236aec99c 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1002,7 +1002,7 @@ class BudgetExceededError(Exception): ): self.current_cost = current_cost self.max_budget = max_budget - self.status_code = 429 + self.status_code = 422 self.llm_provider = llm_provider or "" self.entity_type = entity_type self.entity_id = entity_id diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b0d57cb6228..b0640e4f0dd 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1154,7 +1154,7 @@ class MCPRequestHandler: Failures surface with the status the standard pipeline would give them, mirroring ``UserAPIKeyAuthExceptionHandler``: a disallowed route is the route gate's own 403, an - over-budget identity is a 429, a sub-check that raised its own ``HTTPException``/ + over-budget identity is a 422, a sub-check that raised its own ``HTTPException``/ ``ProxyException`` keeps that status, a transient database outage is a retryable 503, and only a genuinely unresolvable failure (a blocked team/project raises a bare ``Exception``, same as the standard pipeline's fallback) becomes the fail-closed 401. Collapsing every diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 626a13c8061..325052ebda9 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3033,7 +3033,7 @@ def test_get_error_information_budget_exceeded_structured_fields(): assert result["error_budget_entity_id"] == "repro-user" assert result["error_budget_limit"] == 1e-06 assert result["error_budget_spend"] == 3.4e-05 - assert result["error_code"] == "429" + assert result["error_code"] == "422" assert result["error_class"] == "BudgetExceededError" assert result["error_rate_limit_type"] == "budget" @@ -6407,7 +6407,7 @@ def test_get_error_information_keeps_traceback_for_unmapped_provider_4xx(): def test_get_error_information_skips_traceback_for_budget_rejection_with_provider(): - """A key-over-budget 429 is the proxy's own rejection even after the auth + """A key-over-budget 422 is the proxy's own rejection even after the auth handler stamps the requested model's provider onto it, so it stays cheap.""" from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -6416,7 +6416,7 @@ def test_get_error_information_skips_traceback_for_budget_rejection_with_provide litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic") ) result = StandardLoggingPayloadSetup.get_error_information(over_budget) - assert result["error_code"] == "429" + assert result["error_code"] == "422" assert result["llm_provider"] == "anthropic" assert result["traceback"] == "" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 4380df194ed..087c5a03498 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6339,15 +6339,15 @@ class TestMCPDcrBridgeDelegateAdmission: ) return exc_info.value - async def test_over_budget_admission_surfaces_429_not_401(self): - """A validly-authenticated but over-budget identity surfaces the standard pipeline's 429, not + async def test_over_budget_admission_surfaces_422_not_401(self): + """A validly-authenticated but over-budget identity surfaces the standard pipeline's 422, not a misleading 401. Flattening budget to 401 told the caller their credential was invalid, which on a DCR client reads as broken auth and triggers a re-authorize that cannot fix a budget problem. Regression for the status-flattening finding on the live-policy gate.""" import litellm mapped = await self._enforce_with_gate_error(litellm.BudgetExceededError(current_cost=10.0, max_budget=1.0)) - assert mapped.status_code == 429 + assert mapped.status_code == 422 async def test_db_outage_during_policy_surfaces_503_not_401(self): """A transient database outage during the live-policy gate surfaces a retryable 503, not a 401 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 125b8862dfc..3edc57af124 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -448,7 +448,7 @@ async def test_handle_authentication_error_budget_exceeded(): ) assert exc_info.value.type == ProxyErrorTypes.budget_exceeded - assert int(exc_info.value.code) == status.HTTP_429_TOO_MANY_REQUESTS + assert int(exc_info.value.code) == status.HTTP_422_UNPROCESSABLE_CONTENT @pytest.mark.asyncio @@ -687,7 +687,7 @@ def _http_request(client_host: str | None = "10.1.2.3", headers: dict[str, str] {"allow_requests_on_db_unavailable": False}, {}, "10.1.2.3", - id="429_budget_exceeded", + id="422_budget_exceeded", ), ], ) @@ -697,7 +697,7 @@ async def test_auth_failure_logs_requester_ip_address( request_kwargs: dict[str, dict[str, str]], expected_ip: str, ) -> None: - """401s and budget 429s are rejected before `add_litellm_data_to_request` stamps + """401s and budget 422s are rejected before `add_litellm_data_to_request` stamps the caller IP, so without this the failure logs (spend logs, prometheus client_ip) had no IP, and a 401 rarely carries a key or user identity either.""" with ( diff --git a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py index 0f01391b2f5..1c928448bd8 100644 --- a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py +++ b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py @@ -75,7 +75,7 @@ async def test_over_first_window_raises(): await _virtual_key_multi_budget_check(valid_token=token) err = exc_info.value - assert err.status_code == 429 + assert err.status_code == 422 assert "24h" in str(err) assert "Key over" in str(err) @@ -107,7 +107,7 @@ async def test_over_second_window_raises(): await _virtual_key_multi_budget_check(valid_token=token) err = exc_info.value - assert err.status_code == 429 + assert err.status_code == 422 assert "30d" in str(err) 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 e2a68988ee2..b02ff47ed52 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 @@ -8156,7 +8156,7 @@ async def test_reset_key_spend_resets_budget_windows(monkeypatch): counter without also advancing reset_at is not durable either: the very next request would re-sum the unchanged historical spend and put the counter right back above the window's max_budget, so - _virtual_key_multi_budget_check kept raising BudgetExceededError (429) on + _virtual_key_multi_budget_check kept raising BudgetExceededError (422) on every request even though the key's own reported spend read $0. """ mock_prisma_client = MagicMock() @@ -16593,7 +16593,7 @@ async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch): It used to probe a second, provider-stripped key because the counter was written under the request model instead, which is what let a key report zero - usage while being blocked at 429. + usage while being blocked at 422. """ from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index e4ca0b03d59..0b872400be0 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -495,7 +495,7 @@ class TestProxyBaseLLMRequestProcessing: ) assert exc_info.value.type == ProxyErrorTypes.budget_exceeded - assert exc_info.value.code == "429" + assert exc_info.value.code == "422" tag_budget_check.assert_awaited_once() _, call_kwargs = tag_budget_check.call_args assert call_kwargs["tags"] == ("guardrail-tag",) @@ -702,7 +702,7 @@ class TestProxyBaseLLMRequestProcessing: ) assert exc_info.value.type == ProxyErrorTypes.budget_exceeded - assert exc_info.value.code == "429" + assert exc_info.value.code == "422" assert "guardrail-tag" in exc_info.value.message @pytest.mark.asyncio From 92d3a1d87de637efd57acfd31dd8068f5ee0848c Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 06:10:16 +0000 Subject: [PATCH 030/146] test(proxy): expect 422 for per-model budget rejections on cursor route Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/response_api_endpoints/test_endpoints.py | 6 +++--- tests/test_litellm/proxy/test_proxy_server.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index f7abb209015..4153bf7d7ee 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -2099,7 +2099,7 @@ class TestCursorVariantPerModelBudgetEnforcement: response = _post_cursor_with_real_auth(valid_token, attrs, request_model="claude-opus-5-thinking-high") - assert response.status_code == 429, response.text + assert response.status_code == 422, response.text error = response.json()["error"] assert error["type"] == "budget_exceeded" assert "exceeded budget for model=claude-opus-5" in error["message"] @@ -2110,8 +2110,8 @@ class TestCursorVariantPerModelBudgetEnforcement: base_response = _post_cursor_with_real_auth(valid_token, attrs, request_model="claude-opus-5") alias_response = _post_cursor_with_real_auth(valid_token, attrs, request_model="claude-opus-5-fast") - assert base_response.status_code == 429, base_response.text - assert alias_response.status_code == 429, alias_response.text + assert base_response.status_code == 422, base_response.text + assert alias_response.status_code == 422, alias_response.text assert alias_response.json() == base_response.json() diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index f71f9c20f3b..950a6cc3c40 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10872,7 +10872,7 @@ async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reserva """A rate-limit or guardrail rejection happens before route_request, so the relay never runs and no success log can own the reservation. The endpoint must release it on that exit too, or the key stays pinned at the reserved - amount and its next requests 429 with budget_exceeded while /key/info shows + amount and its next requests 422 with budget_exceeded while /key/info shows spend 0 (reproduced live with rpm_limit=1). The client still gets the pre-call error event and the 1011 close it got before.""" reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} From b7e11546b569786bd6733ba7a9c90c54f90984d2 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 06:33:57 +0000 Subject: [PATCH 031/146] test: expect 422 for budget refusals in unification, e2e and integration suites Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/e2e_http.py | 2 +- .../e2e/management/test_key_management_e2e.py | 4 +-- .../budgets/test_budget_enforcement_e2e.py | 30 +++++++++---------- .../budgets/test_multi_window_budget_e2e.py | 2 +- .../test_team_multi_window_budget_e2e.py | 2 +- .../test_partial_update_sequences.py | 4 +-- .../integration/spend/test_cache_and_quota.py | 4 +-- .../test_rate_limit_error_unification.py | 6 ++-- 8 files changed, 27 insertions(+), 27 deletions(-) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 4184b6cbefc..d4978601b20 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -95,7 +95,7 @@ class UnauthorizedError(BaseModel): class RateLimitedError(BaseModel): kind: Literal["rate_limited"] = "rate_limited" retry_after_seconds: int | None = None - # litellm overloads 429 for budget_exceeded too, so keep the body to tell them apart. + # keep the body so callers can tell limiter kinds apart. body: str = "" diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index 353b0f7cf09..39a9e657b8c 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -96,8 +96,8 @@ def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None: for _ in range(40): outcome = client.chat_status(key, SPEND_MODEL, f"spend {unique_marker()}") if _is_budget_block(outcome): - assert outcome.status_code == 429, ( - f"budget refusal must be 429, got {outcome.status_code}: {outcome.body[:200]}" + assert outcome.status_code == 422, ( + f"budget refusal must be 422, got {outcome.status_code}: {outcome.body[:200]}" ) return assert outcome.ok, f"paid call failed before the budget tripped ({outcome.status_code}): {outcome.body[:300]}" diff --git a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py index 918739863ce..8a9be1d1385 100644 --- a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py @@ -46,10 +46,10 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> pytest.fail("budget never enforced within the call budget") -def _assert_blocked_429(client: BudgetClient, key: str) -> StreamingResponse: +def _assert_blocked_422(client: BudgetClient, key: str) -> StreamingResponse: blocked = _assert_budget_blocks(client, key) - assert blocked.status_code == 429, ( - f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + assert blocked.status_code == 422, ( + f"budget refusal must be 422, got {blocked.status_code}: {blocked.body[:200]}" ) return blocked @@ -60,7 +60,7 @@ class TestBudgetBlocksPerLevel: key = client.generate_key(max_budget=TINY_CAP) resources.defer(lambda: client.delete_key(key)) - _assert_blocked_429(client, key) + _assert_blocked_422(client, key) @pytest.mark.covers("quota_management.budget.team.blocks_over_limit") def test_team_budget_blocks_every_team_key(self, client: BudgetClient, resources: ResourceManager) -> None: @@ -71,10 +71,10 @@ class TestBudgetBlocksPerLevel: sibling_key = client.generate_key(team_id=team_id) resources.defer(lambda: client.delete_key(sibling_key)) - _assert_blocked_429(client, spender_key) + _assert_blocked_422(client, spender_key) sibling = _chat(client, sibling_key) - assert is_budget_block(sibling) and sibling.status_code == 429, ( - f"a sibling key on the capped team must get the same 429 budget_exceeded, " + assert is_budget_block(sibling) and sibling.status_code == 422, ( + f"a sibling key on the capped team must get the same 422 budget_exceeded, " f"got {sibling.status_code}: {sibling.body[:200]}" ) @@ -99,10 +99,10 @@ class TestBudgetBlocksPerLevel: team_key = client.generate_key(team_id=team_id, user_id=user_id) resources.defer(lambda: client.delete_key(team_key)) - _assert_blocked_429(client, first_key) + _assert_blocked_422(client, first_key) second = _chat(client, second_key) - assert is_budget_block(second) and second.status_code == 429, ( - f"the second personal key of a user over budget must get the same 429 budget_exceeded, " + assert is_budget_block(second) and second.status_code == 422, ( + f"the second personal key of a user over budget must get the same 422 budget_exceeded, " f"got {second.status_code}: {second.body[:200]}" ) team_result = _chat(client, team_key) @@ -133,7 +133,7 @@ class TestBudgetBlocksPerLevel: key = client.generate_key(team_id=team_id) resources.defer(lambda: client.delete_key(key)) - blocked = _assert_blocked_429(client, key) + blocked = _assert_blocked_422(client, key) assert f"Organization={org_id}" in blocked.body, ( f"refusal must name the org as the blocker, got: {blocked.body[:200]}" ) @@ -155,7 +155,7 @@ class TestBudgetBlocksPerLevel: teammate_key = client.generate_key(team_id=team_id, user_id=teammate_id) resources.defer(lambda: client.delete_key(teammate_key)) - _assert_blocked_429(client, member_key) + _assert_blocked_422(client, member_key) require_successful_call(_chat(client, teammate_key)) @@ -176,7 +176,7 @@ class TestKeyBudgetBlocksAcrossKeyKinds: control_key = client.generate_key(user_id=user_id) resources.defer(lambda: client.delete_key(control_key)) - _assert_blocked_429(client, capped_key) + _assert_blocked_422(client, capped_key) require_successful_call(_chat(client, control_key)) @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") @@ -188,7 +188,7 @@ class TestKeyBudgetBlocksAcrossKeyKinds: control_key = client.generate_key(team_id=team_id) resources.defer(lambda: client.delete_key(control_key)) - _assert_blocked_429(client, capped_key) + _assert_blocked_422(client, capped_key) require_successful_call(_chat(client, control_key)) @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") @@ -205,5 +205,5 @@ class TestKeyBudgetBlocksAcrossKeyKinds: control_key = client.generate_key(team_id=team_id, user_id=member_id) resources.defer(lambda: client.delete_key(control_key)) - _assert_blocked_429(client, capped_key) + _assert_blocked_422(client, capped_key) require_successful_call(_chat(client, control_key)) diff --git a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py index e1cca0c0414..e04f857545d 100644 --- a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py @@ -102,7 +102,7 @@ def test_long_window_blocks_after_short_window_resets(client: BudgetClient, reso # 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error blocked = _drive_to_block(client, key) - assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + assert blocked.status_code == 422, f"budget block was not a 422: {blocked.status_code} {blocked.body[:200]}" # 2. check the reset times of both budget windows after we drove to being blocked blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) diff --git a/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py index 1db68e6afe9..7683132776b 100644 --- a/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py @@ -101,7 +101,7 @@ def test_team_long_window_blocks_after_short_window_resets(client: BudgetClient, # 1. drive the key to being blocked, assert its blocked by budget budget_exceeded blocked = _drive_to_block(client, key) - assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + assert blocked.status_code == 422, f"budget block was not a 422: {blocked.status_code} {blocked.body[:200]}" # 2. check the the teams budget windows blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) diff --git a/tests/integration/management/test_partial_update_sequences.py b/tests/integration/management/test_partial_update_sequences.py index d79c145a685..64d807de39d 100644 --- a/tests/integration/management/test_partial_update_sequences.py +++ b/tests/integration/management/test_partial_update_sequences.py @@ -97,7 +97,7 @@ def test_zero_false_and_empty_values_are_not_treated_as_omission(gateway: Gatewa "POST", "/v1/chat/completions", {"model": models[0], "messages": [{"role": "user", "content": "zero budget"}]}, key=key, ) - assert denied.status_code == 429, denied.text + assert denied.status_code == 422, denied.text assert denied.json()["error"]["type"] == "budget_exceeded" gateway.post("/key/update", {"key": key, "max_budget": 1, "models": [], "metadata": {}}) info: Final = object_value(gateway.get("/key/info", {"key": key})["info"]) @@ -127,7 +127,7 @@ def test_zero_false_and_empty_values_are_not_treated_as_omission(gateway: Gatewa "POST", "/v1/chat/completions", {"model": models[0], "messages": [{"role": "user", "content": "updated zero budget"}]}, key=key, ) - assert zero_after_update.status_code == 429, zero_after_update.text + assert zero_after_update.status_code == 422, zero_after_update.text assert zero_after_update.json()["error"]["type"] == "budget_exceeded" gateway.post("/key/update", {"key": key, "max_budget": None}) assert read_rows( diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 840594c1a96..d32297765f6 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -185,7 +185,7 @@ def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gat {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, key=key, ) - assert denied.status_code == 429 and denied.json()["error"]["type"] == "budget_exceeded", denied.text + assert denied.status_code == 422 and denied.json()["error"]["type"] == "budget_exceeded", denied.text assert upstream.get("/__observations").json()["requests"] == [] assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 gateway.post("/key/update", {"key": key, "spend": 0}) @@ -205,7 +205,7 @@ def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gat {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, key=key, ) - assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", ( + assert denied_again.status_code == 422 and denied_again.json()["error"]["type"] == "budget_exceeded", ( denied_again.text ) assert upstream.get("/__observations").json()["requests"] == [] diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py index 8241b29aff1..256d33845a5 100644 --- a/tests/test_litellm/test_rate_limit_error_unification.py +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -1397,10 +1397,10 @@ class TestBudgetExceededErrorSurfacesUnifiedFields: assert e.llm_provider == "anthropic" def test_should_keep_existing_status_code_and_message(self): - # Backward-compat guard: existing callers depend on `status_code=429` + # Backward-compat guard: existing callers depend on `status_code=422` # and the canonical message format. e = litellm.BudgetExceededError(current_cost=0.000109, max_budget=0.0001) - assert e.status_code == 429 + assert e.status_code == 422 assert "Current cost: 0.000109" in e.message assert "Max budget: 0.0001" in e.message @@ -1424,7 +1424,7 @@ class TestBudgetExceededErrorSurfacesUnifiedFields: info = StandardLoggingPayloadSetup.get_error_information(e) assert info["error_rate_limit_category"] == "litellm_rate_limit" assert info["error_rate_limit_type"] == "budget" - assert info["error_code"] == "429" + assert info["error_code"] == "422" assert info["error_class"] == "BudgetExceededError" def test_should_propagate_llm_provider_to_standard_logging_payload(self): From bf804f51885820a6163ddcd2322ad4494f1caeb8 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 07:00:10 +0000 Subject: [PATCH 032/146] feat(proxy): add budget_exceeded_status_code setting to restore 429 for budget refusals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/exceptions.py | 3 ++- tests/test_litellm/test_rate_limit_error_unification.py | 5 +++++ 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 738dd0cac76..be8f59d210b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -400,6 +400,7 @@ default_redis_batch_cache_expiry: Optional[float] = None model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers +budget_exceeded_status_code: int = 422 # set to 429 to restore the pre-422 budget_exceeded response code budget_duration: Optional[str] = ( None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). ) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 4b236aec99c..c8de2ab12ed 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -16,6 +16,7 @@ from typing import Any, Final import httpx import openai +import litellm from litellm.types.utils import LiteLLMCommonStrings from litellm.types.vector_stores import VectorStoreSearchFailure @@ -1002,7 +1003,7 @@ class BudgetExceededError(Exception): ): self.current_cost = current_cost self.max_budget = max_budget - self.status_code = 422 + self.status_code = litellm.budget_exceeded_status_code self.llm_provider = llm_provider or "" self.entity_type = entity_type self.entity_id = entity_id diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py index 256d33845a5..e5acba938c7 100644 --- a/tests/test_litellm/test_rate_limit_error_unification.py +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -1404,6 +1404,11 @@ class TestBudgetExceededErrorSurfacesUnifiedFields: assert "Current cost: 0.000109" in e.message assert "Max budget: 0.0001" in e.message + def test_should_honor_budget_exceeded_status_code_override(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "budget_exceeded_status_code", 429) + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.status_code == 429 + def test_should_still_be_catchable_as_exception_not_rate_limit_error(self): # Critical: we deliberately did NOT make BudgetExceededError a # RateLimitError subclass. Existing `except BudgetExceededError:` From fac518dbf7dd0449b6cbd8f941c822f695f11233 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 20 Sep 2026 00:08:06 -0700 Subject: [PATCH 033/146] feat(proxy): default to the v2 migration resolver The migrations Job entrypoint (migrations/run.py) has defaulted to v2 with USE_V2_MIGRATION_RESOLVER=false as the opt-out, and the Helm chart documents that knob. Proxy startup still defaulted to v1, so the two paths disagreed about which resolver a deployment runs. Proxy startup now resolves the same way: v2 unless USE_V2_MIGRATION_RESOLVER is false or --use_legacy_migration_resolver is passed. - --use_v2_migration_resolver stays accepted as a no-op that warns, so existing commands and Helm values do not fail on an unknown option. - The dedicated Postgres smoke-test job is repointed at the legacy resolver so v1 keeps real-DB proxy-boot coverage, and the two jobs that deselected it by name are updated to match the rename. #39178 reverted an earlier flip because two replicas sharing a database deadlocked (40P01 / P3018) with neither answering /health/liveliness. That contention is what #40932 coordinates, which is why this builds on it. --- .circleci/config.yml | 8 +- litellm/proxy/proxy_cli.py | 50 ++++++++-- .../test_basic_python_version.py | 10 +- tests/test_litellm/proxy/test_proxy_cli.py | 92 +++++++++++++++++-- 4 files changed, 135 insertions(+), 25 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 2dcedbfac4a..6d5b9a2258e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1508,7 +1508,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: @@ -1532,7 +1532,7 @@ 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: docker: @@ -1561,10 +1561,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: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 464d1141f8d..14d0331c0ff 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -181,6 +181,14 @@ def append_query_params(url: str | None, params: dict) -> str: return modified_url +def resolve_v2_migration_resolver(*, use_legacy_flag: bool) -> bool: + from litellm_proxy_extras.utils import str_to_bool + + if use_legacy_flag: + return False + return bool(str_to_bool(os.getenv("USE_V2_MIGRATION_RESOLVER", "true"))) + + class ProxyInitializationHelpers: @staticmethod def _echo_litellm_version(): @@ -932,12 +940,24 @@ class ProxyInitializationHelpers: is_flag=True, default=False, 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." + "Deprecated and ignored: the v2 migration resolver is now the default, " + "so this flag has no effect. It is still accepted so existing commands " + "keep working. Pass --use_legacy_migration_resolver, or set " + "USE_V2_MIGRATION_RESOLVER=false, to opt back into v1." ), envvar="USE_V2_MIGRATION_RESOLVER", ) +@click.option( + "--use_legacy_migration_resolver", + is_flag=True, + default=False, + help=( + "Fall back to the legacy v1 migration resolver. By default the proxy " + "uses the v2 resolver, which avoids the diff-and-force recovery path " + "that can cause schema thrashing during rolling deploys where two " + "LiteLLM versions contend for the same DB." + ), +) @click.option( "--reload", is_flag=True, @@ -1005,6 +1025,7 @@ def run_server( limit_concurrency: int | None, enforce_prisma_migration_check: bool, use_v2_migration_resolver: bool, + use_legacy_migration_resolver: bool, reload: bool, prometheus_metrics_port: int | None, ): @@ -1346,17 +1367,28 @@ def run_server( if should_update_prisma_schema(general_settings.get("disable_prisma_schema_update")) is False: check_prisma_schema_diff(db_url=None) else: - if not use_v2_migration_resolver: + use_v2_resolver: Final = resolve_v2_migration_resolver( + use_legacy_flag=use_legacy_migration_resolver + ) + if use_v2_migration_resolver and use_v2_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: --use_v2_migration_resolver is " + "deprecated and has no effect \u2014 the v2 migration resolver " + "is now the default. You can safely remove it. To opt back " + "into the legacy v1 resolver, pass " + "--use_legacy_migration_resolver.\033[0m" + ) + if not use_v2_resolver: + print( + "\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration " + "resolver. It performs the diff-and-force recovery that can " + "cause schema thrashing during rolling deploys where two " + "LiteLLM versions contend for the same DB.\033[0m" ) try: setup_ok: Final = PrismaManager.setup_database( use_migrate=not use_prisma_db_push, - use_v2_resolver=use_v2_migration_resolver, + use_v2_resolver=use_v2_resolver, ) except RuntimeError as e: # Raised on unrecoverable migration errors: the v2 diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index fb06ed6b69d..0ce59332417 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -305,14 +305,14 @@ 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 opt-out legacy (v1) migration resolver. Runs in a separate CI job against a local Postgres to avoid collisions - with the v1 variant when they share a database. + with the default variant when they share a database. """ - _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 8cbae859b5c..4b83044b36d 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1995,7 +1995,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 @@ -2010,7 +2010,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("atexit.register") @@ -2070,7 +2070,7 @@ class TestRunServerDbSetup: assert "prisma CLI is neither on PATH" not in capsys.readouterr().out mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=False + use_migrate=True, use_v2_resolver=True ) @patch("subprocess.run") @@ -2137,7 +2137,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") @@ -2204,11 +2204,11 @@ class TestRunServerDbSetup: mock_atexit_register, mock_subprocess_run, ): - """USE_V2_MIGRATION_RESOLVER must select the v2 resolver. + """USE_V2_MIGRATION_RESOLVER=true must select the v2 resolver. The Helm migrations Job runs `python litellm/proxy/prisma_migration.py`, - which calls run_server with a fixed argv, so a deployment has no way to - pass --use_v2_migration_resolver and an env var is the only route in. + which calls run_server with a fixed argv, so a deployment reaches the + resolver through the env var rather than a CLI flag. """ from litellm.proxy.proxy_cli import run_server @@ -2249,6 +2249,84 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=True ) + @pytest.mark.parametrize( + "argv_extra, env_extra, expected_v2", + [ + ([], {}, True), + ([], {"USE_V2_MIGRATION_RESOLVER": "false"}, False), + (["--use_legacy_migration_resolver"], {}, False), + ( + ["--use_legacy_migration_resolver"], + {"USE_V2_MIGRATION_RESOLVER": "true"}, + False, + ), + (["--use_v2_migration_resolver"], {}, True), + ], + ids=[ + "default-is-v2", + "env-false-opts-out", + "legacy-flag-opts-out", + "legacy-flag-beats-env-true", + "deprecated-v2-flag-still-accepted", + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + def test_migration_resolver_selection( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + argv_extra, + env_extra, + expected_v2, + ): + 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" + clean_env.update(env_extra) + + 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_extra], + 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 556c7f6b68ec5e38bc13a3d9ad10b58a300ef7bc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 20 Sep 2026 00:50:05 -0700 Subject: [PATCH 034/146] ci: keep real-database coverage for both migration resolvers The Postgres-backed smoke job previously exercised one resolver. Running the default and the legacy variants in it covers v2 now that it is the default, without losing v1's coverage. --- .circleci/config.yml | 3 ++- tests/local_testing/test_basic_python_version.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 6d5b9a2258e..ba107472b8d 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1561,9 +1561,10 @@ jobs: url: tcp://localhost:5432 timeout: "60" - run: - name: Run legacy migration resolver proxy smoke test + name: Run both migration resolvers against Postgres 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 \ tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver helm_chart_testing: diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index 0ce59332417..ef500fdff42 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -312,7 +312,7 @@ def test_litellm_proxy_server_config_no_general_settings(): def test_litellm_proxy_server_config_no_general_settings_legacy_resolver(): """Exercises the opt-out legacy (v1) migration resolver. - Runs in a separate CI job against a local Postgres to avoid collisions - with the default variant when they share a database. + Runs after the default variant in the CI job that provides a local + Postgres, so both resolvers get real-database proxy-boot coverage. """ _run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"]) From e1d2789d29650fc95533eae7b7b5a824ec5e87c4 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:35:23 +0000 Subject: [PATCH 035/146] test(llms): migrate phase 6 provider unit tests to tests/unit Migrate 18 provider test files from tests/test_litellm/llms to tests/unit/llms. 194 kept tests move as-is after mutation testing; 1 test deleted (test_completion_datarobot_with_environment_variables, env-gated no-assert); the fixture-only fal_ai cost calculator file is removed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/fal_ai/test_cost_calculator.py | 19 --------------- .../llms/chatgpt/chat/test_streaming_utils.py | 0 .../test_chatgpt_responses_transformation.py | 9 +++++++ .../test_cloudflare_transformation.py | 0 .../cohere/chat/test_cohere_transformation.py | 0 .../cohere/embed/test_v1_transformation.py | 0 .../ocr/test_cohere_parse_transformation.py | 0 .../rerank/test_rerank_guardrail_handler.py | 0 .../llms/crusoe/test_crusoe.py | 0 .../test_databricks_chat_transformation.py | 0 ...est_databricks_responses_transformation.py | 0 .../test_datarobot_chat_transformation.py | 0 .../llms/datarobot/test_datarobot.py | 24 ------------------- .../chat/test_deepseek_chat_transformation.py | 0 ...pseek_anthropic_messages_transformation.py | 0 .../deepseek/test_deepseek_cost_calculator.py | 10 ++++++++ ...docker_model_runner_chat_transformation.py | 0 ...levenlabs_text_to_speech_transformation.py | 0 .../fastcrw/search/test_transformation.py | 0 19 files changed, 19 insertions(+), 43 deletions(-) delete mode 100644 tests/test_litellm/llms/fal_ai/test_cost_calculator.py rename tests/{test_litellm => unit}/llms/chatgpt/chat/test_streaming_utils.py (100%) rename tests/{test_litellm => unit}/llms/chatgpt/responses/test_chatgpt_responses_transformation.py (97%) rename tests/{test_litellm => unit}/llms/cloudflare/test_cloudflare_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/chat/test_cohere_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/embed/test_v1_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/ocr/test_cohere_parse_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/rerank/test_rerank_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/crusoe/test_crusoe.py (100%) rename tests/{test_litellm => unit}/llms/databricks/chat/test_databricks_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/databricks/responses/test_databricks_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/datarobot/chat/test_datarobot_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/datarobot/test_datarobot.py (75%) rename tests/{test_litellm => unit}/llms/deepseek/chat/test_deepseek_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/deepseek/test_deepseek_cost_calculator.py (91%) rename tests/{test_litellm => unit}/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fastcrw/search/test_transformation.py (100%) diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py deleted file mode 100644 index 419aff42059..00000000000 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ /dev/null @@ -1,19 +0,0 @@ -import pytest - -import litellm -from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils -from litellm.llms.fal_ai.cost_calculator import cost_calculator -from litellm.types.utils import ImageObject, ImageResponse - - -@pytest.fixture(autouse=True) -def _use_local_model_cost_map(monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() - yield - litellm.get_model_info.cache_clear() - - -def _image_response(num_images: int = 1) -> ImageResponse: - return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)]) diff --git a/tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py b/tests/unit/llms/chatgpt/chat/test_streaming_utils.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py rename to tests/unit/llms/chatgpt/chat/test_streaming_utils.py diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py similarity index 97% rename from tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py rename to tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 9bf3eec61f9..c01ec312796 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -19,6 +19,15 @@ from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class TestChatGPTResponsesAPITransformation: @pytest.mark.parametrize( "model_name", diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py similarity index 100% rename from tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py rename to tests/unit/llms/cloudflare/test_cloudflare_transformation.py diff --git a/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py similarity index 100% rename from tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py rename to tests/unit/llms/cohere/chat/test_cohere_transformation.py diff --git a/tests/test_litellm/llms/cohere/embed/test_v1_transformation.py b/tests/unit/llms/cohere/embed/test_v1_transformation.py similarity index 100% rename from tests/test_litellm/llms/cohere/embed/test_v1_transformation.py rename to tests/unit/llms/cohere/embed/test_v1_transformation.py diff --git a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py b/tests/unit/llms/cohere/ocr/test_cohere_parse_transformation.py similarity index 100% rename from tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py rename to tests/unit/llms/cohere/ocr/test_cohere_parse_transformation.py diff --git a/tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py b/tests/unit/llms/cohere/rerank/test_rerank_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py rename to tests/unit/llms/cohere/rerank/test_rerank_guardrail_handler.py diff --git a/tests/test_litellm/llms/crusoe/test_crusoe.py b/tests/unit/llms/crusoe/test_crusoe.py similarity index 100% rename from tests/test_litellm/llms/crusoe/test_crusoe.py rename to tests/unit/llms/crusoe/test_crusoe.py diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py rename to tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py diff --git a/tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py b/tests/unit/llms/databricks/responses/test_databricks_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py rename to tests/unit/llms/databricks/responses/test_databricks_responses_transformation.py diff --git a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py b/tests/unit/llms/datarobot/chat/test_datarobot_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py rename to tests/unit/llms/datarobot/chat/test_datarobot_chat_transformation.py diff --git a/tests/test_litellm/llms/datarobot/test_datarobot.py b/tests/unit/llms/datarobot/test_datarobot.py similarity index 75% rename from tests/test_litellm/llms/datarobot/test_datarobot.py rename to tests/unit/llms/datarobot/test_datarobot.py index d9f42960601..c98faf0151e 100644 --- a/tests/test_litellm/llms/datarobot/test_datarobot.py +++ b/tests/unit/llms/datarobot/test_datarobot.py @@ -78,27 +78,3 @@ def test_completion_datarobot_with_deployment(): except Exception as e: pytest.fail(f"Error occurred: {e}") - -def test_completion_datarobot_with_environment_variables(): - """Allow the test to run with environment variables if they are set for integrations.""" - # If keys are not set, the test will be skipped - if os.environ.get("DATAROBOT_API_TOKEN") is None: - return - - messages = [ - {"role": "user", "content": "What's the weather like in San Francisco?"} - ] - try: - response = completion( - model="datarobot/vertex_ai/gemini-1.5-flash-002", - messages=messages, - max_tokens=5, - clientId="custom-model", - ) - print(response) - assert response["object"] == "chat.completion" - assert response["model"] == "gemini-1.5-flash-002" - assert len(response["choices"]) == 1 - assert len(response["choices"][0]["message"]["content"]) > 0 - except Exception as e: - pytest.fail(f"Error occurred: {e}") diff --git a/tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py rename to tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py diff --git a/tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py b/tests/unit/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py rename to tests/unit/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py similarity index 91% rename from tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py rename to tests/unit/llms/deepseek/test_deepseek_cost_calculator.py index c3a4cdad0ac..cde4a8a4244 100644 --- a/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py +++ b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py @@ -7,6 +7,16 @@ import litellm from litellm._internal_context import pinned_billing_time from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + PEAK_MOMENTS: Final = ( pytest.param(datetime(2026, 9, 22, 8, 0, tzinfo=timezone.utc), id="tuesday-08:00"), pytest.param(datetime(2026, 9, 25, 9, 59, tzinfo=timezone.utc), id="friday-09:59"), diff --git a/tests/test_litellm/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py b/tests/unit/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py rename to tests/unit/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py diff --git a/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py rename to tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/fastcrw/search/test_transformation.py b/tests/unit/llms/fastcrw/search/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/fastcrw/search/test_transformation.py rename to tests/unit/llms/fastcrw/search/test_transformation.py From 4048062e537329784c8197fc9b16ba77a08b7611 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:48:07 +0000 Subject: [PATCH 036/146] test(llms): annotate local_model_cost_map fixtures in phase 6 tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chatgpt/responses/test_chatgpt_responses_transformation.py | 3 ++- tests/unit/llms/deepseek/test_deepseek_cost_calculator.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index c01ec312796..0b04dd0ed78 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -5,6 +5,7 @@ Source: litellm/llms/chatgpt/responses/transformation.py """ import json +from collections.abc import Generator from unittest.mock import MagicMock, patch import httpx @@ -20,7 +21,7 @@ from litellm.utils import ProviderConfigManager @pytest.fixture -def local_model_cost_map(monkeypatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) litellm.get_model_info.cache_clear() diff --git a/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py index cde4a8a4244..e61c15c3746 100644 --- a/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py +++ b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py @@ -1,3 +1,4 @@ +from collections.abc import Generator from datetime import datetime, timezone from typing import Final @@ -9,7 +10,7 @@ from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.fixture -def local_model_cost_map(monkeypatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) litellm.get_model_info.cache_clear() From b450baa402527c5be47a532508d7708fd1ca9a41 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:47:33 +0000 Subject: [PATCH 037/146] test(llms): migrate phase 5 provider unit tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../image/test_bedrock_image_bearer_token.py | 158 ------------------ .../test_amazon_nova_canvas_transformation.py | 0 .../test_amazon_stability3_transformation.py | 0 .../image/test_bedrock_image_bearer_token.py | 21 +++ .../test_bedrock_image_prepare_request.py | 0 .../test_amazon_nova_canvas_image_edit.py | 0 .../test_bedrock_agent_transformation.py | 0 .../guardrail_translation/test_handler.py | 0 ...test_bedrock_passthrough_transformation.py | 2 - .../realtime/test_bedrock_realtime_handler.py | 0 .../test_bedrock_realtime_transformation.py | 0 .../test_bedrock_rerank_header_forwarding.py | 0 ...est_bedrock_vector_store_transformation.py | 0 ...drock_mantle_passthrough_transformation.py | 0 .../test_bfl_image_edit_transformation.py | 0 ...est_bfl_image_generation_transformation.py | 0 .../test_bfl_common_utils.py | 0 .../chat/test_bytez_chat_transformation.py | 0 .../test_cerebras_chat_transformation.py | 0 .../llms/chat/test_converse_handler.py | 0 .../chatgpt/test_chatgpt_authenticator.py | 0 21 files changed, 21 insertions(+), 160 deletions(-) delete mode 100644 tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py rename tests/{test_litellm => unit}/llms/bedrock/image/test_amazon_nova_canvas_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/image/test_amazon_stability3_transformation.py (100%) create mode 100644 tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py rename tests/{test_litellm => unit}/llms/bedrock/image/test_bedrock_image_prepare_request.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/passthrough/guardrail_translation/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/realtime/test_bedrock_realtime_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/realtime/test_bedrock_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/test_bfl_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/bytez/chat/test_bytez_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cerebras/test_cerebras_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/chat/test_converse_handler.py (100%) rename tests/{test_litellm => unit}/llms/chatgpt/test_chatgpt_authenticator.py (100%) diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py deleted file mode 100644 index 0b11a66c100..00000000000 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ /dev/null @@ -1,158 +0,0 @@ -import json -import os -from unittest.mock import Mock, patch -import pytest - - -import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler - -# Mock response for Bedrock image generation -mock_image_response = {"images": ["base64_encoded_image_data"], "error": None} - - -class TestBedrockImageGeneration: - def test_image_generation_with_api_key_bearer_token(self): - """Test image generation with bearer token authentication""" - test_api_key = "test-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - # Setup mock response - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_env_variable_bearer_token(self, monkeypatch): - """Test image generation with bearer token from environment variable""" - test_api_key = "env-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - # Mock the environment variable - with ( - patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": test_api_key}), - patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen, - ): - - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - @pytest.mark.asyncio - async def test_async_image_generation_with_bearer_token(self): - """Test async image generation with bearer token authentication""" - test_api_key = "async-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.async_image_generation" - ) as mock_async_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_async_bedrock_image_gen.return_value = mock_image_response_obj - - # Call async image generation with api_key parameter - response = await litellm.aimage_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_async_bedrock_image_gen.assert_called_once() - for call in mock_async_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_sigv4(self): - """Test image generation falls back to SigV4 auth when no bearer token is provided""" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - mock_bedrock_image_gen.assert_called_once() - - -def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): - """The deployment's AWS profile does not exist, so resolving SigV4 credentials - raises; a bearer-token deployment must still sign the request with the - bearer token alone.""" - from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration - - monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") - - request = BedrockImageGeneration()._prepare_request( - model="amazon.nova-canvas-v1:0", - prompt="A cute baby sea otter", - optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, - api_base=None, - extra_headers=None, - api_key=None, - logging_obj=Mock(), - ) - - assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py new file mode 100644 index 00000000000..599507da03d --- /dev/null +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -0,0 +1,21 @@ +from unittest.mock import Mock + +def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageGeneration()._prepare_request( + model="amazon.nova-canvas-v1:0", + prompt="A cute baby sea otter", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + api_key=None, + logging_obj=Mock(), + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py rename to tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py rename to tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py diff --git a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py b/tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py rename to tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py rename to tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py rename to tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index f2a9af11af7..854ef92fa4b 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -367,8 +367,6 @@ def test_bedrock_passthrough_region_extraction_from_inference_profile_arn(): assert ( "us-west-2" in api_base ), f"Expected region 'us-west-2' from ARN in base URL, but got: {api_base}" - - def test_bedrock_passthrough_model_id_arn_encoding(): """ Test that model_id ARNs are properly URL-encoded when used in endpoints. diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py similarity index 100% rename from tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py rename to tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py rename to tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py rename to tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py rename to tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py rename to tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py b/tests/unit/llms/black_forest_labs/test_bfl_common_utils.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py rename to tests/unit/llms/black_forest_labs/test_bfl_common_utils.py diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py rename to tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/unit/llms/cerebras/test_cerebras_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py rename to tests/unit/llms/cerebras/test_cerebras_chat_transformation.py diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py similarity index 100% rename from tests/test_litellm/llms/chat/test_converse_handler.py rename to tests/unit/llms/chat/test_converse_handler.py diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py rename to tests/unit/llms/chatgpt/test_chatgpt_authenticator.py From 522c3e3ed6a5ce39b742eed3cbea68303683aeda Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:45:29 +0000 Subject: [PATCH 038/146] test: migrate wave 1 phase 2 legacy unit tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_amazon_nova_chat_completion.py | 195 ------------------ .../integrations/levo/test_levo.py | 42 ---- .../test_litellm_agent_model_resolver.py | 0 .../test_mavvrik_focus_logger.py | 0 .../integrations/opik/test_opik_extractors.py | 0 .../integrations/pointfive/test_logger.py | 0 .../integrations/pointfive/test_payload.py | 0 .../pointfive/test_upload_client.py | 0 .../test_vector_store_pre_call_hook.py | 0 .../audio_utils/test_subtitle_utils.py | 0 .../test_convert_dict_to_response.py | 0 .../test_convert_to_streaming_response.py | 0 .../test_get_formatted_prompt.py | 0 .../test_response_metadata.py | 0 .../test_a2a_guardrail_handler.py | 0 .../chat/test_a2a_chat_streaming_iterator.py | 0 .../a2a/chat/test_a2a_chat_transformation.py | 0 .../llms/a2a/test_common_utils.py | 0 .../llms/anthropic/batches/test_handler.py | 0 .../anthropic/batches/test_transformation.py | 0 20 files changed, 237 deletions(-) delete mode 100644 tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py rename tests/{test_litellm => unit}/integrations/levo/test_levo.py (88%) rename tests/{test_litellm => unit}/integrations/litellm_agent/test_litellm_agent_model_resolver.py (100%) rename tests/{test_litellm => unit}/integrations/mavvrik_focus/test_mavvrik_focus_logger.py (100%) rename tests/{test_litellm => unit}/integrations/opik/test_opik_extractors.py (100%) rename tests/{test_litellm => unit}/integrations/pointfive/test_logger.py (100%) rename tests/{test_litellm => unit}/integrations/pointfive/test_payload.py (100%) rename tests/{test_litellm => unit}/integrations/pointfive/test_upload_client.py (100%) rename tests/{test_litellm => unit}/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/audio_utils/test_subtitle_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_response_metadata.py (100%) rename tests/{test_litellm => unit}/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/a2a/chat/test_a2a_chat_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/llms/a2a/chat/test_a2a_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/a2a/test_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/batches/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/batches/test_transformation.py (100%) diff --git a/tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py b/tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py deleted file mode 100644 index ecdd1b36333..00000000000 --- a/tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py +++ /dev/null @@ -1,195 +0,0 @@ -import os -import pytest - -# Ensure the project root is on the import path - -from litellm import completion -from litellm.types.utils import ModelResponse, Usage, Choices, Message - - -def _has_api_key() -> bool: - """Check if Amazon Nova API key is available""" - return ( - "AMAZON_NOVA_API_KEY" in os.environ - and os.environ["AMAZON_NOVA_API_KEY"] is not None - ) - - -def _create_mock_nova_response(): - """Helper function to create mock Amazon Nova response for testing""" - return ModelResponse( - id="chatcmpl-test-nova-micro", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="I am Amazon Nova Micro. 777 times 9 equals 6993.", - role="assistant", - ), - ) - ], - created=1234567890, - model="amazon-nova/nova-micro-v1", - object="chat.completion", - usage=Usage(prompt_tokens=25, completion_tokens=15, total_tokens=40), - ) - - -def test_amazon_nova_chat_completion_nova_micro(): - if _has_api_key(): - response: ModelResponse = completion( - model="amazon-nova/nova-micro-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? Can you calculate 777 times 9?", - }, - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - else: - # Use mock response when API key is not available - response = _create_mock_nova_response() - # Additional mock-specific assertions for code review reference - assert ( - response.choices[0].message.content - == "I am Amazon Nova Micro. 777 times 9 equals 6993." - ) - assert response.model == "amazon-nova/nova-micro-v1" - assert response.usage.prompt_tokens == 25 - assert response.usage.completion_tokens == 15 - assert response.object == "chat.completion" - assert response.choices[0].finish_reason == "stop" - assert response.choices[0].message.role == "assistant" - - # Common assertions for both real and mock responses - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_nova_lite(): - response: ModelResponse = completion( - model="amazon-nova/nova-lite-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? Please tell me a poem on rain", - }, - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_nova_pro(): - response: ModelResponse = completion( - model="amazon-nova/nova-pro-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? What is MCP server and how does that help in building GenAI applications?", - }, - ], - timeout=30, - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_nova_premier(): - response: ModelResponse = completion( - model="amazon-nova/nova-premier-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? Can you help me understand what Trigonometry is?", - }, - ], - timeout=60, - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - print(response.choices[0].message.content) - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_with_tool_usage(): - response: ModelResponse = completion( - model="amazon-nova/nova-micro-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": "What is the temperature in SFO?"}, - ], - tools=[ - { - "type": "function", - "function": { - "name": "getCurrentWeather", - "description": "Get the current weather in a given city", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "City and country e.g. Bogotá, Colombia", - } - }, - "required": ["location"], - }, - }, - } - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message is not None - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_with_stream_response(): - response = completion( - model="amazon-nova/nova-micro-v1", - stream=True, - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What are MMO games? Can you give me some sample references?", - }, - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - chunks = list(response) - assert chunks is not None - assert len(chunks) > 0 diff --git a/tests/test_litellm/integrations/levo/test_levo.py b/tests/unit/integrations/levo/test_levo.py similarity index 88% rename from tests/test_litellm/integrations/levo/test_levo.py rename to tests/unit/integrations/levo/test_levo.py index 903be644671..647bcb3154e 100644 --- a/tests/test_litellm/integrations/levo/test_levo.py +++ b/tests/unit/integrations/levo/test_levo.py @@ -151,48 +151,6 @@ class TestLevoConfig(unittest.TestCase): class TestLevoIntegration(unittest.TestCase): """Integration tests for LevoLogger.""" - @patch.dict( - "os.environ", - { - "LEVOAI_API_KEY": "test-api-key", - "LEVOAI_ORG_ID": "test-org-id", - "LEVOAI_WORKSPACE_ID": "test-workspace-id", - "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", - }, - ) - @pytest.mark.skipif( - not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed" - ) - @patch( - "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy" - ) - @pytest.mark.asyncio - async def test_levo_logger_health_check_healthy(self, mock_init_proxy): - """Test health check returns healthy status when config is valid.""" - # Mock the proxy initialization to avoid importing proxy code - mock_init_proxy.return_value = None - - config = LevoLogger.get_levo_config() - otel_config = OpenTelemetryConfig( - exporter=config.protocol, - endpoint=config.endpoint, - headers=config.otlp_auth_headers, - ) - - # Create tracer provider with in-memory exporter - tracer_provider = TracerProvider() - tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) - - levo_logger = LevoLogger( - config=otel_config, callback_name="levo", tracer_provider=tracer_provider - ) - - # Run health check - result = await levo_logger.async_health_check() - - self.assertEqual(result["status"], "healthy") - self.assertIn("message", result) - @patch.dict("os.environ", {}, clear=True) def test_levo_logger_health_check_unhealthy(self): """Test health check returns unhealthy status when required vars are missing.""" diff --git a/tests/test_litellm/integrations/litellm_agent/test_litellm_agent_model_resolver.py b/tests/unit/integrations/litellm_agent/test_litellm_agent_model_resolver.py similarity index 100% rename from tests/test_litellm/integrations/litellm_agent/test_litellm_agent_model_resolver.py rename to tests/unit/integrations/litellm_agent/test_litellm_agent_model_resolver.py diff --git a/tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py b/tests/unit/integrations/mavvrik_focus/test_mavvrik_focus_logger.py similarity index 100% rename from tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py rename to tests/unit/integrations/mavvrik_focus/test_mavvrik_focus_logger.py diff --git a/tests/test_litellm/integrations/opik/test_opik_extractors.py b/tests/unit/integrations/opik/test_opik_extractors.py similarity index 100% rename from tests/test_litellm/integrations/opik/test_opik_extractors.py rename to tests/unit/integrations/opik/test_opik_extractors.py diff --git a/tests/test_litellm/integrations/pointfive/test_logger.py b/tests/unit/integrations/pointfive/test_logger.py similarity index 100% rename from tests/test_litellm/integrations/pointfive/test_logger.py rename to tests/unit/integrations/pointfive/test_logger.py diff --git a/tests/test_litellm/integrations/pointfive/test_payload.py b/tests/unit/integrations/pointfive/test_payload.py similarity index 100% rename from tests/test_litellm/integrations/pointfive/test_payload.py rename to tests/unit/integrations/pointfive/test_payload.py diff --git a/tests/test_litellm/integrations/pointfive/test_upload_client.py b/tests/unit/integrations/pointfive/test_upload_client.py similarity index 100% rename from tests/test_litellm/integrations/pointfive/test_upload_client.py rename to tests/unit/integrations/pointfive/test_upload_client.py diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py similarity index 100% rename from tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py rename to tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py b/tests/unit/litellm_core_utils/audio_utils/test_subtitle_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py rename to tests/unit/litellm_core_utils/audio_utils/test_subtitle_utils.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py rename to tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_streaming_iterator.py b/tests/unit/llms/a2a/chat/test_a2a_chat_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/a2a/chat/test_a2a_chat_streaming_iterator.py rename to tests/unit/llms/a2a/chat/test_a2a_chat_streaming_iterator.py diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py rename to tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py diff --git a/tests/test_litellm/llms/a2a/test_common_utils.py b/tests/unit/llms/a2a/test_common_utils.py similarity index 100% rename from tests/test_litellm/llms/a2a/test_common_utils.py rename to tests/unit/llms/a2a/test_common_utils.py diff --git a/tests/test_litellm/llms/anthropic/batches/test_handler.py b/tests/unit/llms/anthropic/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/batches/test_handler.py rename to tests/unit/llms/anthropic/batches/test_handler.py diff --git a/tests/test_litellm/llms/anthropic/batches/test_transformation.py b/tests/unit/llms/anthropic/batches/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/batches/test_transformation.py rename to tests/unit/llms/anthropic/batches/test_transformation.py From 02ccdbae906dfc21edb29421d799692c62d7054d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:50:50 +0000 Subject: [PATCH 039/146] test(llms): migrate bedrock, baseten and base_llm batch tests to tests/unit --- .../files/test_bedrock_files_integration.py | 115 ------------------ .../base_llm/batches/test_transformation.py | 16 --- .../realtime/test_transcription_protocol.py | 0 .../baseten/chat/test_baseten_completions.py | 0 .../test_agentcore_transformation.py | 0 .../test_amazon_moonshot_transformation.py | 0 .../test_amazon_nova_transformation.py | 22 ++++ .../test_amazon_qwen2_transformation.py | 0 .../test_amazon_qwen3_transformation.py | 0 .../test_base_invoke_transformation.py | 0 ...ations_anthropic_claude3_transformation.py | 76 ++++++++++++ .../test_twelvelabs_pegasus_transformation.py | 0 ...test_bedrock_chat_mantle_transformation.py | 43 +++++++ .../test_bedrock_count_tokens_handler.py | 0 ...est_bedrock_count_tokens_transformation.py | 0 .../expected_bedrock_batch_completions.jsonl | 0 .../expected_bedrock_batch_embeddings.jsonl | 0 .../files/input_batch_completions.jsonl | 0 .../files/input_batch_embeddings.jsonl | 0 .../files/test_bedrock_files_handler.py | 0 .../test_bedrock_files_transformation.py | 0 21 files changed, 141 insertions(+), 131 deletions(-) delete mode 100644 tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py rename tests/{test_litellm => unit}/llms/base_llm/batches/test_transformation.py (92%) rename tests/{test_litellm => unit}/llms/base_llm/realtime/test_transcription_protocol.py (100%) rename tests/{test_litellm => unit}/llms/baseten/chat/test_baseten_completions.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/agentcore/test_agentcore_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py (85%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py (92%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py (68%) rename tests/{test_litellm => unit}/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/expected_bedrock_batch_completions.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/input_batch_completions.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/input_batch_embeddings.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/test_bedrock_files_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/test_bedrock_files_transformation.py (100%) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py deleted file mode 100644 index 6d37d43b028..00000000000 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Test Bedrock files integration with main files API -""" - -import base64 -from unittest.mock import MagicMock, patch - -import pytest - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent -from litellm.types.utils import SpecialEnums - - -class TestBedrockFilesIntegration: - """Test integration of Bedrock files with main litellm API""" - - @pytest.mark.asyncio - async def test_litellm_afile_content_bedrock_provider_with_s3_uri(self): - """Test litellm.afile_content with bedrock provider using direct S3 URI""" - file_id = "s3://test-bucket/test-file.jsonl" - expected_content = ( - b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - ) - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="s3://test-bucket/test-file.jsonl"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content since the code - # now routes through ProviderConfigManager -> base_llm_http_handler - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - mock_retrieve.return_value = mock_result - - # Call litellm.afile_content - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called with correct parameters - mock_retrieve.assert_called_once() - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["_is_async"] is True - assert call_kwargs["file_content_request"]["file_id"] == file_id - - @pytest.mark.asyncio - async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self): - """Test litellm.afile_content with bedrock provider using unified file ID""" - # Create a unified file ID - s3_uri = "s3://test-bucket/batch-outputs/output.jsonl" - unified_id = "test-unified-id-123" - model_id = "test-model-id-456" - - unified_file_id_str = f"litellm_proxy:application/json;unified_id,{unified_id};target_model_names,;llm_output_file_id,{s3_uri};llm_output_file_model_id,{model_id}" - encoded_file_id = ( - base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=") - ) - - expected_content = ( - b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - ) - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=s3_uri), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - mock_retrieve.return_value = mock_result - - # Call litellm.afile_content with unified file ID - result = await litellm.afile_content( - file_id=encoded_file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called - mock_retrieve.assert_called_once() - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["_is_async"] is True - # The handler passes the encoded file_id as-is - assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id diff --git a/tests/test_litellm/llms/base_llm/batches/test_transformation.py b/tests/unit/llms/base_llm/batches/test_transformation.py similarity index 92% rename from tests/test_litellm/llms/base_llm/batches/test_transformation.py rename to tests/unit/llms/base_llm/batches/test_transformation.py index d84c820228f..0c360ce2ed9 100644 --- a/tests/test_litellm/llms/base_llm/batches/test_transformation.py +++ b/tests/unit/llms/base_llm/batches/test_transformation.py @@ -129,22 +129,6 @@ def test_subclass_missing_any_abstract_member_cannot_instantiate(missing_member) Incomplete() -def test_concrete_instance_methods_run(): - """Sanity: the trivial overrides actually execute through the base contract.""" - instance = _ConcreteBatchesConfig() - assert instance.custom_llm_provider == LlmProviders.OPENAI - assert instance.validate_environment( - headers={"x": "1"}, - model="m", - messages=[], - optional_params={}, - litellm_params={}, - ) == {"x": "1"} - assert instance.transform_retrieve_batch_request( - batch_id="b-1", optional_params={}, litellm_params={} - ) == {"batch_id": "b-1"} - - # =========================================================================== # # get_config() # =========================================================================== # diff --git a/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py b/tests/unit/llms/base_llm/realtime/test_transcription_protocol.py similarity index 100% rename from tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py rename to tests/unit/llms/base_llm/realtime/test_transcription_protocol.py diff --git a/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py b/tests/unit/llms/baseten/chat/test_baseten_completions.py similarity index 100% rename from tests/test_litellm/llms/baseten/chat/test_baseten_completions.py rename to tests/unit/llms/baseten/chat/test_baseten_completions.py diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py rename to tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py similarity index 85% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py index 6c370344ae7..f0f0f9160fb 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py @@ -1,5 +1,8 @@ import json +import pytest + +import litellm from litellm.llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import ( AmazonInvokeNovaConfig, ) @@ -13,6 +16,25 @@ TOOL_CALL = {"id": "call_1", "type": "function", "function": {"name": "f", "argu PNG_DATA_URL = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled in-repo cost map so capability and pricing assertions do not + depend on the network-fetched ``main`` copy, which lags this branch until merge. + + ``get_model_info`` is lru_cached, so swapping ``model_cost`` is not enough on its + own; clear on the way in and out so entries warmed against either map never leak + across tests.""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _transform_request(messages, optional_params, litellm_params=None): return AmazonInvokeNovaConfig().transform_request( model=MODEL, diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py similarity index 92% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 84db0733227..2c74d23a6a2 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,6 +1,8 @@ import asyncio +import base64 import json import uuid +from types import SimpleNamespace from typing import Final from unittest.mock import patch @@ -17,6 +19,80 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled in-repo cost map so capability and pricing assertions do not + depend on the network-fetched ``main`` copy, which lags this branch until merge. + + ``get_model_info`` is lru_cached, so swapping ``model_cost`` is not enough on its + own; clear on the way in and out so entries warmed against either map never leak + across tests.""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.fixture +def local_beta_headers_config(monkeypatch): + """Pin the bundled ``anthropic_beta_headers_config.json`` so beta header assertions + do not depend on the network-fetched copy or on what earlier tests left cached.""" + from litellm.anthropic_beta_headers_manager import reload_beta_headers_config + + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + reload_beta_headers_config() + try: + yield + finally: + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + reload_beta_headers_config() + + def test_get_supported_params_thinking(): config = AmazonAnthropicClaudeConfig() params = config.get_supported_openai_params( diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py similarity index 68% rename from tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py rename to tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index a8448f5fa7a..cb892b1ea11 100644 --- a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -1,12 +1,55 @@ +import base64 import json import uuid +from types import SimpleNamespace import httpx +import pytest import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch + + async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): image_url = f"http://img.example/{uuid.uuid4()}.png" captured = {} diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py rename to tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py rename to tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_completions.jsonl b/tests/unit/llms/bedrock/files/expected_bedrock_batch_completions.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_completions.jsonl rename to tests/unit/llms/bedrock/files/expected_bedrock_batch_completions.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl b/tests/unit/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl rename to tests/unit/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/input_batch_completions.jsonl b/tests/unit/llms/bedrock/files/input_batch_completions.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/input_batch_completions.jsonl rename to tests/unit/llms/bedrock/files/input_batch_completions.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl b/tests/unit/llms/bedrock/files/input_batch_embeddings.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl rename to tests/unit/llms/bedrock/files/input_batch_embeddings.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/unit/llms/bedrock/files/test_bedrock_files_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py rename to tests/unit/llms/bedrock/files/test_bedrock_files_handler.py diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py rename to tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py From 9aec964bace897dd4713ee5820da357b35d19eaa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 20 Sep 2026 01:02:19 -0700 Subject: [PATCH 040/146] Merge remote-tracking branch 'origin/main' into litellm_flip_v2_migration_resolver_default Drops the TQ008 suppressions the new test carried; main removed that rule. --- tests/test_litellm/proxy/test_proxy_cli.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4b83044b36d..d2835142194 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2272,9 +2272,9 @@ class TestRunServerDbSetup: ) @patch("subprocess.run") @patch("atexit.register") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @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_selection( self, mock_should_update_schema, From 4defed7f2e7eaacdf8e130857eddc6018c7bc3f7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:03:54 +0000 Subject: [PATCH 041/146] test: migrate wave 1 phase 8 legacy llm tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_hosted_vllm_ssl_verify.py | 147 ------------------ .../test_hosted_vllm_embedding_ssl_verify.py | 135 ---------------- ..._github_copilot_messages_transformation.py | 7 - ...github_copilot_responses_transformation.py | 116 +++++--------- .../test_gradient_ai_chat_transformation.py | 0 .../chat/test_groq_chat_transformation.py | 2 - .../llms/groq/test_groq_cost_calculator.py | 0 .../test_hosted_vllm_chat_transformation.py | 71 +-------- ...st_hosted_vllm_embedding_transformation.py | 8 +- ...t_hosted_vllm_image_edit_transformation.py | 0 .../responses/test_hosted_vllm_responses.py | 9 +- .../test_hosted_vllm_rerank_transformation.py | 0 .../test_hosted_vllm_video_transformation.py | 0 .../test_huggingface_rerank_transformation.py | 40 +---- .../test_inception_chat_transformation.py | 18 +-- ...est_inception_completion_transformation.py | 18 +-- .../test_jina_embedding_transformation.py | 0 .../chat/test_langflow_chat_transformation.py | 31 +--- .../litellm_proxy/test_sandbox_executor.py | 25 +-- .../litellm_proxy/test_skills_ownership.py | 73 ++------- 20 files changed, 84 insertions(+), 616 deletions(-) delete mode 100644 tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py delete mode 100644 tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py rename tests/{test_litellm => unit}/llms/github_copilot/messages/test_github_copilot_messages_transformation.py (98%) rename tests/{test_litellm => unit}/llms/github_copilot/responses/test_github_copilot_responses_transformation.py (89%) rename tests/{test_litellm => unit}/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/groq/chat/test_groq_chat_transformation.py (99%) rename tests/{test_litellm => unit}/llms/groq/test_groq_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py (82%) rename tests/{test_litellm => unit}/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py (97%) rename tests/{test_litellm => unit}/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/responses/test_hosted_vllm_responses.py (96%) rename tests/{test_litellm => unit}/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/huggingface/rerank/test_huggingface_rerank_transformation.py (91%) rename tests/{test_litellm => unit}/llms/inception/test_inception_chat_transformation.py (96%) rename tests/{test_litellm => unit}/llms/inception/test_inception_completion_transformation.py (95%) rename tests/{test_litellm => unit}/llms/jina_ai/embedding/test_jina_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/langflow/chat/test_langflow_chat_transformation.py (93%) rename tests/{test_litellm => unit}/llms/litellm_proxy/test_sandbox_executor.py (84%) rename tests/{test_litellm => unit}/llms/litellm_proxy/test_skills_ownership.py (88%) diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py deleted file mode 100644 index 2364468efe1..00000000000 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py +++ /dev/null @@ -1,147 +0,0 @@ -""" -Test SSL verification for hosted_vllm provider. - -This test ensures that the ssl_verify parameter is properly passed through -to the HTTP client when using the hosted_vllm provider. - -Issue: ssl_verify parameter was being ignored because hosted_vllm fell through -to the OpenAI catch-all path in main.py, which doesn't pass ssl_verify to the HTTP client. -""" - -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm - - -class TestHostedVLLMSSLVerify: - """Test suite for SSL verification in hosted_vllm provider.""" - - @patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") - def test_hosted_vllm_ssl_verify_false_sync(self, mock_get_httpx_client): - """Test that ssl_verify=False is passed to the HTTP client for sync calls.""" - # Setup mock client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Test response", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - mock_response.text = '{"id": "chatcmpl-test", "object": "chat.completion", "created": 1234567890, "model": "test-model", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Test response"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}' - mock_client.post.return_value = mock_response - mock_get_httpx_client.return_value = mock_client - - try: - litellm.completion( - model="hosted_vllm/test-model", - messages=[{"role": "user", "content": "Hello"}], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify _get_httpx_client was called with ssl_verify=False - mock_get_httpx_client.assert_called() - call_args = mock_get_httpx_client.call_args - - # Check that params contains ssl_verify=False - if call_args[0]: - # Positional argument - params = call_args[0][0] - else: - # Keyword argument - params = call_args[1].get("params", {}) - - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - @patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") - @pytest.mark.asyncio - async def test_hosted_vllm_ssl_verify_false_async( - self, mock_get_async_httpx_client - ): - """Test that ssl_verify=False is passed to the HTTP client for async calls.""" - # Setup mock async client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Test response", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - mock_response.text = '{"id": "chatcmpl-test", "object": "chat.completion", "created": 1234567890, "model": "test-model", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Test response"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}' - - async def mock_post(*args, **kwargs): - return mock_response - - mock_client.post = mock_post - mock_get_async_httpx_client.return_value = mock_client - - try: - await litellm.acompletion( - model="hosted_vllm/test-model", - messages=[{"role": "user", "content": "Hello"}], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify get_async_httpx_client was called with ssl_verify=False - mock_get_async_httpx_client.assert_called() - call_kwargs = mock_get_async_httpx_client.call_args[1] - - # Check that params contains ssl_verify=False - params = call_kwargs.get("params", {}) - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py b/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py deleted file mode 100644 index de94da49384..00000000000 --- a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py +++ /dev/null @@ -1,135 +0,0 @@ -""" -Test SSL verification for hosted_vllm provider embeddings. - -This test ensures that the ssl_verify parameter is properly passed through -to the HTTP client when using the hosted_vllm provider for embeddings. - -Issue: ssl_verify parameter was being ignored because hosted_vllm fell through -to the openai_like catch-all path in main.py, which doesn't pass ssl_verify to the HTTP client. -""" - -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm - - -class TestHostedVLLMEmbeddingSSLVerify: - """Test suite for SSL verification in hosted_vllm provider embeddings.""" - - @patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") - def test_hosted_vllm_embedding_ssl_verify_false_sync(self, mock_get_httpx_client): - """Test that ssl_verify=False is passed to the HTTP client for sync embedding calls.""" - # Setup mock client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], - } - ], - "model": "text-embedding-model", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5, - }, - } - mock_response.text = '{"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}], "model": "text-embedding-model", "usage": {"prompt_tokens": 5, "total_tokens": 5}}' - mock_client.post.return_value = mock_response - mock_get_httpx_client.return_value = mock_client - - try: - litellm.embedding( - model="hosted_vllm/text-embedding-model", - input=["hello world"], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify _get_httpx_client was called with ssl_verify=False - mock_get_httpx_client.assert_called() - call_args = mock_get_httpx_client.call_args - - # Check that params contains ssl_verify=False - if call_args[0]: - # Positional argument - params = call_args[0][0] - else: - # Keyword argument - params = call_args[1].get("params", {}) - - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - @patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") - @pytest.mark.asyncio - async def test_hosted_vllm_embedding_ssl_verify_false_async( - self, mock_get_async_httpx_client - ): - """Test that ssl_verify=False is passed to the HTTP client for async embedding calls.""" - # Setup mock async client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], - } - ], - "model": "text-embedding-model", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5, - }, - } - mock_response.text = '{"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}], "model": "text-embedding-model", "usage": {"prompt_tokens": 5, "total_tokens": 5}}' - - async def mock_post(*args, **kwargs): - return mock_response - - mock_client.post = mock_post - mock_get_async_httpx_client.return_value = mock_client - - try: - await litellm.aembedding( - model="hosted_vllm/text-embedding-model", - input=["hello world"], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify get_async_httpx_client was called with ssl_verify=False - mock_get_async_httpx_client.assert_called() - call_kwargs = mock_get_async_httpx_client.call_args[1] - - # Check that params contains ssl_verify=False - params = call_kwargs.get("params", {}) - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py similarity index 98% rename from tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py rename to tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 8039e744f46..9e9760650cf 100644 --- a/tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -10,13 +10,6 @@ from litellm.llms.github_copilot.messages.transformation import ( ) -def test_github_copilot_anthropic_messages_config_init(): - """Test GithubCopilotAnthropicMessagesConfig initialization.""" - config = GithubCopilotAnthropicMessagesConfig() - assert config is not None - assert hasattr(config, "authenticator") - - def test_github_copilot_anthropic_messages_get_complete_url(): """get_complete_url builds the /v1/messages URL from the base it is handed. diff --git a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py b/tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py similarity index 89% rename from tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py rename to tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py index 0174465b0cc..b8380b7adb4 100644 --- a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py +++ b/tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py @@ -26,9 +26,7 @@ def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): """Pin litellm.model_cost to the bundled local backup so tests don't depend on remote catalog fetches (and don't change behavior across remote refreshes).""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr( - litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url) - ) + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) litellm.add_known_models(model_cost_map=litellm.model_cost) @@ -44,49 +42,35 @@ class TestGithubCopilotResponsesAPITransformation: provider=LlmProviders.GITHUB_COPILOT, ) - assert ( - config is not None - ), "Config should not be None for GitHub Copilot provider" - assert isinstance( - config, GithubCopilotResponsesAPIConfig - ), f"Expected GithubCopilotResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.GITHUB_COPILOT - ), "custom_llm_provider should be GITHUB_COPILOT" + assert config is not None, "Config should not be None for GitHub Copilot provider" + assert isinstance(config, GithubCopilotResponsesAPIConfig), ( + f"Expected GithubCopilotResponsesAPIConfig, got {type(config)}" + ) + assert config.custom_llm_provider == LlmProviders.GITHUB_COPILOT, "custom_llm_provider should be GITHUB_COPILOT" @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_github_copilot_responses_endpoint_url(self, mock_authenticator_class): """Test that get_complete_url returns correct GitHub Copilot endpoint""" # Mock authenticator to return default base mock_auth_instance = MagicMock() - mock_auth_instance.get_api_base.return_value = ( - "https://api.individual.githubcopilot.com" - ) + mock_auth_instance.get_api_base.return_value = "https://api.individual.githubcopilot.com" mock_authenticator_class.return_value = mock_auth_instance config = GithubCopilotResponsesAPIConfig() # Test with default GitHub Copilot API base (from authenticator) url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.individual.githubcopilot.com/responses" - ), f"Expected GitHub Copilot responses endpoint, got {url}" + assert url == "https://api.individual.githubcopilot.com/responses", ( + f"Expected GitHub Copilot responses endpoint, got {url}" + ) # Test with custom api_base (overrides authenticator) - custom_url = config.get_complete_url( - api_base="https://custom.githubcopilot.com", litellm_params={} - ) - assert ( - custom_url == "https://custom.githubcopilot.com/responses" - ), f"Expected custom endpoint, got {custom_url}" + custom_url = config.get_complete_url(api_base="https://custom.githubcopilot.com", litellm_params={}) + assert custom_url == "https://custom.githubcopilot.com/responses", f"Expected custom endpoint, got {custom_url}" # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.githubcopilot.com/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.githubcopilot.com/responses" - ), "Should handle trailing slash" + url_with_slash = config.get_complete_url(api_base="https://api.githubcopilot.com/", litellm_params={}) + assert url_with_slash == "https://api.githubcopilot.com/responses", "Should handle trailing slash" @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_validate_environment_default_headers(self, mock_authenticator_class): @@ -98,9 +82,7 @@ class TestGithubCopilotResponsesAPITransformation: config = GithubCopilotResponsesAPIConfig() - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params={} - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params={}) # Check required headers assert headers["Authorization"] == "Bearer test-api-key-123" @@ -127,9 +109,7 @@ class TestGithubCopilotResponsesAPITransformation: "custom-header": "custom-value", } - headers = config.validate_environment( - headers=custom_headers, model="gpt-5.1-codex", litellm_params={} - ) + headers = config.validate_environment(headers=custom_headers, model="gpt-5.1-codex", litellm_params={}) # User header should override default assert headers["editor-version"] == "custom/2.0.0" @@ -182,9 +162,7 @@ class TestGithubCopilotResponsesAPITransformation: """Test _has_vision_input detects input_image type""" config = GithubCopilotResponsesAPIConfig() - input_with_vision = [ - {"role": "user", "content": [{"type": "input_image", "data": "base64..."}]} - ] + input_with_vision = [{"role": "user", "content": [{"type": "input_image", "data": "base64..."}]}] has_vision = config._has_vision_input(input_with_vision) assert has_vision is True, "Should detect input_image type" @@ -246,13 +224,11 @@ class TestGithubCopilotResponsesAPITransformation: } ] - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params) - assert ( - headers.get("copilot-vision-request") == "true" - ), "Should add copilot-vision-request header for vision input" + assert headers.get("copilot-vision-request") == "true", ( + "Should add copilot-vision-request header for vision input" + ) @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_validate_environment_with_x_initiator(self, mock_authenticator_class): @@ -270,21 +246,15 @@ class TestGithubCopilotResponsesAPITransformation: {"role": "assistant", "content": "Hi"}, ] - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params) - assert ( - headers.get("X-Initiator") == "agent" - ), "Should set X-Initiator to 'agent' for assistant role" + assert headers.get("X-Initiator") == "agent", "Should set X-Initiator to 'agent' for assistant role" def test_map_openai_params_no_transformation(self): """Test that map_openai_params passes through parameters unchanged""" config = GithubCopilotResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - temperature=0.7, max_output_tokens=1000, stream=False - ) + params = ResponsesAPIOptionalRequestParams(temperature=0.7, max_output_tokens=1000, stream=False) result = config.map_openai_params( response_api_optional_params=params, @@ -338,9 +308,9 @@ class TestGithubCopilotResponsesAPITransformation: result = config._handle_reasoning_item(reasoning_item) # encrypted_content should be preserved - assert ( - result.get("encrypted_content") == "encrypted-blob-abc123" - ), "encrypted_content must be preserved for GitHub Copilot multi-turn conversations" + assert result.get("encrypted_content") == "encrypted-blob-abc123", ( + "encrypted_content must be preserved for GitHub Copilot multi-turn conversations" + ) # status=None should be filtered out assert "status" not in result, "status=None should be filtered out" # content=None should be filtered out @@ -393,9 +363,7 @@ class TestGithubCopilotResponsesAPIRouting: in the (already-merged) model info; otherwise returns None so the dispatcher routes through the chat-completions translation bridge.""" - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_config_when_mode_is_responses(self, mock_get_info): """``mode=responses`` returns native config.""" mock_get_info.return_value = {"mode": "responses"} @@ -405,9 +373,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_mode_is_chat(self, mock_get_info): """``mode=chat`` returns None so dispatcher uses bridge.""" mock_get_info.return_value = {"mode": "chat"} @@ -417,9 +383,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_mode_is_unset_and_no_endpoints(self, mock_get_info): """Entry without ``mode`` and without ``supported_endpoints`` returns None (conservative default).""" @@ -499,9 +463,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_get_model_info_raises(self, mock_get_info): """Catalog lookup failure (model not registered) returns None (conservative default; bridge handles unknown models safely).""" @@ -512,9 +474,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_user_override_via_register_model(self, mock_get_info): """User-supplied per-deployment ``model_info`` flows through ``litellm.register_model`` (called by the router) into the merged @@ -528,9 +488,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_realistic_chat_only_entry_returns_none(self, mock_get_info): """Realistic ``model_prices_and_context_window.json`` shape for a chat-only Copilot model (e.g. github_copilot/gemini-3.1-pro-preview) @@ -554,9 +512,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_realistic_responses_only_entry_returns_config(self, mock_get_info): """Realistic catalog entry for a Responses-only Copilot model (e.g. github_copilot/gpt-5.5) returns the native config.""" @@ -592,9 +548,7 @@ class TestGithubCopilotReasoningStreamItemIdNormalization: output_index group to the id from its output_item.added.""" def _config(self): - with patch( - "litellm.llms.github_copilot.responses.transformation.Authenticator" - ): + with patch("litellm.llms.github_copilot.responses.transformation.Authenticator"): return GithubCopilotResponsesAPIConfig() def _transform(self, config, chunk): diff --git a/tests/test_litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py b/tests/unit/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py rename to tests/unit/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py similarity index 99% rename from tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py rename to tests/unit/llms/groq/chat/test_groq_chat_transformation.py index f605958b979..f5a7a920124 100644 --- a/tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -202,5 +202,3 @@ class TestGroqWebSearchUsageSignal: model_response = litellm.ModelResponse() GroqChatConfig()._add_web_search_usage(model_response=model_response) assert getattr(model_response, "usage", None) is None - - diff --git a/tests/test_litellm/llms/groq/test_groq_cost_calculator.py b/tests/unit/llms/groq/test_groq_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/groq/test_groq_cost_calculator.py rename to tests/unit/llms/groq/test_groq_cost_calculator.py diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py similarity index 82% rename from tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py rename to tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index 82b05601a85..1cc6a1457fc 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -41,74 +41,9 @@ def test_hosted_vllm_chat_transformation_file_url(): ] -def test_hosted_vllm_chat_transformation_with_audio_url(): - from litellm import completion - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "llama-3.1-70b-instruct", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - } - mock_response.text = json.dumps(mock_response.json.return_value) - mock_client.post.return_value = mock_response - - with patch( - "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", - return_value=mock_client, - ): - try: - completion( - model="hosted_vllm/llama-3.1-70b-instruct", - messages=[ - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3"}, - }, - ], - }, - ], - api_base="https://test-vllm.example.com/v1", - ) - except Exception: - pass - - mock_client.post.assert_called_once() - call_kwargs = mock_client.post.call_args[1] - request_data = json.loads(call_kwargs["data"]) - assert request_data["messages"] == [ - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3"}, - } - ], - } - ] - - def test_hosted_vllm_supports_reasoning_effort(): config = HostedVLLMChatConfig() - supported_params = config.get_supported_openai_params( - model="hosted_vllm/gpt-oss-120b" - ) + supported_params = config.get_supported_openai_params(model="hosted_vllm/gpt-oss-120b") assert "reasoning_effort" in supported_params optional_params = config.map_openai_params( non_default_params={"reasoning_effort": "high"}, @@ -129,9 +64,7 @@ def test_hosted_vllm_supports_thinking(): Related issue: https://github.com/BerriAI/litellm/issues/19761 """ config = HostedVLLMChatConfig() - supported_params = config.get_supported_openai_params( - model="hosted_vllm/GLM-4.6-FP8" - ) + supported_params = config.get_supported_openai_params(model="hosted_vllm/GLM-4.6-FP8") assert "thinking" in supported_params # Test thinking below the low threshold -> "minimal" diff --git a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py b/tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py similarity index 97% rename from tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py rename to tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py index 34be3e12abd..5854b1596b4 100644 --- a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py +++ b/tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py @@ -87,9 +87,7 @@ class TestHostedVLLMEmbeddingTransformation: headers={}, ) - assert ( - "encoding_format" not in result - ), "encoding_format should not be in request when not provided" + assert "encoding_format" not in result, "encoding_format should not be in request when not provided" def test_encoding_format_not_included_when_none(self): """ @@ -278,9 +276,7 @@ class TestHostedVLLMEmbeddingTransformation: sent_data = json.loads(call_kwargs["data"]) # Assert that encoding_format is NOT in the sent data - assert ( - "encoding_format" not in sent_data - ), "encoding_format should not be in request when not provided" + assert "encoding_format" not in sent_data, "encoding_format should not be in request when not provided" assert sent_data["model"] == "BAAI/bge-small-en-v1.5" assert sent_data["input"] == ["Hello world"] diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/unit/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py rename to tests/unit/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py diff --git a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py similarity index 96% rename from tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py rename to tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py index e81bf0c4f1f..55d0ce1e68e 100644 --- a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py +++ b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py @@ -68,9 +68,7 @@ def test_hosted_vllm_responses_create_with_string_input(): Test that hosted_vllm routes directly to the native /v1/responses endpoint when the Responses API config is registered, and correctly parses the response. """ - mock_client = _make_mock_http_client( - _make_mock_responses_api_response("I'm doing well, thanks!") - ) + mock_client = _make_mock_http_client(_make_mock_responses_api_response("I'm doing well, thanks!")) with patch( "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", @@ -109,10 +107,7 @@ def test_hosted_vllm_responses_create_with_explicit_none_extra_body(): ) # extra_body=None should be normalized to an empty dict (or absent) - assert ( - optional_params.get("extra_body") is not None - or "extra_body" not in optional_params - ) + assert optional_params.get("extra_body") is not None or "extra_body" not in optional_params def test_hosted_vllm_provider_config_registration(): diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py rename to tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py diff --git a/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py b/tests/unit/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py rename to tests/unit/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py diff --git a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py similarity index 91% rename from tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py rename to tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index 9d6b7290eb6..6fd2b006fef 100644 --- a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -219,29 +219,6 @@ def test_huggingface_rerank_return_documents(mock_post): assert "text" in result["document"] -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_huggingface_rerank_error_handling(mock_post): - """Test HuggingFace rerank error handling.""" - - def return_val(): - return {"error": "Unauthorized"} - - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.json = return_val - mock_response.text = "Unauthorized" - mock_post.return_value = mock_response - - with pytest.raises(litellm.APIConnectionError): - litellm.rerank( - model="huggingface/BAAI/bge-reranker-base", - query="hello", - documents=["hello", "world"], - top_n=2, - api_key="invalid_key", - ) - - def test_huggingface_rerank_config(): """Test HuggingFaceRerankConfig class functionality.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig @@ -249,10 +226,7 @@ def test_huggingface_rerank_config(): config = HuggingFaceRerankConfig() # Test complete URL generation - assert ( - config.get_complete_url(None, "test") - == "https://api-inference.huggingface.co/rerank" - ) + assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" # Test custom API base custom_url = config.get_complete_url("https://custom.huggingface.co", "test") @@ -292,13 +266,9 @@ def test_request_transformation(): config = HuggingFaceRerankConfig() - optional_params = OptionalRerankParams( - query="hello", texts=["hello", "world"], top_n=2, return_text=True - ) + optional_params = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) - request_body = config.transform_rerank_request( - model="test", optional_rerank_params=optional_params, headers={} - ) + request_body = config.transform_rerank_request(model="test", optional_rerank_params=optional_params, headers={}) assert request_body["query"] == "hello" assert request_body["texts"] == ["hello", "world"] @@ -368,9 +338,7 @@ def test_validate_environment(): # Test headers override custom_headers = {"custom": "header"} - headers = config.validate_environment( - headers=custom_headers, model="test", api_key="test_key" - ) + headers = config.validate_environment(headers=custom_headers, model="test", api_key="test_key") assert "custom" in headers assert headers["custom"] == "header" diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/unit/llms/inception/test_inception_chat_transformation.py similarity index 96% rename from tests/test_litellm/llms/inception/test_inception_chat_transformation.py rename to tests/unit/llms/inception/test_inception_chat_transformation.py index 1d12be2adee..c4c023077fc 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/unit/llms/inception/test_inception_chat_transformation.py @@ -188,21 +188,15 @@ def test_inception_does_not_leak_key_to_caller_api_base(): caller also supplies their own key. """ config = InceptionChatConfig() - with mock.patch.dict( - os.environ, {"INCEPTION_API_KEY": "server-secret"}, clear=True - ): + with mock.patch.dict(os.environ, {"INCEPTION_API_KEY": "server-secret"}, clear=True): with mock.patch.object(litellm, "inception_key", "module-secret"): # caller overrides api_base without a key -> server key withheld - api_base, api_key = config._get_openai_compatible_provider_info( - "https://attacker.example/v1", None - ) + api_base, api_key = config._get_openai_compatible_provider_info("https://attacker.example/v1", None) assert api_base == "https://attacker.example/v1" assert api_key is None # caller overrides api_base AND supplies their own key -> used as-is - _, api_key = config._get_openai_compatible_provider_info( - "https://attacker.example/v1", "caller-key" - ) + _, api_key = config._get_openai_compatible_provider_info("https://attacker.example/v1", "caller-key") assert api_key == "caller-key" # default/server base -> server-managed key resolved @@ -217,9 +211,7 @@ def test_get_llm_provider_inception(): assert model == "mercury-2" assert provider == "inception" - model, provider, _, api_base = get_llm_provider( - "mercury-2", api_base="https://api.inceptionlabs.ai/v1" - ) + model, provider, _, api_base = get_llm_provider("mercury-2", api_base="https://api.inceptionlabs.ai/v1") assert model == "mercury-2" assert provider == "inception" assert api_base == "https://api.inceptionlabs.ai/v1" @@ -293,5 +285,3 @@ def test_inception_completion_targets_inception_endpoint(): assert captured["body"]["model"] == "mercury-2" assert captured["body"]["tool_choice"] == "auto" assert response.choices[0].message.content == "hi" - - diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/unit/llms/inception/test_inception_completion_transformation.py similarity index 95% rename from tests/test_litellm/llms/inception/test_inception_completion_transformation.py rename to tests/unit/llms/inception/test_inception_completion_transformation.py index ed3f34fc744..84923229e20 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/unit/llms/inception/test_inception_completion_transformation.py @@ -22,9 +22,7 @@ def _fim_response_bytes(): "object": "text_completion", "created": 1, "model": "mercury-edit-2", - "choices": [ - {"text": "a + b", "index": 0, "finish_reason": "stop", "logprobs": None} - ], + "choices": [{"text": "a + b", "index": 0, "finish_reason": "stop", "logprobs": None}], "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, } ).encode() @@ -47,9 +45,7 @@ def test_inception_fim_supports_suffix_param(): def test_inception_fim_supported_params_match_schema(): """FIM exposes the OpenAI subset of Inception's FIMCompletionRequest only""" - params = InceptionTextCompletionConfig().get_supported_openai_params( - "mercury-edit-2" - ) + params = InceptionTextCompletionConfig().get_supported_openai_params("mercury-edit-2") for p in ("suffix", "top_p", "frequency_penalty", "presence_penalty", "stop"): assert p in params # Chat-only sampling controls are not part of Inception's FIM schema @@ -75,11 +71,7 @@ def test_inception_get_supported_openai_params_dispatch(): @pytest.mark.parametrize("provider", ["inception", "text-completion-inception"]) def test_inception_validate_environment(provider): - model = ( - "inception/mercury-2" - if provider == "inception" - else "text-completion-inception/mercury-edit-2" - ) + model = "inception/mercury-2" if provider == "inception" else "text-completion-inception/mercury-edit-2" with mock.patch.dict(os.environ, {}, clear=True): result = litellm.validate_environment(model) @@ -217,9 +209,7 @@ def test_inception_fim_does_not_leak_global_api_key(): content=_fim_response_bytes(), ) - with mock.patch.dict( - os.environ, {"INCEPTION_API_KEY": "sk-inception-correct"}, clear=True - ): + with mock.patch.dict(os.environ, {"INCEPTION_API_KEY": "sk-inception-correct"}, clear=True): with mock.patch.object(litellm, "inception_key", None): with mock.patch.object(litellm, "api_key", "sk-global-should-not-leak"): with mock.patch("httpx.Client.send", new=fake_send): diff --git a/tests/test_litellm/llms/jina_ai/embedding/test_jina_embedding_transformation.py b/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/jina_ai/embedding/test_jina_embedding_transformation.py rename to tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py diff --git a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py similarity index 93% rename from tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py rename to tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py index 383a7afbe93..179a6cad4aa 100644 --- a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py @@ -46,7 +46,7 @@ def test_langflow_config_get_complete_url(): def test_langflow_config_get_complete_url_requires_api_base(): config = LangFlowConfig() - with pytest.raises(ValueError, match='api_base is required for LangFlow\\. Set it via'): + with pytest.raises(ValueError, match="api_base is required for LangFlow\\. Set it via"): config.get_complete_url( api_base=None, api_key=None, @@ -225,9 +225,7 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload(): posted_bodies.append(json.loads(body) if isinstance(body, str) else body) resp = MagicMock(spec=httpx.Response) resp.status_code = 200 - resp.json.return_value = { - "outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}] - } + resp.json.return_value = {"outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}]} resp.headers = {} resp.text = "{}" return resp @@ -275,9 +273,7 @@ def test_langflow_config_extract_response_from_outputs_dict(): "outputs": [ { "results": {}, - "outputs": { - "message": {"message": {"text": "via outputs dict"}} - }, + "outputs": {"message": {"message": {"text": "via outputs dict"}}}, } ] } @@ -292,14 +288,9 @@ def test_langflow_extract_response_returns_none_when_no_message(): assert config._extract_content_from_response({"outputs": []}) is None assert config._extract_content_from_response({"detail": "flow failed"}) is None assert config._extract_content_from_response({"outputs": ["not-a-dict"]}) is None + assert config._extract_content_from_response({"outputs": [{"outputs": ["bad"]}]}) is None assert ( - config._extract_content_from_response({"outputs": [{"outputs": ["bad"]}]}) - is None - ) - assert ( - config._extract_content_from_response( - {"outputs": [{"outputs": [{"results": {"message": {"text": ""}}}]}]} - ) + config._extract_content_from_response({"outputs": [{"outputs": [{"results": {"message": {"text": ""}}}]}]}) is None ) @@ -310,9 +301,7 @@ def test_langflow_transform_response_builds_model_response_with_usage(): status_code=200, json={ "session_id": "sess-abc", - "outputs": [ - {"outputs": [{"results": {"message": {"text": "Hello from LangFlow"}}}]} - ], + "outputs": [{"outputs": [{"results": {"message": {"text": "Hello from LangFlow"}}}]}], }, ) @@ -332,9 +321,7 @@ def test_langflow_transform_response_builds_model_response_with_usage(): assert result.choices[0].finish_reason == "stop" assert result.model == "langflow/my-flow-id" assert result.usage.completion_tokens > 0 - assert result.usage.total_tokens == ( - result.usage.prompt_tokens + result.usage.completion_tokens - ) + assert result.usage.total_tokens == (result.usage.prompt_tokens + result.usage.completion_tokens) def test_langflow_transform_response_raises_on_unparseable_body(): @@ -357,9 +344,7 @@ def test_langflow_transform_response_raises_on_unparseable_body(): def test_langflow_transform_response_raises_on_non_json_body(): config = LangFlowConfig() - raw_response = httpx.Response( - status_code=200, content=b"not json", headers={"content-type": "text/plain"} - ) + raw_response = httpx.Response(status_code=200, content=b"not json", headers={"content-type": "text/plain"}) with pytest.raises(LangFlowError): config.transform_response( diff --git a/tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py b/tests/unit/llms/litellm_proxy/test_sandbox_executor.py similarity index 84% rename from tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py rename to tests/unit/llms/litellm_proxy/test_sandbox_executor.py index 422e7a3cf4d..e7a03b9231a 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py +++ b/tests/unit/llms/litellm_proxy/test_sandbox_executor.py @@ -55,9 +55,7 @@ def _install_fake_sandbox(monkeypatch, session_cls=_FakeSandboxSession): def test_execute_installs_inline_requirements_file(monkeypatch): _install_fake_sandbox(monkeypatch) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) requirements = "git+https://example.com/repo.git#egg=foo\n-r extra.txt\n-e ./pkg\n" result = executor.execute( @@ -69,22 +67,15 @@ def test_execute_installs_inline_requirements_file(monkeypatch): assert result["success"] is True created_session = _FakeSandboxSession.last_instance - assert created_session.copied_contents[ - "/sandbox/.litellm_requirements.txt" - ] == requirements.encode("utf-8") - assert ( - "pip', 'install', '-r', '.litellm_requirements.txt'" - in created_session.run_calls[0] - ) + assert created_session.copied_contents["/sandbox/.litellm_requirements.txt"] == requirements.encode("utf-8") + assert "pip', 'install', '-r', '.litellm_requirements.txt'" in created_session.run_calls[0] assert "os.chdir('/sandbox')" in created_session.run_calls[1] def test_execute_uses_skill_requirements_txt(monkeypatch): _install_fake_sandbox(monkeypatch) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) result = executor.execute( code="print('hello')", @@ -97,9 +88,7 @@ def test_execute_uses_skill_requirements_txt(monkeypatch): assert result["success"] is True created_session = _FakeSandboxSession.last_instance - copied_paths = { - sandbox_path for _, sandbox_path in created_session.copy_to_runtime_calls - } + copied_paths = {sandbox_path for _, sandbox_path in created_session.copy_to_runtime_calls} assert "/sandbox/requirements.txt" in copied_paths assert "/sandbox/.litellm_requirements.txt" not in copied_paths assert "pip', 'install', '-r', 'requirements.txt'" in created_session.run_calls[0] @@ -118,9 +107,7 @@ def test_execute_returns_install_failure(monkeypatch): _install_fake_sandbox(monkeypatch, session_cls=_FailingSandboxSession) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) result = executor.execute( code="print('hello')", diff --git a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py b/tests/unit/llms/litellm_proxy/test_skills_ownership.py similarity index 88% rename from tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py rename to tests/unit/llms/litellm_proxy/test_skills_ownership.py index e538c50cde8..6caa2da3169 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/unit/llms/litellm_proxy/test_skills_ownership.py @@ -37,12 +37,7 @@ def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: def test_should_extract_skill_auth_from_supported_metadata_fields(): auth = UserAPIKeyAuth(user_id="user-1") - assert ( - skills_main._get_user_api_key_auth_from_kwargs( - {"metadata": {"user_api_key_auth": auth}} - ) - is auth - ) + assert skills_main._get_user_api_key_auth_from_kwargs({"metadata": {"user_api_key_auth": auth}}) is auth assert ( skills_main._get_user_api_key_auth_from_kwargs( {"metadata": {}, "litellm_metadata": {"user_api_key_auth": auth}} @@ -122,9 +117,7 @@ def test_should_forward_skill_auth_through_sdk_entrypoints(monkeypatch): == "deleted" ) - assert handler.create_skill_handler.call_args.kwargs["metadata"] == { - "source": "request" - } + assert handler.create_skill_handler.call_args.kwargs["metadata"] == {"source": "request"} assert handler.create_skill_handler.call_args.kwargs["user_api_key_dict"] is auth assert handler.list_skills_handler.call_args.kwargs["user_api_key_dict"] is auth assert handler.get_skill_handler.call_args.kwargs["user_api_key_dict"] is auth @@ -149,9 +142,7 @@ def test_should_build_resource_owner_scopes_for_auth_context(): ] assert resource_ownership.get_primary_resource_owner_scope(auth) == "user-1" assert resource_ownership.user_can_access_resource_owner("team:team-1", auth) - assert resource_ownership.get_resource_owner_scopes( - UserAPIKeyAuth(token="token-hash") - ) == ["key:token-hash"] + assert resource_ownership.get_resource_owner_scopes(UserAPIKeyAuth(token="token-hash")) == ["key:token-hash"] # Identity-less callers get an empty scope set — sharing a sentinel # would collapse every identity-less caller into the same logical # owner, which is a cross-tenant data-access primitive. @@ -165,9 +156,7 @@ def test_should_allow_admin_and_anonymous_resource_owner_paths(): assert resource_ownership.is_proxy_admin(admin) assert resource_ownership.user_can_access_resource_owner(None, admin) assert resource_ownership.user_can_access_resource_owner(None, None) - assert not resource_ownership.user_can_access_resource_owner( - None, UserAPIKeyAuth(user_id="user-1") - ) + assert not resource_ownership.user_can_access_resource_owner(None, UserAPIKeyAuth(user_id="user-1")) @pytest.mark.asyncio @@ -218,9 +207,7 @@ async def test_should_forward_skill_auth_through_transformation_handler(monkeypa async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -242,9 +229,7 @@ async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): async def test_should_store_token_owner_for_keys_without_user_team_or_org(monkeypatch): table = AsyncMock() table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -268,9 +253,7 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc sentinel as ``created_by`` would let any two such callers see each other's skills via the resulting shared owner scope.""" table = AsyncMock() - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -291,9 +274,7 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypatch): table = AsyncMock() table.find_many.return_value = [_skill("litellm_skill_owner", "user-1")] - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -318,9 +299,7 @@ async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypat async def test_should_hide_skill_from_different_owner(monkeypatch): table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_other", "user-2") - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -340,9 +319,7 @@ async def test_should_hide_skill_from_different_owner(monkeypatch): async def test_should_hide_unowned_skill_by_default(monkeypatch): table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_unowned", None) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -364,9 +341,7 @@ async def test_list_skills_excludes_unowned_for_non_admin(monkeypatch): with ``created_by IS NULL`` are excluded — admin-only.""" table = AsyncMock() table.find_many.return_value = [] - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -422,9 +397,7 @@ async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") table = AsyncMock() table.find_unique = AsyncMock(return_value=fake_skill) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -432,10 +405,7 @@ async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): ) for _ in range(3): - assert ( - await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - is fake_skill - ) + assert await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") is fake_skill assert table.find_unique.await_count == 1 @@ -445,9 +415,7 @@ async def test_load_skill_caches_negative_lookups(monkeypatch): the DB and the caller still sees ``None``.""" table = AsyncMock() table.find_unique = AsyncMock(return_value=None) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -466,9 +434,7 @@ async def test_delete_skill_invalidates_cache(monkeypatch): table = AsyncMock() table.find_unique = AsyncMock(return_value=fake_skill) table.delete = AsyncMock() - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -480,12 +446,7 @@ async def test_delete_skill_invalidates_cache(monkeypatch): assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") is fake_skill auth = UserAPIKeyAuth(user_id="user-1") - await skills_handler.LiteLLMSkillsHandler.delete_skill( - "litellm_skill_a", user_api_key_dict=auth - ) + await skills_handler.LiteLLMSkillsHandler.delete_skill("litellm_skill_a", user_api_key_dict=auth) # Post-delete, the cache holds the negative sentinel — not the stale row. - assert ( - skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") - == skills_handler._NEGATIVE_SKILL_SENTINEL - ) + assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") == skills_handler._NEGATIVE_SKILL_SENTINEL From 99c2ef4d73efbfa657da009ef0a6a70205b03073 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:07:50 +0000 Subject: [PATCH 042/146] test(unit): block external sockets at import time and add a socket policy regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++------ tests/unit/test_socket_policy.py | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_socket_policy.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3bdab1d231a..017e63ed1b8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,9 +1,11 @@ -from collections.abc import Iterator +import os from typing import Final import pytest from pytest_socket import enable_socket, socket_allow_hosts +os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -11,13 +13,13 @@ def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) -@pytest.fixture(autouse=True, scope="session") -def block_external_sockets() -> Iterator[None]: - _allow_loopback_only() - yield - enable_socket() +_allow_loopback_only() @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() + + +def pytest_sessionfinish() -> None: + enable_socket() diff --git a/tests/unit/test_socket_policy.py b/tests/unit/test_socket_policy.py new file mode 100644 index 00000000000..f93794d1ba8 --- /dev/null +++ b/tests/unit/test_socket_policy.py @@ -0,0 +1,17 @@ +import socket + +import pytest +from pytest_socket import SocketConnectBlockedError + + +def test_external_connect_is_refused_before_a_packet_leaves() -> None: + with pytest.raises(SocketConnectBlockedError): + socket.create_connection(("192.0.2.1", 9), timeout=1) + + +def test_loopback_connect_is_allowed() -> None: + with socket.socket() as server: + server.bind(("127.0.0.1", 0)) + server.listen() + with socket.create_connection(server.getsockname(), timeout=1) as client: + assert client.getpeername() == server.getsockname() From d47008129c175a068b23c9e0afe306463e39974f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:08:12 +0000 Subject: [PATCH 043/146] test(llms): migrate phase 7 provider unit tests to tests/unit Move the wave 1 phase 7 batch (fireworks_ai, gemini, gigachat, github_copilot; 20 files) from tests/test_litellm to tests/unit after judging every test function under a behaviour mutation. Seven wiring or mock-echo tests that stayed green are deleted. The fireworks cost calculator tests get a local model_cost save/restore fixture since the tests/unit tree has no shared conftest for it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_fireworks_ai_chat_transformation.py | 0 ...test_fireworks_ai_rerank_transformation.py | 0 ...t_fireworks_ai_responses_transformation.py | 15 -------- .../test_fireworks_ai_cache_pricing.py | 0 .../test_fireworks_ai_common_utils.py | 0 .../test_fireworks_ai_cost_calculator.py | 22 ++++++++--- ...mini_audio_transcription_transformation.py | 0 .../files/test_gemini_files_transformation.py | 0 .../test_google_genai_guardrail_handler.py | 0 .../test_gemini_image_edit_transformation.py | 0 .../test_gemini_realtime_transformation.py | 0 .../test_gemini_video_transformation.py | 0 .../chat/test_gigachat_chat_streaming.py | 0 .../chat/test_gigachat_chat_transformation.py | 28 -------------- .../test_gigachat_embedding_transformation.py | 30 --------------- ...est_gigachat_passthrough_transformation.py | 0 .../llms/gigachat/test_authenticator.py | 0 .../llms/gigachat/test_file_handler.py | 38 ------------------- .../llms/gigachat/test_utils.py | 0 ...github_copilot_embedding_transformation.py | 0 20 files changed, 16 insertions(+), 117 deletions(-) rename tests/{test_litellm => unit}/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py (97%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py (90%) rename tests/{test_litellm => unit}/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/files/test_gemini_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/gemini/image_edit/test_gemini_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/realtime/test_gemini_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/videos/test_gemini_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/chat/test_gigachat_chat_streaming.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/chat/test_gigachat_chat_transformation.py (95%) rename tests/{test_litellm => unit}/llms/gigachat/embedding/test_gigachat_embedding_transformation.py (91%) rename tests/{test_litellm => unit}/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/test_authenticator.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/test_file_handler.py (91%) rename tests/{test_litellm => unit}/llms/gigachat/test_utils.py (100%) rename tests/{test_litellm => unit}/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py (100%) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py rename to tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py rename to tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py similarity index 97% rename from tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py rename to tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index d0697ca9b0e..05e3812152e 100644 --- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -406,21 +406,6 @@ def test_responses_call_sends_session_affinity_for_caller_session_id() -> None: assert headers["x-session-affinity"] == "sess-42" -def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) - pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"}) - with patch(HTTPX_CLIENT_FACTORY, return_value=client): - litellm.responses( - model="fireworks_ai/kimi-k3", - input="hi", - api_key="fw-test-key", - litellm_session_id="sess-42", - extra_headers=pinned, - ) - _, headers, _ = _sent_request(client) - assert headers["x-session-affinity"] == "explicit-node" - - def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None: client: Final = MagicMock() request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py similarity index 90% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index 52222f22a51..c6096ba2745 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -1,4 +1,5 @@ import math +from collections.abc import Generator from datetime import datetime, timezone from typing import Final @@ -24,6 +25,15 @@ CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="firew OUTPUT_COST = 4.4e-06 +@pytest.fixture(autouse=True) +def restore_model_cost() -> Generator[None, None, None]: + original: Final = litellm.model_cost + litellm.get_model_info.cache_clear() + yield + litellm.model_cost = original + litellm.get_model_info.cache_clear() + + def _usage(prompt_tokens: int, cached_tokens: int, completion_tokens: int) -> Usage: return Usage( prompt_tokens=prompt_tokens, @@ -57,7 +67,7 @@ def _register_off_peak_model( cache_read_cost: float | None = STANDARD_CACHE_READ_COST, model: str = OFF_PEAK_MODEL, ) -> None: - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -151,7 +161,7 @@ def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documente """Fireworks documents a default 50% cached-token discount for serverless models: https://docs.fireworks.ai/guides/prompt-caching, accessed 2026-09-19.""" model = "accounts/fireworks/models/default-cache-read-test" - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -171,7 +181,7 @@ def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documente def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): model = "accounts/fireworks/models/breakdown-cache-read-test" - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -204,7 +214,7 @@ def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): def test_generic_cost_per_token_applies_fireworks_cache_read_default_with_or_without_model_info(): model = "accounts/fireworks/models/generic-cache-read-test" - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -257,7 +267,7 @@ COMPONENT_AUDIO_OUT_COST = 6e-06 def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_rates(): - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, f"fireworks_ai/{COMPONENT_MODEL}": { "litellm_provider": "fireworks_ai", @@ -302,7 +312,7 @@ def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_ra def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback(): - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, # pyright: ignore[reportUnknownMemberType] # the SDK types model_cost as dict[Unknown, Unknown] "fireworks_ai/accounts/fireworks/models/no-input-rate-test": { "litellm_provider": "fireworks_ai", diff --git a/tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py b/tests/unit/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py rename to tests/unit/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/unit/llms/gemini/files/test_gemini_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py rename to tests/unit/llms/gemini/files/test_gemini_files_transformation.py diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py rename to tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py rename to tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py rename to tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py diff --git a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py b/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py rename to tests/unit/llms/gemini/videos/test_gemini_video_transformation.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py similarity index 100% rename from tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py rename to tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py similarity index 95% rename from tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py rename to tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index 2f9511e642c..8e84072e549 100644 --- a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -141,22 +141,6 @@ class TestValidateEnvironment: assert self.config._current_credentials == "my-creds" assert self.config._current_api_base == "https://my-api.example.com" - @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") - @patch(f"{TRANSFORM_MODULE}.get_secret_str") - def test_falls_back_to_env_for_credentials( # test-quality-ok: mock-echo of internal wiring - self, mock_get_secret, mock_get_token - ): - mock_get_secret.return_value = "env-creds" - self.config.validate_environment( - headers={}, - model="GigaChat", - messages=[], - optional_params={}, - litellm_params={}, - api_key=None, - api_base=None, - ) - mock_get_secret.assert_any_call("GIGACHAT_CREDENTIALS") # test-quality-ok: mock-echo of internal wiring class TestGetSupportedOpenAiParams: @@ -865,18 +849,6 @@ class TestUploadImage: def setup_method(self): self.config = GigaChatConfig() - @patch(f"{TRANSFORM_MODULE}.upload_file_sync", return_value="file-uploaded") - def test_upload_image_success(self, mock_upload): - self.config._current_credentials = "creds" - self.config._current_api_base = "https://api.example.com" - result = self.config._upload_image("https://example.com/img.jpg") - assert result == "file-uploaded" - mock_upload.assert_called_once_with( - image_url="https://example.com/img.jpg", - credentials="creds", - api_base="https://api.example.com", - ) - @patch(f"{TRANSFORM_MODULE}.upload_file_sync", side_effect=Exception("fail")) def test_upload_image_failure_returns_none(self, mock_upload): result = self.config._upload_image("https://example.com/img.jpg") diff --git a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py similarity index 91% rename from tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py rename to tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py index 8537793ea72..01fe66ca4c7 100644 --- a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py +++ b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -37,17 +37,6 @@ def _make_httpx_response(body: dict, status_code: int = 200) -> httpx.Response: # --------------------------------------------------------------------------- -class TestGetConfig: - def setup_method(self): - self.config = GigaChatEmbeddingConfig() - - def test_contains_only_abc_impl(self): - """get_config returns ABC internal data due to inheritance.""" - result = self.config.get_config() - # The only key should be _abc_impl from ABC base class - assert set(result.keys()) == {"_abc_impl"} - - class TestGetSupportedOpenAiParams: def setup_method(self): self.config = GigaChatEmbeddingConfig() @@ -287,25 +276,6 @@ class TestTransformEmbeddingResponse: ) assert result.model == "Embeddings" - def test_calls_logging_post_call(self): - raw = self._make_gigachat_response([ - {"object": "embedding", "embedding": [0.1], "index": 0}, - ]) - model_response = EmbeddingResponse() - self.config.transform_embedding_response( - model="gigachat/Embeddings", - raw_response=raw, - model_response=model_response, - logging_obj=self.logging_obj, - api_key="test-api-key", - request_data={"input": ["hello"]}, - optional_params={}, - litellm_params={}, - ) - self.logging_obj.post_call.assert_called_once() - args = self.logging_obj.post_call.call_args.kwargs - assert args["api_key"] == "test-api-key" - assert args["input"] == ["hello"] class TestValidateEnvironment: diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/unit/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py rename to tests/unit/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py diff --git a/tests/test_litellm/llms/gigachat/test_authenticator.py b/tests/unit/llms/gigachat/test_authenticator.py similarity index 100% rename from tests/test_litellm/llms/gigachat/test_authenticator.py rename to tests/unit/llms/gigachat/test_authenticator.py diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/unit/llms/gigachat/test_file_handler.py similarity index 91% rename from tests/test_litellm/llms/gigachat/test_file_handler.py rename to tests/unit/llms/gigachat/test_file_handler.py index ce9505f11f2..de83b2ddf5f 100644 --- a/tests/test_litellm/llms/gigachat/test_file_handler.py +++ b/tests/unit/llms/gigachat/test_file_handler.py @@ -344,25 +344,6 @@ class TestUploadFileSync: assert result is None - @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") - @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_uploads_without_optional_args( - self, mock_http_handler_cls, mock_get_token, mock_get_api_base - ): - """Verify that credentials, api_base, and litellm_params are optional.""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.json.return_value = {"id": "file-no-args"} - mock_response.raise_for_status = MagicMock() - mock_client.post.return_value = mock_response - mock_http_handler_cls.return_value = mock_client - - result = upload_file_sync(image_url=_RED_PNG_DATA_URL) - - assert result == "file-no-args" - # Should still have called get_access_token without args - mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) # --------------------------------------------------------------------------- @@ -483,22 +464,3 @@ class TestUploadFileAsync: ) assert result is None - - @pytest.mark.asyncio - @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") - @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") - @patch(f"{FILE_MODULE}.get_async_httpx_client") - async def test_uploads_without_optional_args( - self, mock_get_client, mock_get_token, mock_get_api_base - ): - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.json = MagicMock(return_value={"id": "async-no-args"}) - mock_response.raise_for_status = MagicMock() - mock_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_client - - result = await upload_file_async(image_url=_RED_PNG_DATA_URL) - - assert result == "async-no-args" - mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) \ No newline at end of file diff --git a/tests/test_litellm/llms/gigachat/test_utils.py b/tests/unit/llms/gigachat/test_utils.py similarity index 100% rename from tests/test_litellm/llms/gigachat/test_utils.py rename to tests/unit/llms/gigachat/test_utils.py diff --git a/tests/test_litellm/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py rename to tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py From 443c9f838533027f9d09d60afc5021c6561e8197 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:09:49 +0000 Subject: [PATCH 044/146] test: migrate nvidia, oci, ocr, oobabooga and openai legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../audio_transcription/__init__.py | 0 tests/test_litellm/llms/oci/embed/__init__.py | 0 .../ocr/guardrail_translation/__init__.py | 0 .../test_litellm/llms/openai/chat/__init__.py | 0 .../chat/guardrail_translation/__init__.py | 0 ...t_nvidia_nim_passthrough_transformation.py | 0 .../test_nvidia_nim_rerank_transformation.py | 0 .../audio_transcription/test_audio_utils.py | 13 -- .../audio_transcription/test_handler.py | 0 .../test_transformation.py | 0 .../oci/chat/test_oci_chat_transformation.py | 150 ------------------ .../test_oci_chat_transformation_for_14158.py | 0 .../oci/chat/test_oci_cohere_tool_calls.py | 20 --- .../llms/oci/chat/test_oci_generic_chat.py | 12 -- .../llms/oci/chat/test_oci_sse_splitter.py | 0 .../oci/chat/test_oci_streaming_tool_calls.py | 0 .../embed/test_oci_embed_transformation.py | 22 --- .../llms/oci/embed/test_oci_embedding.py | 0 .../test_ocr_guardrail_handler.py | 0 .../llms/oobabooga/chat/test_oobabooga.py | 0 .../test_openai_guardrail_handler.py | 19 --- .../chat/test_openai_gpt_transformation.py | 0 .../completion/test_completion_handler.py | 0 .../test_text_completion_guardrail_handler.py | 0 .../test_text_completion_token_ids.py | 0 25 files changed, 236 deletions(-) delete mode 100644 tests/test_litellm/llms/nvidia_riva/audio_transcription/__init__.py delete mode 100644 tests/test_litellm/llms/oci/embed/__init__.py delete mode 100644 tests/test_litellm/llms/ocr/guardrail_translation/__init__.py delete mode 100644 tests/test_litellm/llms/openai/chat/__init__.py delete mode 100644 tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py rename tests/{test_litellm => unit}/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/nvidia_riva/audio_transcription/test_audio_utils.py (90%) rename tests/{test_litellm => unit}/llms/nvidia_riva/audio_transcription/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/nvidia_riva/audio_transcription/test_transformation.py (100%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_chat_transformation.py (91%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_chat_transformation_for_14158.py (100%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_cohere_tool_calls.py (97%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_generic_chat.py (97%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_sse_splitter.py (100%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_streaming_tool_calls.py (100%) rename tests/{test_litellm => unit}/llms/oci/embed/test_oci_embed_transformation.py (95%) rename tests/{test_litellm => unit}/llms/oci/embed/test_oci_embedding.py (100%) rename tests/{test_litellm => unit}/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/oobabooga/chat/test_oobabooga.py (100%) rename tests/{test_litellm => unit}/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py (99%) rename tests/{test_litellm => unit}/llms/openai/chat/test_openai_gpt_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/completion/test_completion_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/completion/test_text_completion_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/completion/test_text_completion_token_ids.py (100%) diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/__init__.py b/tests/test_litellm/llms/nvidia_riva/audio_transcription/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/oci/embed/__init__.py b/tests/test_litellm/llms/oci/embed/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/ocr/guardrail_translation/__init__.py b/tests/test_litellm/llms/ocr/guardrail_translation/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/chat/__init__.py b/tests/test_litellm/llms/openai/chat/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py b/tests/unit/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py rename to tests/unit/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py diff --git a/tests/test_litellm/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py b/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py rename to tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_audio_utils.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_audio_utils.py similarity index 90% rename from tests/test_litellm/llms/nvidia_riva/audio_transcription/test_audio_utils.py rename to tests/unit/llms/nvidia_riva/audio_transcription/test_audio_utils.py index 63a53c2c97b..54fc30e6f2d 100644 --- a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_audio_utils.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_audio_utils.py @@ -62,19 +62,6 @@ def test_resample_16khz_mono_passes_through_int16_bytes_match_length(): assert resampled.duration_seconds == pytest.approx(1.0, abs=0.001) -def test_resample_preserves_int16_clip_range(): - sample_rate = 16000 - samples = np.array([2.0, -2.0, 0.0, 1.0], dtype=np.float32) - wav_in = _wav_bytes(samples, sample_rate) - - resampled = resample_to_riva_pcm(wav_in) - - decoded = np.frombuffer(resampled.pcm_bytes, dtype="= -32767 - - def test_unknown_format_raises_clear_error(): # 4 random bytes are not valid audio in any container we can decode. with pytest.raises(NvidiaRivaException) as excinfo: diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_handler.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py similarity index 100% rename from tests/test_litellm/llms/nvidia_riva/audio_transcription/test_handler.py rename to tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_transformation.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/nvidia_riva/audio_transcription/test_transformation.py rename to tests/unit/llms/nvidia_riva/audio_transcription/test_transformation.py diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py similarity index 91% rename from tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py rename to tests/unit/llms/oci/chat/test_oci_chat_transformation.py index 4c9bd29b337..708187b8ae1 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py @@ -922,91 +922,6 @@ class TestOCISignerSupport: assert wrapper.path_url == "/api/v1/chat" -class TestOCISplitChunks: - """ - Unit tests for the SSE split_chunks helpers used in sync and async streaming. - - These validate the fix for: - - Sync: JSONDecodeError when iter_text() returns chunks spanning multiple events - - Async: whitespace-only chunks being yielded before stripping (Greptile P2) - """ - - def _run_sync_split(self, raw_chunks): - """Invoke the sync split_chunks logic directly (extracted for testability).""" - results = [] - for item in raw_chunks: - for chunk in item.split("\n\n"): - stripped = chunk.strip() - if stripped: - results.append(stripped) - return results - - async def _run_async_split(self, raw_chunks): - """Invoke the async split_chunks logic directly.""" - results = [] - - async def _gen(): - for c in raw_chunks: - yield c - - async for item in _gen(): - for chunk in item.split("\n\n"): - stripped = chunk.strip() - if stripped: - results.append(stripped) - return results - - def test_sync_single_event_per_chunk(self): - """Normal case: one SSE event per iter_text() chunk.""" - chunks = ['data: {"text":"hello"}', 'data: {"text":"world"}'] - assert self._run_sync_split(chunks) == [ - 'data: {"text":"hello"}', - 'data: {"text":"world"}', - ] - - def test_sync_multiple_events_in_one_chunk(self): - """iter_text() returns two SSE events concatenated — must be split.""" - chunks = ['data: {"text":"a"}\n\ndata: {"text":"b"}'] - assert self._run_sync_split(chunks) == [ - 'data: {"text":"a"}', - 'data: {"text":"b"}', - ] - - def test_sync_whitespace_only_chunks_discarded(self): - """Whitespace between events must not be yielded.""" - chunks = ["data: {}\n\n \n\ndata: {}"] - result = self._run_sync_split(chunks) - assert result == ["data: {}", "data: {}"] - - def test_sync_empty_string_discarded(self): - """Empty string produced by splitting trailing \\n\\n must be discarded.""" - chunks = ["data: {}\n\n"] - assert self._run_sync_split(chunks) == ["data: {}"] - - @pytest.mark.asyncio - async def test_async_whitespace_only_chunks_discarded(self): - """ - Regression test for Greptile P2: async version was checking `if not chunk` - BEFORE stripping, so '\\n ' would pass the guard and yield '' downstream, - causing ValueError in chunk_creator ('Chunk does not start with data:'). - """ - chunks = ["data: {}\n\n \n\ndata: {}"] - result = await self._run_async_split(chunks) - assert result == ["data: {}", "data: {}"] - - @pytest.mark.asyncio - async def test_async_empty_string_discarded(self): - """Trailing \\n\\n must not produce an empty yielded chunk in async path.""" - chunks = ["data: {}\n\n"] - result = await self._run_async_split(chunks) - assert result == ["data: {}"] - - @pytest.mark.asyncio - async def test_async_multiple_events_in_one_chunk(self): - """Async path must split concatenated SSE events just like sync.""" - chunks = ['data: {"text":"x"}\n\ndata: {"text":"y"}'] - result = await self._run_async_split(chunks) - assert result == ['data: {"text":"x"}', 'data: {"text":"y"}'] class TestOCIProviderEmbeddingConfig: @@ -1026,21 +941,6 @@ class TestOCIProviderEmbeddingConfig: ) assert isinstance(config, OCIEmbedConfig) - def test_no_duplicate_oci_branch(self): - """ - Ensure utils.py does not contain two separate OCI embedding branches. - The dead code was removed in commit 64dfbe2b; this test guards against - regression (e.g. a future merge re-introducing it). - """ - import inspect - from litellm.utils import ProviderConfigManager - - source = inspect.getsource(ProviderConfigManager.get_provider_embedding_config) - oci_count = source.count("LlmProviders.OCI") - assert oci_count == 1, ( - f"Expected exactly 1 OCI branch in get_provider_embedding_config, found {oci_count}. " - "A duplicate dead-code branch may have been reintroduced." - ) class TestOCICohereParamMapping: @@ -1586,57 +1486,7 @@ def config(): class TestOCIKeyNormalization: """Tests for OCI private key content normalization.""" - def test_oci_key_with_escaped_newlines(self, config): - """Test that escaped newlines (\\n) are converted to actual newlines.""" - # Simulate PEM content with escaped newlines (as would come from JSON/UI input) - escaped_pem = "-----BEGIN RSA PRIVATE KEY-----\\nMIIEowIBAAKCAQEA...\\n-----END RSA PRIVATE KEY-----" - optional_params = { - "oci_user": "ocid1.user.oc1..test", - "oci_fingerprint": "aa:bb:cc:dd", - "oci_tenancy": "ocid1.tenancy.oc1..test", - "oci_region": "us-ashburn-1", - "oci_key": escaped_pem, - } - - # We can't fully test signing without a real key, but we can verify - # the error message indicates the key was processed (not a type error) - with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: - sign_with_manual_credentials( - headers={}, - optional_params=optional_params, - request_data={"test": "data"}, - api_base="https://test.oci.oraclecloud.com/api", - ) - - # The error should be about key format/loading, not about type - # This confirms the string was processed and newlines were normalized - error_message = str(exc_info.value) - assert "must be a string" not in error_message.lower() - - def test_oci_key_with_crlf_newlines(self, config): - """Test that Windows-style CRLF newlines are normalized to LF.""" - # Simulate PEM content with CRLF newlines - crlf_pem = "-----BEGIN RSA PRIVATE KEY-----\r\nMIIEowIBAAKCAQEA...\r\n-----END RSA PRIVATE KEY-----" - - optional_params = { - "oci_user": "ocid1.user.oc1..test", - "oci_fingerprint": "aa:bb:cc:dd", - "oci_tenancy": "ocid1.tenancy.oc1..test", - "oci_region": "us-ashburn-1", - "oci_key": crlf_pem, - } - - with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: - sign_with_manual_credentials( - headers={}, - optional_params=optional_params, - request_data={"test": "data"}, - api_base="https://test.oci.oraclecloud.com/api", - ) - - error_message = str(exc_info.value) - assert "must be a string" not in error_message.lower() def test_oci_key_rejects_non_string_type(self, config): """Test that non-string oci_key values raise OCIError.""" diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation_for_14158.py similarity index 100% rename from tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py rename to tests/unit/llms/oci/chat/test_oci_chat_transformation_for_14158.py diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py similarity index 97% rename from tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py rename to tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py index 729a2d25f41..9b06b01aa00 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -966,13 +966,6 @@ class TestOCICohereStreaming: completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging ) - def test_cohere_streaming_wrapper_initialization(self): - """Test OCIStreamWrapper initialization""" - stream_wrapper = self._create_stream_wrapper() - - # chunk_creator is the public dispatch entry point - assert hasattr(stream_wrapper, "chunk_creator") - assert callable(stream_wrapper.chunk_creator) def test_cohere_streaming_chunk_parsing(self): """Test parsing of Cohere streaming chunks""" @@ -1003,16 +996,3 @@ class TestOCICohereStreaming: # Test non-JSON chunk with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): stream_wrapper.chunk_creator("data: invalid json") - - def test_cohere_streaming_generic_chunk_fallback(self): - """Test fallback to generic chunk handling for non-Cohere chunks""" - stream_wrapper = self._create_stream_wrapper() - - # Test generic chunk (no apiFormat or different apiFormat) - generic_chunk = {"apiFormat": "GEMINI", "text": "Hello from Gemini"} - chunk_data = f"data: {json.dumps(generic_chunk)}" - - # This should fall back to generic handling - result = stream_wrapper.chunk_creator(chunk_data) - # The exact structure depends on the generic handler implementation - assert hasattr(result, "choices") diff --git a/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py b/tests/unit/llms/oci/chat/test_oci_generic_chat.py similarity index 97% rename from tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py rename to tests/unit/llms/oci/chat/test_oci_generic_chat.py index 0a47852d085..9ec5ab9aed4 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py +++ b/tests/unit/llms/oci/chat/test_oci_generic_chat.py @@ -450,15 +450,3 @@ class TestGpt5MaxCompletionTokens: ) assert out.get("maxTokens") == 64 assert "maxCompletionTokens" not in out - - def test_payload_serializes_max_completion_tokens(self): - from litellm.types.llms.oci import OCIChatRequestPayload - - payload = OCIChatRequestPayload( - apiFormat="GENERIC", - messages=[], - maxCompletionTokens=64, - ) - dumped = payload.model_dump(exclude_none=True) - assert dumped["maxCompletionTokens"] == 64 - assert "maxTokens" not in dumped diff --git a/tests/test_litellm/llms/oci/chat/test_oci_sse_splitter.py b/tests/unit/llms/oci/chat/test_oci_sse_splitter.py similarity index 100% rename from tests/test_litellm/llms/oci/chat/test_oci_sse_splitter.py rename to tests/unit/llms/oci/chat/test_oci_sse_splitter.py diff --git a/tests/test_litellm/llms/oci/chat/test_oci_streaming_tool_calls.py b/tests/unit/llms/oci/chat/test_oci_streaming_tool_calls.py similarity index 100% rename from tests/test_litellm/llms/oci/chat/test_oci_streaming_tool_calls.py rename to tests/unit/llms/oci/chat/test_oci_streaming_tool_calls.py diff --git a/tests/test_litellm/llms/oci/embed/test_oci_embed_transformation.py b/tests/unit/llms/oci/embed/test_oci_embed_transformation.py similarity index 95% rename from tests/test_litellm/llms/oci/embed/test_oci_embed_transformation.py rename to tests/unit/llms/oci/embed/test_oci_embed_transformation.py index 363c0b46809..4ffd79ff147 100644 --- a/tests/test_litellm/llms/oci/embed/test_oci_embed_transformation.py +++ b/tests/unit/llms/oci/embed/test_oci_embed_transformation.py @@ -269,28 +269,6 @@ class TestOCIEmbedConfig: assert result.model == "cohere.embed-v3.0" assert result.usage.prompt_tokens == 10 - def test_transform_response_no_usage(self): - cfg = self._config() - model_response = EmbeddingResponse() - raw = self._mock_response( - 200, - { - "embeddings": [[0.1]], - "modelId": "cohere.embed-v3.0", - "modelVersion": "3.0.0", - }, - ) - result = cfg.transform_embedding_response( - model="cohere.embed-v3.0", - raw_response=raw, - model_response=model_response, - logging_obj=MagicMock(), - api_key=None, - request_data={}, - optional_params={}, - litellm_params={}, - ) - assert len(result.data) == 1 def test_transform_response_http_error_raises(self): cfg = self._config() diff --git a/tests/test_litellm/llms/oci/embed/test_oci_embedding.py b/tests/unit/llms/oci/embed/test_oci_embedding.py similarity index 100% rename from tests/test_litellm/llms/oci/embed/test_oci_embedding.py rename to tests/unit/llms/oci/embed/test_oci_embedding.py diff --git a/tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py b/tests/unit/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py rename to tests/unit/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py diff --git a/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py b/tests/unit/llms/oobabooga/chat/test_oobabooga.py similarity index 100% rename from tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py rename to tests/unit/llms/oobabooga/chat/test_oobabooga.py diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py similarity index 99% rename from tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py rename to tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 258226ae22c..5c85faa5e13 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -545,25 +545,6 @@ class TestOpenAIChatCompletionsHandlerToolCallsInput: assert data["messages"][0]["content"] == "HELLO" assert data["messages"][1]["content"] == "HI THERE!" - @pytest.mark.asyncio - async def test_empty_tool_calls_list(self): - """Test that empty tool_calls list is handled correctly""" - handler = OpenAIChatCompletionsHandler() - guardrail = MockGuardrail() - - data = { - "messages": [ - {"role": "assistant", "content": "Hello", "tool_calls": []}, - ] - } - - # Process the input - await handler.process_input_messages(data, guardrail) - - # Verify empty tool_calls doesn't cause issues - assert guardrail.last_inputs is not None - tool_calls = guardrail.last_inputs.get("tool_calls", []) - assert len(tool_calls) == 0 class TestOpenAIChatCompletionsHandlerToolCallsOutput: diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py rename to tests/unit/llms/openai/chat/test_openai_gpt_transformation.py diff --git a/tests/test_litellm/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/completion/test_completion_handler.py rename to tests/unit/llms/openai/completion/test_completion_handler.py diff --git a/tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py rename to tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/completion/test_text_completion_token_ids.py b/tests/unit/llms/openai/completion/test_text_completion_token_ids.py similarity index 100% rename from tests/test_litellm/llms/openai/completion/test_text_completion_token_ids.py rename to tests/unit/llms/openai/completion/test_text_completion_token_ids.py From feca00248cf78fda0da5a1fa5b7f1a164fbe5608 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:11:56 +0000 Subject: [PATCH 045/146] test(bedrock): isolate host AWS config in realtime and rerank unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/realtime/test_bedrock_realtime_handler.py | 9 +++++++++ .../rerank/test_bedrock_rerank_header_forwarding.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index a7f0f64ef68..73a78a94e9f 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -18,6 +18,15 @@ from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +@pytest.fixture(autouse=True) +def _isolate_host_aws_config(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(env_var, raising=False) + + class FakePayloadPart: def __init__(self, bytes_): self.bytes_ = bytes_ diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 2ea61b5e978..aa93ddb21b8 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -15,6 +15,15 @@ from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + +@pytest.fixture(autouse=True) +def _isolate_host_aws_config(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(env_var, raising=False) + # Mock response for Bedrock rerank # Format based on Bedrock rerank API response structure bedrock_rerank_response = { From 65a4a009585973a6d328cd981866b4d8e5188c1a Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:12:39 +0000 Subject: [PATCH 046/146] fix(ci): excuse retired test-quality rules in the budget ratchet Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/budget_ratchet_check.py | 43 +++++++++++++++---- scripts/check_test_quality.py | 4 ++ .../test_litellm/test_budget_ratchet_check.py | 30 +++++++++++++ tests/test_litellm/test_check_test_quality.py | 15 +++++++ 4 files changed, 84 insertions(+), 8 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 485e118efd2..470adddfcd3 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -7,13 +7,17 @@ driven DOWN over time. This check compares every budget file against its own content at the merge-base with the target branch and fails (exits 1, red) if: * a rule's `limit` went up, - * a rule was dropped from a budget (its ceiling effectively became infinite), or + * a rule was dropped from a budget (its ceiling effectively became infinite) while + its checker still emits it, or * an entire budget file was deleted. New rules and lowered/equal limits are fine. So is a rule that graduated: once a paired config (ruff.toml for the ruff-strict budget) selects the rule outright it hard-fails at the first violation, which is stricter than any ceiling the budget could hold, so dropping its entry tightens the guard rather than removing it. +Likewise a retired rule: once the paired checker (check_test_quality.py for the +test-quality budget) no longer emits a code, its entry has no ceiling left to +loosen. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -29,11 +33,12 @@ Usage: from __future__ import annotations import argparse +import importlib.util import json import subprocess import sys from pathlib import Path -from types import MappingProxyType +from types import MappingProxyType, ModuleType from typing import Final, NamedTuple if sys.version_info >= (3, 11): @@ -49,6 +54,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "test-quality-budget.json", ) GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"}) +RETIREMENT_SOURCES = MappingProxyType({"test-quality-budget.json": "check_test_quality"}) class Regression(NamedTuple): @@ -139,20 +145,40 @@ def graduated_selectors(rel: str) -> tuple[str, ...]: ) +def _load_script(name: str) -> ModuleType: + if name in sys.modules: + return sys.modules[name] + spec: Final = importlib.util.spec_from_file_location(name, REPO_ROOT / "scripts" / f"{name}.py") + assert spec is not None and spec.loader is not None + module: Final = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def retired_rules(rel: str, base: dict) -> frozenset[str]: + """Rules in the base budget that the paired checker can no longer emit, so there is no ceiling to loosen.""" + source: Final = RETIREMENT_SOURCES.get(rel) + if source is None: + return frozenset() + return frozenset(_limits(base)) - _load_script(source).RULE_CODES + + def _regression_detail( rule: str, base_limits: dict[str, int], head_limits: dict[str, int], graduated: tuple[str, ...], + retired: frozenset[str] = frozenset(), ) -> str | None: - """Why `rule` regressed vs base, or None when it held flat, fell, or graduated. + """Why `rule` regressed vs base, or None when it held flat, fell, or left the budget legitimately. - A dropped rule is terminal unless it graduated; otherwise the only loosening - left is a raised limit. + A dropped rule is terminal unless it graduated or retired; otherwise the only + loosening left is a raised limit. """ base_limit = base_limits[rule] if rule not in head_limits: - if graduated and rule.startswith(graduated): + if rule in retired or (graduated and rule.startswith(graduated)): return None return f"rule dropped (limit {base_limit} -> removed)" if head_limits[rule] > base_limit: @@ -165,6 +191,7 @@ def regressions_for( base: dict | None, head: dict | None, graduated: tuple[str, ...] = (), + retired: frozenset[str] = frozenset(), ) -> list[Regression]: if base is None: return [] # new budget file: nothing to ratchet against yet @@ -175,7 +202,7 @@ def regressions_for( return [ Regression(rel, rule, detail) for rule in sorted(base_limits) - if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None + if (detail := _regression_detail(rule, base_limits, head_limits, graduated, retired)) is not None ] @@ -209,7 +236,7 @@ def main() -> int: print(f"skip {rel}: new file (no base at {base_ref} to ratchet against)") continue checked.append(rel) - regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) + regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel), retired_rules(rel, base))) if regressions: print( diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index dddc9d61982..9f93023cd53 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -146,6 +146,10 @@ SDK_MODULE: Final = "litellm" SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_call", "call")) INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P")) +RULE_CODES: Final = frozenset(( + "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", +)) + CREDENTIAL_NAME_RE: Final = re.compile( r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$" ) diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 22d05f4d00d..d6809b22161 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -92,6 +92,36 @@ def test_graduation_never_excuses_a_raised_limit(): assert "0 -> 7" in regs[0].detail +def test_dropped_rule_the_checker_retired_is_clean(): + base = {"TQ008": _spec_of(10993)} + assert ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) == [] + + +def test_dropped_rule_the_checker_still_emits_is_a_regression(): + base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + regs = ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) + assert [r.rule for r in regs] == ["TQ001"] + assert "dropped" in regs[0].detail + + +def test_retirement_never_excuses_a_raised_limit(): + base = {"TQ008": _spec_of(0)} + regs = ratchet.regressions_for("b.json", base, {"TQ008": _spec_of(7)}, retired=frozenset({"TQ008"})) + assert [r.rule for r in regs] == ["TQ008"] + assert "0 -> 7" in regs[0].detail + + +def test_retired_rules_come_from_the_paired_checker(): + base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + assert ratchet.retired_rules("test-quality-budget.json", base) == frozenset({"TQ008"}) + + +def test_budgets_without_a_paired_checker_never_retire(): + base = {"TQ008": _spec_of(1)} + for rel in ("ruff-strict-budget.json", "type-discipline-budget.json", "basedpyright-code-budget.json"): + assert ratchet.retired_rules(rel, base) == frozenset() + + def test_graduated_selectors_come_from_the_paired_ruff_config(): selectors = ratchet.graduated_selectors("ruff-strict-budget.json") assert "UP006" in selectors diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 05c25fb19fb..5a5c53fc31c 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -7,8 +7,10 @@ produced against tests/e2e, where the assertions live in a shared helper rather in the test body. """ +import ast import importlib.util import os +import re import subprocess import sys from pathlib import Path @@ -612,6 +614,19 @@ def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): assert all(" TQ001 " in line for line in reported) +def test_rule_codes_match_every_code_the_checker_emits(): + source = _MODULE_PATH.read_text(encoding="utf-8") + tree = ast.parse(source) + definition = next( + node + for node in tree.body + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "RULE_CODES" + ) + lines = source.splitlines() + outside = "\n".join(lines[: definition.lineno - 1] + lines[definition.end_lineno :]) + assert frozenset(re.findall(r'"(TQ\d{3})"', outside)) == checker.RULE_CODES + + def test_sys_executable_child_without_isolation_flag_is_flagged(tmp_path): source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n' assert _codes(tmp_path, source) == ["TQ009"] From ad523edb254fdb58e3a965339a0e44362d263320 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:13:11 +0000 Subject: [PATCH 047/146] test: migrate phase 9 legacy llm provider tests to tests/unit --- .../test_litellm_proxy_chat_transformation.py | 0 .../skills/test_code_execution.py | 0 .../litellm_proxy/skills/test_skill_search.py | 0 .../test_llamafile_chat_transformation.py | 0 tests/unit/llms/manus/__init__.py | 0 tests/unit/llms/manus/responses/__init__.py | 0 .../test_manus_responses_transformation.py | 0 .../test_meta_realtime_transformation.py | 0 .../test_meta_llama_chat_transformation.py | 0 tests/unit/llms/minimax/__init__.py | 0 tests/unit/llms/minimax/chat/__init__.py | 0 .../llms/minimax/chat/test_transformation.py | 96 ------------------- tests/unit/llms/minimax/messages/__init__.py | 0 .../minimax/messages/test_transformation.py | 74 -------------- tests/unit/llms/mistral/__init__.py | 0 ...est_mistral_audio_speech_transformation.py | 0 tests/unit/llms/mistral/batches/__init__.py | 0 .../test_mistral_batches_transformation.py | 0 tests/unit/llms/mistral/files/__init__.py | 0 .../test_mistral_files_transformation.py | 0 tests/unit/llms/mistral/ocr/__init__.py | 0 .../ocr/test_mistral_ocr_transformation.py | 0 ...est_modelscope_image_gen_transformation.py | 0 .../test_mongodb_transformation.py | 0 .../test_moonshot_chat_transformation.py | 25 ----- .../llms/neosantara/test_neosantara.py | 0 .../test_nimble_search_transformation.py | 0 .../chat/test_novita_chat_transformation.py | 9 -- .../chat/test_nscale_chat_transformation.py | 0 29 files changed, 204 deletions(-) rename tests/{test_litellm => unit}/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/litellm_proxy/skills/test_code_execution.py (100%) rename tests/{test_litellm => unit}/llms/litellm_proxy/skills/test_skill_search.py (100%) rename tests/{test_litellm => unit}/llms/llamafile/chat/test_llamafile_chat_transformation.py (100%) create mode 100644 tests/unit/llms/manus/__init__.py create mode 100644 tests/unit/llms/manus/responses/__init__.py rename tests/{test_litellm => unit}/llms/manus/responses/test_manus_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/meta/realtime/test_meta_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/meta_llama/test_meta_llama_chat_transformation.py (100%) create mode 100644 tests/unit/llms/minimax/__init__.py create mode 100644 tests/unit/llms/minimax/chat/__init__.py rename tests/{test_litellm => unit}/llms/minimax/chat/test_transformation.py (54%) create mode 100644 tests/unit/llms/minimax/messages/__init__.py rename tests/{test_litellm => unit}/llms/minimax/messages/test_transformation.py (57%) create mode 100644 tests/unit/llms/mistral/__init__.py rename tests/{test_litellm => unit}/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py (100%) create mode 100644 tests/unit/llms/mistral/batches/__init__.py rename tests/{test_litellm => unit}/llms/mistral/batches/test_mistral_batches_transformation.py (100%) create mode 100644 tests/unit/llms/mistral/files/__init__.py rename tests/{test_litellm => unit}/llms/mistral/files/test_mistral_files_transformation.py (100%) create mode 100644 tests/unit/llms/mistral/ocr/__init__.py rename tests/{test_litellm => unit}/llms/mistral/ocr/test_mistral_ocr_transformation.py (100%) rename tests/{test_litellm => unit}/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/mongodb/vector_stores/test_mongodb_transformation.py (100%) rename tests/{test_litellm => unit}/llms/moonshot/test_moonshot_chat_transformation.py (96%) rename tests/{test_litellm => unit}/llms/neosantara/test_neosantara.py (100%) rename tests/{test_litellm => unit}/llms/nimble/search/test_nimble_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/novita/chat/test_novita_chat_transformation.py (85%) rename tests/{test_litellm => unit}/llms/nscale/chat/test_nscale_chat_transformation.py (100%) diff --git a/tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py rename to tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_code_execution.py b/tests/unit/llms/litellm_proxy/skills/test_code_execution.py similarity index 100% rename from tests/test_litellm/llms/litellm_proxy/skills/test_code_execution.py rename to tests/unit/llms/litellm_proxy/skills/test_code_execution.py diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py similarity index 100% rename from tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py rename to tests/unit/llms/litellm_proxy/skills/test_skill_search.py diff --git a/tests/test_litellm/llms/llamafile/chat/test_llamafile_chat_transformation.py b/tests/unit/llms/llamafile/chat/test_llamafile_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/llamafile/chat/test_llamafile_chat_transformation.py rename to tests/unit/llms/llamafile/chat/test_llamafile_chat_transformation.py diff --git a/tests/unit/llms/manus/__init__.py b/tests/unit/llms/manus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/manus/responses/__init__.py b/tests/unit/llms/manus/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py b/tests/unit/llms/manus/responses/test_manus_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py rename to tests/unit/llms/manus/responses/test_manus_responses_transformation.py diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py rename to tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py diff --git a/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py b/tests/unit/llms/meta_llama/test_meta_llama_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py rename to tests/unit/llms/meta_llama/test_meta_llama_chat_transformation.py diff --git a/tests/unit/llms/minimax/__init__.py b/tests/unit/llms/minimax/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/minimax/chat/__init__.py b/tests/unit/llms/minimax/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/minimax/chat/test_transformation.py b/tests/unit/llms/minimax/chat/test_transformation.py similarity index 54% rename from tests/test_litellm/llms/minimax/chat/test_transformation.py rename to tests/unit/llms/minimax/chat/test_transformation.py index 9d51b556500..2645b2832aa 100644 --- a/tests/test_litellm/llms/minimax/chat/test_transformation.py +++ b/tests/unit/llms/minimax/chat/test_transformation.py @@ -2,14 +2,9 @@ Test MiniMax OpenAI-compatible API support """ -import os from unittest.mock import MagicMock, patch -import pytest - - import litellm -from litellm import completion from litellm.llms.minimax.chat.transformation import MinimaxChatConfig @@ -107,97 +102,6 @@ def test_minimax_provider_config_manager(): assert isinstance(config, MinimaxChatConfig) -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_basic(): - """Test basic chat completion with MiniMax OpenAI-compatible API""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello, how are you?"}, - ], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_with_reasoning_split(): - """Test completion with reasoning_split parameter (MiniMax M2.1 feature)""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Solve this problem: 2+2=?"}, - ], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - extra_body={"reasoning_split": True}, - ) - - assert response is not None - # Check if reasoning_details is present in response - if hasattr(response.choices[0].message, "reasoning_details"): - assert response.choices[0].message.reasoning_details is not None - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_with_tools(): - """Test completion with tool calling (function calling)""" - tools = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather in a location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - } - }, - "required": ["location"], - }, - }, - } - ] - - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "What's the weather in San Francisco?"}], - tools=tools, - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - ) - - assert response is not None - assert hasattr(response, "choices") - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_streaming(): - """Test streaming completion""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "Count to 5"}], - stream=True, - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - ) - - chunks = [] - for chunk in response: - chunks.append(chunk) - - assert len(chunks) > 0 - - if __name__ == "__main__": # Run basic tests that don't require API key print("Testing MiniMax Chat Config...") diff --git a/tests/unit/llms/minimax/messages/__init__.py b/tests/unit/llms/minimax/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/minimax/messages/test_transformation.py b/tests/unit/llms/minimax/messages/test_transformation.py similarity index 57% rename from tests/test_litellm/llms/minimax/messages/test_transformation.py rename to tests/unit/llms/minimax/messages/test_transformation.py index c7435a52890..a4b075414e3 100644 --- a/tests/test_litellm/llms/minimax/messages/test_transformation.py +++ b/tests/unit/llms/minimax/messages/test_transformation.py @@ -2,14 +2,9 @@ Test MiniMax Anthropic-compatible API support """ -import os from unittest.mock import MagicMock, patch -import pytest - - import litellm -from litellm import completion from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig @@ -58,75 +53,6 @@ def test_minimax_provider_config_manager(): assert config.custom_llm_provider == "minimax" -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_completion_basic(): - """Test basic completion with MiniMax Anthropic-compatible API""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "Hello, how are you?"}], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/anthropic/v1/messages", - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_completion_with_thinking(): - """Test completion with thinking parameter (MiniMax M2.1 feature)""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "Solve this problem: 2+2=?"}], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/anthropic/v1/messages", - thinking={"type": "enabled", "budget_tokens": 1000}, - ) - - assert response is not None - # Check if thinking content is present in response - for choice in response.choices: - if hasattr(choice.message, "content"): - # MiniMax returns thinking blocks similar to Anthropic - assert choice.message.content is not None - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_completion_with_tools(): - """Test completion with tool calling (function calling)""" - tools = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather in a location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - } - }, - "required": ["location"], - }, - }, - } - ] - - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "What's the weather in San Francisco?"}], - tools=tools, - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/anthropic/v1/messages", - ) - - assert response is not None - assert hasattr(response, "choices") - - if __name__ == "__main__": # Run basic tests that don't require API key print("Testing MiniMax Anthropic Config...") diff --git a/tests/unit/llms/mistral/__init__.py b/tests/unit/llms/mistral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py rename to tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py diff --git a/tests/unit/llms/mistral/batches/__init__.py b/tests/unit/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/unit/llms/mistral/batches/test_mistral_batches_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py rename to tests/unit/llms/mistral/batches/test_mistral_batches_transformation.py diff --git a/tests/unit/llms/mistral/files/__init__.py b/tests/unit/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/unit/llms/mistral/files/test_mistral_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py rename to tests/unit/llms/mistral/files/test_mistral_files_transformation.py diff --git a/tests/unit/llms/mistral/ocr/__init__.py b/tests/unit/llms/mistral/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py b/tests/unit/llms/mistral/ocr/test_mistral_ocr_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py rename to tests/unit/llms/mistral/ocr/test_mistral_ocr_transformation.py diff --git a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py rename to tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py diff --git a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py b/tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py similarity index 100% rename from tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py rename to tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py similarity index 96% rename from tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py rename to tests/unit/llms/moonshot/test_moonshot_chat_transformation.py index f94ea5e3db2..c39affc18a8 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py @@ -19,31 +19,6 @@ from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig class TestMoonshotConfig: """Test class for Moonshot AI functionality""" - def test_default_api_base(self): - """Test that default API base is used when none is provided""" - config = MoonshotChatConfig() - headers = {} - api_key = "fake-moonshot-key" - - # Call validate_environment without specifying api_base - result = config.validate_environment( - headers=headers, - model="moonshot-v1-8k", - messages=[{"role": "user", "content": "Hey"}], - optional_params={}, - litellm_params={}, - api_key=api_key, - api_base=None, # Not providing api_base - ) - - # Verify headers are still set correctly - assert result["Authorization"] == f"Bearer {api_key}" - assert result["Content-Type"] == "application/json" - - # We can't directly test the api_base value here since validate_environment - # only returns the headers, but we can verify it doesn't raise an exception - # which would happen if api_base handling was incorrect - def test_get_supported_openai_params(self): """Test that get_supported_openai_params returns correct params""" config = MoonshotChatConfig() diff --git a/tests/test_litellm/llms/neosantara/test_neosantara.py b/tests/unit/llms/neosantara/test_neosantara.py similarity index 100% rename from tests/test_litellm/llms/neosantara/test_neosantara.py rename to tests/unit/llms/neosantara/test_neosantara.py diff --git a/tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py b/tests/unit/llms/nimble/search/test_nimble_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py rename to tests/unit/llms/nimble/search/test_nimble_search_transformation.py diff --git a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py b/tests/unit/llms/novita/chat/test_novita_chat_transformation.py similarity index 85% rename from tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py rename to tests/unit/llms/novita/chat/test_novita_chat_transformation.py index 3f2a3f77c41..1381cf95a5d 100644 --- a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py +++ b/tests/unit/llms/novita/chat/test_novita_chat_transformation.py @@ -54,12 +54,3 @@ class TestNovitaConfig: ) assert "Missing Novita AI API Key" in str(excinfo.value) - - def test_inheritance(self): - """Test proper inheritance from OpenAIGPTConfig""" - config = NovitaConfig() - - from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig - - assert isinstance(config, OpenAIGPTConfig) - assert hasattr(config, "get_supported_openai_params") diff --git a/tests/test_litellm/llms/nscale/chat/test_nscale_chat_transformation.py b/tests/unit/llms/nscale/chat/test_nscale_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/nscale/chat/test_nscale_chat_transformation.py rename to tests/unit/llms/nscale/chat/test_nscale_chat_transformation.py From fcabb626acdcf240c43ba78dfb92d7822e239cfe Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:19:05 +0000 Subject: [PATCH 048/146] test(unit): migrate wave 1 phase 3 anthropic, apiserpent, azure and azure_ai legacy tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 11 +++ .../test_reasoning_effort_fields.py | 0 .../test_anthropic_files_transformation.py | 0 .../messages/test_advisor_orchestration.py | 0 .../llms/apiserpent/test_apiserpent_search.py | 0 .../test_azure_image_edit_transformation.py | 0 .../test_azure_image_generation_init.py | 82 ------------------- .../test_azure_passthrough_transformation.py | 0 .../realtime/test_azure_realtime_handler.py | 35 -------- .../response/test_azure_transformation.py | 0 .../foundry_responses_web_search_fixture.json | 0 ...st_bing_grounding_search_transformation.py | 0 .../test_azure_tts_transformation.py | 0 ...test_azure_vector_stores_transformation.py | 0 .../chat/test_azure_ai_transformation.py | 15 ---- .../embed/test_azure_ai_embed_handler.py | 0 ...test_azure_ai_image_edit_transformation.py | 0 .../test_mai_image_edit_transformation.py | 0 ...st_azure_ai_cohere_parse_transformation.py | 0 ...est_azure_ai_passthrough_transformation.py | 0 .../test_azure_ai_rerank_transformation.py | 0 .../test_azure_ai_responses_transformation.py | 0 22 files changed, 11 insertions(+), 132 deletions(-) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/files/test_anthropic_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/messages/test_advisor_orchestration.py (100%) rename tests/{test_litellm => unit}/llms/apiserpent/test_apiserpent_search.py (100%) rename tests/{test_litellm => unit}/llms/azure/image_edit/test_azure_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/image_generation/test_azure_image_generation_init.py (91%) rename tests/{test_litellm => unit}/llms/azure/passthrough/test_azure_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/realtime/test_azure_realtime_handler.py (94%) rename tests/{test_litellm => unit}/llms/azure/response/test_azure_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/search/foundry_responses_web_search_fixture.json (100%) rename tests/{test_litellm => unit}/llms/azure/search/test_bing_grounding_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/text_to_speech/test_azure_tts_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/vector_stores/test_azure_vector_stores_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/chat/test_azure_ai_transformation.py (97%) rename tests/{test_litellm => unit}/llms/azure_ai/embed/test_azure_ai_embed_handler.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/responses/test_azure_ai_responses_transformation.py (100%) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3bdab1d231a..d002452d14c 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,6 +2,8 @@ from collections.abc import Iterator from typing import Final import pytest + +import litellm from pytest_socket import enable_socket, socket_allow_hosts LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -21,3 +23,12 @@ def block_external_sockets() -> Iterator[None]: @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() + + +@pytest.fixture +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py b/tests/unit/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py rename to tests/unit/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py rename to tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py diff --git a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py similarity index 100% rename from tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py rename to tests/unit/llms/anthropic/messages/test_advisor_orchestration.py diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/unit/llms/apiserpent/test_apiserpent_search.py similarity index 100% rename from tests/test_litellm/llms/apiserpent/test_apiserpent_search.py rename to tests/unit/llms/apiserpent/test_apiserpent_search.py diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/unit/llms/azure/image_edit/test_azure_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py rename to tests/unit/llms/azure/image_edit/test_azure_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py similarity index 91% rename from tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py rename to tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py index cfde1760389..eabd5c8427d 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py @@ -133,88 +133,6 @@ def test_azure_image_generation_flattens_extra_body(): assert data["size"] == "1024x1024" -def test_azure_image_generation_creates_token_provider_from_credentials(): - """ - Test that azure_ad_token_provider is created from tenant_id, client_id, client_secret. - - This test verifies the fix in images/main.py where we now create the - azure_ad_token_provider from credentials in litellm_params if it's not already provided. - """ - # Simulate the fix in images/main.py - litellm_params_dict = { - "tenant_id": "test-tenant-id", - "client_id": "test-client-id", - "client_secret": "test-client-secret", - "azure_scope": None, - } - - azure_ad_token_provider = None - - # This is the logic we added in images/main.py - if azure_ad_token_provider is None: - tenant_id = litellm_params_dict.get("tenant_id") - client_id = litellm_params_dict.get("client_id") - client_secret = litellm_params_dict.get("client_secret") - azure_scope = ( - litellm_params_dict.get("azure_scope") - or "https://cognitiveservices.azure.com/.default" - ) - - # Verify the credentials are extracted correctly - assert tenant_id == "test-tenant-id" - assert client_id == "test-client-id" - assert client_secret == "test-client-secret" - assert azure_scope == "https://cognitiveservices.azure.com/.default" - - # Verify the condition to create token provider is met - assert ( - tenant_id and client_id and client_secret - ), "Credentials should be present to create token provider" - - -def test_azure_image_generation_headers_without_api_key(): - """ - Test that when api_key is None, the api-key header is not added to headers. - - This prevents the httpx TypeError: "Header value must be str or bytes, not " - that was occurring when api_key was None and being set in headers. - - This is a unit test for the fix in images/main.py where we now check: - if api_key is not None: - default_headers["api-key"] = api_key - """ - from litellm.images.main import image_generation - - # Test the header building logic directly - api_key = None - - default_headers = { - "Content-Type": "application/json", - } - - # This is the fix: only add api-key if it's not None - if api_key is not None: - default_headers["api-key"] = api_key - - # Verify api-key is not in headers when api_key is None - assert "api-key" not in default_headers - - # Verify Content-Type is still there - assert default_headers["Content-Type"] == "application/json" - - # Test with a valid api_key - api_key = "valid-key-123" - default_headers_with_key = { - "Content-Type": "application/json", - } - if api_key is not None: - default_headers_with_key["api-key"] = api_key - - # Verify api-key is added when api_key is valid - assert "api-key" in default_headers_with_key - assert default_headers_with_key["api-key"] == "valid-key-123" - - def test_azure_image_generation_drop_params_response_format(): """ Test that unsupported params like response_format are dropped when drop_params=True. diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py rename to tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py similarity index 94% rename from tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py rename to tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index 7d24e604569..c1ba286f8c0 100644 --- a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -426,41 +426,6 @@ async def test_async_realtime_beta_without_api_version_raises(): ) -@pytest.mark.asyncio -async def test_realtime_protocol_env_var_fallback(): - """ - Test that LITELLM_AZURE_REALTIME_PROTOCOL env var is used as fallback. - Fixes #22127: no way to set realtime_protocol from config. - """ - from litellm.realtime_api.main import _arealtime - from litellm.types.router import GenericLiteLLMParams - - with patch.dict(os.environ, {"LITELLM_AZURE_REALTIME_PROTOCOL": "v1"}): - # Create a GenericLiteLLMParams without realtime_protocol - litellm_params = GenericLiteLLMParams() - # The env var should be picked up as fallback - realtime_protocol = ( - {}.get("realtime_protocol") - or litellm_params.get("realtime_protocol") - or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") - or "beta" - ) - assert realtime_protocol == "v1" - - -@pytest.mark.asyncio -async def test_realtime_protocol_from_litellm_params(): - """ - Test that realtime_protocol is read from litellm_params (config.yaml extra field). - Fixes #22127: realtime_protocol in litellm_params was not used. - """ - from litellm.types.router import GenericLiteLLMParams - - # Simulate config.yaml with realtime_protocol as an extra field - litellm_params = GenericLiteLLMParams(realtime_protocol="GA") - assert litellm_params.get("realtime_protocol") == "GA" - - @pytest.mark.asyncio async def test_arealtime_transcription_intent_defaults_to_ga(monkeypatch): """ diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/response/test_azure_transformation.py rename to tests/unit/llms/azure/response/test_azure_transformation.py diff --git a/tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json b/tests/unit/llms/azure/search/foundry_responses_web_search_fixture.json similarity index 100% rename from tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json rename to tests/unit/llms/azure/search/foundry_responses_web_search_fixture.json diff --git a/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py rename to tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py diff --git a/tests/test_litellm/llms/azure/text_to_speech/test_azure_tts_transformation.py b/tests/unit/llms/azure/text_to_speech/test_azure_tts_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/text_to_speech/test_azure_tts_transformation.py rename to tests/unit/llms/azure/text_to_speech/test_azure_tts_transformation.py diff --git a/tests/test_litellm/llms/azure/vector_stores/test_azure_vector_stores_transformation.py b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/vector_stores/test_azure_vector_stores_transformation.py rename to tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py similarity index 97% rename from tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py rename to tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index f8cc0b5071e..e4a33d5772c 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -352,21 +352,6 @@ def test_azure_model_router_stamps_selected_model_on_hidden_params(): ) -def test_azure_model_router_stamp_does_not_leak_across_responses(): - """ - ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written - as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. - """ - from litellm.llms.azure_ai.common_utils import ( - AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, - ) - from litellm.types.utils import ModelResponse - - untouched = ModelResponse() - - assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {}) - - def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): """ Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name. diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/unit/llms/azure_ai/embed/test_azure_ai_embed_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py rename to tests/unit/llms/azure_ai/embed/test_azure_ai_embed_handler.py diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py rename to tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py rename to tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py b/tests/unit/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py rename to tests/unit/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py rename to tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py rename to tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py rename to tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py From 6d327fff6f812b9c7dbbbb1ecc9b997a7807146e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 08:21:28 +0000 Subject: [PATCH 049/146] test(bedrock): keep beta headers fixture teardown off the network --- ...oke_transformations_anthropic_claude3_transformation.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 2c74d23a6a2..cf2fd78a896 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -86,11 +86,8 @@ def local_beta_headers_config(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") reload_beta_headers_config() - try: - yield - finally: - monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) - reload_beta_headers_config() + yield + reload_beta_headers_config() def test_get_supported_params_thinking(): From cf2a9b372cffe0e00e44ee252168384c8f5c058d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:21:51 +0000 Subject: [PATCH 050/146] test(gigachat): cover env credential fallback by its resulting auth header Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_gigachat_chat_transformation.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index 8e84072e549..b1307f56336 100644 --- a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -141,6 +141,25 @@ class TestValidateEnvironment: assert self.config._current_credentials == "my-creds" assert self.config._current_api_base == "https://my-api.example.com" + @patch( + f"{TRANSFORM_MODULE}.get_access_token", + side_effect=lambda credentials, litellm_params: f"token-for-{credentials}", + ) + def test_falls_back_to_env_credentials_when_api_key_missing( + self, mock_get_token, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setenv("GIGACHAT_CREDENTIALS", "env-creds") + result = self.config.validate_environment( + headers={}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + assert result["Authorization"] == "Bearer token-for-env-creds" + assert self.config._current_credentials == "env-creds" class TestGetSupportedOpenAiParams: From 129c4a703b7224965b7d0c6ff402a4db2b85a635 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 20 Sep 2026 01:21:01 -0700 Subject: [PATCH 051/146] fix(proxy): only warn about the deprecated flag when it came from the CLI USE_V2_MIGRATION_RESOLVER=true is a supported way to select v2, but click sets the same parameter from that env var, so the deprecation notice fired for environment-based config that is not deprecated. The notice now keys off click's parameter source. Also drops an em dash from the notice, and moves the resolver decision under mock-free tests by making it take the env value as an argument. --- litellm/proxy/proxy_cli.py | 25 ++++++--- tests/test_litellm/proxy/test_proxy_cli.py | 63 ++++++++++++++-------- 2 files changed, 58 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 14d0331c0ff..78885461724 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final import click import httpx +from click.core import ParameterSource from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict @@ -181,12 +182,21 @@ def append_query_params(url: str | None, params: dict) -> str: return modified_url -def resolve_v2_migration_resolver(*, use_legacy_flag: bool) -> bool: +def resolve_v2_migration_resolver(*, use_legacy_flag: bool, env_value: str | None) -> bool: from litellm_proxy_extras.utils import str_to_bool if use_legacy_flag: return False - return bool(str_to_bool(os.getenv("USE_V2_MIGRATION_RESOLVER", "true"))) + if env_value is None: + return True + return bool(str_to_bool(env_value)) + + +def deprecated_v2_flag_passed_on_cli() -> bool: + ctx: Final = click.get_current_context(silent=True) + if ctx is None: + return False + return ctx.get_parameter_source("use_v2_migration_resolver") is ParameterSource.COMMANDLINE class ProxyInitializationHelpers: @@ -1368,14 +1378,15 @@ def run_server( check_prisma_schema_diff(db_url=None) else: use_v2_resolver: Final = resolve_v2_migration_resolver( - use_legacy_flag=use_legacy_migration_resolver + use_legacy_flag=use_legacy_migration_resolver, + env_value=os.getenv("USE_V2_MIGRATION_RESOLVER"), ) - if use_v2_migration_resolver and use_v2_resolver: + if deprecated_v2_flag_passed_on_cli() and use_v2_resolver: print( "\033[1;33mLiteLLM Proxy: --use_v2_migration_resolver is " - "deprecated and has no effect \u2014 the v2 migration resolver " - "is now the default. You can safely remove it. To opt back " - "into the legacy v1 resolver, pass " + "deprecated and has no effect, because the v2 migration " + "resolver is now the default. You can safely remove it. To " + "opt back into the legacy v1 resolver, pass " "--use_legacy_migration_resolver.\033[0m" ) if not use_v2_resolver: diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index d2835142194..a38470d1fdf 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2203,6 +2203,7 @@ class TestRunServerDbSetup: mock_setup_database, mock_atexit_register, mock_subprocess_run, + capsys, ): """USE_V2_MIGRATION_RESOLVER=true must select the v2 resolver. @@ -2248,44 +2249,58 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_once_with( use_migrate=True, use_v2_resolver=True ) + assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out @pytest.mark.parametrize( - "argv_extra, env_extra, expected_v2", + "use_legacy_flag, env_value, expected", [ - ([], {}, True), - ([], {"USE_V2_MIGRATION_RESOLVER": "false"}, False), - (["--use_legacy_migration_resolver"], {}, False), - ( - ["--use_legacy_migration_resolver"], - {"USE_V2_MIGRATION_RESOLVER": "true"}, - False, - ), - (["--use_v2_migration_resolver"], {}, True), + (False, None, True), + (False, "true", True), + (False, "false", False), + (True, None, False), + (True, "true", False), ], ids=[ - "default-is-v2", - "env-false-opts-out", - "legacy-flag-opts-out", + "unset-env-defaults-to-v2", + "env-true-selects-v2", + "env-false-selects-v1", + "legacy-flag-selects-v1", "legacy-flag-beats-env-true", - "deprecated-v2-flag-still-accepted", ], ) + def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected): + from litellm.proxy.proxy_cli import resolve_v2_migration_resolver + + assert ( + resolve_v2_migration_resolver( + use_legacy_flag=use_legacy_flag, env_value=env_value + ) + is expected + ) + + def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self): + from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli + + assert deprecated_v2_flag_passed_on_cli() is False + @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_selection( + def test_legacy_resolver_flag_reaches_database_setup( self, mock_should_update_schema, mock_check_schema_diff, mock_setup_database, mock_atexit_register, mock_subprocess_run, - argv_extra, - env_extra, - expected_v2, ): + """--use_legacy_migration_resolver must reach the database setup call. + + The resolver decision itself is covered mock-free above; this is the + one wiring check that the flag is threaded through run_server. + """ from litellm.proxy.proxy_cli import run_server mock_subprocess_run.return_value = MagicMock(returncode=0) @@ -2302,11 +2317,9 @@ class TestRunServerDbSetup: clean_env = { k: v for k, v in os.environ.items() - if k - not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") } clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" - clean_env.update(env_extra) with ( patch.dict(os.environ, clean_env, clear=True), @@ -2319,12 +2332,16 @@ class TestRunServerDbSetup: ), ): run_server.main( - ["--local", "--skip_server_startup", *argv_extra], + [ + "--local", + "--skip_server_startup", + "--use_legacy_migration_resolver", + ], standalone_mode=False, ) mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=expected_v2 + use_migrate=True, use_v2_resolver=False ) From 9850cd14f7e72bbe99f0b025cbb722d9a18cb523 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:27:02 +0000 Subject: [PATCH 052/146] test(ci): prove RULE_CODES by running every checker rule Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/budget_ratchet_check.py | 2 +- tests/test_litellm/test_check_test_quality.py | 49 ++++++++++++++----- 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 470adddfcd3..3ca5e9f3e9d 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -156,7 +156,7 @@ def _load_script(name: str) -> ModuleType: return module -def retired_rules(rel: str, base: dict) -> frozenset[str]: +def retired_rules(rel: str, base: dict[str, object]) -> frozenset[str]: """Rules in the base budget that the paired checker can no longer emit, so there is no ceiling to loosen.""" source: Final = RETIREMENT_SOURCES.get(rel) if source is None: diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 5a5c53fc31c..2a0e32d4de7 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -7,13 +7,13 @@ produced against tests/e2e, where the assertions live in a shared helper rather in the test body. """ -import ast import importlib.util import os -import re import subprocess import sys from pathlib import Path +from types import MappingProxyType +from typing import Final import pytest @@ -614,17 +614,42 @@ def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): assert all(" TQ001 " in line for line in reported) -def test_rule_codes_match_every_code_the_checker_emits(): - source = _MODULE_PATH.read_text(encoding="utf-8") - tree = ast.parse(source) - definition = next( - node - for node in tree.body - if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "RULE_CODES" +_VIOLATING_SNIPPETS: Final = MappingProxyType( + { + "TQ000": ("test_snippet.py", "def test_broken(:\n pass\n"), + "TQ001": ("test_snippet.py", "def test_nothing():\n compute()\n"), + "TQ002": ( + "test_snippet.py", + "from unittest.mock import patch\n" + "\n" + "\n" + "def test_echo():\n" + " with patch('litellm.completion') as mock_completion:\n" + " run()\n" + " mock_completion.assert_called_once()\n", + ), + "TQ003": ("test_snippet.py", "import sys\n\nsys.path.insert(0, '..')\n"), + "TQ004": ("test_snippet.py", "import os\n\nos.environ['KEY'] = 'v'\n"), + "TQ005": ("test_snippet.py", "import litellm\n\nlitellm.drop_params = True\n"), + "TQ006": ("test_snippet.py", _DIRECT_GATE), + "TQ007": ("conftest.py", _SNAPSHOT_CONFTEST), + "TQ009": ( + "test_snippet.py", + 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n', + ), + } +) + + +def test_rule_codes_match_every_code_the_checker_emits(tmp_path): + emitted = frozenset( + v.code + for name, source in _VIOLATING_SNIPPETS.values() + for v in checker.check_file(_written(tmp_path, source, name)) ) - lines = source.splitlines() - outside = "\n".join(lines[: definition.lineno - 1] + lines[definition.end_lineno :]) - assert frozenset(re.findall(r'"(TQ\d{3})"', outside)) == checker.RULE_CODES + for code, (name, source) in _VIOLATING_SNIPPETS.items(): + assert code in [v.code for v in checker.check_file(_written(tmp_path, source, name))], code + assert emitted == checker.RULE_CODES def test_sys_executable_child_without_isolation_flag_is_flagged(tmp_path): From decb28f5b575378cff5e8ce941c999da24df3875 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:34:55 +0000 Subject: [PATCH 053/146] test(unit): restore live router and runtime model cost state between unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index d002452d14c..d47f71b21a8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,9 +2,11 @@ from collections.abc import Iterator from typing import Final import pytest +from pytest_socket import enable_socket, socket_allow_hosts import litellm -from pytest_socket import enable_socket, socket_allow_hosts +import litellm.router as litellm_router_module +import litellm.utils as litellm_utils_module LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -25,6 +27,24 @@ def pytest_runtest_setup() -> None: _allow_loopback_only() +@pytest.fixture(autouse=True) +def isolate_router_model_cost_state() -> Iterator[None]: + original_live_routers: Final = frozenset(litellm_router_module._live_routers) + original_runtime_registered_model_cost: Final = { + model_key: dict(model_value) + for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() + } + yield + for router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(router) + for router in original_live_routers: + litellm_router_module._live_routers.add(router) + litellm_utils_module._runtime_registered_model_cost.clear() + litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + litellm_utils_module._invalidate_model_cost_lowercase_map() + litellm.get_model_info.cache_clear() + + @pytest.fixture def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") From bbfa853d844b8c54c66746a416cafeef93f15b8f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:07:51 +0000 Subject: [PATCH 054/146] test(ci): annotate new test locals as Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_budget_ratchet_check.py | 15 ++++++++------- tests/test_litellm/test_check_test_quality.py | 2 +- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index d6809b22161..b359b8b42e5 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -9,6 +9,7 @@ import importlib.util import subprocess import sys from pathlib import Path +from typing import Final _MODULE_PATH = ( Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" @@ -93,31 +94,31 @@ def test_graduation_never_excuses_a_raised_limit(): def test_dropped_rule_the_checker_retired_is_clean(): - base = {"TQ008": _spec_of(10993)} + base: Final = {"TQ008": _spec_of(10993)} assert ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) == [] def test_dropped_rule_the_checker_still_emits_is_a_regression(): - base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} - regs = ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) + base: Final = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + regs: Final = ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) assert [r.rule for r in regs] == ["TQ001"] assert "dropped" in regs[0].detail def test_retirement_never_excuses_a_raised_limit(): - base = {"TQ008": _spec_of(0)} - regs = ratchet.regressions_for("b.json", base, {"TQ008": _spec_of(7)}, retired=frozenset({"TQ008"})) + base: Final = {"TQ008": _spec_of(0)} + regs: Final = ratchet.regressions_for("b.json", base, {"TQ008": _spec_of(7)}, retired=frozenset({"TQ008"})) assert [r.rule for r in regs] == ["TQ008"] assert "0 -> 7" in regs[0].detail def test_retired_rules_come_from_the_paired_checker(): - base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + base: Final = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} assert ratchet.retired_rules("test-quality-budget.json", base) == frozenset({"TQ008"}) def test_budgets_without_a_paired_checker_never_retire(): - base = {"TQ008": _spec_of(1)} + base: Final = {"TQ008": _spec_of(1)} for rel in ("ruff-strict-budget.json", "type-discipline-budget.json", "basedpyright-code-budget.json"): assert ratchet.retired_rules(rel, base) == frozenset() diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 2a0e32d4de7..bf05775d09d 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -642,7 +642,7 @@ _VIOLATING_SNIPPETS: Final = MappingProxyType( def test_rule_codes_match_every_code_the_checker_emits(tmp_path): - emitted = frozenset( + emitted: Final = frozenset( v.code for name, source in _VIOLATING_SNIPPETS.values() for v in checker.check_file(_written(tmp_path, source, name)) From 03a650c8a95802633c792d290a95bb8cbbb1159c Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:17:57 +0000 Subject: [PATCH 055/146] test: migrate wave 1 phase 1 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 153 ------------------ .../test_pydantic_ai_agent_headers.py | 0 .../test_pydantic_ai_agent_transformation.py | 0 ...test_watsonx_orchestrate_transformation.py | 18 --- .../test_exception_mapping_utils.py | 0 .../batches/test_batch_utils.py | 0 .../batches/test_main.py | 0 .../batches/test_responses_batch_cost.py | 25 +-- .../chat_completions/test_dispatch.py | 13 +- ...responses_transformation_transformation.py | 0 .../compression/test_compress.py | 0 .../test_transformation.py | 0 .../test_callback_controls.py | 0 .../enterprise_callbacks/test_llm_guard.py | 0 .../test_secret_detection.py | 0 .../test_compression_interception_handler.py | 0 .../gcs_bucket/test_gcs_bucket_base.py | 0 .../integrations/gcs_pubsub/test_pub_sub.py | 0 .../helicone/test_helicone_gemini.py | 34 ---- 19 files changed, 14 insertions(+), 229 deletions(-) delete mode 100644 tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py rename tests/{test_litellm => unit}/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py (100%) rename tests/{test_litellm => unit}/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py (100%) rename tests/{test_litellm => unit}/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py (95%) rename tests/{test_litellm => unit}/anthropic_interface/exceptions/test_exception_mapping_utils.py (100%) rename tests/{test_litellm => unit}/batches/test_batch_utils.py (100%) rename tests/{test_litellm => unit}/batches/test_main.py (100%) rename tests/{test_litellm => unit}/batches/test_responses_batch_cost.py (87%) rename tests/{test_litellm => unit}/chat_completions/test_dispatch.py (93%) rename tests/{test_litellm => unit}/completion_extras/test_litellm_responses_transformation_transformation.py (100%) rename tests/{test_litellm => unit}/compression/test_compress.py (100%) rename tests/{test_litellm => unit}/endpoints/speech/speech_to_completion_bridge/test_transformation.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_callback_controls.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_llm_guard.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_secret_detection.py (100%) rename tests/{test_litellm => unit}/integrations/compression_interception/test_compression_interception_handler.py (100%) rename tests/{test_litellm => unit}/integrations/gcs_bucket/test_gcs_bucket_base.py (100%) rename tests/{test_litellm => unit}/integrations/gcs_pubsub/test_pub_sub.py (100%) rename tests/{test_litellm => unit}/integrations/helicone/test_helicone_gemini.py (73%) diff --git a/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py b/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py deleted file mode 100644 index 8036c72679e..00000000000 --- a/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py +++ /dev/null @@ -1,153 +0,0 @@ -""" -Regression test for https://github.com/BerriAI/litellm/issues/28505 - -the Responses API bridge double-strips the provider prefix from the -model name when a Chat Completions request has both `tools` and -`reasoning_effort`. - -Root cause: the bridge handler called `litellm.responses()` / -`litellm.aresponses()` without passing the already-resolved -`custom_llm_provider`. The downstream call then re-invoked -`get_llm_provider()` with `custom_llm_provider=None`, which stripped -a second provider prefix from a `provider/provider/model` deployment -string. - -This test pins both the sync and async bridge handler call sites: -the resolved `custom_llm_provider` must be forwarded to the underlying -`responses` / `aresponses` call so the provider isn't re-detected. -""" - -from unittest.mock import MagicMock, patch - -import pytest - -from litellm.completion_extras.litellm_responses_transformation.handler import ( - ResponsesToCompletionBridgeHandler, -) - - -def _validated_kwargs(): - return { - "model": "openai/openai/openai/gpt-5.5", - "messages": [{"role": "user", "content": "hi"}], - "optional_params": {}, - "litellm_params": {}, - "headers": {}, - "model_response": MagicMock(), - "logging_obj": MagicMock(), - "custom_llm_provider": "openai", - } - - -def test_sync_completion_forwards_custom_llm_provider(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai/openai/openai/gpt-5.5", - "input": [], - # `_build_sanitized_litellm_params` spreads `custom_llm_provider` from - # `litellm_params` into request_data on the real bridge path. Seed - # it here so the test exercises the overwrite (not an explicit kwarg - # that would TypeError against an already-present key). - "custom_llm_provider": "should-be-overwritten", - } - handler.transformation_handler.transform_response.return_value = ( - _validated_kwargs()["model_response"] - ) - with ( - patch.object( - handler, "validate_input_kwargs", return_value=_validated_kwargs() - ), - patch( - "litellm.responses", - return_value=MagicMock(spec=[]), - ) as mock_responses, - ): - # The handler routes ResponsesAPIResponse through transform_response. - # We just want to verify the kwargs going INTO responses(). - try: - handler.completion(acompletion=False) - except Exception: - # Downstream handling (transform_response, type checks) is not - # the subject of this test. - pass - assert mock_responses.called - kwargs = mock_responses.call_args.kwargs - assert kwargs.get("custom_llm_provider") == "openai", ( - "sync bridge must forward custom_llm_provider to litellm.responses() " - "so the downstream get_llm_provider() call does not re-strip the " - "provider prefix on a provider/provider/model deployment string" - ) - - -@pytest.mark.asyncio -async def test_async_completion_forwards_custom_llm_provider(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai/openai/openai/gpt-5.5", - "input": [], - # `_build_sanitized_litellm_params` spreads `custom_llm_provider` from - # `litellm_params` into request_data on the real bridge path. Seed - # it here so the test exercises the overwrite (not an explicit kwarg - # that would TypeError against an already-present key). - "custom_llm_provider": "should-be-overwritten", - } - - async def _fake_aresponses(**kwargs): - _fake_aresponses.kwargs = kwargs - return MagicMock(spec=[]) - - _fake_aresponses.kwargs = {} - - with ( - patch.object( - handler, "validate_input_kwargs", return_value=_validated_kwargs() - ), - patch("litellm.aresponses", _fake_aresponses), - ): - try: - await handler.acompletion() - except Exception: - pass - assert _fake_aresponses.kwargs.get("custom_llm_provider") == "openai", ( - "async bridge must forward custom_llm_provider to litellm.aresponses() " - "so the downstream get_llm_provider() call does not re-strip the " - "provider prefix on a provider/provider/model deployment string" - ) - - -@pytest.mark.asyncio -async def test_async_completion_forwards_aws_region_name(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai.gpt-5.5", - "input": [], - "aws_region_name": "us-east-2", - "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", - "custom_llm_provider": "bedrock_mantle", - } - - async def _fake_aresponses(**kwargs): - _fake_aresponses.kwargs = kwargs - return MagicMock(spec=[]) - - _fake_aresponses.kwargs = {} - - validated = _validated_kwargs() - validated["custom_llm_provider"] = "bedrock_mantle" - validated["litellm_params"] = { - "aws_region_name": "us-east-2", - "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", - "custom_llm_provider": "bedrock_mantle", - } - - with ( - patch.object(handler, "validate_input_kwargs", return_value=validated), - patch("litellm.aresponses", _fake_aresponses), - ): - try: - await handler.acompletion() - except Exception: - pass - assert _fake_aresponses.kwargs.get("aws_region_name") == "us-east-2" diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py rename to tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py rename to tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py similarity index 95% rename from tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py rename to tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 43dfdaba02d..a300560ae9d 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,7 +1,6 @@ import asyncio import json import time -from pathlib import Path import httpx import pytest @@ -571,20 +570,3 @@ def test_config_manager_returns_wxo_provider(): ) assert config is not None assert config.__class__.__name__ == "WatsonxOrchestrateA2AConfig" - - -def test_wxo_dashboard_auth_fields(): - fields_path = ( - Path(__file__).resolve().parents[5] - / "litellm/proxy/public_endpoints/agent_create_fields.json" - ) - agent_fields = json.loads(fields_path.read_text()) - wxo_agent = next( - agent for agent in agent_fields if agent["agent_type"] == "watsonx_orchestrate" - ) - fields_by_key = {field["key"]: field for field in wxo_agent["credential_fields"]} - - assert fields_by_key["auth_mode"]["default_value"] == "cp4d" - # Username is CP4D-only; UI does not require it so ibm_cloud users are not blocked. - assert fields_by_key["username"]["required"] is False - assert "cp4d" in fields_by_key["username"]["tooltip"].lower() diff --git a/tests/test_litellm/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py similarity index 100% rename from tests/test_litellm/anthropic_interface/exceptions/test_exception_mapping_utils.py rename to tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py similarity index 100% rename from tests/test_litellm/batches/test_batch_utils.py rename to tests/unit/batches/test_batch_utils.py diff --git a/tests/test_litellm/batches/test_main.py b/tests/unit/batches/test_main.py similarity index 100% rename from tests/test_litellm/batches/test_main.py rename to tests/unit/batches/test_main.py diff --git a/tests/test_litellm/batches/test_responses_batch_cost.py b/tests/unit/batches/test_responses_batch_cost.py similarity index 87% rename from tests/test_litellm/batches/test_responses_batch_cost.py rename to tests/unit/batches/test_responses_batch_cost.py index b634f5f73db..63b28fb3b42 100644 --- a/tests/test_litellm/batches/test_responses_batch_cost.py +++ b/tests/unit/batches/test_responses_batch_cost.py @@ -12,17 +12,28 @@ Line shape decides the parse, not the batch's declared endpoint, so an output file mixing Responses-shaped and chat-shaped lines sums across both. """ -from typing import Literal, get_args, get_type_hints import pytest import litellm import litellm.batches.batch_utils as bu -from litellm.types.llms.openai import CreateBatchRequest MODEL = "gpt-5.6" +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _responses_line(input_tokens: int, output_tokens: int) -> dict: return { "response": { @@ -107,13 +118,3 @@ async def test_mixed_shape_batch_output_sums_across_both_line_shapes(local_model assert result.cost == pytest.approx( 133 * model_info["input_cost_per_token_batches"] + 107 * model_info["output_cost_per_token_batches"] ) - - -def test_create_batch_endpoint_accepts_v1_responses(): - """A type-checked caller can pass endpoint="/v1/responses", which the runtime - already forwarded correctly.""" - endpoint_annotation = get_type_hints(CreateBatchRequest)["endpoint"] - assert "/v1/responses" in get_args(endpoint_annotation) - - for create_fn in (litellm.create_batch, litellm.acreate_batch): - assert "/v1/responses" in get_args(get_type_hints(create_fn)["endpoint"]) diff --git a/tests/test_litellm/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py similarity index 93% rename from tests/test_litellm/chat_completions/test_dispatch.py rename to tests/unit/chat_completions/test_dispatch.py index d4bfeaf8d70..63821c74208 100644 --- a/tests/test_litellm/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -1,11 +1,9 @@ -import inspect from collections.abc import Awaitable, Callable, Mapping -from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect +from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures import pytest import litellm -from litellm import main as python_chat from litellm.chat_completions.dispatch import ( _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch @@ -40,15 +38,6 @@ def acompletion_binding(native: NativeAcompletion | None) -> NativeBinding[Nativ return binding -def test_public_signature_is_the_legacy_signature() -> None: - public_completion: Final = cast(Callable[..., object], litellm.completion) - legacy_completion: Final = cast(Callable[..., object], python_chat.completion) - public_acompletion: Final = cast(Callable[..., object], litellm.acompletion) - legacy_acompletion: Final = cast(Callable[..., object], python_chat.acompletion) - assert inspect.signature(public_completion) == inspect.signature(legacy_completion) - assert inspect.signature(public_acompletion) == inspect.signature(legacy_acompletion) - - def test_python_route_forwards_original_call_shape() -> None: metadata: Final = {"user_id": "u"} args: Final[tuple[object, ...]] = ("gpt-4o", MESSAGES) diff --git a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/test_litellm_responses_transformation_transformation.py similarity index 100% rename from tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py rename to tests/unit/completion_extras/test_litellm_responses_transformation_transformation.py diff --git a/tests/test_litellm/compression/test_compress.py b/tests/unit/compression/test_compress.py similarity index 100% rename from tests/test_litellm/compression/test_compress.py rename to tests/unit/compression/test_compress.py diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py similarity index 100% rename from tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py rename to tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py b/tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py rename to tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py rename to tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py rename to tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py similarity index 100% rename from tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py rename to tests/unit/integrations/compression_interception/test_compression_interception_handler.py diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py similarity index 100% rename from tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py rename to tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py diff --git a/tests/test_litellm/integrations/gcs_pubsub/test_pub_sub.py b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py similarity index 100% rename from tests/test_litellm/integrations/gcs_pubsub/test_pub_sub.py rename to tests/unit/integrations/gcs_pubsub/test_pub_sub.py diff --git a/tests/test_litellm/integrations/helicone/test_helicone_gemini.py b/tests/unit/integrations/helicone/test_helicone_gemini.py similarity index 73% rename from tests/test_litellm/integrations/helicone/test_helicone_gemini.py rename to tests/unit/integrations/helicone/test_helicone_gemini.py index 8ce02784345..667b16a48a1 100644 --- a/tests/test_litellm/integrations/helicone/test_helicone_gemini.py +++ b/tests/unit/integrations/helicone/test_helicone_gemini.py @@ -3,7 +3,6 @@ Test HeliconeLogger Gemini/Vertex AI support. Fixes: https://github.com/BerriAI/litellm/issues/19093 """ -import pytest def test_helicone_gemini_model_in_list(): @@ -36,39 +35,6 @@ def test_helicone_gemini_models_recognized(): assert is_recognized, f"{model} should be recognized by helicone_model_list" -def test_helicone_vertex_ai_models_recognized(): - """ - Test that Vertex AI models (GLM, DeepSeek, etc.) are recognized via custom_llm_provider. - """ - # Test models that don't contain "gemini" but are vertex_ai - test_models = [ - "vertex_ai/zai-org/glm-4.7-maas", - "vertex_ai/deepseek-ai/deepseek-v3", - "vertex_ai/meta/llama-3.1-405b", - ] - for model in test_models: - is_vertex_ai = model.startswith("vertex_ai/") - assert is_vertex_ai, f"{model} should be recognized as vertex_ai model" - - -def test_helicone_vertex_ai_via_custom_llm_provider(): - """ - Test that vertex_ai models are recognized when custom_llm_provider is set. - """ - # Models without vertex_ai/ prefix but with custom_llm_provider="vertex_ai" - test_cases = [ - ("zai-org/glm-4.7-maas", "vertex_ai"), - ("deepseek-ai/deepseek-v3", "vertex_ai"), - ] - for model, custom_llm_provider in test_cases: - is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith( - "vertex_ai/" - ) - assert ( - is_vertex_ai - ), f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai" - - def test_helicone_vertex_gemini_gets_vertex_provider_url(): """ Test that vertex_ai/gemini-* models route to aiplatform.googleapis.com, From ba629f2537a73f74f5d466e8b7d2703902b74014 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:37:17 +0000 Subject: [PATCH 056/146] test(bedrock): wrap long lines flagged by review in migrated unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../image/test_bedrock_image_prepare_request.py | 9 ++++++--- .../test_bedrock_passthrough_transformation.py | 13 ++++++++++--- .../realtime/test_bedrock_realtime_handler.py | 8 +++++++- .../rerank/test_bedrock_rerank_header_forwarding.py | 11 +++++++++-- .../test_bedrock_vector_store_transformation.py | 3 ++- 5 files changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py index 1575ccb5739..b010db3a840 100644 --- a/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py @@ -11,7 +11,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None: with ( patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params" + "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration." + "_get_boto_credentials_from_optional_params" ), patch( "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers" @@ -31,7 +32,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None: assert ( request.endpoint_url - == "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012%3Aapplication-inference-profile%2Fabcdefghi123/invoke" + == "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012" + "%3Aapplication-inference-profile%2Fabcdefghi123/invoke" ) @@ -41,7 +43,8 @@ def test_bedrock_image_prepare_request_without_arn() -> None: with ( patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params" + "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration." + "_get_boto_credentials_from_optional_params" ), patch( "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers" diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 854ef92fa4b..d1d636a15f7 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -419,7 +419,9 @@ def test_bedrock_passthrough_model_id_arn_encoding(): ), f"ARN slash should be encoded, but found unencoded version in: {url_str}" # Verify the complete expected URL structure - expected_encoded_model_id = "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7" + expected_encoded_model_id = ( + "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7" + ) expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{expected_encoded_model_id}/converse" assert url_str == expected_url, f"Expected {expected_url}, but got: {url_str}" @@ -515,7 +517,10 @@ def test_bedrock_passthrough_model_id_without_arn(): def _event_frame(event_type: str, payload: dict) -> bytes: def header(name: str, value: str) -> bytes: name_b, value_b = name.encode(), value.encode() - return struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + return ( + struct.pack("!B", len(name_b)) + name_b + + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + ) payload_b = json.dumps(payload, separators=(",", ":")).encode() headers_b = ( @@ -589,7 +594,9 @@ def _feed(collector: PassthroughStreamCollector, stream: bytes, chunk_size: int def test_converse_stream_collector_keeps_usage_without_retaining_the_stream(): texts = [f"tok{i} " for i in range(4000)] - stream = _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + stream = ( + _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + ) _feed(_converse_stream_collector(), stream) tracemalloc.start() diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 73a78a94e9f..3aa827beb80 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -23,7 +23,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") - for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + for env_var in ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", + ): monkeypatch.delenv(env_var, raising=False) diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index aa93ddb21b8..c40830b238f 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -21,7 +21,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") - for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + for env_var in ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", + ): monkeypatch.delenv(env_var, raising=False) # Mock response for Bedrock rerank @@ -39,7 +45,8 @@ bedrock_rerank_response = { test_query = "What is the capital of the United States?" test_documents = [ "Carson City is the capital city of the American state of Nevada.", - "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.", + "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. " + "Its capital is Saipan.", "Washington, D.C. is the capital of the United States.", ] diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index ab5a2531461..b45e70e31d3 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -46,7 +46,8 @@ def test_transform_search_request_encodes_vector_store_id(): assert ( url - == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother%3Fx%3D1%23frag/retrieve" + == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother" + "%3Fx%3D1%23frag/retrieve" ) assert body["retrievalQuery"].get("text") == "hello" From f61abe00a6bd3eff086e0fc5fc503e4418850d3e Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:58:23 +0000 Subject: [PATCH 057/146] test(llms): wrap remaining lines over 120 chars in migrated unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_translation/test_handler.py | 12 ++++-- ...drock_mantle_passthrough_transformation.py | 8 +++- .../chat/test_bytez_chat_transformation.py | 42 +++++++++++++++---- tests/unit/llms/chat/test_converse_handler.py | 5 ++- 4 files changed, 53 insertions(+), 14 deletions(-) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index dee8366ce2d..da7ed635dcb 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -1072,7 +1072,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_reasoning_text_delta_de_anonymized(self): - """Reasoning deltas carry model output; their text must be guardrailed while the reasoning signature is left untouched.""" + """Reasoning deltas carry model output; their text must be guardrailed while the + reasoning signature is left untouched.""" stream_bytes = ( _build_event_stream_frame("messageStart", {"role": "assistant"}) + _build_event_stream_frame( @@ -1105,7 +1106,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_tool_use_input_delta_de_anonymized(self): - """toolUse.input deltas carry model-generated tool arguments and must be guardrailed instead of being forwarded raw.""" + """toolUse.input deltas carry model-generated tool arguments and must be + guardrailed instead of being forwarded raw.""" stream_bytes = _build_event_stream_frame( "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": '{"q":""}'}}}, @@ -1154,7 +1156,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_text_and_reasoning_deltas_de_anonymized_independently(self): - """Distinct delta kinds must each be guardrailed and written back into their own field without bleeding the de-anonymized text across kinds.""" + """Distinct delta kinds must each be guardrailed and written back into their own + field without bleeding the de-anonymized text across kinds.""" captured = {} async def mock_hook(data, user_api_key_dict, response): @@ -1192,7 +1195,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_reasoning_signature_only_frame_left_unmodified(self): - """A reasoning delta carrying only a signature has no guardrailable text; it must be forwarded untouched and the guardrail must not run.""" + """A reasoning delta carrying only a signature has no guardrailable text; it must + be forwarded untouched and the guardrail must not run.""" stream_bytes = _build_event_stream_frame( "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig"}}}, diff --git a/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py index 090de0a9d3e..27f6c9a9140 100644 --- a/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py +++ b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -109,7 +109,13 @@ def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"), ], ) -def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer): +def test_sign_request_uses_the_deployment_bearer_token( + no_ambient_aws, + monkeypatch, + litellm_params, + env, + expected_bearer, +): for name, value in env.items(): monkeypatch.setenv(name, value) headers, body = BedrockMantlePassthroughConfig().sign_request( diff --git a/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py index 440304aeac1..157e2e51175 100644 --- a/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py @@ -8,6 +8,32 @@ from litellm.llms.bytez.chat.transformation import BytezChatConfig, API_BASE, ve TEST_API_KEY = "MOCK_BYTEZ_API_KEY" TEST_MODEL_NAME = "google/gemma-3-4b-it" TEST_MODEL = f"bytez/{TEST_MODEL_NAME}" +CAT_IMAGE_URL = ( + "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUX" + "VRLHI/male-orange-tabby-cat.jpg" +) +KAGGLE_AUDIO_URL = ( + "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_" + "SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-1616" + "07.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&" + "X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf" + "81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc39" + "0679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250" + "f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817" + "000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468" + "adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3" +) +KAGGLE_VIDEO_URL = ( + "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG" + "4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F202507" + "11%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-Signed" + "Headers=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5f" + "c6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72" + "084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb" + "90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d9" + "99f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189" + "c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947" +) TEST_MESSAGES = [{"role": "user", "content": "Hello"}] @@ -148,7 +174,7 @@ class TestBytezChatConfig: "What color is this cat?", { "type": "image_url", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -160,7 +186,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -174,7 +200,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image_url", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -186,7 +212,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -200,7 +226,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of cat meow is this?"}, { "type": "input_audio", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc390679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3", + "url": KAGGLE_AUDIO_URL, }, ], } @@ -212,7 +238,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of cat meow is this?"}, { "type": "audio", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc390679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3", + "url": KAGGLE_AUDIO_URL, }, ], } @@ -226,7 +252,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of dog is this?"}, { "type": "video_url", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5fc6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d999f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947", + "url": KAGGLE_VIDEO_URL, }, ], } @@ -238,7 +264,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of dog is this?"}, { "type": "video", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5fc6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d999f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947", + "url": KAGGLE_VIDEO_URL, }, ], } diff --git a/tests/unit/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py index 12b5f03aedc..05debee0602 100644 --- a/tests/unit/llms/chat/test_converse_handler.py +++ b/tests/unit/llms/chat/test_converse_handler.py @@ -106,7 +106,10 @@ class TestBedrockRegionInModelPath: ), f"modelId mismatch for {model!r}: got {model_id!r}, expected {expected_model_id!r}" assert ( optional_params.get("aws_region_name") == expected_region - ), f"region mismatch for {model!r}: got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}" + ), ( + f"region mismatch for {model!r}: " + f"got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}" + ) def test_explicit_aws_region_name_not_overridden(self): """ From ccd8997b0846469ff0e624b2cd837c4d0f8a3da6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:16:48 +0000 Subject: [PATCH 058/146] refactor(types): replace Any with proven types in 34 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_agentcore/transformation.py | 2 +- .../providers/watsonx_orchestrate/config.py | 4 ++-- litellm/a2a_protocol/utils.py | 6 ++--- .../gitlab/gitlab_prompt_manager.py | 4 ++-- litellm/integrations/otel/presets/agentops.py | 7 ++++-- .../integrations/vantage/vantage_logger.py | 4 ++-- litellm/interactions/agents/http_handler.py | 22 +++++++++---------- litellm/litellm_core_utils/logging_utils.py | 4 +++- litellm/litellm_core_utils/url_utils.py | 4 ++-- litellm/llms/anthropic/chat/handler.py | 2 +- litellm/llms/azure/completion/handler.py | 9 ++++---- .../document_intelligence/transformation.py | 8 +++---- .../guardrail_translation/base_translation.py | 2 +- .../llms/bedrock/batches/transformation.py | 10 ++++----- ...mazon_twelvelabs_pegasus_transformation.py | 2 +- litellm/llms/bytez/chat/transformation.py | 6 ++--- litellm/llms/custom_httpx/aiohttp_handler.py | 4 ++-- .../llms/deprecated_providers/aleph_alpha.py | 2 +- litellm/llms/lemonade/chat/transformation.py | 2 +- .../llms/openai/image_edit/transformation.py | 2 +- .../runwayml/text_to_speech/transformation.py | 2 +- .../llms/vertex_ai/files/transformation.py | 6 ++--- .../batch_embed_content_handler.py | 4 ++-- .../llms/vertex_ai/vertex_ai_non_gemini.py | 2 +- .../audio_transcription/transformation.py | 4 ++-- litellm/proxy/a2a/agent_card.py | 16 +++++++------- .../proxy/agent_endpoints/a2a_endpoints.py | 2 +- litellm/proxy/client/credentials.py | 5 +++-- .../guardrails_ai/guardrails_ai.py | 4 ++-- .../guardrail_hooks/singulr/singulr.py | 6 ++--- .../object_permission_utils.py | 13 ++++++----- .../proxy/policy_engine/pipeline_executor.py | 12 +++++----- .../router_strategy/adaptive_router/hooks.py | 2 +- .../auto_router/auto_router.py | 4 ++-- 34 files changed, 98 insertions(+), 90 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 9fa9db48af8..c486f1f6d95 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -85,7 +85,7 @@ def _filter_reserved_headers( def _request_scoped_runtime_session_id( - params: Mapping[str, Any], + params: Mapping[str, object], litellm_params: Mapping[str, Any], ) -> str | None: context_id: Final = get_session_id_from_a2a_params(params) diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index ca84d3e07b4..44873edf271 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -20,7 +20,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Handle a non-streaming A2A request via WXO runs API.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: @@ -40,7 +40,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """Handle a streaming A2A request via WXO streaming runs API.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 47f561068cd..7400844bf28 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -17,7 +17,7 @@ class A2ARequestUtils: """Utility class for A2A request/response processing.""" @staticmethod - def extract_text_from_message(message: Any) -> str: + def extract_text_from_message(message: object) -> str: """ Extract text content from A2A message parts. @@ -142,7 +142,7 @@ class A2ARequestUtils: return prompt_tokens, completion_tokens, total_tokens -def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None: +def get_session_id_from_a2a_params(params: Mapping[str, object]) -> str | None: message: Final = params.get("message", {}) if isinstance(message, dict): return message.get("contextId") @@ -166,7 +166,7 @@ def scope_session_to_principal(session_id: str, principal: str | None) -> str: # Backwards compatibility aliases -def extract_text_from_a2a_message(message: Any) -> str: +def extract_text_from_a2a_message(message: object) -> str: return A2ARequestUtils.extract_text_from_message(message) diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index d4602176650..817d280074f 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -200,8 +200,8 @@ class GitLabTemplateManager: metadata=metadata, ) - def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]: - result: Final[dict[str, Any]] = {} + def _parse_yaml_basic(self, yaml_str: str) -> dict[str, bool | int | float | str]: + result: Final[dict[str, bool | int | float | str]] = {} for line in yaml_str.split("\n"): line = line.strip() if ":" in line and not line.startswith("#"): diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 965213f2ee4..58123656caa 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,9 +9,12 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ +from collections.abc import Sequence from typing import Any, Final import httpx +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -71,7 +74,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> Any: +def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -106,7 +109,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> Any: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Any) -> Any: + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: self._ensure_authenticated() return super().export(spans) diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index c219ba392ab..48a492fdd72 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -59,7 +59,7 @@ class VantageLogger(FocusLogger): raw_interval, ) - destination_config: Final[dict[str, Any]] = {} + destination_config: Final[dict[str, str]] = {} if resolved_api_key: destination_config["api_key"] = resolved_api_key if resolved_token: @@ -93,7 +93,7 @@ class VantageLogger(FocusLogger): pod_lock_manager = None if proxy_logging_obj is not None: - writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) + writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: pod_lock_manager = getattr(writer, "pod_lock_manager", None) diff --git a/litellm/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index ec9df0fb488..2afedc34d36 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -7,7 +7,7 @@ duplicated. BaseAgentsAPIConfig stays as pure transform code. """ from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final import httpx @@ -38,7 +38,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, @@ -93,7 +93,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, @@ -141,7 +141,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): agents_api_config: BaseAgentsAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -181,7 +181,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): agents_api_config: BaseAgentsAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentListResponse: @@ -216,7 +216,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -259,7 +259,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentCreateResponse: @@ -295,7 +295,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -338,7 +338,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentDeleteResult: @@ -374,7 +374,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -417,7 +417,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentVersionsResponse: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 5be9dd7be2f..38a501ecaae 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -67,7 +67,9 @@ def _truncate_base64_in_string(value: str) -> str: return _DATA_URI_RE.sub(_base64_data_uri_replacer, value) -def _truncate_base64_in_value(value: Any) -> Any: +def _truncate_base64_in_value( + value: str | dict[str, object] | list[object] | None, +) -> str | dict[str, object] | list[object] | None: """Iteratively truncate base64 data URIs in a JSON-like value (str/list/dict). Uses an explicit stack instead of recursion to satisfy the project's diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index fa070a648f5..b94d5a6886d 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -418,7 +418,7 @@ def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: +def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -461,7 +461,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: +async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Any) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 359b8bb08c9..ef0f45d8f8b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -596,7 +596,7 @@ class ModelResponseIterator: self.reasoning_content_chunks: list[str] = [] # Track server tool use inputs and results for code_interpreter_results - self._server_tool_inputs: dict[str, Any] = {} + self._server_tool_inputs: dict[str, object] = {} self.tool_results: list[dict[str, Any]] = [] self._current_server_tool_id: str | None = None self._container_id: str | None = None diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 80934e994f6..23eef51e7ee 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -1,6 +1,7 @@ from collections.abc import Callable -from typing import Any, Final +from typing import Final +import httpx from openai import AsyncAzureOpenAI, AzureOpenAI from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -191,7 +192,7 @@ class AzureTextCompletion(BaseAzureLLM): model: str, api_base: str, data: dict, - timeout: Any, + timeout: float | httpx.Timeout | None, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, max_retries: int, @@ -253,7 +254,7 @@ class AzureTextCompletion(BaseAzureLLM): api_version: str, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout | None, azure_ad_token: str | None = None, client=None, litellm_params: dict = {}, @@ -306,7 +307,7 @@ class AzureTextCompletion(BaseAzureLLM): api_version: str, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout | None, azure_ad_token: str | None = None, client=None, litellm_params: dict = {}, diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 3a2af8a5aba..23f532be757 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -12,7 +12,7 @@ import asyncio import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from urllib.parse import quote import httpx @@ -127,7 +127,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def map_ocr_params( self, - non_default_params: dict, + non_default_params: Mapping[str, object], optional_params: dict, model: str, ) -> dict: @@ -164,7 +164,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise UnsupportedParamsError(message=f"{e}", model=model, llm_provider="azure_ai") from e @staticmethod - def _normalize_pages_param(pages: Any) -> str: + def _normalize_pages_param(pages: object) -> str: """ Convert a caller-provided `pages` value to Azure DI's query-string form. Azure expects 1-based page numbers, grammar: `^(\\d+(-\\d+)?)(,\\s*(\\d+(-\\d+)?))*$`. @@ -412,7 +412,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise ValueError("Document URL is required") # Build Azure DI request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} # Check if it's a data URI (base64) if document_url.startswith("data:"): diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 89ad67f0485..ace5af8124f 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -81,7 +81,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Any | None, + user_api_key_dict: Optional["UserAPIKeyAuth"], ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index ae0f8c5935b..973388ca5bd 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -2,7 +2,7 @@ import os import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError @@ -170,7 +170,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform the batch creation request to Bedrock format. @@ -354,7 +354,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) @staticmethod - def _get_openai_compatible_batch_metadata(metadata: Any) -> dict[str, str]: + def _get_openai_compatible_batch_metadata(metadata: object) -> dict[str, str]: """ OpenAI Batch metadata only accepts string values. """ @@ -379,7 +379,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): batch_id: str, optional_params: dict, litellm_params: dict, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform batch retrieval request for Bedrock. @@ -523,7 +523,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Enrich metadata with useful Bedrock fields - enriched_metadata_raw: Final[dict[str, Any]] = { + enriched_metadata_raw: Final[dict[str, object]] = { "jobName": response_data.get("jobName"), "clientRequestToken": response_data.get("clientRequestToken"), "modelId": response_data.get("modelId"), diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index d12c8aee48c..39cded4ed64 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -110,7 +110,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): headers: dict, ) -> dict: input_prompt: Final = self._convert_messages_to_prompt(messages=messages) - request_data: Final[dict[str, Any]] = {"inputPrompt": input_prompt} + request_data: Final[dict[str, object]] = {"inputPrompt": input_prompt} media_source: Final = self._build_media_source(optional_params) if media_source is not None: diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index d9a0c98b6db..7977db0f056 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -335,10 +335,10 @@ class BytezChatConfig(BaseConfig): class BytezCustomStreamWrapper(CustomStreamWrapper): - def chunk_creator(self, chunk: Any): + def chunk_creator(self, chunk: object): try: model_response: Final = self.model_response_creator() - response_obj: dict[str, Any] = {} + response_obj: dict[str, object] = {} response_obj = { "text": chunk, @@ -346,7 +346,7 @@ class BytezCustomStreamWrapper(CustomStreamWrapper): "finish_reason": "", } - completion_obj: Final[dict[str, Any]] = {"content": chunk} + completion_obj: Final[dict[str, object]] = {"content": chunk} return self.return_processed_chunk_logic( completion_obj=completion_obj, diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 7035ce58ae1..0809ef5274f 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,5 +1,5 @@ import ssl -from collections.abc import Callable +from collections.abc import AsyncIterable, Callable, Iterable from typing import TYPE_CHECKING, Any, Final, cast import aiohttp @@ -212,7 +212,7 @@ class BaseLLMAIOHTTPHandler: litellm_params: dict, stream: bool = False, files: dict | None = None, - content: Any = None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, params: dict | None = None, ) -> httpx.Response: max_retry_on_unprocessable_entity_error: Final = provider_config.max_retry_on_unprocessable_entity_error diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index 4a29549b6aa..2ad9ce4edc8 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -146,7 +146,7 @@ class AlephAlphaConfig: setattr(self.__class__, key, value) @classmethod - def get_config(cls): + def get_config(cls) -> dict[str, object]: return { k: v for k, v in cls.__dict__.items() diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index 553478aec16..c01ad2a0edd 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -170,7 +170,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> Any: + ) -> dict[str, object]: if model.startswith("lemonade/"): model = model.split("/", 1)[1] diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index f55084adbde..d54522597a0 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -66,7 +66,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): def _add_image_to_files( self, files_list: list[tuple[str, Any]], - image: Any, + image: object, field_name: str, ) -> None: """Add an image to the files list with appropriate content type""" diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 19e6d8ff494..6769accc1d6 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -78,7 +78,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", Coroutine[object, object, "HttpxBinaryResponseContent"], diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 85ec2911464..80d32289c94 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -651,7 +651,7 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], ) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. @@ -774,7 +774,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -948,7 +948,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index f81d4ca777e..c6ac87d646b 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Final, Literal import httpx @@ -210,7 +210,7 @@ class GoogleBatchEmbeddings(VertexLLM): ) ### TRANSFORMATION (sync path) ### - request_data: Any + request_data: VertexAIBatchEmbeddingsRequestBody | dict[str, object] if use_embed_content: resolved_files = {} if api_key: diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index 1c582c7c376..a7a1ea8d88d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -64,7 +64,7 @@ def _get_client_from_cache(client_cache_key: str): return litellm.in_memory_llm_clients_cache.get_cache(client_cache_key) -def _set_client_in_cache(client_cache_key: str, vertex_llm_model: Any): +def _set_client_in_cache(client_cache_key: str, vertex_llm_model: object): litellm.in_memory_llm_clients_cache.set_cache( key=client_cache_key, value=vertex_llm_model, diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 7d1aba63428..2169b9bf49a 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -4,7 +4,7 @@ Translates from OpenAI's `/v1/audio/transcriptions` to IBM WatsonX's `/ml/v1/aud WatsonX follows the OpenAI spec for audio transcription. """ -from typing import Any, Final +from typing import Final from httpx import Response @@ -124,7 +124,7 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran } # Convert TypedDict to regular dict for AudioTranscriptionRequestData - form_data_dict: Final[dict[str, Any]] = dict(form_data) + form_data_dict: Final[dict[str, object]] = dict(form_data) return AudioTranscriptionRequestData(data=form_data_dict, files=files) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index 5bec5158bcc..3718359fbb6 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -10,7 +10,7 @@ and uses LiteLLM auth. import re from collections.abc import Mapping from copy import deepcopy -from typing import Any, Final, Literal +from typing import Final, Literal SupportedA2AVersion = Literal["0.3", "1.0"] @@ -44,7 +44,7 @@ def normalize_protocol_version(version: object) -> SupportedA2AVersion | None: return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None) -def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: +def resolve_served_protocol_version(card: Mapping[str, object] | None) -> str: """Return the validated protocol version an agent card pins, else the default.""" normalized: Final = normalize_protocol_version(card.get("protocolVersion") if card else None) return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION @@ -53,7 +53,7 @@ def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: # Security scheme exposed by the LiteLLM-fronted agent card. Always replaces # whatever upstream advertised — the client must authenticate to the proxy, # not the upstream agent. -LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, Any]]] = { +LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, str]]] = { "LiteLLMKey": { "type": "http", "scheme": "bearer", @@ -112,7 +112,7 @@ _ALLOWED_TOP_LEVEL_KEYS: Final = { "url", } -_DEFAULT_SKILLS: Final[list[dict[str, Any]]] = [ +_DEFAULT_SKILLS: Final[list[dict[str, str | list[str]]]] = [ { "id": "chat", "name": "Chat", @@ -129,7 +129,7 @@ _DEFAULT_MODES: Final[list[str]] = ["text"] _DEFAULT_AGENT_VERSION: Final = "1.0.0" -def _filter_capabilities(upstream_capabilities: Any) -> dict[str, Any]: +def _filter_capabilities(upstream_capabilities: object) -> dict[str, object]: """Return a capabilities dict containing only allowlisted, truthy keys.""" if not isinstance(upstream_capabilities, dict): return {} @@ -143,13 +143,13 @@ def _default_litellm_provider(proxy_base_url: str) -> dict[str, str]: def merge_agent_card( - upstream_card: Mapping[str, Any] | None, + upstream_card: Mapping[str, object] | None, *, proxy_url: str, proxy_base_url: str, name: str | None = None, description: str | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build the LiteLLM-fronted agent card. @@ -169,7 +169,7 @@ def merge_agent_card( A dict suitable for serving as the proxy's agent card. Only keys in the v1.0 AgentCard schema (plus ``supportedInterfaces``) are emitted. """ - base: Final[dict[str, Any]] = deepcopy(dict(upstream_card)) if upstream_card else {} + base: Final[dict[str, object]] = deepcopy(dict(upstream_card)) if upstream_card else {} # Keep the upstream ``url`` on the stored card: the runtime A2A # invocation path reads it from ``agent_card_params`` to know where to diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 834c16ba6dc..faf3e98a3a7 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -880,7 +880,7 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None _enqueue_fn() - response_dict: Final[dict[str, Any]] = ( + response_dict: Final[dict[str, object]] = ( response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index a9bff67b1c5..d9edecd2eb7 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import requests @@ -69,8 +70,8 @@ class CredentialsManagementClient: def create( self, credential_name: str, - credential_info: dict[str, Any], - credential_values: dict[str, Any], + credential_info: Mapping[str, object], + credential_values: Mapping[str, object], return_request: bool = False, ) -> dict[str, Any] | requests.Request: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index 47324471650..18451df574f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -7,7 +7,7 @@ import json import os -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict +from typing import TYPE_CHECKING, Final, Literal, TypedDict from fastapi import HTTPException @@ -181,7 +181,7 @@ class GuardrailsAI(CustomGuardrail): ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm return await self.process_input(data=data, call_type=call_type) - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: if call_type == "acompletion" or call_type == "completion": kwargs = await self.process_input(data=kwargs, call_type=call_type) diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 06d4b39f5f6..bd5b18e368d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, ) @@ -36,7 +36,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( ToolCall, ToolCallFunction, ) -from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, GenericGuardrailAPIInputs _DEFAULT_API_BASE: Final = "http://localhost:8003" _GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2" @@ -339,7 +339,7 @@ class SingulrGuardrail(CustomGuardrail): return inputs @staticmethod - def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None": + def _build_tool_call(tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall) -> "ToolCall | None": tool_call_id: Final = tool_call.get("id") fun: Final = tool_call.get("function") if not tool_call_id or not fun: diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index daab38d3662..61e432daa16 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -156,12 +156,13 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - new_object_permission: dict | str | None = data_json.pop("object_permission", None) - if new_object_permission is None: + raw_object_permission: Final[dict | str | None] = data_json.pop("object_permission", None) + if raw_object_permission is None: return None - if isinstance(new_object_permission, str): - new_object_permission = json.loads(new_object_permission) + new_object_permission: Final[object] = ( + json.loads(raw_object_permission) if isinstance(raw_object_permission, str) else raw_object_permission + ) upsert: Final = await prepare_object_permission_upsert( new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, @@ -230,7 +231,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 0b81e7af84d..e9d23436b59 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -8,7 +8,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. import copy import time from collections.abc import Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar +from typing import TYPE_CHECKING, Final, Literal, TypeVar from pydantic import BaseModel @@ -314,11 +314,11 @@ class PipelineExecutor: steps: list[PipelineStep], mode: str, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ @@ -490,10 +490,10 @@ class PipelineExecutor: step: PipelineStep, mode: str, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], @@ -722,7 +722,7 @@ def _extract_error_message(e: Exception) -> str: if isinstance(e, ModifyResponseException): return str(e) if HTTPException is not None and isinstance(e, HTTPException): - detail: Final = getattr(e, "detail", None) + detail: Final[object] = getattr(e, "detail", None) if detail: return str(detail) return str(e) diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 709910753f2..c4a2eae1ef9 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -86,7 +86,7 @@ def _resolve_session_key(kwargs: dict[str, Any]) -> str | None: return hashlib.sha256(payload.encode("utf-8")).hexdigest() -def _last_user_content(messages: list[dict[str, Any]] | None) -> str | None: +def _last_user_content(messages: Sequence[Mapping[str, object]] | None) -> str | None: if not messages: return None for msg in reversed(messages): diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index d08afa8c1f6..250201a46d1 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -3,7 +3,7 @@ Auto-Routing Strategy that works with a Semantic Router Config """ import asyncio -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel, ConfigDict @@ -158,7 +158,7 @@ class AutoRouter(CustomLogger): return await asyncio.shield(build_task) @staticmethod - def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str: + def _extract_text_from_messages(messages: Sequence[Mapping[str, object]]) -> str: """ Extract text content from the last user message for routing. From dbf566873ea32a82a9915876615dad67febd2e57 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:28:06 +0000 Subject: [PATCH 059/146] refactor(types): keep agentops preset imports optional Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/presets/agentops.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 58123656caa..965213f2ee4 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,12 +9,9 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ -from collections.abc import Sequence from typing import Any, Final import httpx -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -74,7 +71,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: +def _build_agentops_exporter(spec: ExporterSpec) -> Any: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -109,7 +106,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + def export(self, spans: Any) -> Any: self._ensure_authenticated() return super().export(spans) From 729a96e4eaff1d6fe4e68d5a1137ba83b07c0fe0 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:40:27 +0000 Subject: [PATCH 060/146] test: migrate openai, openai_like and openrouter legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/llms/openai/evals/__init__.py | 1 - tests/test_litellm/llms/openai_like/embedding/__init__.py | 1 - tests/test_litellm/llms/openai_like/messages/__init__.py | 0 tests/test_litellm/llms/openrouter/image_edit/__init__.py | 0 .../guardrail_translation/test_embeddings_guardrail_handler.py | 0 .../llms/openai/evals/test_openai_evals_transformation.py | 0 .../llms/openai/image_generation/test_gpt_transformation.py | 0 .../image_generation/test_image_generation_guardrail_handler.py | 0 .../test_openai_image_generation_extra_headers.py | 0 .../llms/openai/speech/test_text_to_speech_guardrail_handler.py | 0 .../transcriptions/test_audio_transcription_guardrail_handler.py | 0 .../openai/transcriptions/test_transcription_duration_hidden.py | 0 .../llms/openai/transcriptions/test_whisper_transformation.py | 0 .../test_openai_vector_store_files_transformation.py | 0 .../vector_stores/test_openai_vector_stores_transformation.py | 0 .../llms/openai/videos/test_openai_video_transformation.py | 0 .../openai_like/chat/test_openai_like_chat_transformation.py | 0 .../llms/openai_like/embedding/test_openai_like_embedding.py | 0 .../test_openai_like_anthropic_messages_transformation.py | 0 .../llms/openrouter/chat/test_openrouter_chat_transformation.py | 0 .../image_edit/test_openrouter_image_edit_transformation.py | 0 .../image_generation/test_openrouter_image_gen_transformation.py | 0 .../llms/openrouter/test_openrouter_embedding_transformation.py | 0 .../llms/openrouter/test_openrouter_provider_routing.py | 0 24 files changed, 2 deletions(-) delete mode 100644 tests/test_litellm/llms/openai/evals/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/embedding/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/messages/__init__.py delete mode 100644 tests/test_litellm/llms/openrouter/image_edit/__init__.py rename tests/{test_litellm => unit}/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/evals/test_openai_evals_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_gpt_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_image_generation_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_openai_image_generation_extra_headers.py (100%) rename tests/{test_litellm => unit}/llms/openai/speech/test_text_to_speech_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_transcription_duration_hidden.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_whisper_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/vector_stores/test_openai_vector_stores_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/videos/test_openai_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/chat/test_openai_like_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/embedding/test_openai_like_embedding.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/chat/test_openrouter_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/test_openrouter_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/test_openrouter_provider_routing.py (100%) diff --git a/tests/test_litellm/llms/openai/evals/__init__.py b/tests/test_litellm/llms/openai/evals/__init__.py deleted file mode 100644 index 47a8a2f0aed..00000000000 --- a/tests/test_litellm/llms/openai/evals/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""OpenAI Evals API tests""" diff --git a/tests/test_litellm/llms/openai_like/embedding/__init__.py b/tests/test_litellm/llms/openai_like/embedding/__init__.py deleted file mode 100644 index 2cb77227ed0..00000000000 --- a/tests/test_litellm/llms/openai_like/embedding/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Test module for OpenAI-like embedding handler diff --git a/tests/test_litellm/llms/openai_like/messages/__init__.py b/tests/test_litellm/llms/openai_like/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openrouter/image_edit/__init__.py b/tests/test_litellm/llms/openrouter/image_edit/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py b/tests/unit/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py rename to tests/unit/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py rename to tests/unit/llms/openai/evals/test_openai_evals_transformation.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py b/tests/unit/llms/openai/image_generation/test_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py rename to tests/unit/llms/openai/image_generation/test_gpt_transformation.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py b/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py rename to tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py rename to tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py diff --git a/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py b/tests/unit/llms/openai/speech/test_text_to_speech_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py rename to tests/unit/llms/openai/speech/test_text_to_speech_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py b/tests/unit/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py rename to tests/unit/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_transcription_duration_hidden.py b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_transcription_duration_hidden.py rename to tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_whisper_transformation.py b/tests/unit/llms/openai/transcriptions/test_whisper_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_whisper_transformation.py rename to tests/unit/llms/openai/transcriptions/test_whisper_transformation.py diff --git a/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py b/tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py rename to tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py diff --git a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py b/tests/unit/llms/openai/vector_stores/test_openai_vector_stores_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py rename to tests/unit/llms/openai/vector_stores/test_openai_vector_stores_transformation.py diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/unit/llms/openai/videos/test_openai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py rename to tests/unit/llms/openai/videos/test_openai_video_transformation.py diff --git a/tests/test_litellm/llms/openai_like/chat/test_openai_like_chat_transformation.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai_like/chat/test_openai_like_chat_transformation.py rename to tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py diff --git a/tests/test_litellm/llms/openai_like/embedding/test_openai_like_embedding.py b/tests/unit/llms/openai_like/embedding/test_openai_like_embedding.py similarity index 100% rename from tests/test_litellm/llms/openai_like/embedding/test_openai_like_embedding.py rename to tests/unit/llms/openai_like/embedding/test_openai_like_embedding.py diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py rename to tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py rename to tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py diff --git a/tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py b/tests/unit/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py rename to tests/unit/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py diff --git a/tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py rename to tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py rename to tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_provider_routing.py b/tests/unit/llms/openrouter/test_openrouter_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/openrouter/test_openrouter_provider_routing.py rename to tests/unit/llms/openrouter/test_openrouter_provider_routing.py From 7cf9a3035ce130fb4a07a02c7e139df6915fd378 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:02 +0000 Subject: [PATCH 061/146] test: migrate phase 16 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/rust_bridge/__init__.py | 0 tests/unit/rust_bridge/ocr/__init__.py | 0 tests/unit/rust_bridge/ocr/test_route_host.py | 85 ++ tests/unit/rust_bridge/responses/__init__.py | 0 .../rust_bridge/responses/test_route_host.py | 57 + tests/unit/sandbox/test_e2b_sandbox.py | 318 +++++ .../unit/sandbox/test_opensandbox_sandbox.py | 647 ++++++++++ tests/unit/sandbox/test_sandbox_tools.py | 181 +++ tests/unit/skills/test_skills_main.py | 57 + .../test_enforce_model_rate_limits.py | 468 ++++++++ .../test_router/test_io_token_rate_limits.py | 1041 +++++++++++++++++ .../types/llms/test_types_llms_bedrock.py | 46 + .../unit/types/llms/test_types_llms_openai.py | 591 ++++++++++ .../types/proxy/policy_engine/__init__.py | 0 .../policy_engine/test_pipeline_types.py | 168 +++ .../proxy/policy_engine/test_policy_types.py | 15 + .../policy_engine/test_resolver_types.py | 115 ++ tests/unit/videos/__init__.py | 0 tests/unit/videos/test_main.py | 455 +++++++ tests/unit/videos/test_utils.py | 181 +++ 20 files changed, 4425 insertions(+) create mode 100644 tests/unit/rust_bridge/__init__.py create mode 100644 tests/unit/rust_bridge/ocr/__init__.py create mode 100644 tests/unit/rust_bridge/ocr/test_route_host.py create mode 100644 tests/unit/rust_bridge/responses/__init__.py create mode 100644 tests/unit/rust_bridge/responses/test_route_host.py create mode 100644 tests/unit/sandbox/test_e2b_sandbox.py create mode 100644 tests/unit/sandbox/test_opensandbox_sandbox.py create mode 100644 tests/unit/sandbox/test_sandbox_tools.py create mode 100644 tests/unit/skills/test_skills_main.py create mode 100644 tests/unit/test_router/test_enforce_model_rate_limits.py create mode 100644 tests/unit/test_router/test_io_token_rate_limits.py create mode 100644 tests/unit/types/llms/test_types_llms_bedrock.py create mode 100644 tests/unit/types/llms/test_types_llms_openai.py create mode 100644 tests/unit/types/proxy/policy_engine/__init__.py create mode 100644 tests/unit/types/proxy/policy_engine/test_pipeline_types.py create mode 100644 tests/unit/types/proxy/policy_engine/test_policy_types.py create mode 100644 tests/unit/types/proxy/policy_engine/test_resolver_types.py create mode 100644 tests/unit/videos/__init__.py create mode 100644 tests/unit/videos/test_main.py create mode 100644 tests/unit/videos/test_utils.py diff --git a/tests/unit/rust_bridge/__init__.py b/tests/unit/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/ocr/__init__.py b/tests/unit/rust_bridge/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/ocr/test_route_host.py b/tests/unit/rust_bridge/ocr/test_route_host.py new file mode 100644 index 00000000000..699492e4424 --- /dev/null +++ b/tests/unit/rust_bridge/ocr/test_route_host.py @@ -0,0 +1,85 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure +from litellm.rust_bridge.ocr.route_host import response as build_ocr_response +from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest + +REQUEST: Final = LiteLLMOcrRequest( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key="test-key", + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"req_format": "markdown"}, +) + + +class RustUpstreamError(Exception): + def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: + super().__init__(status, body) + self.headers: Final = list(headers) + + +class RustFormatError(Exception): + ocr_request_format_error: Final = True + + +def test_rust_ocr_response_retains_provider_native_response(): + provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} + response = build_ocr_response( + { + "pages": [], + "model": "prebuilt-layout", + "document_annotation": None, + "usage_info": {"pages_processed": 0}, + "object": "ocr", + "provider_native_response": provider_response, + } + ) + + assert response.get_provider_native_response() == provider_response + assert response.model_dump().get("provider_native_response") is None + + +def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: + error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.RateLimitError) + assert public_error.status_code == 429 + assert public_error.response.headers["retry-after"] == "7" + assert public_error.response.text == '{"message": "slow down"}' + assert public_error.__context__ is error + assert public_error.llm_provider == "mistral" + + +def test_map_failure_maps_upstream_401_to_authentication_error() -> None: + error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.AuthenticationError) + assert public_error.status_code == 401 + assert public_error.response.text == '{"message": "Unauthorized"}' + assert public_error.__context__ is error + + +def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: + error: Final = RuntimeError("bridge exploded") + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert not isinstance(public_error, UpstreamFailure) + assert isinstance(public_error, litellm.APIConnectionError) + assert "bridge exploded" in str(public_error) + + +def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): + raise map_failure(RustFormatError(), REQUEST, "mistral") diff --git a/tests/unit/rust_bridge/responses/__init__.py b/tests/unit/rust_bridge/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/responses/test_route_host.py b/tests/unit/rust_bridge/responses/test_route_host.py new file mode 100644 index 00000000000..49bf19e7d8a --- /dev/null +++ b/tests/unit/rust_bridge/responses/test_route_host.py @@ -0,0 +1,57 @@ +from types import MappingProxyType +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.rust_bridge.responses.route_host import arguments, response +from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest +from litellm.types.llms.openai import ResponsesAPIResponse + + +def test_response_validates_into_the_public_responses_model() -> None: + built: Final = response( + MappingProxyType( + { + "id": "resp_native", + "object": "response", + "created_at": 1, + "model": "gpt-4o", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_native", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "native", "annotations": []}], + } + ], + } + ) + ) + + assert isinstance(built, ResponsesAPIResponse) + assert built.id == "resp_native" + assert built.output[0].content[0].text == "native" + + +def test_response_rejects_a_payload_missing_required_fields() -> None: + with pytest.raises(ValidationError): + response(MappingProxyType({"object": "response"})) + + +def test_arguments_are_the_public_kwargs_view() -> None: + kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) + request: Final = LiteLLMResponsesRequest( + model="gpt-4o", + input="hi", + stream=None, + api_key=None, + api_base=None, + custom_llm_provider="openai", + extra_headers=None, + kwargs=kwargs, + ) + + assert arguments(request) is kwargs diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py new file mode 100644 index 00000000000..e01b9120416 --- /dev/null +++ b/tests/unit/sandbox/test_e2b_sandbox.py @@ -0,0 +1,318 @@ +""" +Tests for the e2b code execution sandbox primitive. + +Unit tests inject a fake async HTTP client (dependency injection, no +monkeypatching) and assert request shapes and result mapping. Real-network +integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. +""" + +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.e2b.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + E2BSandboxConfig, +) + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + """Records outbound requests and returns canned responses keyed by URL.""" + + def __init__( + self, + *, + create_json=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "sandboxID": "sbx_123", + "domain": "e2b.app", + "envdAccessToken": "tok_abc", + } + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json)) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/execute"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None)) + if not (200 <= self.delete_status < 300): + raise httpx.HTTPStatusError( + f"status {self.delete_status}", + request=httpx.Request("DELETE", url), + response=httpx.Response(self.delete_status), + ) + return FakeResponse(status_code=self.delete_status) + + +# ---------- pure parser ---------- + + +def test_parse_lines_stdout_and_count(): + lines = [ + json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}), + json.dumps({"type": "number_of_executions", "execution_count": 1}), + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.stdout == "6\n" + assert result.execution_count == 1 + assert result.error is None + + +def test_parse_lines_error_surfaces_name_and_traceback(): + lines = [ + json.dumps( + { + "type": "error", + "name": "ZeroDivisionError", + "value": "division by zero", + "traceback": "Traceback (most recent call last): ...", + } + ) + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.error["name"] == "ZeroDivisionError" + assert "Traceback" in result.error["traceback"] + + +def test_parse_lines_result_carries_png(): + lines = [ + json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True}) + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.results and result.results[0]["png"] == "BASE64DATA" + assert "type" not in result.results[0] + + +# ---------- request shapes ---------- + + +@pytest.mark.asyncio +async def test_template_flows_into_create_request_as_templateID(): + client = FakeHTTPClient() + cfg = E2BSandboxConfig() + handle = await cfg.acreate_sandbox( + template="my-custom-template", api_key="e2b_key", client=client + ) + + method, url, headers, body = client.calls[0] + assert method == "POST" + assert url.endswith("/sandboxes") + assert body["templateID"] == "my-custom-template" # not "template" + assert body["secure"] is True + assert headers["X-API-Key"] == "e2b_key" + assert handle.id == "sbx_123" + assert handle._hidden_params["envd_access_token"] == "tok_abc" + + +@pytest.mark.asyncio +async def test_create_defaults_template_when_omitted(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client) + _, _, _, body = client.calls[0] + assert body["templateID"] == "code-interpreter-v1" + + +@pytest.mark.asyncio +async def test_run_code_targets_jupyter_host_with_access_token(): + client = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})] + ) + handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok_run"} + + result = await E2BSandboxConfig().arun_code( + container=handle, code="print(6*7)", client=client + ) + + method, url, headers, body = client.calls[0] + assert url == "https://49999-sbx_xyz.e2b.app/execute" + assert headers["X-Access-Token"] == "tok_run" + assert body["code"] == "print(6*7)" + assert result.stdout.strip() == "42" + + +@pytest.mark.asyncio +async def test_delete_issues_delete_to_sandbox_id(): + client = FakeHTTPClient(delete_status=204) + handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + + ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + + method, url, headers, _ = client.calls[0] + assert method == "DELETE" + assert url.endswith("/sandboxes/sbx_del") + assert ok is True + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + assert ok is False + + +# ---------- ephemeral teardown ---------- + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="e2b", code="1/0", api_key="e2b_key", client=client + ) + + methods = [c[0] for c in client.calls] + urls = [c[1] for c in client.calls] + assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete + assert urls[0].endswith("/sandboxes") + assert urls[1].endswith("/execute") + assert urls[2].endswith("/sandboxes/sbx_123") + + +# ---------- correctness guards ---------- + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + with pytest.raises(httpx.HTTPStatusError): + await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + + +@pytest.mark.asyncio +async def test_create_preserves_explicit_zero_timeout(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + timeout=0, api_key="e2b_key", client=client + ) + _, _, _, body = client.calls[0] + assert body["timeout"] == 0 + + +@pytest.mark.asyncio +async def test_run_code_rejects_bare_id_without_access_token(): + client = FakeHTTPClient() + with pytest.raises(ValueError, match="access token"): + await E2BSandboxConfig().arun_code( + container="sbx_no_token", code="print(1)", client=client + ) + assert client.calls == [] # never reached the network + + +def test_parse_lines_skips_non_json_lines(): + lines = [ + "not-json-heartbeat", + json.dumps({"type": "stdout", "text": "ok\n"}), + "", + "{partial", + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.stdout == "ok\n" + assert result.error is None + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + big_line = "x" * (MAX_OUTPUT_BYTES + 1) + client = FakeHTTPClient(execute_lines=[big_line]) + handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok"} + with pytest.raises(ValueError, match="exceeded"): + await E2BSandboxConfig().arun_code( + container=handle, code="print('x'*999)", client=client + ) + + +# ---------- public entrypoints ---------- + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})] + ) + container = await litellm.acreate_sandbox( + provider="e2b", api_key="e2b_key", client=client + ) + assert container.id == "sbx_123" + + result = await litellm.arun_code( + provider="e2b", + container=container, + api_key="e2b_key", + code="print(6*7)", + client=client, + ) + assert result.stdout.strip() == "42" + + assert ( + await litellm.adelete_sandbox( + provider="e2b", container=container, api_key="e2b_key", client=client + ) + is True + ) + + +@pytest.mark.asyncio +async def test_unsupported_provider_raises(): + with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): + await litellm.acreate_sandbox(provider="not-a-provider") + + +# ---------- api_base override ---------- + + +@pytest.mark.asyncio +async def test_create_uses_api_base_override(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + api_base="http://my-sandbox:8080", api_key="k", client=client + ) + _, url, _, _ = client.calls[0] + assert url == "http://my-sandbox:8080/sandboxes" + + +@pytest.mark.asyncio +async def test_create_defaults_to_e2b_api_base(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) + _, url, _, _ = client.calls[0] + assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py new file mode 100644 index 00000000000..2928dea100e --- /dev/null +++ b/tests/unit/sandbox/test_opensandbox_sandbox.py @@ -0,0 +1,647 @@ +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.opensandbox.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OpenSandboxSandboxConfig, +) +from litellm.utils import ProviderConfigManager + +TEST_API_BASE = "https://sandbox.test/v1" + + +def http_status_error(status_code, url="http://test"): + return httpx.HTTPStatusError( + f"status {status_code}", + request=httpx.Request("GET", url), + response=httpx.Response(status_code), + ) + + +def sse(data): + return f"data: {json.dumps(data)}" + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + def raise_for_status(self): + if self.status_code >= 400: + raise http_status_error(self.status_code) + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + def __init__( + self, + *, + create_json=None, + sandbox_states=None, + endpoint_json=None, + endpoint_responses=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + self.sandbox_states = list( + sandbox_states + or [ + { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + ] + ) + self.endpoint_json = endpoint_json or { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + } + self.endpoint_responses = ( + list(endpoint_responses) if endpoint_responses is not None else None + ) + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json, {"stream": stream})) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/code"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def get(self, url, headers=None, params=None, **kwargs): + self.calls.append(("GET", url, headers, None, params)) + if "/endpoints/44772" in url: + if self.endpoint_responses is not None and self.endpoint_responses: + response = self.endpoint_responses.pop(0) + if isinstance(response, Exception): + raise response + if isinstance(response, FakeResponse): + return response + return FakeResponse(json_data=response) + return FakeResponse(json_data=self.endpoint_json) + if "/sandboxes/" in url: + state = self.sandbox_states.pop(0) + return FakeResponse(json_data=state) + raise AssertionError(f"unexpected GET {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None, None)) + if not (200 <= self.delete_status < 300): + raise http_status_error(self.delete_status, url) + return FakeResponse(status_code=self.delete_status) + + +def test_parse_sse_lines_maps_output_result_count_and_error(): + lines = [ + sse({"type": "stdout", "text": "hello\n"}), + sse({"type": "stderr", "text": "warn\n"}), + sse({"type": "result", "results": {"text/plain": "4"}}), + sse({"type": "execution_count", "execution_count": 7}), + sse( + { + "type": "error", + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["Traceback"], + }, + } + ), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "hello\n" + assert result.stderr == "warn\n" + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 7 + assert result.error == { + "name": "ValueError", + "value": "bad", + "traceback": ["Traceback"], + } + + +def test_parse_sse_lines_skips_non_json_and_control_lines(): + lines = [ + "event: message", + "not-json", + "", + sse({"type": "stdout", "text": "ok\n"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "ok\n" + assert result.error is None + + +def test_parse_sse_lines_maps_fallback_shapes(): + lines = [ + "data:", + sse(["not-a-dict"]), + sse({"code": "BadRequest", "message": "nope"}), + sse({"type": "result", "text/plain": "4"}), + sse({"type": "error", "name": "RuntimeError", "text": "boom"}), + sse({"type": "execution_count", "execution_count": "8"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 8 + assert result.error == { + "name": "BadRequest", + "value": "nope", + "traceback": [], + } + fallback_error = OpenSandboxSandboxConfig._parse_lines( + [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] + ) + assert fallback_error.error == { + "name": "RuntimeError", + "value": "boom", + "traceback": [], + } + empty_string_error = OpenSandboxSandboxConfig._parse_lines( + [ + sse( + { + "type": "error", + "error": { + "ename": "", + "name": "FallbackName", + "evalue": "", + "value": "fallback value", + "traceback": [], + }, + } + ) + ] + ) + assert empty_string_error.error == { + "name": "", + "value": "", + "traceback": [], + } + + +def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): + def fake_secret(key): + if key == "OPEN_SANDBOX_API_KEY": + return "env-key" + if key == "OPEN_SANDBOX_API_BASE": + return TEST_API_BASE + return None + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + fake_secret, + ) + config = OpenSandboxSandboxConfig() + handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") + + assert config.validate_environment() == "env-key" + assert config.validate_environment(api_key="") == "" + assert config._api_key(api_key=None, handle=handle) == "env-key" + + handle._hidden_params = {"api_key": "stored-key"} + assert config._api_key(api_key=None, handle=handle) == "stored-key" + assert config._http(None) is not None + + body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy={"egress": [{"domain": "example.com"}]}, + secure_access=True, + ) + assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} + assert body["secureAccess"] is True + + other_body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy=None, + secure_access=False, + ) + assert body["resourceLimits"] is not other_body["resourceLimits"] + + assert config._sandbox_state(None) is None + assert config._sandbox_state({"status": "Running"}) is None + assert config._as_str_dict(None) == {} + assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( + "http://execd.local" + ) + assert config._api_base(None) == TEST_API_BASE + assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" + assert config._as_int("9") == 9 + assert config._as_int("nope") is None + assert config._as_int(None) is None + assert isinstance( + ProviderConfigManager.get_provider_sandbox_config("opensandbox"), + OpenSandboxSandboxConfig, + ) + + +def test_api_base_requires_kwarg_or_env(monkeypatch): + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + lambda key: None, + ) + + with pytest.raises(ValueError, match="api_base is required"): + OpenSandboxSandboxConfig._api_base(None) + + +@pytest.mark.asyncio +async def test_create_posts_default_body_and_omits_empty_api_key(): + client = FakeHTTPClient() + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + method, url, headers, body, _ = client.calls[0] + assert method == "POST" + assert url == f"{TEST_API_BASE}/sandboxes" + assert "OPEN-SANDBOX-API-KEY" not in headers + assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} + assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] + assert body["timeout"] == 300 + assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert handle.id == "osb_123" + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + + +@pytest.mark.asyncio +async def test_create_can_opt_into_internet_access(): + client = FakeHTTPClient() + + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + allow_internet_access=True, + client=client, + ) + + _, _, _, body, _ = client.calls[0] + assert "networkPolicy" not in body + + +@pytest.mark.asyncio +async def test_create_custom_options_poll_and_endpoint_resolution(): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + }, + sandbox_states=[ + { + "id": "osb_pending", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + } + ], + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + template="custom/image:latest", + timeout=600, + allow_internet_access=False, + api_key="osb-key", + api_base="https://sandbox.example/v1", + metadata={"suite": "unit"}, + env_vars={"PYTHONUNBUFFERED": "1"}, + resource_limits={"cpu": "500m", "memory": "512Mi"}, + resource_requests={"cpu": "250m", "memory": "256Mi"}, + entrypoint=["/bin/sh", "-lc", "sleep 3600"], + use_server_proxy=True, + client=client, + ) + + _, create_url, create_headers, body, _ = client.calls[0] + _, poll_url, poll_headers, _, _ = client.calls[1] + _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] + + assert create_url == "https://sandbox.example/v1/sandboxes" + assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert body["image"] == {"uri": "custom/image:latest"} + assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] + assert body["metadata"] == {"suite": "unit"} + assert body["env"] == {"PYTHONUNBUFFERED": "1"} + assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} + assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" + assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") + assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_params == {"use_server_proxy": True} + assert handle.id == "osb_pending" + + +@pytest.mark.asyncio +async def test_create_waits_across_pending_state(monkeypatch): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + }, + sandbox_states=[ + {"id": "osb_pending", "status": {"state": "Pending"}}, + {"id": "osb_pending", "status": {"state": "Running"}}, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + assert handle.id == "osb_pending" + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_for_terminal_state(): + client = FakeHTTPClient( + create_json={"id": "osb_failed", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_failed", "status": {"state": "Failed"}}, + ], + ) + + with pytest.raises(ValueError, match="entered Failed"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_create_times_out_waiting_for_running(): + client = FakeHTTPClient( + create_json={"id": "osb_slow", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_slow", "status": {"state": "Pending"}}, + ], + ) + + with pytest.raises(TimeoutError, match="was not Running"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=0, + poll_interval=0, + client=client, + ) + + +@pytest.mark.asyncio +async def test_create_waits_for_endpoint_resolution(monkeypatch): + client = FakeHTTPClient( + endpoint_responses=[ + http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), + { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + }, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + assert len(endpoint_calls) == 2 + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_when_endpoint_is_missing(): + client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) + + with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client + ) + + +@pytest.mark.asyncio +async def test_create_reraises_non_404_endpoint_error(): + client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_run_code_resolves_bare_id_and_posts_sse_request(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + result = await OpenSandboxSandboxConfig().arun_code( + container="osb_bare", + code="print(6*7)", + language="python", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + endpoint_call = client.calls[0] + run_call = client.calls[1] + assert endpoint_call[0] == "GET" + assert ( + endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" + ) + assert run_call[0] == "POST" + assert run_call[1] == "http://execd.local:44772/code" + assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" + assert run_call[3] == { + "code": "print(6*7)", + "context": {"language": "python"}, + } + assert run_call[4] == {"stream": True} + assert result.stdout == "42\n" + + +@pytest.mark.asyncio +async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): + client = FakeHTTPClient() + handle = ContainerHandle( + id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" + ) + handle._hidden_params = { + "execd_endpoint": "execd.example/route/44772", + "execd_headers": {}, + } + + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print(1)", client=client + ) + + assert client.calls[0][1] == "https://execd.example/route/44772/code" + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) + handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") + handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} + + with pytest.raises(ValueError, match="exceeded"): + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print('x')", client=client + ) + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + + ok = await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_gone", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + assert ok is False + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_err", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + container = await litellm.acreate_sandbox( + provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client + ) + result = await litellm.arun_code( + provider="opensandbox", + container=container, + code="print(6*7)", + api_key="", + client=client, + ) + ok = await litellm.adelete_sandbox( + provider="opensandbox", + container=container, + api_key="", + client=client, + ) + + assert container.id == "osb_123" + assert result.stdout == "42\n" + assert ok is True + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="opensandbox", + code="1/0", + api_key="", + api_base=TEST_API_BASE, + client=client, + ) + + assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] + assert client.calls[0][1].endswith("/sandboxes") + assert client.calls[1][1].endswith("/endpoints/44772") + assert client.calls[2][1].endswith("/code") + assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/unit/sandbox/test_sandbox_tools.py b/tests/unit/sandbox/test_sandbox_tools.py new file mode 100644 index 00000000000..06136534b13 --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_tools.py @@ -0,0 +1,181 @@ +"""Unit tests for the sandbox-tool registry.""" + +from litellm.sandbox import sandbox_tools + + +def _reset(): + sandbox_tools.clear_sandbox_tools() + + +def test_register_resolves_provider_key_and_base(): + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved == { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + } + _reset() + + +def test_register_clears_stale_entries_on_reload(): + """A tool removed from the config must not survive a re-registration.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "old", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("old") is not None + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "new", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("new") is not None + assert ( + sandbox_tools.resolve_sandbox_tool("old") is None + ), "stale tool must be gone after the config is reloaded" + _reset() + + +def test_register_empty_list_clears_removed_tools(): + """Reloading a config with sandbox_tools removed (the proxy passes an empty + list) must drop previously registered credentials from the process.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None + + sandbox_tools.register_sandbox_tools([]) + + assert ( + sandbox_tools.resolve_sandbox_tool("e2b_default") is None + ), "removing sandbox_tools from config must clear stale credentials" + _reset() + + +def test_register_resolves_secret_from_env(monkeypatch): + _reset() + monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "os.environ/MY_SANDBOX_KEY", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved is not None + assert resolved["api_key"] == "sk-from-env" + assert resolved["api_base"] is None + _reset() + + +def test_resolve_unknown_returns_none(): + _reset() + assert sandbox_tools.resolve_sandbox_tool("nope") is None + + +def test_register_skips_malformed_entries_without_crashing(): + """A single malformed entry (missing sandbox_tool_name, or not a dict) must + not crash registration during proxy startup/hot-reload; valid entries in the + same list must still register.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name + "not-a-dict", # wrong type + {"sandbox_tool_name": "", "litellm_params": {}}, # empty name + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("good") is not None + assert sandbox_tools.resolve_sandbox_tool("") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_skips_entry_missing_sandbox_provider(): + """An entry with a name but no sandbox_provider must be skipped at + registration so it cannot later resolve and call acreate_sandbox(provider=None), + which fails with a cryptic runtime error instead of a clear startup warning.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, + { + "sandbox_tool_name": "null_provider", + "litellm_params": {"sandbox_provider": None}, + }, + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("no_provider") is None + assert sandbox_tools.resolve_sandbox_tool("null_provider") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_swaps_registry_atomically(): + """register_sandbox_tools must replace the registry in one rebind so a + concurrent resolve never observes a half-populated or transiently empty + registry between clearing and repopulating.""" + _reset() + sandbox_tools.register_sandbox_tools( + [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] + ) + before = sandbox_tools._SANDBOX_TOOL_REGISTRY + + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, + {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, + ] + ) + after = sandbox_tools._SANDBOX_TOOL_REGISTRY + + assert after is not before, "the registry must be replaced, not mutated in place" + assert set(after) == {"b", "c"} + assert "a" not in after + _reset() diff --git a/tests/unit/skills/test_skills_main.py b/tests/unit/skills/test_skills_main.py new file mode 100644 index 00000000000..e1c66c8d9ea --- /dev/null +++ b/tests/unit/skills/test_skills_main.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import litellm.skills.main as skills_main +from litellm.types.utils import LlmProviders + + +def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( + monkeypatch, +) -> None: + """The REST /v1/skills form endpoint passes description/instructions as top-level + kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch + branch of create_skill() dropped both, so every LiteLLM-hosted skill was created + with description=None and instructions=None regardless of what the caller sent.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Converts files from one language into another" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( + "Take an uploaded document and produce it in the target language" + ) + + +def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: + """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Warehouse SQL Analyst", + extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Runs SQL against the inventory database" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" + + +def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) + + assert handler.create_skill_handler.call_args.kwargs["description"] is None + assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py new file mode 100644 index 00000000000..7577064b7f9 --- /dev/null +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -0,0 +1,468 @@ +""" +Tests for enforce_model_rate_limits feature. + +This feature allows users to enforce TPM/RPM limits set on model deployments +regardless of the routing strategy being used. +""" + +import asyncio +from datetime import timedelta +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) + +TPM_DEPLOYMENT = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "replica-test-id"}, + "model_name": "test-model", +} + + +def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: + dual_cache = DualCache(redis_cache=redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache) + now = litellm.utils.get_utc_datetime() + for minute in (now, now + timedelta(minutes=1)): + tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) + dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) + return dual_cache + + +class TestModelRateLimitingCheck: + """Test the ModelRateLimitingCheck class directly.""" + + def test_get_deployment_limits_from_top_level(self): + """Test extracting limits from top-level deployment config.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "tpm": 1000, + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 1000 + assert rpm == 10 + + def test_get_deployment_limits_from_litellm_params(self): + """Test extracting limits from litellm_params.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 2000 + assert rpm == 20 + + def test_get_deployment_limits_from_model_info(self): + """Test extracting limits from model_info.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 3000 + assert rpm == 30 + + def test_get_deployment_limits_none_when_not_set(self): + """Test that None is returned when limits are not set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm is None + assert rpm is None + + def test_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 11 # Over limit after increment + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + assert "current usage=11" in str(exc_info.value) + + def test_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit.""" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 6 + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 1000 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.get_cache.return_value = 1000 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.parametrize( + "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] + ) + def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): + redis_cache = MagicMock() + redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() + redis_cache.increment_cache.return_value = 2 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert check.pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + + def test_log_success_event_increments_cache(self): + """Test that log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + check.log_success_event(kwargs, None, None, None) + + # Verify increment_cache was called + mock_cache.increment_cache.assert_called_once() + _, kwarg_params = mock_cache.increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestModelRateLimitingCheckAsync: + """Test async methods of ModelRateLimitingCheck.""" + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache = AsyncMock(return_value=6) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=1000) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] + ) + async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.async_get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( + self, + ): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) + redis_cache.async_increment = AsyncMock(return_value=2) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert await check.async_pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_log_success_event_increments_cache(self): + """Test that async_log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + await check.async_log_success_event(kwargs, None, None, None) + + # Verify async_increment_cache was called + mock_cache.async_increment_cache.assert_called_once() + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestRouterWithEnforceModelRateLimits: + """Test Router integration with enforce_model_rate_limits.""" + + def test_router_initializes_with_enforce_model_rate_limits(self): + """Test that Router properly initializes the ModelRateLimitingCheck.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + router = Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Check that the callback was added + assert router.optional_callbacks is not None + assert len(router.optional_callbacks) == 1 + assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) + + def test_router_optional_callbacks_contains_model_rate_limiting(self): + """Test that ModelRateLimitingCheck is in the callbacks list.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Find the ModelRateLimitingCheck in litellm.callbacks + found = False + for callback in litellm.callbacks: + if isinstance(callback, ModelRateLimitingCheck): + found = True + break + + assert found, "ModelRateLimitingCheck should be in litellm.callbacks" + + +class TestModelRateLimitConcurrency: + """Test that RPM rate limiting is atomic under concurrent requests.""" + + @pytest.mark.asyncio + async def test_concurrent_requests_respect_rpm_limit(self): + """ + Fire 4 concurrent async requests with RPM limit of 2. + Exactly 2 should succeed and 2 should raise RateLimitError. + + This test validates the atomic increment-first pattern: + the old check-then-increment pattern would let 3+ through + due to a race condition on the local cache read. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + + deployment = { + "rpm": 2, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "concurrent-test-id"}, + "model_name": "test-model", + } + + async def attempt_request(): + return await check.async_pre_call_check(deployment) + + results = await asyncio.gather( + *[attempt_request() for _ in range(4)], + return_exceptions=True, + ) + + successes = [r for r in results if not isinstance(r, Exception)] + failures = [r for r in results if isinstance(r, litellm.RateLimitError)] + + assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" + assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py new file mode 100644 index 00000000000..a5a68271111 --- /dev/null +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -0,0 +1,1041 @@ +""" +Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). +""" + +import asyncio + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, + _reservation_value, + _resolve_max_tokens, + async_io_token_pre_call_check, + async_io_token_reconcile_success, + build_io_token_rate_limit_headers, + deployment_has_io_token_limits, + get_io_token_rate_limit_request_kwargs, + io_token_reconcile_success, + io_token_refund_failure, + refund_stale_reservation_before_retry, + set_io_token_rate_limit_request_kwargs, +) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) +from litellm.types.utils import ModelResponse, Usage + + +class TestIOTokenRateLimitHelpers: + def test_deployment_has_io_token_limits(self): + assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) + assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) + + def test_reservation_value_minimal_when_estimate_fails(self): + # A failed/empty estimate (0) must reserve a minimal slot, not the + # entire limit - otherwise one request whose estimate failed fills + # the whole bucket and blocks every concurrent request until it + # completes and reconciles. + assert _reservation_value(0, 100) == 1 + assert _reservation_value(0, 1) == 1 + # A real non-zero estimate is reserved as-is. + assert _reservation_value(42, 100) == 42 + + def test_resolve_max_tokens_respects_explicit_zero(self): + deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} + # An explicit max_tokens=0 is honored, not replaced by the model default. + assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 + # max_completion_tokens is the fallback only when max_tokens is absent. + assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 + assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 + + def test_build_io_token_rate_limit_headers(self): + headers = build_io_token_rate_limit_headers( + itpm_limit=200, + otpm_limit=40, + current_itpm=15, + current_otpm=4, + ) + assert headers["x-ratelimit-limit-input-tokens"] == 200 + assert headers["x-ratelimit-remaining-input-tokens"] == 185 + assert headers["x-ratelimit-limit-output-tokens"] == 40 + assert headers["x-ratelimit-remaining-output-tokens"] == 36 + + +class TestModelRateLimitingCheckIOTokens: + @pytest.mark.asyncio + async def test_itpm_reservation_and_reconcile(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-test-id"}, + "model_name": "opus", + } + + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + + kwargs = { + "standard_logging_object": { + "model_id": "io-test-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "index": 0, + "finish_reason": "stop", + } + ], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + await check.async_log_success_event(kwargs, response, None, None) + + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + # ITPM tracks input tokens only (billable prompt tokens), not output. + assert current_itpm == 5 + assert current_otpm == 3 + + @pytest.mark.asyncio + async def test_itpm_limit_raises_429(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 5, + }, + "model_info": {"id": "io-limit-id"}, + "model_name": "opus", + } + + # ITPM enforces input tokens only; the prompt alone must exceed the limit, + # a large max_tokens must not contribute to the ITPM reservation. + set_io_token_rate_limit_request_kwargs( + { + "messages": [ + { + "role": "user", + "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", + } + ], + "max_tokens": 10, + "metadata": {}, + } + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "ITPM limit=5" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + otpm_limit = 10 + max_tokens = 4 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": otpm_limit, + }, + "model_info": {"id": "io-otpm-race-id"}, + "model_name": "opus", + } + + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": max_tokens, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + + # Atomic reservation must never let concurrent requests overshoot the limit. + assert current_otpm is not None + assert current_otpm <= otpm_limit + assert successes == otpm_limit // max_tokens + assert current_otpm == successes * max_tokens + + @pytest.mark.asyncio + async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): + """ + When input-token estimation yields 0 (no messages/prompt/input field, + unsupported model, tokenizer error), the reservation must be a + minimal 1 token, not the entire itpm limit. Otherwise the first + request whose estimate fails fills the whole bucket and every + concurrent request is rejected until it completes - effectively + serializing traffic to the deployment. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_limit = 5 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": itpm_limit, + }, + "model_info": {"id": "io-itpm-estimate-fail-id"}, + "model_name": "opus", + } + + # No messages/prompt/input field -> _estimate_input_tokens returns 0. + set_io_token_rate_limit_request_kwargs( + { + "max_tokens": 5, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + + # A minimal 1-token reservation per request lets itpm_limit concurrent + # requests through, instead of a single request starving the rest. + assert current_itpm is not None + assert current_itpm <= itpm_limit + assert successes == itpm_limit + + @pytest.mark.asyncio + async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + # Production kwargs commonly carry litellm_params.metadata; the stashed + # reservation lives in the top-level metadata and must still be found. + kwargs = { + "standard_logging_object": { + "model_id": "io-lp-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, + "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + @pytest.mark.asyncio + async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + }, + "model_info": {"id": "io-zero-est-id"}, + "model_name": "opus", + } + + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A failed/zero estimate reserves a minimal 1 token, not the full + # itpm limit, so it doesn't starve concurrent requests. + assert await dual_cache.async_get_cache(key=itpm_key) == 1 + + kwargs = { + "standard_logging_object": { + "model_id": "io-zero-est-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + + @pytest.mark.asyncio + async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): + """ + A zero/failed estimate reserves a minimal 1 token rather than the + full itpm limit, so up to itpm_limit such calls are allowed + concurrently instead of the first one claiming the entire bucket. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 2, + }, + "model_info": {"id": "io-zero-cap-id"}, + "model_name": "opus", + } + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + # Second zero-estimate call still fits within the itpm=2 limit. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + await check.async_pre_call_check(deployment) + + # A third exceeds the limit and is rejected. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + @pytest.mark.asyncio + async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": 5, + }, + "model_info": {"id": "io-zero-output-id"}, + "model_name": "opus", + } + zero_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 0, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + await check.async_pre_call_check(deployment) + + zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 + assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 + + normal_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + await check.async_pre_call_check(deployment) + + normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + + def test_sync_io_pre_call_reserves_and_reconciles(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-sync-id"}, + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + check.pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + kwargs = { + "standard_logging_object": { + "model_id": "io-sync-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 5 + assert dual_cache.get_cache(key=otpm_key) == 3 + + @pytest.mark.asyncio + async def test_reconcile_runs_via_success_event_without_model_id(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + # standard_logging_object has no model_id (only the TPM path needs it); + # IO reconciliation must still run off the stashed cache key. + kwargs = { + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + + @pytest.mark.asyncio + async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + + # Shared request metadata carrying the first (IO) deployment's reservation. + metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} + fail_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "io-first", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + } + await check.async_log_failure_event(fail_kwargs, None, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + # Retry succeeds on a non-IO fallback deployment reusing the same metadata. + retry_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "non-io-second", + "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, + "metadata": {}, + "total_tokens": 12, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), + ) + await check.async_log_success_event(retry_kwargs, response, None, None) + + # The first deployment's ITPM counter is not driven negative... + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + # ...and the non-IO deployment's TPM usage is tracked normally. + tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" + assert await dual_cache.async_get_cache(key=tpm_key) == 12 + + @pytest.mark.asyncio + async def test_stale_reservation_refunded_before_retry_overwrites_it(self): + """ + A retry reuses the same mutable kwargs dict for the next deployment. + If deployment A's failure event hasn't run yet (e.g. it was scheduled + as a background task) when the retry calls + set_io_token_rate_limit_request_kwargs for deployment B, the router + must first synchronously refund + clear A's reservation via + refund_stale_reservation_before_retry - otherwise A's counter stays + elevated by the reservation until its TTL expires, and the + now-orphaned sentinels must not leak into B's accounting either. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + minute = get_utc_datetime().strftime("%H-%M") + itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) + + # Deployment A's still-unreconciled reservation, stashed on the shared + # kwargs dict the retry loop reuses. + shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} + + # Router calls this before overwriting kwargs for deployment B's attempt - + # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. + refund_stale_reservation_before_retry(dual_cache, shared_kwargs) + + # A's reservation is refunded immediately, not left stranded for a + # background failure task that may run arbitrarily later (or never, + # if the sentinels get cleared out from under it first). + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] + assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] + + # A's own (now-late) failure event finds nothing left to refund and + # is a safe no-op, since the sentinels were already cleared above. + io_token_refund_failure(dual_cache, shared_kwargs) + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + + # The retry proceeds to stash deployment B's own reservation on the + # same dict; it starts clean, unaffected by A's cleared sentinels. + set_io_token_rate_limit_request_kwargs(shared_kwargs) + itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" + shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 + shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b + await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) + assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 + + @pytest.mark.asyncio + async def test_client_supplied_reservation_keys_are_stripped(self): + # metadata is caller-controlled; the server-only reservation sentinels + # must be removed before the router captures the request kwargs. + forged = { + "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, + "litellm_metadata": {OTPM_RESERVED_KEY: 7}, + "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, + } + set_io_token_rate_limit_request_kwargs(forged) + stored = get_io_token_rate_limit_request_kwargs() + + assert ITPM_RESERVED_KEY not in stored["metadata"] + assert ITPM_CACHE_KEY not in stored["metadata"] + assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] + assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] + + @pytest.mark.asyncio + async def test_forged_reservation_cannot_decrement_counter(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + victim_key = "global_router:victim:model:itpm:00-00" + await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) + + # A caller forges a reservation pointing at another deployment's counter. + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, + "standard_logging_object": { + "model_id": "m", + "hidden_params": {"litellm_model_name": "model"}, + "metadata": {}, + "total_tokens": 2, + }, + } + # The router sanitizes the request kwargs before the call runs. + set_io_token_rate_limit_request_kwargs(kwargs) + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + await check.async_log_success_event(kwargs, response, None, None) + + # The forged reservation was stripped, so the victim counter is untouched. + assert await dual_cache.async_get_cache(key=victim_key) == 100 + + @pytest.mark.asyncio + async def test_otpm_reservation_error_rolls_back_itpm(self): + from litellm.utils import get_utc_datetime + + class _OtpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":otpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _OtpmFailCache() + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 1000, + "otpm": 1000, + }, + "model_info": {"id": "io-rollback-id"}, + "model_name": "opus", + } + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + ) + + with pytest.raises(RuntimeError): + await async_io_token_pre_call_check(dual_cache, deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A transient OTPM error must release the ITPM reservation, not leak it. + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + + @pytest.mark.asyncio + async def test_reconcile_clears_stash_even_when_increment_errors(self): + class _ItpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":itpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _ItpmFailCache() + metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} + kwargs = {"metadata": metadata} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + with pytest.raises(RuntimeError): + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + # The stash is cleared even though reconciliation raised, so a duplicate + # success event can't re-process it. + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + @pytest.mark.asyncio + async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): + import logging + + check = ModelRateLimitingCheck(dual_cache=DualCache()) + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, + "model_info": {}, + } + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + + warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] + # id-less deployments are not collapsed onto a single dedup key. + assert len(warnings) == 2 + + @pytest.mark.asyncio + async def test_missing_deployment_id_skips_io_reservation(self): + dual_cache = DualCache() + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, + "model_info": {}, # no id -> cannot build a per-deployment cache key + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + result = await async_io_token_pre_call_check(dual_cache, deployment) + + assert result is deployment + # No reservation is stashed, so nothing lands in a shared None:None bucket. + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_uses_reservation_minute_key(self): + dual_cache = DualCache() + # Reservation was made on a fixed minute key; a call that finishes in a + # later minute must reconcile against that same key, never a key built + # from the response-time minute. + itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), + ) + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + @pytest.mark.asyncio + async def test_reconcile_missing_usage_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + ) + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_total_tokens_only_keeps_reservation(self): + """ + A response usage object with only total_tokens (no prompt/completion + breakdown) can't be split into input/output, so it must be treated the + same as missing usage: keep the reservation instead of resolving to + (0, 0) and refunding it in full. + """ + dual_cache = DualCache() + itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = {"type": "message", "usage": {"total_tokens": 13}} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + + def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" + dual_cache.set_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": {"total_tokens": 4}, + } + response = {"type": "message", "role": "assistant", "content": []} + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 10 + + @pytest.mark.asyncio + async def test_reconcile_falls_back_to_standard_logging_object(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "prompt_tokens": 4, + "completion_tokens": 0, + "total_tokens": 4, + }, + } + response = {"type": "message", "role": "assistant", "content": []} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + def test_sync_reconcile_anthropic_dict_usage(self): + dual_cache = DualCache() + itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" + dual_cache.set_cache(key=itpm_key, value=6, ttl=60) + dual_cache.set_cache(key=otpm_key, value=4, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 6, + OTPM_RESERVED_KEY: 4, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = { + "type": "message", + "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, + } + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 2 + assert dual_cache.get_cache(key=otpm_key) == 2 + + @pytest.mark.asyncio + async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): + import logging + + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + deployment = { + "litellm_params": { + "model": deployment_name, + "itpm": 100, + "rpm": 1, + }, + "model_info": {"id": model_id}, + "model_name": "opus", + } + + minute = get_utc_datetime().strftime("%H-%M") + rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) + + request_kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + assert await dual_cache.async_get_cache(key=rpm_key) == 6 + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + assert any("both limit types are enforced" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): + """ + A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on + success, otherwise the tpm_key the pre-call check reads is never written + and the tpm_limit can never be enforced. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + await check.async_log_success_event(kwargs, response, None, None) + + # ITPM reconciled down from the 5-token reservation to actual usage (3). + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. + assert await dual_cache.async_get_cache(key=tpm_key) == 7 + + def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-sync-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + dual_cache.set_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 3 + assert dual_cache.get_cache(key=tpm_key) == 7 + + @pytest.mark.asyncio + async def test_failure_refunds_itpm_reservation(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} + kwargs = { + "standard_logging_object": { + "model_id": "io-refund-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": dict(reservation), + }, + "metadata": dict(reservation), + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + +class TestRouterIOTokenIntegration: + @pytest.mark.asyncio + async def test_model_group_info_aggregates_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "opus", + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 20, + }, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + info = router.get_model_group_info("opus") + assert info is not None + assert info.itpm == 100 + assert info.otpm == 20 + + +class TestContextSlotRetention: + def test_setter_stores_kwargs_only_for_io_limited_deployments(self): + """ + The context slot pins the entire request kwargs (messages included) + for the lifetime of the surrounding asyncio context, and pooled + resources created mid-request (e.g. redis connections) capture that + context, extending the pin far past the request. Only ITPM/OTPM + pre-call checks read the slot, so the setter must store None for + deployments without io token limits and still clear reservation + sentinels from kwargs either way. + """ + kwargs = { + "messages": [{"role": "user", "content": "x" * 1000}], + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + } + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + assert ITPM_CACHE_KEY not in kwargs["metadata"] + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + assert get_io_token_rate_limit_request_kwargs() is kwargs + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_does_not_pin_kwargs_without_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + } + ] + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("plain") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_pins_kwargs_for_io_limited_deployment(self): + router = Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("limited") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is kwargs diff --git a/tests/unit/types/llms/test_types_llms_bedrock.py b/tests/unit/types/llms/test_types_llms_bedrock.py new file mode 100644 index 00000000000..a5ad882e775 --- /dev/null +++ b/tests/unit/types/llms/test_types_llms_bedrock.py @@ -0,0 +1,46 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams + + +def test_model_validate_keeps_auth_params_and_ignores_request_params(): + auth_params = AwsAuthParams.model_validate( + { + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-role", + "aws_session_name": "litellm-session", + "aws_external_id": "litellm-external-id", + "aws_region_name": "us-west-2", + "aws_bedrock_runtime_endpoint": "https://bedrock.example.com", + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "temperature": 0.1, + "messages": [{"role": "user", "content": "hi"}], + } + ) + + assert auth_params.aws_role_name == "arn:aws:iam::999999999999:role/litellm-role" + assert auth_params.aws_session_name == "litellm-session" + assert auth_params.aws_external_id == "litellm-external-id" + assert auth_params.aws_access_key_id is None + assert set(auth_params.model_dump()) == set(AWS_AUTH_PARAM_KEYS) + assert not set(AWS_AUTH_PARAM_KEYS) & {"aws_region_name", "aws_bedrock_runtime_endpoint", "model", "temperature"} + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("aws_role_name", 1234), + ("aws_session_name", ["litellm-session"]), + ("aws_external_id", {"id": "x"}), + ], +) +def test_model_validate_rejects_non_string_credentials(field, value): + with pytest.raises(ValidationError): + AwsAuthParams.model_validate({field: value}) + + +def test_frozen_struct_rejects_field_assignment(): + auth_params = AwsAuthParams(aws_role_name="arn:aws:iam::999999999999:role/litellm-role") + + with pytest.raises(ValidationError): + auth_params.aws_role_name = "arn:aws:iam::999999999999:role/other-role" diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py new file mode 100644 index 00000000000..64ec09838e8 --- /dev/null +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -0,0 +1,591 @@ +import asyncio +from typing import Optional +from unittest.mock import AsyncMock, patch + +import pytest + +import json + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent + + +@pytest.mark.parametrize("stream", (False, True)) +def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: + from typing import Final + + from litellm.types.llms.openai import ( + ChatCompletionReasoningItem, + ChatCompletionReasoningSummaryTextBlock, + ) + from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + ) + + reasoning_item: Final = ChatCompletionReasoningItem( + type="reasoning", + id="rs_123", + encrypted_content="encrypted", + summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], + ) + response: Final = ( + ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) + if stream + else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) + ) + message_key: Final = "delta" if stream else "message" + assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + restored: Final = type(response).model_validate_json(response.model_dump_json()) + assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + +def test_generic_event(): + from litellm.types.llms.openai import GenericEvent + + event = {"type": "test", "test": "test"} + event = GenericEvent(**event) + assert event.type == "test" + assert event.test == "test" + + +def test_output_item_added_event(): + from litellm.types.llms.openai import OutputItemAddedEvent + + event = { + "type": "response.output_item.added", + "sequence_number": 4, + "output_index": 1, + "item": None, + } + event = OutputItemAddedEvent(**event) + assert event.type == "response.output_item.added" + assert event.sequence_number == 4 + assert event.output_index == 1 + assert event.item is None + + +class TestResponsesAPIResponseOutputText: + """Tests for the output_text property on ResponsesAPIResponse""" + + def test_output_text_with_single_message(self): + """Test output_text with a single message containing text output""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello, world!", + } + ], + } + ], + ) + + assert response.output_text == "Hello, world!" + + def test_output_text_with_multiple_messages(self): + """Test output_text with multiple messages aggregates all text""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "First part. ", + } + ], + }, + { + "type": "message", + "id": "msg_2", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Second part.", + } + ], + }, + ], + ) + + assert response.output_text == "First part. Second part." + + def test_output_text_with_no_text_content(self): + """Test output_text returns empty string when no output_text content exists""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + } + ], + ) + + assert response.output_text == "" + + def test_output_text_with_mixed_content(self): + """Test output_text only aggregates output_text type content""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "The weather is sunny. ", + }, + { + "type": "refusal", + "refusal": "I cannot do that.", + }, + ], + }, + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + }, + ], + ) + + assert response.output_text == "The weather is sunny. " + + def test_output_text_with_empty_output(self): + """Test output_text returns empty string with empty output list""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[], + ) + + assert response.output_text == "" + + +class TestAssistantMessageImageUrlContent: + """ + Regression tests for image_url blocks in assistant message content. + + Bug: ChatCompletionAssistantMessage.content did not include + ChatCompletionImageObject in its union, so Pydantic v2 silently dropped + image_url blocks (content → []) when serialising via AllMessageValues. + This affects users who store conversation history as JSON (e.g. in a DB) + and read it back typed as list[AllMessageValues]. + """ + + ASSISTANT_MESSAGE_WITH_IMAGE = { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here is the image you requested:"}, + { + "type": "image_url", + "image_url": { + "url": ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" + "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + }, + }, + ], + } + + def test_assistant_message_image_url_preserved_single(self): + """ + TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive + validate_python → dump_python without being dropped or raising an error. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import ChatCompletionAssistantMessage + + adapter = TypeAdapter(ChatCompletionAssistantMessage) + validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) + dumped = adapter.dump_python(validated) + + raw_content = dumped.get("content") + # Pydantic may return a lazy SerializationIterator for Iterable fields; + # convert to list to consume it — this must not raise ValidationError. + content_blocks = list(raw_content) if raw_content is not None else [] + + assert ( + len(content_blocks) == 2 + ), f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" + types = [b.get("type") for b in content_blocks if isinstance(b, dict)] + assert ( + "image_url" in types + ), f"image_url block was silently dropped; blocks: {content_blocks}" + + def test_assistant_message_image_url_preserved_in_all_message_values(self): + """ + TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an + assistant message must not be silently dropped during dump_python(mode='json'). + + This is the primary failing path: conversation history stored as JSON in a + database and read back typed as list[AllMessageValues]. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import AllMessageValues + + conversation = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a LiteLLM costume", + }, + self.ASSISTANT_MESSAGE_WITH_IMAGE, + ] + + adapter = TypeAdapter(List[AllMessageValues]) + validated = adapter.validate_python(conversation) + dumped = adapter.dump_python(validated, mode="json") + + assistant = next((m for m in dumped if m.get("role") == "assistant"), None) + assert assistant is not None, "Assistant message missing after serialisation" + + content = assistant.get("content", []) + assert isinstance( + content, list + ), f"content should be a list, got {type(content)}" + assert ( + len(content) == 2 + ), f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" + types = [b.get("type") for b in content if isinstance(b, dict)] + assert ( + "image_url" in types + ), f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" + + +class TestResponsesAPIReasoningNullFields: + """ + Tests for issue #16824: reasoning output items should not include null + status/content/encrypted_content fields. + + When a provider returns reasoning items without these fields, LiteLLM's + Pydantic parsing adds them as Optional defaults (None). Serializing them + as null breaks downstream SDKs (e.g., the OpenAI C# SDK crashes on + status=null). + + The fix uses a field_serializer on ResponsesAPIResponse.output that + mirrors the request-side filtering in + OpenAIResponsesAPIConfig._handle_reasoning_item(). + """ + + def _make_response(self, output): + from litellm.types.llms.openai import ResponsesAPIResponse + + return ResponsesAPIResponse( + id="resp_test", + created_at=1741476542, + model="gpt-5-mini", + object="response", + status="completed", + output=output, + ) + + def test_reasoning_item_null_fields_removed_model_dump(self): + """Null status/content/encrypted_content should be absent from model_dump.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert "encrypted_content" not in reasoning + + def test_reasoning_item_null_fields_removed_model_dump_json(self): + """Null fields should also be absent from model_dump_json.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + parsed = json.loads(response.model_dump_json()) + reasoning = parsed["output"][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert "encrypted_content" not in reasoning + + def test_reasoning_item_non_null_values_preserved(self): + """Non-null values on reasoning items should be kept.""" + response = self._make_response( + output=[ + { + "id": "rs_abc", + "type": "reasoning", + "summary": [], + "status": "completed", + "encrypted_content": "gAAAA...", + } + ] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert reasoning["status"] == "completed" + assert reasoning["encrypted_content"] == "gAAAA..." + + def test_message_item_not_affected(self): + """Non-reasoning output items should keep all their fields.""" + response = self._make_response( + output=[ + { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Hello!", + "annotations": [], + } + ], + } + ] + ) + dumped = response.model_dump() + message = dumped["output"][0] + assert message["status"] == "completed" + assert message["type"] == "message" + assert len(message["content"]) == 1 + + def test_mixed_output_reasoning_and_message(self): + """Reasoning items cleaned, message items untouched in same response.""" + response = self._make_response( + output=[ + {"id": "rs_abc", "type": "reasoning", "summary": []}, + { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Answer", + "annotations": [], + } + ], + }, + ] + ) + dumped = response.model_dump() + reasoning = [ + o + for o in dumped["output"] + if isinstance(o, dict) and o.get("type") == "reasoning" + ][0] + message = [ + o + for o in dumped["output"] + if isinstance(o, dict) and o.get("type") == "message" + ][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert message["status"] == "completed" + assert len(message["content"]) == 1 + + def test_reasoning_core_fields_preserved(self): + """id, type, summary should always be present on reasoning items.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": ["thinking..."]}] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert reasoning["id"] == "rs_abc" + assert reasoning["type"] == "reasoning" + assert reasoning["summary"] == ["thinking..."] + + def test_top_level_null_fields_unaffected(self): + """Top-level response fields with None should not be affected.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + dumped = response.model_dump() + assert "error" in dumped + assert dumped["error"] is None + assert "instructions" in dumped + assert dumped["instructions"] is None + + +def test_normalize_fine_tuning_job_dict_maps_azure_pending(): + from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict + + out = _normalize_fine_tuning_job_dict( + {"organization_id": None, "result_files": None, "status": "pending"}, + is_azure=True, + ) + assert out["organization_id"] == "" + assert out["result_files"] == [] + assert out["status"] == "queued" + + +def test_normalize_fine_tuning_job_dict_openai_unchanged(): + from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict + + data = {"organization_id": None, "result_files": None, "status": "pending"} + out = _normalize_fine_tuning_job_dict(data, is_azure=False) + assert out is data + + +def test_openai_file_object_accepts_pending_status(): + from litellm.types.llms.openai import OpenAIFileObject + + file_obj = OpenAIFileObject( + id="file-123", + bytes=1024, + created_at=1677610602, + filename="train.jsonl", + object="file", + purpose="fine-tune", + status="pending", + ) + assert file_obj.status == "pending" + + +class TestOpenAIFileObjectBatchGuardrailSerialization: + """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" + + @staticmethod + def _file_object(**overrides): + from litellm.types.llms.openai import OpenAIFileObject + + return OpenAIFileObject( + id="file-123", + object="file", + bytes=1024, + created_at=1677610602, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + **overrides, + ) + + @staticmethod + def _report(): + from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport + + return BatchGuardrailReport( + submitted_records=3, + modified_records=(BatchGuardrailRecord(line=2, custom_id="dirty", action="redacted"),), + ) + + @pytest.mark.parametrize("mode", ["python", "json"]) + def test_key_absent_when_unset(self, mode): + assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode=mode) + + @pytest.mark.parametrize("mode", ["python", "json"]) + def test_key_present_when_set(self, mode): + dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode=mode) + assert dumped["litellm_batch_guardrail"]["submitted_records"] == 3 + + def test_nested_nulls_of_a_set_report_survive(self): + """`exclude_none=True` was rejected as the fix because it would strip these.""" + dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode="json") + assert dumped["litellm_batch_guardrail"]["modified_records"] == [ + {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} + ] + + def test_by_alias_dump_also_omits_the_key(self): + """Tripwire: the serializer filters a literal key name, which an added alias would bypass.""" + assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode="json", by_alias=True) + + def test_other_optional_fields_still_serialize_as_null(self): + dumped = self._file_object().model_dump(mode="json") + assert dumped["expires_at"] is None + assert dumped["status_details"] is None + + def test_round_trip_of_a_set_report_is_lossless(self): + from litellm.types.llms.openai import OpenAIFileObject + + original = self._file_object(litellm_batch_guardrail=self._report()) + assert OpenAIFileObject(**original.model_dump()) == original + + def test_serialization_json_schema_still_describes_the_model(self): + """A return annotation on the wrap serializer would collapse this to a bare object.""" + from litellm.types.llms.openai import OpenAIFileObject + + schema = OpenAIFileObject.model_json_schema(mode="serialization") + assert "litellm_batch_guardrail" in schema["properties"] + + def test_key_omitted_inside_a_file_list_page(self): + from litellm.types.llms.openai import FileListPage + + page = FileListPage(object="list", data=[self._file_object()], has_more=False) + assert "litellm_batch_guardrail" not in page.model_dump(mode="json")["data"][0] + + +def _binary_content(payload: bytes) -> HttpxBinaryResponseContent: + import httpx + + return HttpxBinaryResponseContent(httpx.Response(200, content=payload)) + + +def test_httpx_binary_response_content_hidden_params_are_per_instance(): + first = _binary_content(b"first") + second = _binary_content(b"second") + + first._hidden_params["response_cost"] = 0.5 + + assert second._hidden_params == {} + + +def test_set_response_cost_none_leaves_hidden_params_empty(): + binary_response = _binary_content(b"audio") + + binary_response.set_response_cost(None) + + assert "response_cost" not in binary_response._hidden_params + + binary_response.set_response_cost(0.25) + + assert binary_response._hidden_params["response_cost"] == 0.25 + + binary_response.set_response_cost(None) + + assert "response_cost" not in binary_response._hidden_params diff --git a/tests/unit/types/proxy/policy_engine/__init__.py b/tests/unit/types/proxy/policy_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/policy_engine/test_pipeline_types.py b/tests/unit/types/proxy/policy_engine/test_pipeline_types.py new file mode 100644 index 00000000000..2e5986d3ef8 --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_pipeline_types.py @@ -0,0 +1,168 @@ +""" +Tests for pipeline type definitions. +""" + +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.pipeline_types import ( + GuardrailPipeline, + PipelineExecutionResult, + PipelineStep, + PipelineStepResult, +) +from litellm.types.proxy.policy_engine.policy_types import ( + Policy, + PolicyGuardrails, +) + + +def test_pipeline_step_defaults(): + step = PipelineStep(guardrail="my-guard") + assert step.on_fail == "block" + assert step.on_pass == "allow" + assert step.on_error is None + assert step.pass_data is False + assert step.modify_response_message is None + + +def test_pipeline_step_valid_actions(): + step = PipelineStep(guardrail="my-guard", on_fail="next", on_pass="next") + assert step.on_fail == "next" + assert step.on_pass == "next" + + +def test_pipeline_step_all_action_types(): + for action in ("allow", "block", "next", "modify_response"): + step = PipelineStep( + guardrail="g", on_fail=action, on_pass=action, on_error=action + ) + assert step.on_fail == action + assert step.on_pass == action + assert step.on_error == action + + +def test_pipeline_step_invalid_action_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_fail="invalid_action") + + +def test_pipeline_step_invalid_on_pass_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_pass="skip") + + +def test_pipeline_step_on_error_valid(): + step = PipelineStep( + guardrail="g", on_error="next", on_fail="block", on_pass="allow" + ) + assert step.on_error == "next" + + +def test_pipeline_step_invalid_on_error_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_error="invalid") + + +def test_pipeline_requires_at_least_one_step(): + with pytest.raises(ValidationError): + GuardrailPipeline(mode="pre_call", steps=[]) + + +def test_pipeline_invalid_mode_rejected(): + with pytest.raises(ValidationError): + GuardrailPipeline( + mode="during_call", + steps=[PipelineStep(guardrail="g")], + ) + + +def test_pipeline_valid_modes(): + for mode in ("pre_call", "post_call"): + pipeline = GuardrailPipeline( + mode=mode, + steps=[PipelineStep(guardrail="g")], + ) + assert pipeline.mode == mode + + +def test_pipeline_with_multiple_steps(): + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="g1", on_fail="next", on_pass="allow"), + PipelineStep(guardrail="g2", on_fail="block", on_pass="allow"), + ], + ) + assert len(pipeline.steps) == 2 + assert pipeline.steps[0].guardrail == "g1" + assert pipeline.steps[1].guardrail == "g2" + + +def test_policy_with_pipeline_parses(): + policy = Policy( + guardrails=PolicyGuardrails(add=["g1", "g2"]), + pipeline=GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="g1", on_fail="next"), + PipelineStep(guardrail="g2"), + ], + ), + ) + assert policy.pipeline is not None + assert len(policy.pipeline.steps) == 2 + + +def test_policy_without_pipeline(): + policy = Policy( + guardrails=PolicyGuardrails(add=["g1"]), + ) + assert policy.pipeline is None + + +def test_pipeline_step_result(): + result = PipelineStepResult( + guardrail_name="g1", + outcome="fail", + action_taken="next", + error_detail="Content policy violation", + duration_seconds=0.05, + ) + assert result.outcome == "fail" + assert result.action_taken == "next" + + +def test_pipeline_execution_result(): + result = PipelineExecutionResult( + terminal_action="block", + step_results=[ + PipelineStepResult( + guardrail_name="g1", + outcome="fail", + action_taken="next", + ), + PipelineStepResult( + guardrail_name="g2", + outcome="fail", + action_taken="block", + ), + ], + error_message="Content blocked", + ) + assert result.terminal_action == "block" + assert len(result.step_results) == 2 + + +def test_pipeline_step_extra_fields_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="g", unknown_field="value") + + +def test_pipeline_extra_fields_rejected(): + with pytest.raises(ValidationError): + GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="g")], + unknown="value", + ) diff --git a/tests/unit/types/proxy/policy_engine/test_policy_types.py b/tests/unit/types/proxy/policy_engine/test_policy_types.py new file mode 100644 index 00000000000..bcd6d39aa4d --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_policy_types.py @@ -0,0 +1,15 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment + + +@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) +def test_policy_attachment_accepts_int32_priority(priority: int): + assert PolicyAttachment(policy="p", priority=priority).priority == priority + + +@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) +def test_policy_attachment_rejects_priority_outside_int32(priority: int): + with pytest.raises(ValidationError): + PolicyAttachment(policy="p", priority=priority) diff --git a/tests/unit/types/proxy/policy_engine/test_resolver_types.py b/tests/unit/types/proxy/policy_engine/test_resolver_types.py new file mode 100644 index 00000000000..f31b9d7e873 --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_resolver_types.py @@ -0,0 +1,115 @@ +""" +Tests for pipeline field on policy CRUD types (resolver_types.py). +""" + +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.resolver_types import ( + PolicyAttachmentCreateRequest, + PolicyCreateRequest, + PolicyDBResponse, + PolicyUpdateRequest, +) + + +def test_policy_create_request_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, + {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, + ], + } + req = PolicyCreateRequest( + policy_name="test-policy", + guardrails_add=["g1", "g2"], + pipeline=pipeline_data, + ) + assert req.pipeline is not None + assert req.pipeline["mode"] == "pre_call" + assert len(req.pipeline["steps"]) == 2 + + +def test_policy_create_request_without_pipeline(): + req = PolicyCreateRequest( + policy_name="test-policy", + guardrails_add=["g1"], + ) + assert req.pipeline is None + + +def test_policy_update_request_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "block", "on_pass": "allow"}, + ], + } + req = PolicyUpdateRequest(pipeline=pipeline_data) + assert req.pipeline is not None + assert req.pipeline["steps"][0]["guardrail"] == "g1" + + +def test_policy_db_response_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, + {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, + ], + } + resp = PolicyDBResponse( + policy_id="test-id", + policy_name="test-policy", + guardrails_add=["g1", "g2"], + pipeline=pipeline_data, + ) + assert resp.pipeline is not None + assert resp.pipeline["mode"] == "pre_call" + dumped = resp.model_dump() + assert dumped["pipeline"]["steps"][0]["guardrail"] == "g1" + + +def test_policy_db_response_without_pipeline(): + resp = PolicyDBResponse( + policy_id="test-id", + policy_name="test-policy", + ) + assert resp.pipeline is None + dumped = resp.model_dump() + assert dumped["pipeline"] is None + + +def test_policy_create_request_roundtrip(): + pipeline_data = { + "mode": "post_call", + "steps": [ + { + "guardrail": "g1", + "on_fail": "modify_response", + "on_pass": "next", + "pass_data": True, + "modify_response_message": "custom msg", + }, + ], + } + req = PolicyCreateRequest( + policy_name="roundtrip-test", + guardrails_add=["g1"], + pipeline=pipeline_data, + ) + dumped = req.model_dump() + restored = PolicyCreateRequest(**dumped) + assert restored.pipeline == pipeline_data + + +@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) +def test_policy_attachment_create_request_accepts_int32_priority(priority: int): + assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority + + +@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) +def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int): + with pytest.raises(ValidationError): + PolicyAttachmentCreateRequest(policy_name="p", priority=priority) diff --git a/tests/unit/videos/__init__.py b/tests/unit/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py new file mode 100644 index 00000000000..22e1e5c05eb --- /dev/null +++ b/tests/unit/videos/test_main.py @@ -0,0 +1,455 @@ +""" +Dispatch-contract tests for litellm/videos/main.py + +Each public video operation is a pair: a sync `video_*` worker (decorated with +@client) that resolves the provider, fetches the provider config, logs, and then +forwards to exactly one `base_llm_http_handler.video_*_handler`; and an async +`avideo_*` wrapper that delegates to the sync worker in an executor. + +This file locks the contract of that layer so a regression fails loudly: + + 1. DISPATCH - the one correct handler fired and every sibling video handler + asserted NOT called. A copy-paste that calls the wrong handler + (e.g. remix -> edit) flips this. + 2. RESULT - the handler's return value is propagated by identity. + 3. PROVIDER - custom_llm_provider is decoded from an encoded video id when not + passed (status/content/remix/edit/extension), or defaults to + "openai" (list/create_character/get_character). This is the exact + surface of the historical "content defaulted to openai" bug. + 4. PAYLOAD - the provider config object and the operation's identifying args + (video_id/prompt/name/...) reach the handler; _is_async is False + on the sync path. + 5. SHORT-CIRCUIT - mock_response returns a typed object without any handler call. + 6. UNSUPPORTED - a None provider config raises before any handler fires. + 7. DELEGATION - avideo_* returns the sync worker's result untouched, sets + async_call=True, and pre-resolves the provider where it must. + +Seams mocked: the http handler (network), the provider-config registry lookup, +get_llm_provider, and the video-generation optional-param builders. The id decode +helper runs for real against genuinely-encoded ids, so the provider assertions +reflect production. +""" + +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Any, Dict +from unittest.mock import MagicMock, patch + +import pytest + + +import litellm +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.videos.main import CharacterObject, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider +from litellm.videos import main as videos_main + +# A real model-encoded video id: decodes (for real) to provider "azure". Used to +# prove the sync workers derive custom_llm_provider from the id, not a hardcode. +AZURE_VIDEO_ID = encode_video_id_with_provider("video_raw", "azure", "deployment-1") + +# The nine sync handlers on base_llm_http_handler. Dispatch tests assert exactly +# one fired and the other eight did not. +SYNC_HANDLERS = ( + "video_generation_handler", + "video_content_handler", + "video_remix_handler", + "video_create_character_handler", + "video_get_character_handler", + "video_edit_handler", + "video_extension_handler", + "video_list_handler", + "video_status_handler", +) + +GEN_OPTIONAL_PARAMS = {"seconds": "8", "size": "720x1280"} + + +@dataclass +class Seams: + handler: MagicMock + get_config: MagicMock + config: MagicMock + + def kwargs_of(self, handler_name: str) -> Dict[str, Any]: + method = getattr(self.handler, handler_name) + assert method.call_count == 1 + return dict(method.call_args.kwargs) + + def assert_only(self, handler_name: str) -> None: + for name in SYNC_HANDLERS: + method = getattr(self.handler, name) + if name == handler_name: + method.assert_called_once() + else: + method.assert_not_called() + + +@pytest.fixture +def seams(): + handler = MagicMock(spec=BaseLLMHTTPHandler) + config = MagicMock(name="provider_video_config") + get_config = MagicMock(return_value=config) + + with ExitStack() as stack: + stack.enter_context(patch.object(videos_main, "base_llm_http_handler", handler)) + stack.enter_context( + patch.object( + videos_main.ProviderConfigManager, + "get_provider_video_config", + get_config, + ) + ) + # video_generation resolves model+provider through get_llm_provider and + # builds optional params; mock those so the dispatch payload is deterministic. + stack.enter_context( + patch.object( + videos_main, + "get_llm_provider", + MagicMock(return_value=("sora-2", "openai", None, None)), + ) + ) + stack.enter_context( + patch.object( + videos_main.VideoGenerationRequestUtils, + "get_requested_video_generation_optional_param", + MagicMock(return_value={"seconds": "8"}), + ) + ) + stack.enter_context( + patch.object( + videos_main.VideoGenerationRequestUtils, + "get_optional_params_video_generation", + MagicMock(return_value=dict(GEN_OPTIONAL_PARAMS)), + ) + ) + yield Seams(handler=handler, get_config=get_config, config=config) + + +# =========================================================================== # +# Dispatch contract - one rich test per sync worker. +# =========================================================================== # + + +def test_video_generation__dispatch(seams): + result = videos_main.video_generation(prompt="a sunset", model="sora-2") + + seams.assert_only("video_generation_handler") + assert result is seams.handler.video_generation_handler.return_value + kw = seams.kwargs_of("video_generation_handler") + assert kw["model"] == "sora-2" + assert kw["prompt"] == "a sunset" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_generation_provider_config"] is seams.config + assert kw["video_generation_optional_request_params"] == GEN_OPTIONAL_PARAMS + assert kw["_is_async"] is False + + +def test_video_status__dispatch_and_provider_from_id(seams): + result = videos_main.video_status(video_id=AZURE_VIDEO_ID) + + seams.assert_only("video_status_handler") + assert result is seams.handler.video_status_handler.return_value + kw = seams.kwargs_of("video_status_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["custom_llm_provider"] == "azure" # decoded from the id, not openai + assert kw["video_status_provider_config"] is seams.config + assert kw["_is_async"] is False + # provider config requested for the decoded provider, not a hardcode. + assert seams.get_config.call_args.kwargs["provider"] == litellm.LlmProviders.AZURE + + +def test_video_content__dispatch_and_provider_from_id(seams): + result = videos_main.video_content(video_id=AZURE_VIDEO_ID, variant="thumbnail") + + seams.assert_only("video_content_handler") + assert result is seams.handler.video_content_handler.return_value + kw = seams.kwargs_of("video_content_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["custom_llm_provider"] == "azure" + assert kw["variant"] == "thumbnail" + assert kw["video_content_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_content__plain_id_defaults_to_openai(seams): + videos_main.video_content(video_id="video_plain") + + assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai" + + +def test_video_remix__dispatch_and_provider_from_id(seams): + result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors") + + seams.assert_only("video_remix_handler") + assert result is seams.handler.video_remix_handler.return_value + kw = seams.kwargs_of("video_remix_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "new colors" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_remix_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_edit__dispatch_and_provider_from_id(seams): + result = videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="brighter") + + seams.assert_only("video_edit_handler") + assert result is seams.handler.video_edit_handler.return_value + kw = seams.kwargs_of("video_edit_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "brighter" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_extension__dispatch_and_provider_from_id(seams): + result = videos_main.video_extension( + video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" + ) + + seams.assert_only("video_extension_handler") + assert result is seams.handler.video_extension_handler.return_value + kw = seams.kwargs_of("video_extension_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "continue" + assert kw["seconds"] == "5" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_list__dispatch_defaults_to_openai(seams): + result = videos_main.video_list(after="cur", limit=5, order="desc") + + seams.assert_only("video_list_handler") + assert result is seams.handler.video_list_handler.return_value + kw = seams.kwargs_of("video_list_handler") + assert kw["after"] == "cur" + assert kw["limit"] == 5 + assert kw["order"] == "desc" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_list_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_create_character__dispatch_defaults_to_openai(seams): + video = MagicMock(name="video_upload") + result = videos_main.video_create_character(name="hero", video=video) + + seams.assert_only("video_create_character_handler") + assert result is seams.handler.video_create_character_handler.return_value + kw = seams.kwargs_of("video_create_character_handler") + assert kw["name"] == "hero" + assert kw["video"] is video + assert kw["custom_llm_provider"] == "openai" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_get_character__dispatch_defaults_to_openai(seams): + result = videos_main.video_get_character(character_id="char_1") + + seams.assert_only("video_get_character_handler") + assert result is seams.handler.video_get_character_handler.return_value + kw = seams.kwargs_of("video_get_character_handler") + assert kw["character_id"] == "char_1" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_explicit_provider_beats_decoded_id(seams): + """An explicit custom_llm_provider wins over the one encoded in the id.""" + videos_main.video_status(video_id=AZURE_VIDEO_ID, custom_llm_provider="vertex_ai") + + assert seams.kwargs_of("video_status_handler")["custom_llm_provider"] == "vertex_ai" + + +# =========================================================================== # +# mock_response short-circuit - returns a typed object, no handler call. +# =========================================================================== # + + +def test_generation__mock_response_short_circuits(seams): + resp = videos_main.video_generation( + prompt="x", + model="sora-2", + mock_response={"id": "v1", "object": "video", "status": "queued"}, + ) + + assert isinstance(resp, VideoObject) + assert resp.id == "v1" + seams.handler.video_generation_handler.assert_not_called() + + +def test_list__mock_response_short_circuits(seams): + resp = videos_main.video_list( + mock_response=[{"id": "v1", "object": "video", "status": "completed"}] + ) + + assert isinstance(resp, list) + assert resp[0].id == "v1" + seams.handler.video_list_handler.assert_not_called() + + +def test_get_character__mock_response_short_circuits(seams): + resp = videos_main.video_get_character( + character_id="char_1", + mock_response={ + "id": "char_1", + "object": "character", + "created_at": 1, + "name": "hero", + }, + ) + + assert isinstance(resp, CharacterObject) + assert resp.id == "char_1" + seams.handler.video_get_character_handler.assert_not_called() + + +# =========================================================================== # +# Unsupported provider - a None provider config raises before any dispatch. +# =========================================================================== # + + +def test_unsupported_provider_raises_without_dispatch(seams): + seams.get_config.return_value = None + + with pytest.raises(litellm.APIConnectionError): + videos_main.video_status(video_id=AZURE_VIDEO_ID) + + seams.handler.video_status_handler.assert_not_called() + + +# =========================================================================== # +# Async-wrapper delegation - representative coverage. +# =========================================================================== # + + +@pytest.mark.asyncio +async def test_avideo_generation__delegates_with_async_flag(): + sentinel = VideoObject(id="v-async", object="video", status="queued") + with ( + patch.object( + videos_main, "video_generation", MagicMock(return_value=sentinel) + ) as sync, + patch.object( + litellm, + "get_llm_provider", + MagicMock(return_value=("sora-2", "openai", None, None)), + ), + ): + result = await videos_main.avideo_generation(prompt="x", model="sora-2") + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_avideo_status__delegates_untouched(): + sentinel = VideoObject(id="v-async", object="video", status="queued") + with patch.object( + videos_main, "video_status", MagicMock(return_value=sentinel) + ) as sync: + result = await videos_main.avideo_status(video_id="video_plain") + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["video_id"] == "video_plain" + + +@pytest.mark.asyncio +async def test_avideo_content__pre_decodes_provider_before_delegating(): + """avideo_content resolves the provider from the encoded id itself before + handing off, so the sync worker receives the decoded provider, not None.""" + sentinel = b"mp4-bytes" + with patch.object( + videos_main, "video_content", MagicMock(return_value=sentinel) + ) as sync: + result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["custom_llm_provider"] == "azure" + + +# =========================================================================== # +# Credential passthrough - DB/YAML model-config credentials the router injects +# via kwargs must reach the provider call for EVERY video handler, carried in +# litellm_params. Distinct per-field values catch a cross-wired field. +# =========================================================================== # + +DB_YAML_CREDS = { + "api_key": "sk-db-credential", + "api_base": "https://db-resource.test", + "api_version": "2024-12-31", + "vertex_project": "db-project-xyz", +} + +CREDENTIAL_OPERATIONS = [ + ( + "video_generation_handler", + lambda: videos_main.video_generation( + prompt="p", model="sora-2", **DB_YAML_CREDS + ), + ), + ( + "video_status_handler", + lambda: videos_main.video_status(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), + ), + ( + "video_content_handler", + lambda: videos_main.video_content(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), + ), + ( + "video_remix_handler", + lambda: videos_main.video_remix( + video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS + ), + ), + ( + "video_edit_handler", + lambda: videos_main.video_edit( + video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS + ), + ), + ( + "video_extension_handler", + lambda: videos_main.video_extension( + video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS + ), + ), + ( + "video_list_handler", + lambda: videos_main.video_list(**DB_YAML_CREDS), + ), + ( + "video_create_character_handler", + lambda: videos_main.video_create_character( + name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS + ), + ), + ( + "video_get_character_handler", + lambda: videos_main.video_get_character(character_id="char_1", **DB_YAML_CREDS), + ), +] + + +@pytest.mark.parametrize( + "handler_name,invoke", + CREDENTIAL_OPERATIONS, + ids=[op[0] for op in CREDENTIAL_OPERATIONS], +) +def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): + invoke() + + litellm_params = seams.kwargs_of(handler_name)["litellm_params"] + assert litellm_params.get("api_key") == DB_YAML_CREDS["api_key"] + assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] + assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] + assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] diff --git a/tests/unit/videos/test_utils.py b/tests/unit/videos/test_utils.py new file mode 100644 index 00000000000..728644cdda5 --- /dev/null +++ b/tests/unit/videos/test_utils.py @@ -0,0 +1,181 @@ +""" +Pure-logic contract tests for litellm/videos/main.py's request utils +(litellm/videos/utils.py: VideoGenerationRequestUtils). + +These lock the exact param-shaping behavior so a mutation that drops a filter, +flips a precedence, or stops removing a key fails loudly. The only seam is the +provider config's map_openai_params (a provider boundary); filter_out_litellm_params +runs for real, so the "litellm-internal params get stripped" assertions reflect +production. Every test asserts the exact resulting dict, never "ran without error". +""" + +from unittest.mock import MagicMock + + + +import litellm +from litellm.videos.utils import VideoGenerationRequestUtils + +get_requested = ( + VideoGenerationRequestUtils.get_requested_video_generation_optional_param +) +get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation + + +# =========================================================================== # +# get_requested_video_generation_optional_param +# +# Receives the caller's full local_vars; must return only the API-bound optional +# params. filter_out_litellm_params strips known internal keys for real; the +# values used below were chosen against the live set: seconds/size/user/foo_param/ +# vertex_project/extra/a/b survive, api_key/metadata/litellm_* are stripped. +# =========================================================================== # + + +def test_requested__drops_none_and_excluded_keys(): + result = get_requested( + { + "seconds": "8", + "size": None, # None -> dropped + "prompt": "a sunset", # excluded + "model": "sora-2", # excluded + "user": "u1", + } + ) + assert result == {"seconds": "8", "user": "u1"} + + +def test_requested__strips_litellm_internal_params(): + result = get_requested( + { + "seconds": "8", + "api_key": "sk-secret", + "metadata": {"x": 1}, + "litellm_call_id": "id-123", + } + ) + assert result == {"seconds": "8"} + + +def test_requested__timeout_always_removed(): + # timeout is NOT a litellm-internal param, so only the explicit pop removes it. + result = get_requested({"seconds": "8", "timeout": 30}) + assert result == {"seconds": "8"} + + +def test_requested__nested_kwargs_merge_and_override_base(): + result = get_requested( + {"seconds": "8", "kwargs": {"size": "720x1280", "seconds": "override"}} + ) + # nested kwargs win over the top-level base params on collision. + assert result == {"seconds": "override", "size": "720x1280"} + + +def test_requested__non_dict_kwargs_treated_as_empty(): + result = get_requested({"seconds": "8", "kwargs": "not-a-dict"}) + assert result == {"seconds": "8"} + + +def test_requested__none_input_returns_empty(): + assert get_requested(None) == {} + + +def test_requested__top_level_extra_body_spread_and_preserved(): + result = get_requested( + {"seconds": "8", "extra_body": {"vertex_project": "proj", "foo_param": "bar"}} + ) + # extra_body keys are both spread at top level AND kept under "extra_body". + assert result == { + "seconds": "8", + "vertex_project": "proj", + "foo_param": "bar", + "extra_body": {"vertex_project": "proj", "foo_param": "bar"}, + } + + +def test_requested__extra_body_kwargs_overrides_top_level(): + result = get_requested( + { + "extra_body": {"a": "top", "b": "top_b"}, + "kwargs": {"extra_body": {"a": "kw"}}, + } + ) + # kwargs' extra_body wins over the top-level extra_body on collision; the + # non-colliding top-level key survives. + assert result == { + "a": "kw", + "b": "top_b", + "extra_body": {"a": "kw", "b": "top_b"}, + } + + +def test_requested__extra_body_strips_litellm_internal_params(): + result = get_requested({"extra_body": {"api_key": "sk", "foo_param": "bar"}}) + # api_key filtered out of extra_body; only foo_param remains (and is spread). + assert result == {"foo_param": "bar", "extra_body": {"foo_param": "bar"}} + + +def test_requested__empty_extra_body_not_added(): + result = get_requested({"seconds": "8", "extra_body": {}}) + assert result == {"seconds": "8"} + assert "extra_body" not in result + + +# =========================================================================== # +# get_optional_params_video_generation +# +# Delegates mapping to the provider config (the seam) then folds extra_body in. +# =========================================================================== # + + +def _config(map_return): + config = MagicMock() + config.map_openai_params.return_value = map_return + return config + + +def test_optional__delegates_to_map_openai_params_with_drop_params(): + config = _config({"seconds": "8"}) + optional_params = {"seconds": "8"} + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params=optional_params, + ) + + assert result == {"seconds": "8"} + config.map_openai_params.assert_called_once_with( + video_create_optional_params=optional_params, + model="sora-2", + drop_params=litellm.drop_params, + ) + + +def test_optional__extra_body_overrides_mapped_and_is_removed(): + # mapped output carries a leftover extra_body that must be popped; the input + # extra_body overrides a colliding mapped key and is spread in. + config = _config({"seconds": "8", "size": "mapped", "extra_body": {"leftover": 1}}) + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params={ + "extra_body": {"size": "override", "extra": "x"} + }, + ) + + assert result == {"seconds": "8", "size": "override", "extra": "x"} + assert "extra_body" not in result + + +def test_optional__non_dict_extra_body_ignored(): + config = _config({"seconds": "8"}) + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params={"seconds": "8", "extra_body": None}, + ) + + assert result == {"seconds": "8"} From 924ad6e57118b7e42f2f483c9a028c41f90f25f3 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:22 +0000 Subject: [PATCH 062/146] test: remove phase 16 legacy test files from tests/test_litellm Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../rust_bridge/ocr/test_route_host.py | 85 -- .../rust_bridge/responses/test_route_host.py | 57 - .../test_litellm/sandbox/test_e2b_sandbox.py | 318 ----- .../sandbox/test_opensandbox_sandbox.py | 647 ---------- .../sandbox/test_sandbox_tools.py | 181 --- tests/test_litellm/skills/test_skills_main.py | 57 - .../test_enforce_model_rate_limits.py | 468 -------- .../test_router/test_io_token_rate_limits.py | 1069 ----------------- .../types/llms/test_types_llms_bedrock.py | 46 - .../types/llms/test_types_llms_openai.py | 591 --------- .../policy_engine/test_pipeline_types.py | 168 --- .../proxy/policy_engine/test_policy_types.py | 15 - .../policy_engine/test_resolver_types.py | 115 -- tests/test_litellm/videos/test_main.py | 455 ------- tests/test_litellm/videos/test_utils.py | 193 --- 15 files changed, 4465 deletions(-) delete mode 100644 tests/test_litellm/rust_bridge/ocr/test_route_host.py delete mode 100644 tests/test_litellm/rust_bridge/responses/test_route_host.py delete mode 100644 tests/test_litellm/sandbox/test_e2b_sandbox.py delete mode 100644 tests/test_litellm/sandbox/test_opensandbox_sandbox.py delete mode 100644 tests/test_litellm/sandbox/test_sandbox_tools.py delete mode 100644 tests/test_litellm/skills/test_skills_main.py delete mode 100644 tests/test_litellm/test_router/test_enforce_model_rate_limits.py delete mode 100644 tests/test_litellm/test_router/test_io_token_rate_limits.py delete mode 100644 tests/test_litellm/types/llms/test_types_llms_bedrock.py delete mode 100644 tests/test_litellm/types/llms/test_types_llms_openai.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_policy_types.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py delete mode 100644 tests/test_litellm/videos/test_main.py delete mode 100644 tests/test_litellm/videos/test_utils.py diff --git a/tests/test_litellm/rust_bridge/ocr/test_route_host.py b/tests/test_litellm/rust_bridge/ocr/test_route_host.py deleted file mode 100644 index 699492e4424..00000000000 --- a/tests/test_litellm/rust_bridge/ocr/test_route_host.py +++ /dev/null @@ -1,85 +0,0 @@ -from typing import Final - -import pytest - -import litellm -from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure -from litellm.rust_bridge.ocr.route_host import response as build_ocr_response -from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest - -REQUEST: Final = LiteLLMOcrRequest( - model="mistral/mistral-ocr-latest", - document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, - api_key="test-key", - api_base=None, - timeout=None, - custom_llm_provider=None, - extra_headers=None, - kwargs={"req_format": "markdown"}, -) - - -class RustUpstreamError(Exception): - def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: - super().__init__(status, body) - self.headers: Final = list(headers) - - -class RustFormatError(Exception): - ocr_request_format_error: Final = True - - -def test_rust_ocr_response_retains_provider_native_response(): - provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} - response = build_ocr_response( - { - "pages": [], - "model": "prebuilt-layout", - "document_annotation": None, - "usage_info": {"pages_processed": 0}, - "object": "ocr", - "provider_native_response": provider_response, - } - ) - - assert response.get_provider_native_response() == provider_response - assert response.model_dump().get("provider_native_response") is None - - -def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: - error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert isinstance(public_error, litellm.RateLimitError) - assert public_error.status_code == 429 - assert public_error.response.headers["retry-after"] == "7" - assert public_error.response.text == '{"message": "slow down"}' - assert public_error.__context__ is error - assert public_error.llm_provider == "mistral" - - -def test_map_failure_maps_upstream_401_to_authentication_error() -> None: - error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert isinstance(public_error, litellm.AuthenticationError) - assert public_error.status_code == 401 - assert public_error.response.text == '{"message": "Unauthorized"}' - assert public_error.__context__ is error - - -def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: - error: Final = RuntimeError("bridge exploded") - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert not isinstance(public_error, UpstreamFailure) - assert isinstance(public_error, litellm.APIConnectionError) - assert "bridge exploded" in str(public_error) - - -def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: - with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): - raise map_failure(RustFormatError(), REQUEST, "mistral") diff --git a/tests/test_litellm/rust_bridge/responses/test_route_host.py b/tests/test_litellm/rust_bridge/responses/test_route_host.py deleted file mode 100644 index 49bf19e7d8a..00000000000 --- a/tests/test_litellm/rust_bridge/responses/test_route_host.py +++ /dev/null @@ -1,57 +0,0 @@ -from types import MappingProxyType -from typing import Final - -import pytest -from pydantic import ValidationError - -from litellm.rust_bridge.responses.route_host import arguments, response -from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest -from litellm.types.llms.openai import ResponsesAPIResponse - - -def test_response_validates_into_the_public_responses_model() -> None: - built: Final = response( - MappingProxyType( - { - "id": "resp_native", - "object": "response", - "created_at": 1, - "model": "gpt-4o", - "status": "completed", - "output": [ - { - "type": "message", - "id": "msg_native", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "native", "annotations": []}], - } - ], - } - ) - ) - - assert isinstance(built, ResponsesAPIResponse) - assert built.id == "resp_native" - assert built.output[0].content[0].text == "native" - - -def test_response_rejects_a_payload_missing_required_fields() -> None: - with pytest.raises(ValidationError): - response(MappingProxyType({"object": "response"})) - - -def test_arguments_are_the_public_kwargs_view() -> None: - kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) - request: Final = LiteLLMResponsesRequest( - model="gpt-4o", - input="hi", - stream=None, - api_key=None, - api_base=None, - custom_llm_provider="openai", - extra_headers=None, - kwargs=kwargs, - ) - - assert arguments(request) is kwargs diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py deleted file mode 100644 index e01b9120416..00000000000 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ /dev/null @@ -1,318 +0,0 @@ -""" -Tests for the e2b code execution sandbox primitive. - -Unit tests inject a fake async HTTP client (dependency injection, no -monkeypatching) and assert request shapes and result mapping. Real-network -integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. -""" - -import json - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.sandbox.transformation import ContainerHandle -from litellm.llms.e2b.sandbox.transformation import ( - MAX_OUTPUT_BYTES, - E2BSandboxConfig, -) - - -class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): - self._json = json_data - self._lines = lines or [] - self.status_code = status_code - - def json(self): - return self._json - - async def aiter_lines(self): - for line in self._lines: - yield line - - -class FakeHTTPClient: - """Records outbound requests and returns canned responses keyed by URL.""" - - def __init__( - self, - *, - create_json=None, - execute_lines=None, - delete_status=204, - execute_raises=None, - ): - self.create_json = create_json or { - "sandboxID": "sbx_123", - "domain": "e2b.app", - "envdAccessToken": "tok_abc", - } - self.execute_lines = execute_lines or [] - self.delete_status = delete_status - self.execute_raises = execute_raises - self.calls = [] - - async def post(self, url, headers=None, json=None, stream=False, **kwargs): - self.calls.append(("POST", url, headers, json)) - if url.endswith("/sandboxes"): - return FakeResponse(json_data=self.create_json) - if url.endswith("/execute"): - if self.execute_raises is not None: - raise self.execute_raises - return FakeResponse(lines=self.execute_lines) - raise AssertionError(f"unexpected POST {url}") - - async def delete(self, url, headers=None, **kwargs): - self.calls.append(("DELETE", url, headers, None)) - if not (200 <= self.delete_status < 300): - raise httpx.HTTPStatusError( - f"status {self.delete_status}", - request=httpx.Request("DELETE", url), - response=httpx.Response(self.delete_status), - ) - return FakeResponse(status_code=self.delete_status) - - -# ---------- pure parser ---------- - - -def test_parse_lines_stdout_and_count(): - lines = [ - json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}), - json.dumps({"type": "number_of_executions", "execution_count": 1}), - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.stdout == "6\n" - assert result.execution_count == 1 - assert result.error is None - - -def test_parse_lines_error_surfaces_name_and_traceback(): - lines = [ - json.dumps( - { - "type": "error", - "name": "ZeroDivisionError", - "value": "division by zero", - "traceback": "Traceback (most recent call last): ...", - } - ) - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.error["name"] == "ZeroDivisionError" - assert "Traceback" in result.error["traceback"] - - -def test_parse_lines_result_carries_png(): - lines = [ - json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True}) - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.results and result.results[0]["png"] == "BASE64DATA" - assert "type" not in result.results[0] - - -# ---------- request shapes ---------- - - -@pytest.mark.asyncio -async def test_template_flows_into_create_request_as_templateID(): - client = FakeHTTPClient() - cfg = E2BSandboxConfig() - handle = await cfg.acreate_sandbox( - template="my-custom-template", api_key="e2b_key", client=client - ) - - method, url, headers, body = client.calls[0] - assert method == "POST" - assert url.endswith("/sandboxes") - assert body["templateID"] == "my-custom-template" # not "template" - assert body["secure"] is True - assert headers["X-API-Key"] == "e2b_key" - assert handle.id == "sbx_123" - assert handle._hidden_params["envd_access_token"] == "tok_abc" - - -@pytest.mark.asyncio -async def test_create_defaults_template_when_omitted(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client) - _, _, _, body = client.calls[0] - assert body["templateID"] == "code-interpreter-v1" - - -@pytest.mark.asyncio -async def test_run_code_targets_jupyter_host_with_access_token(): - client = FakeHTTPClient( - execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})] - ) - handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app") - handle._hidden_params = {"envd_access_token": "tok_run"} - - result = await E2BSandboxConfig().arun_code( - container=handle, code="print(6*7)", client=client - ) - - method, url, headers, body = client.calls[0] - assert url == "https://49999-sbx_xyz.e2b.app/execute" - assert headers["X-Access-Token"] == "tok_run" - assert body["code"] == "print(6*7)" - assert result.stdout.strip() == "42" - - -@pytest.mark.asyncio -async def test_delete_issues_delete_to_sandbox_id(): - client = FakeHTTPClient(delete_status=204) - handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - - ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - - method, url, headers, _ = client.calls[0] - assert method == "DELETE" - assert url.endswith("/sandboxes/sbx_del") - assert ok is True - - -@pytest.mark.asyncio -async def test_delete_returns_false_on_404(): - client = FakeHTTPClient(delete_status=404) - handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - assert ok is False - - -# ---------- ephemeral teardown ---------- - - -@pytest.mark.asyncio -async def test_code_interpreter_tool_deletes_even_when_run_raises(): - client = FakeHTTPClient(execute_raises=RuntimeError("boom")) - - with pytest.raises(RuntimeError, match="boom"): - await litellm.acode_interpreter_tool( - provider="e2b", code="1/0", api_key="e2b_key", client=client - ) - - methods = [c[0] for c in client.calls] - urls = [c[1] for c in client.calls] - assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete - assert urls[0].endswith("/sandboxes") - assert urls[1].endswith("/execute") - assert urls[2].endswith("/sandboxes/sbx_123") - - -# ---------- correctness guards ---------- - - -@pytest.mark.asyncio -async def test_delete_reraises_non_404_http_error(): - client = FakeHTTPClient(delete_status=500) - handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - with pytest.raises(httpx.HTTPStatusError): - await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - - -@pytest.mark.asyncio -async def test_create_preserves_explicit_zero_timeout(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox( - timeout=0, api_key="e2b_key", client=client - ) - _, _, _, body = client.calls[0] - assert body["timeout"] == 0 - - -@pytest.mark.asyncio -async def test_run_code_rejects_bare_id_without_access_token(): - client = FakeHTTPClient() - with pytest.raises(ValueError, match="access token"): - await E2BSandboxConfig().arun_code( - container="sbx_no_token", code="print(1)", client=client - ) - assert client.calls == [] # never reached the network - - -def test_parse_lines_skips_non_json_lines(): - lines = [ - "not-json-heartbeat", - json.dumps({"type": "stdout", "text": "ok\n"}), - "", - "{partial", - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.stdout == "ok\n" - assert result.error is None - - -@pytest.mark.asyncio -async def test_run_code_aborts_on_output_over_cap(): - big_line = "x" * (MAX_OUTPUT_BYTES + 1) - client = FakeHTTPClient(execute_lines=[big_line]) - handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app") - handle._hidden_params = {"envd_access_token": "tok"} - with pytest.raises(ValueError, match="exceeded"): - await E2BSandboxConfig().arun_code( - container=handle, code="print('x'*999)", client=client - ) - - -# ---------- public entrypoints ---------- - - -@pytest.mark.asyncio -async def test_public_lifecycle_create_run_delete(): - client = FakeHTTPClient( - execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})] - ) - container = await litellm.acreate_sandbox( - provider="e2b", api_key="e2b_key", client=client - ) - assert container.id == "sbx_123" - - result = await litellm.arun_code( - provider="e2b", - container=container, - api_key="e2b_key", - code="print(6*7)", - client=client, - ) - assert result.stdout.strip() == "42" - - assert ( - await litellm.adelete_sandbox( - provider="e2b", container=container, api_key="e2b_key", client=client - ) - is True - ) - - -@pytest.mark.asyncio -async def test_unsupported_provider_raises(): - with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): - await litellm.acreate_sandbox(provider="not-a-provider") - - -# ---------- api_base override ---------- - - -@pytest.mark.asyncio -async def test_create_uses_api_base_override(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox( - api_base="http://my-sandbox:8080", api_key="k", client=client - ) - _, url, _, _ = client.calls[0] - assert url == "http://my-sandbox:8080/sandboxes" - - -@pytest.mark.asyncio -async def test_create_defaults_to_e2b_api_base(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) - _, url, _, _ = client.calls[0] - assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py deleted file mode 100644 index 2928dea100e..00000000000 --- a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py +++ /dev/null @@ -1,647 +0,0 @@ -import json - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.sandbox.transformation import ContainerHandle -from litellm.llms.opensandbox.sandbox.transformation import ( - MAX_OUTPUT_BYTES, - OPEN_SANDBOX_DEFAULT_TEMPLATE, - OpenSandboxSandboxConfig, -) -from litellm.utils import ProviderConfigManager - -TEST_API_BASE = "https://sandbox.test/v1" - - -def http_status_error(status_code, url="http://test"): - return httpx.HTTPStatusError( - f"status {status_code}", - request=httpx.Request("GET", url), - response=httpx.Response(status_code), - ) - - -def sse(data): - return f"data: {json.dumps(data)}" - - -class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): - self._json = json_data - self._lines = lines or [] - self.status_code = status_code - - def json(self): - return self._json - - def raise_for_status(self): - if self.status_code >= 400: - raise http_status_error(self.status_code) - - async def aiter_lines(self): - for line in self._lines: - yield line - - -class FakeHTTPClient: - def __init__( - self, - *, - create_json=None, - sandbox_states=None, - endpoint_json=None, - endpoint_responses=None, - execute_lines=None, - delete_status=204, - execute_raises=None, - ): - self.create_json = create_json or { - "id": "osb_123", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], - } - self.sandbox_states = list( - sandbox_states - or [ - { - "id": "osb_123", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], - } - ] - ) - self.endpoint_json = endpoint_json or { - "endpoint": "execd.local:44772", - "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, - } - self.endpoint_responses = ( - list(endpoint_responses) if endpoint_responses is not None else None - ) - self.execute_lines = execute_lines or [] - self.delete_status = delete_status - self.execute_raises = execute_raises - self.calls = [] - - async def post(self, url, headers=None, json=None, stream=False, **kwargs): - self.calls.append(("POST", url, headers, json, {"stream": stream})) - if url.endswith("/sandboxes"): - return FakeResponse(json_data=self.create_json) - if url.endswith("/code"): - if self.execute_raises is not None: - raise self.execute_raises - return FakeResponse(lines=self.execute_lines) - raise AssertionError(f"unexpected POST {url}") - - async def get(self, url, headers=None, params=None, **kwargs): - self.calls.append(("GET", url, headers, None, params)) - if "/endpoints/44772" in url: - if self.endpoint_responses is not None and self.endpoint_responses: - response = self.endpoint_responses.pop(0) - if isinstance(response, Exception): - raise response - if isinstance(response, FakeResponse): - return response - return FakeResponse(json_data=response) - return FakeResponse(json_data=self.endpoint_json) - if "/sandboxes/" in url: - state = self.sandbox_states.pop(0) - return FakeResponse(json_data=state) - raise AssertionError(f"unexpected GET {url}") - - async def delete(self, url, headers=None, **kwargs): - self.calls.append(("DELETE", url, headers, None, None)) - if not (200 <= self.delete_status < 300): - raise http_status_error(self.delete_status, url) - return FakeResponse(status_code=self.delete_status) - - -def test_parse_sse_lines_maps_output_result_count_and_error(): - lines = [ - sse({"type": "stdout", "text": "hello\n"}), - sse({"type": "stderr", "text": "warn\n"}), - sse({"type": "result", "results": {"text/plain": "4"}}), - sse({"type": "execution_count", "execution_count": 7}), - sse( - { - "type": "error", - "error": { - "ename": "ValueError", - "evalue": "bad", - "traceback": ["Traceback"], - }, - } - ), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.stdout == "hello\n" - assert result.stderr == "warn\n" - assert result.results == [{"text/plain": "4"}] - assert result.execution_count == 7 - assert result.error == { - "name": "ValueError", - "value": "bad", - "traceback": ["Traceback"], - } - - -def test_parse_sse_lines_skips_non_json_and_control_lines(): - lines = [ - "event: message", - "not-json", - "", - sse({"type": "stdout", "text": "ok\n"}), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.stdout == "ok\n" - assert result.error is None - - -def test_parse_sse_lines_maps_fallback_shapes(): - lines = [ - "data:", - sse(["not-a-dict"]), - sse({"code": "BadRequest", "message": "nope"}), - sse({"type": "result", "text/plain": "4"}), - sse({"type": "error", "name": "RuntimeError", "text": "boom"}), - sse({"type": "execution_count", "execution_count": "8"}), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.results == [{"text/plain": "4"}] - assert result.execution_count == 8 - assert result.error == { - "name": "BadRequest", - "value": "nope", - "traceback": [], - } - fallback_error = OpenSandboxSandboxConfig._parse_lines( - [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] - ) - assert fallback_error.error == { - "name": "RuntimeError", - "value": "boom", - "traceback": [], - } - empty_string_error = OpenSandboxSandboxConfig._parse_lines( - [ - sse( - { - "type": "error", - "error": { - "ename": "", - "name": "FallbackName", - "evalue": "", - "value": "fallback value", - "traceback": [], - }, - } - ) - ] - ) - assert empty_string_error.error == { - "name": "", - "value": "", - "traceback": [], - } - - -def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): - def fake_secret(key): - if key == "OPEN_SANDBOX_API_KEY": - return "env-key" - if key == "OPEN_SANDBOX_API_BASE": - return TEST_API_BASE - return None - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", - fake_secret, - ) - config = OpenSandboxSandboxConfig() - handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") - - assert config.validate_environment() == "env-key" - assert config.validate_environment(api_key="") == "" - assert config._api_key(api_key=None, handle=handle) == "env-key" - - handle._hidden_params = {"api_key": "stored-key"} - assert config._api_key(api_key=None, handle=handle) == "stored-key" - assert config._http(None) is not None - - body = config._create_body( - template=None, - timeout=None, - allow_internet_access=False, - metadata=None, - env_vars=None, - resource_limits=None, - resource_requests=None, - entrypoint=None, - network_policy={"egress": [{"domain": "example.com"}]}, - secure_access=True, - ) - assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} - assert body["secureAccess"] is True - - other_body = config._create_body( - template=None, - timeout=None, - allow_internet_access=False, - metadata=None, - env_vars=None, - resource_limits=None, - resource_requests=None, - entrypoint=None, - network_policy=None, - secure_access=False, - ) - assert body["resourceLimits"] is not other_body["resourceLimits"] - - assert config._sandbox_state(None) is None - assert config._sandbox_state({"status": "Running"}) is None - assert config._as_str_dict(None) == {} - assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( - "http://execd.local" - ) - assert config._api_base(None) == TEST_API_BASE - assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" - assert config._as_int("9") == 9 - assert config._as_int("nope") is None - assert config._as_int(None) is None - assert isinstance( - ProviderConfigManager.get_provider_sandbox_config("opensandbox"), - OpenSandboxSandboxConfig, - ) - - -def test_api_base_requires_kwarg_or_env(monkeypatch): - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", - lambda key: None, - ) - - with pytest.raises(ValueError, match="api_base is required"): - OpenSandboxSandboxConfig._api_base(None) - - -@pytest.mark.asyncio -async def test_create_posts_default_body_and_omits_empty_api_key(): - client = FakeHTTPClient() - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - method, url, headers, body, _ = client.calls[0] - assert method == "POST" - assert url == f"{TEST_API_BASE}/sandboxes" - assert "OPEN-SANDBOX-API-KEY" not in headers - assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} - assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] - assert body["timeout"] == 300 - assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} - assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} - assert handle.id == "osb_123" - assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" - - -@pytest.mark.asyncio -async def test_create_can_opt_into_internet_access(): - client = FakeHTTPClient() - - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - allow_internet_access=True, - client=client, - ) - - _, _, _, body, _ = client.calls[0] - assert "networkPolicy" not in body - - -@pytest.mark.asyncio -async def test_create_custom_options_poll_and_endpoint_resolution(): - client = FakeHTTPClient( - create_json={ - "id": "osb_pending", - "status": {"state": "Pending"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/bin/sh"], - }, - sandbox_states=[ - { - "id": "osb_pending", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/bin/sh"], - } - ], - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - template="custom/image:latest", - timeout=600, - allow_internet_access=False, - api_key="osb-key", - api_base="https://sandbox.example/v1", - metadata={"suite": "unit"}, - env_vars={"PYTHONUNBUFFERED": "1"}, - resource_limits={"cpu": "500m", "memory": "512Mi"}, - resource_requests={"cpu": "250m", "memory": "256Mi"}, - entrypoint=["/bin/sh", "-lc", "sleep 3600"], - use_server_proxy=True, - client=client, - ) - - _, create_url, create_headers, body, _ = client.calls[0] - _, poll_url, poll_headers, _, _ = client.calls[1] - _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] - - assert create_url == "https://sandbox.example/v1/sandboxes" - assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert body["image"] == {"uri": "custom/image:latest"} - assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] - assert body["metadata"] == {"suite": "unit"} - assert body["env"] == {"PYTHONUNBUFFERED": "1"} - assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} - assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} - assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} - assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" - assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") - assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert endpoint_params == {"use_server_proxy": True} - assert handle.id == "osb_pending" - - -@pytest.mark.asyncio -async def test_create_waits_across_pending_state(monkeypatch): - client = FakeHTTPClient( - create_json={ - "id": "osb_pending", - "status": {"state": "Pending"}, - "createdAt": "2026-01-01T00:00:00Z", - }, - sandbox_states=[ - {"id": "osb_pending", "status": {"state": "Pending"}}, - {"id": "osb_pending", "status": {"state": "Running"}}, - ], - ) - sleeps = [] - - async def fake_sleep(interval): - sleeps.append(interval) - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=1, - poll_interval=0.01, - client=client, - ) - - assert handle.id == "osb_pending" - assert sleeps == [0.01] - - -@pytest.mark.asyncio -async def test_create_raises_for_terminal_state(): - client = FakeHTTPClient( - create_json={"id": "osb_failed", "status": {"state": "Pending"}}, - sandbox_states=[ - {"id": "osb_failed", "status": {"state": "Failed"}}, - ], - ) - - with pytest.raises(ValueError, match="entered Failed"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - -@pytest.mark.asyncio -async def test_create_times_out_waiting_for_running(): - client = FakeHTTPClient( - create_json={"id": "osb_slow", "status": {"state": "Pending"}}, - sandbox_states=[ - {"id": "osb_slow", "status": {"state": "Pending"}}, - ], - ) - - with pytest.raises(TimeoutError, match="was not Running"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=0, - poll_interval=0, - client=client, - ) - - -@pytest.mark.asyncio -async def test_create_waits_for_endpoint_resolution(monkeypatch): - client = FakeHTTPClient( - endpoint_responses=[ - http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), - { - "endpoint": "execd.local:44772", - "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, - }, - ], - ) - sleeps = [] - - async def fake_sleep(interval): - sleeps.append(interval) - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=1, - poll_interval=0.01, - client=client, - ) - - endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] - assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" - assert len(endpoint_calls) == 2 - assert sleeps == [0.01] - - -@pytest.mark.asyncio -async def test_create_raises_when_endpoint_is_missing(): - client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) - - with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client - ) - - -@pytest.mark.asyncio -async def test_create_reraises_non_404_endpoint_error(): - client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) - - with pytest.raises(httpx.HTTPStatusError): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - -@pytest.mark.asyncio -async def test_run_code_resolves_bare_id_and_posts_sse_request(): - client = FakeHTTPClient( - execute_lines=[ - sse({"type": "stdout", "text": "42\n"}), - ] - ) - - result = await OpenSandboxSandboxConfig().arun_code( - container="osb_bare", - code="print(6*7)", - language="python", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - endpoint_call = client.calls[0] - run_call = client.calls[1] - assert endpoint_call[0] == "GET" - assert ( - endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" - ) - assert run_call[0] == "POST" - assert run_call[1] == "http://execd.local:44772/code" - assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" - assert run_call[3] == { - "code": "print(6*7)", - "context": {"language": "python"}, - } - assert run_call[4] == {"stream": True} - assert result.stdout == "42\n" - - -@pytest.mark.asyncio -async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): - client = FakeHTTPClient() - handle = ContainerHandle( - id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" - ) - handle._hidden_params = { - "execd_endpoint": "execd.example/route/44772", - "execd_headers": {}, - } - - await OpenSandboxSandboxConfig().arun_code( - container=handle, code="print(1)", client=client - ) - - assert client.calls[0][1] == "https://execd.example/route/44772/code" - - -@pytest.mark.asyncio -async def test_run_code_aborts_on_output_over_cap(): - client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) - handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") - handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} - - with pytest.raises(ValueError, match="exceeded"): - await OpenSandboxSandboxConfig().arun_code( - container=handle, code="print('x')", client=client - ) - - -@pytest.mark.asyncio -async def test_delete_returns_false_on_404(): - client = FakeHTTPClient(delete_status=404) - - ok = await OpenSandboxSandboxConfig().adelete_sandbox( - container="osb_gone", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - assert ok is False - - -@pytest.mark.asyncio -async def test_delete_reraises_non_404_http_error(): - client = FakeHTTPClient(delete_status=500) - - with pytest.raises(httpx.HTTPStatusError): - await OpenSandboxSandboxConfig().adelete_sandbox( - container="osb_err", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - -@pytest.mark.asyncio -async def test_public_lifecycle_create_run_delete(): - client = FakeHTTPClient( - execute_lines=[ - sse({"type": "stdout", "text": "42\n"}), - ] - ) - - container = await litellm.acreate_sandbox( - provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client - ) - result = await litellm.arun_code( - provider="opensandbox", - container=container, - code="print(6*7)", - api_key="", - client=client, - ) - ok = await litellm.adelete_sandbox( - provider="opensandbox", - container=container, - api_key="", - client=client, - ) - - assert container.id == "osb_123" - assert result.stdout == "42\n" - assert ok is True - - -@pytest.mark.asyncio -async def test_code_interpreter_tool_deletes_even_when_run_raises(): - client = FakeHTTPClient(execute_raises=RuntimeError("boom")) - - with pytest.raises(RuntimeError, match="boom"): - await litellm.acode_interpreter_tool( - provider="opensandbox", - code="1/0", - api_key="", - api_base=TEST_API_BASE, - client=client, - ) - - assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] - assert client.calls[0][1].endswith("/sandboxes") - assert client.calls[1][1].endswith("/endpoints/44772") - assert client.calls[2][1].endswith("/code") - assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/test_litellm/sandbox/test_sandbox_tools.py b/tests/test_litellm/sandbox/test_sandbox_tools.py deleted file mode 100644 index 06136534b13..00000000000 --- a/tests/test_litellm/sandbox/test_sandbox_tools.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Unit tests for the sandbox-tool registry.""" - -from litellm.sandbox import sandbox_tools - - -def _reset(): - sandbox_tools.clear_sandbox_tools() - - -def test_register_resolves_provider_key_and_base(): - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": { - "sandbox_provider": "e2b", - "api_key": "sk-literal", - "api_base": "https://sandbox.internal", - }, - } - ] - ) - - resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") - assert resolved == { - "sandbox_provider": "e2b", - "api_key": "sk-literal", - "api_base": "https://sandbox.internal", - } - _reset() - - -def test_register_clears_stale_entries_on_reload(): - """A tool removed from the config must not survive a re-registration.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "old", - "litellm_params": {"sandbox_provider": "e2b"}, - } - ] - ) - assert sandbox_tools.resolve_sandbox_tool("old") is not None - - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "new", - "litellm_params": {"sandbox_provider": "e2b"}, - } - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("new") is not None - assert ( - sandbox_tools.resolve_sandbox_tool("old") is None - ), "stale tool must be gone after the config is reloaded" - _reset() - - -def test_register_empty_list_clears_removed_tools(): - """Reloading a config with sandbox_tools removed (the proxy passes an empty - list) must drop previously registered credentials from the process.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, - } - ] - ) - assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None - - sandbox_tools.register_sandbox_tools([]) - - assert ( - sandbox_tools.resolve_sandbox_tool("e2b_default") is None - ), "removing sandbox_tools from config must clear stale credentials" - _reset() - - -def test_register_resolves_secret_from_env(monkeypatch): - _reset() - monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": { - "sandbox_provider": "e2b", - "api_key": "os.environ/MY_SANDBOX_KEY", - }, - } - ] - ) - - resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") - assert resolved is not None - assert resolved["api_key"] == "sk-from-env" - assert resolved["api_base"] is None - _reset() - - -def test_resolve_unknown_returns_none(): - _reset() - assert sandbox_tools.resolve_sandbox_tool("nope") is None - - -def test_register_skips_malformed_entries_without_crashing(): - """A single malformed entry (missing sandbox_tool_name, or not a dict) must - not crash registration during proxy startup/hot-reload; valid entries in the - same list must still register.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name - "not-a-dict", # wrong type - {"sandbox_tool_name": "", "litellm_params": {}}, # empty name - { - "sandbox_tool_name": "good", - "litellm_params": {"sandbox_provider": "e2b"}, - }, - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("good") is not None - assert sandbox_tools.resolve_sandbox_tool("") is None - assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} - _reset() - - -def test_register_skips_entry_missing_sandbox_provider(): - """An entry with a name but no sandbox_provider must be skipped at - registration so it cannot later resolve and call acreate_sandbox(provider=None), - which fails with a cryptic runtime error instead of a clear startup warning.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, - { - "sandbox_tool_name": "null_provider", - "litellm_params": {"sandbox_provider": None}, - }, - { - "sandbox_tool_name": "good", - "litellm_params": {"sandbox_provider": "e2b"}, - }, - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("no_provider") is None - assert sandbox_tools.resolve_sandbox_tool("null_provider") is None - assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} - _reset() - - -def test_register_swaps_registry_atomically(): - """register_sandbox_tools must replace the registry in one rebind so a - concurrent resolve never observes a half-populated or transiently empty - registry between clearing and repopulating.""" - _reset() - sandbox_tools.register_sandbox_tools( - [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] - ) - before = sandbox_tools._SANDBOX_TOOL_REGISTRY - - sandbox_tools.register_sandbox_tools( - [ - {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, - {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, - ] - ) - after = sandbox_tools._SANDBOX_TOOL_REGISTRY - - assert after is not before, "the registry must be replaced, not mutated in place" - assert set(after) == {"b", "c"} - assert "a" not in after - _reset() diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py deleted file mode 100644 index e1c66c8d9ea..00000000000 --- a/tests/test_litellm/skills/test_skills_main.py +++ /dev/null @@ -1,57 +0,0 @@ -from unittest.mock import MagicMock - -import litellm.skills.main as skills_main -from litellm.types.utils import LlmProviders - - -def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( - monkeypatch, -) -> None: - """The REST /v1/skills form endpoint passes description/instructions as top-level - kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch - branch of create_skill() dropped both, so every LiteLLM-hosted skill was created - with description=None and instructions=None regardless of what the caller sent.""" - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill( - display_title="Document Translator", - description="Converts files from one language into another", - instructions="Take an uploaded document and produce it in the target language", - custom_llm_provider=LlmProviders.LITELLM_PROXY.value, - ) - - assert handler.create_skill_handler.call_args.kwargs["description"] == ( - "Converts files from one language into another" - ) - assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( - "Take an uploaded document and produce it in the target language" - ) - - -def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: - """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under - extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill( - display_title="Warehouse SQL Analyst", - extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, - custom_llm_provider=LlmProviders.LITELLM_PROXY.value, - ) - - assert handler.create_skill_handler.call_args.kwargs["description"] == ( - "Runs SQL against the inventory database" - ) - assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" - - -def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) - - assert handler.create_skill_handler.call_args.kwargs["description"] is None - assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py deleted file mode 100644 index 7577064b7f9..00000000000 --- a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py +++ /dev/null @@ -1,468 +0,0 @@ -""" -Tests for enforce_model_rate_limits feature. - -This feature allows users to enforce TPM/RPM limits set on model deployments -regardless of the routing strategy being used. -""" - -import asyncio -from datetime import timedelta -from unittest.mock import AsyncMock, MagicMock - -import pytest - -import litellm -from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) - -TPM_DEPLOYMENT = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "replica-test-id"}, - "model_name": "test-model", -} - - -def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: - dual_cache = DualCache(redis_cache=redis_cache) - check = ModelRateLimitingCheck(dual_cache=dual_cache) - now = litellm.utils.get_utc_datetime() - for minute in (now, now + timedelta(minutes=1)): - tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) - dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) - return dual_cache - - -class TestModelRateLimitingCheck: - """Test the ModelRateLimitingCheck class directly.""" - - def test_get_deployment_limits_from_top_level(self): - """Test extracting limits from top-level deployment config.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "tpm": 1000, - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 1000 - assert rpm == 10 - - def test_get_deployment_limits_from_litellm_params(self): - """Test extracting limits from litellm_params.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 2000 - assert rpm == 20 - - def test_get_deployment_limits_from_model_info(self): - """Test extracting limits from model_info.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 3000 - assert rpm == 30 - - def test_get_deployment_limits_none_when_not_set(self): - """Test that None is returned when limits are not set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm is None - assert rpm is None - - def test_pre_call_check_allows_request_when_no_limits(self): - """Test that requests are allowed when no limits are set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - result = check.pre_call_check(deployment) - assert result == deployment - - def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 11 # Over limit after increment - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "RPM limit=10" in str(exc_info.value) - assert "current usage=11" in str(exc_info.value) - - def test_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 6 - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = check.pre_call_check(deployment) - assert result == deployment - - def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): - """Test that RateLimitError is raised when TPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = 1000 # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "TPM limit=1000" in str(exc_info.value) - assert "current usage=1000" in str(exc_info.value) - - def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 1000 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.parametrize( - "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] - ) - def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): - redis_cache = MagicMock() - redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() - redis_cache.increment_cache.return_value = 2 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) - - assert "RPM limit=1" in str(exc_info.value) - - def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) - deployment = {**TPM_DEPLOYMENT, "rpm": 1} - - assert check.pre_call_check(deployment) == deployment - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "RPM limit=1" in str(exc_info.value) - - def test_log_success_event_increments_cache(self): - """Test that log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } - - check.log_success_event(kwargs, None, None, None) - - # Verify increment_cache was called - mock_cache.increment_cache.assert_called_once() - _, kwarg_params = mock_cache.increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 - - -class TestModelRateLimitingCheckAsync: - """Test async methods of ModelRateLimitingCheck.""" - - @pytest.mark.asyncio - async def test_async_pre_call_check_allows_request_when_no_limits(self): - """Test that requests are allowed when no limits are set (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment - - @pytest.mark.asyncio - async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "RPM limit=10" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=6) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment - - @pytest.mark.asyncio - async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): - """Test that RateLimitError is raised when TPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "TPM limit=1000" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=1000) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] - ) - async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.async_get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( - self, - ): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment = AsyncMock(return_value=2) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) - - assert "RPM limit=1" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) - deployment = {**TPM_DEPLOYMENT, "rpm": 1} - - assert await check.async_pre_call_check(deployment) == deployment - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "RPM limit=1" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_log_success_event_increments_cache(self): - """Test that async_log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - mock_cache.async_increment_cache = AsyncMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } - - await check.async_log_success_event(kwargs, None, None, None) - - # Verify async_increment_cache was called - mock_cache.async_increment_cache.assert_called_once() - _, kwarg_params = mock_cache.async_increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 - - -class TestRouterWithEnforceModelRateLimits: - """Test Router integration with enforce_model_rate_limits.""" - - def test_router_initializes_with_enforce_model_rate_limits(self): - """Test that Router properly initializes the ModelRateLimitingCheck.""" - model_list = [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "test"}, - "rpm": 10, - } - ] - - router = Router( - model_list=model_list, - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - - # Check that the callback was added - assert router.optional_callbacks is not None - assert len(router.optional_callbacks) == 1 - assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) - - def test_router_optional_callbacks_contains_model_rate_limiting(self): - """Test that ModelRateLimitingCheck is in the callbacks list.""" - model_list = [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "test"}, - "rpm": 10, - } - ] - - Router( - model_list=model_list, - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - - # Find the ModelRateLimitingCheck in litellm.callbacks - found = False - for callback in litellm.callbacks: - if isinstance(callback, ModelRateLimitingCheck): - found = True - break - - assert found, "ModelRateLimitingCheck should be in litellm.callbacks" - - -class TestModelRateLimitConcurrency: - """Test that RPM rate limiting is atomic under concurrent requests.""" - - @pytest.mark.asyncio - async def test_concurrent_requests_respect_rpm_limit(self): - """ - Fire 4 concurrent async requests with RPM limit of 2. - Exactly 2 should succeed and 2 should raise RateLimitError. - - This test validates the atomic increment-first pattern: - the old check-then-increment pattern would let 3+ through - due to a race condition on the local cache read. - """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - - deployment = { - "rpm": 2, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "concurrent-test-id"}, - "model_name": "test-model", - } - - async def attempt_request(): - return await check.async_pre_call_check(deployment) - - results = await asyncio.gather( - *[attempt_request() for _ in range(4)], - return_exceptions=True, - ) - - successes = [r for r in results if not isinstance(r, Exception)] - failures = [r for r in results if isinstance(r, litellm.RateLimitError)] - - assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" - assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py deleted file mode 100644 index 3cef1c7bb63..00000000000 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ /dev/null @@ -1,1069 +0,0 @@ -""" -Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). -""" - -import asyncio - -import pytest - -import litellm -from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_CACHE_KEY, - ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, - OTPM_RESERVED_KEY, - _reservation_value, - _resolve_max_tokens, - async_io_token_pre_call_check, - async_io_token_reconcile_success, - build_io_token_rate_limit_headers, - deployment_has_io_token_limits, - get_io_token_rate_limit_request_kwargs, - io_token_reconcile_success, - io_token_refund_failure, - refund_stale_reservation_before_retry, - set_io_token_rate_limit_request_kwargs, -) -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) -from litellm.types.utils import ModelResponse, Usage - - -class TestIOTokenRateLimitHelpers: - def test_deployment_has_io_token_limits(self): - assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) - assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) - - def test_reservation_value_minimal_when_estimate_fails(self): - # A failed/empty estimate (0) must reserve a minimal slot, not the - # entire limit - otherwise one request whose estimate failed fills - # the whole bucket and blocks every concurrent request until it - # completes and reconciles. - assert _reservation_value(0, 100) == 1 - assert _reservation_value(0, 1) == 1 - # A real non-zero estimate is reserved as-is. - assert _reservation_value(42, 100) == 42 - - def test_resolve_max_tokens_respects_explicit_zero(self): - deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} - # An explicit max_tokens=0 is honored, not replaced by the model default. - assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 - # max_completion_tokens is the fallback only when max_tokens is absent. - assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 - assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 - - def test_build_io_token_rate_limit_headers(self): - headers = build_io_token_rate_limit_headers( - itpm_limit=200, - otpm_limit=40, - current_itpm=15, - current_otpm=4, - ) - assert headers["x-ratelimit-limit-input-tokens"] == 200 - assert headers["x-ratelimit-remaining-input-tokens"] == 185 - assert headers["x-ratelimit-limit-output-tokens"] == 40 - assert headers["x-ratelimit-remaining-output-tokens"] == 36 - - -class TestModelRateLimitingCheckIOTokens: - @pytest.mark.asyncio - async def test_itpm_reservation_and_reconcile(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 50, - }, - "model_info": {"id": "io-test-id"}, - "model_name": "opus", - } - - request_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 10, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - - kwargs = { - "standard_logging_object": { - "model_id": "io-test-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[ - { - "message": {"role": "assistant", "content": "hi"}, - "index": 0, - "finish_reason": "stop", - } - ], - usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - await check.async_log_success_event(kwargs, response, None, None) - - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - current_otpm = await dual_cache.async_get_cache(key=otpm_key) - # ITPM tracks input tokens only (billable prompt tokens), not output. - assert current_itpm == 5 - assert current_otpm == 3 - - @pytest.mark.asyncio - async def test_itpm_limit_raises_429(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 5, - }, - "model_info": {"id": "io-limit-id"}, - "model_name": "opus", - } - - # ITPM enforces input tokens only; the prompt alone must exceed the limit, - # a large max_tokens must not contribute to the ITPM reservation. - set_io_token_rate_limit_request_kwargs( - { - "messages": [ - { - "role": "user", - "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", - } - ], - "max_tokens": 10, - "metadata": {}, - } - ) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "ITPM limit=5" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - otpm_limit = 10 - max_tokens = 4 - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "otpm": otpm_limit, - }, - "model_info": {"id": "io-otpm-race-id"}, - "model_name": "opus", - } - - set_io_token_rate_limit_request_kwargs( - { - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": max_tokens, - "metadata": {}, - } - ) - - async def _attempt(): - try: - await check.async_pre_call_check(deployment) - return True - except litellm.RateLimitError: - return False - - results = await asyncio.gather(*[_attempt() for _ in range(8)]) - successes = sum(1 for r in results if r) - - minute = get_utc_datetime().strftime("%H-%M") - otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - current_otpm = await dual_cache.async_get_cache(key=otpm_key) - - # Atomic reservation must never let concurrent requests overshoot the limit. - assert current_otpm is not None - assert current_otpm <= otpm_limit - assert successes == otpm_limit // max_tokens - assert current_otpm == successes * max_tokens - - @pytest.mark.asyncio - async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): - """ - When input-token estimation yields 0 (no messages/prompt/input field, - unsupported model, tokenizer error), the reservation must be a - minimal 1 token, not the entire itpm limit. Otherwise the first - request whose estimate fails fills the whole bucket and every - concurrent request is rejected until it completes - effectively - serializing traffic to the deployment. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_limit = 5 - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": itpm_limit, - }, - "model_info": {"id": "io-itpm-estimate-fail-id"}, - "model_name": "opus", - } - - # No messages/prompt/input field -> _estimate_input_tokens returns 0. - set_io_token_rate_limit_request_kwargs( - { - "max_tokens": 5, - "metadata": {}, - } - ) - - async def _attempt(): - try: - await check.async_pre_call_check(deployment) - return True - except litellm.RateLimitError: - return False - - results = await asyncio.gather(*[_attempt() for _ in range(8)]) - successes = sum(1 for r in results if r) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - - # A minimal 1-token reservation per request lets itpm_limit concurrent - # requests through, instead of a single request starving the rest. - assert current_itpm is not None - assert current_itpm <= itpm_limit - assert successes == itpm_limit - - @pytest.mark.asyncio - async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) - - # Production kwargs commonly carry litellm_params.metadata; the stashed - # reservation lives in the top-level metadata and must still be found. - kwargs = { - "standard_logging_object": { - "model_id": "io-lp-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, - "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, - } - await check.async_log_failure_event(kwargs, None, None, None) - - current = await dual_cache.async_get_cache(key=itpm_key) - assert current == 0 - - @pytest.mark.asyncio - async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - }, - "model_info": {"id": "io-zero-est-id"}, - "model_name": "opus", - } - - request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - # A failed/zero estimate reserves a minimal 1 token, not the full - # itpm limit, so it doesn't starve concurrent requests. - assert await dual_cache.async_get_cache(key=itpm_key) == 1 - - kwargs = { - "standard_logging_object": { - "model_id": "io-zero-est-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), - ) - await check.async_log_success_event(kwargs, response, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 7 - - @pytest.mark.asyncio - async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): - """ - A zero/failed estimate reserves a minimal 1 token rather than the - full itpm limit, so up to itpm_limit such calls are allowed - concurrently instead of the first one claiming the entire bucket. - """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 2, - }, - "model_info": {"id": "io-zero-cap-id"}, - "model_name": "opus", - } - request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - # Second zero-estimate call still fits within the itpm=2 limit. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) - await check.async_pre_call_check(deployment) - - # A third exceeds the limit and is rejected. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) - with pytest.raises(litellm.RateLimitError): - await check.async_pre_call_check(deployment) - - @pytest.mark.asyncio - async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "otpm": 5, - }, - "model_info": {"id": "io-zero-output-id"}, - "model_name": "opus", - } - zero_output_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 0, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(zero_output_kwargs) - await check.async_pre_call_check(deployment) - - zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 - assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 - - normal_output_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(normal_output_kwargs) - await check.async_pre_call_check(deployment) - - normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 - - def test_sync_io_pre_call_reserves_and_reconciles(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 50, - }, - "model_info": {"id": "io-sync-id"}, - "model_name": "opus", - } - request_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 10, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - check.pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - kwargs = { - "standard_logging_object": { - "model_id": "io-sync-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - check.log_success_event(kwargs, response, None, None) - - assert dual_cache.get_cache(key=itpm_key) == 5 - assert dual_cache.get_cache(key=otpm_key) == 3 - - @pytest.mark.asyncio - async def test_reconcile_runs_via_success_event_without_model_id(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - # standard_logging_object has no model_id (only the TPM path needs it); - # IO reconciliation must still run off the stashed cache key. - kwargs = { - "standard_logging_object": { - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - await check.async_log_success_event(kwargs, response, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 3 - - @pytest.mark.asyncio - async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - - # Shared request metadata carrying the first (IO) deployment's reservation. - metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} - fail_kwargs = { - "metadata": metadata, - "standard_logging_object": { - "model_id": "io-first", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - } - await check.async_log_failure_event(fail_kwargs, None, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 0 - assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata - - # Retry succeeds on a non-IO fallback deployment reusing the same metadata. - retry_kwargs = { - "metadata": metadata, - "standard_logging_object": { - "model_id": "non-io-second", - "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, - "metadata": {}, - "total_tokens": 12, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), - ) - await check.async_log_success_event(retry_kwargs, response, None, None) - - # The first deployment's ITPM counter is not driven negative... - assert await dual_cache.async_get_cache(key=itpm_key) == 0 - # ...and the non-IO deployment's TPM usage is tracked normally. - tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" - assert await dual_cache.async_get_cache(key=tpm_key) == 12 - - @pytest.mark.asyncio - async def test_stale_reservation_refunded_before_retry_overwrites_it(self): - """ - A retry reuses the same mutable kwargs dict for the next deployment. - If deployment A's failure event hasn't run yet (e.g. it was scheduled - as a background task) when the retry calls - set_io_token_rate_limit_request_kwargs for deployment B, the router - must first synchronously refund + clear A's reservation via - refund_stale_reservation_before_retry - otherwise A's counter stays - elevated by the reservation until its TTL expires, and the - now-orphaned sentinels must not leak into B's accounting either. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - minute = get_utc_datetime().strftime("%H-%M") - itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) - - # Deployment A's still-unreconciled reservation, stashed on the shared - # kwargs dict the retry loop reuses. - shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} - - # Router calls this before overwriting kwargs for deployment B's attempt - - # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. - refund_stale_reservation_before_retry(dual_cache, shared_kwargs) - - # A's reservation is refunded immediately, not left stranded for a - # background failure task that may run arbitrarily later (or never, - # if the sentinels get cleared out from under it first). - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 - assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] - assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] - - # A's own (now-late) failure event finds nothing left to refund and - # is a safe no-op, since the sentinels were already cleared above. - io_token_refund_failure(dual_cache, shared_kwargs) - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 - - # The retry proceeds to stash deployment B's own reservation on the - # same dict; it starts clean, unaffected by A's cleared sentinels. - set_io_token_rate_limit_request_kwargs(shared_kwargs) - itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" - shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 - shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b - await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) - assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 - - @pytest.mark.asyncio - async def test_client_supplied_reservation_keys_are_stripped(self): - # metadata is caller-controlled; the server-only reservation sentinels - # must be removed before the router captures the request kwargs. - forged = { - "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, - "litellm_metadata": {OTPM_RESERVED_KEY: 7}, - "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, - } - set_io_token_rate_limit_request_kwargs(forged) - stored = get_io_token_rate_limit_request_kwargs() - - assert ITPM_RESERVED_KEY not in stored["metadata"] - assert ITPM_CACHE_KEY not in stored["metadata"] - assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] - assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] - - @pytest.mark.asyncio - async def test_forged_reservation_cannot_decrement_counter(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - victim_key = "global_router:victim:model:itpm:00-00" - await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) - - # A caller forges a reservation pointing at another deployment's counter. - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, - "standard_logging_object": { - "model_id": "m", - "hidden_params": {"litellm_model_name": "model"}, - "metadata": {}, - "total_tokens": 2, - }, - } - # The router sanitizes the request kwargs before the call runs. - set_io_token_rate_limit_request_kwargs(kwargs) - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), - ) - await check.async_log_success_event(kwargs, response, None, None) - - # The forged reservation was stripped, so the victim counter is untouched. - assert await dual_cache.async_get_cache(key=victim_key) == 100 - - @pytest.mark.asyncio - async def test_otpm_reservation_error_rolls_back_itpm(self): - from litellm.utils import get_utc_datetime - - class _OtpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":otpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _OtpmFailCache() - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 1000, - "otpm": 1000, - }, - "model_info": {"id": "io-rollback-id"}, - "model_name": "opus", - } - set_io_token_rate_limit_request_kwargs( - { - "messages": [{"role": "user", "content": "hello world"}], - "max_tokens": 5, - "metadata": {}, - } - ) - - with pytest.raises(RuntimeError): - await async_io_token_pre_call_check(dual_cache, deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - # A transient OTPM error must release the ITPM reservation, not leak it. - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 - - @pytest.mark.asyncio - async def test_reconcile_clears_stash_even_when_increment_errors(self): - class _ItpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":itpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _ItpmFailCache() - metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} - kwargs = {"metadata": metadata} - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - with pytest.raises(RuntimeError): - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - # The stash is cleared even though reconciliation raised, so a duplicate - # success event can't re-process it. - assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata - - @pytest.mark.asyncio - async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): - import logging - - check = ModelRateLimitingCheck(dual_cache=DualCache()) - deployment = { - "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, - "model_info": {}, - } - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - check._warn_io_token_and_tpm_rpm_coexist_once(deployment) - check._warn_io_token_and_tpm_rpm_coexist_once(deployment) - - warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] - # id-less deployments are not collapsed onto a single dedup key. - assert len(warnings) == 2 - - @pytest.mark.asyncio - async def test_missing_deployment_id_skips_io_reservation(self): - dual_cache = DualCache() - deployment = { - "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, - "model_info": {}, # no id -> cannot build a per-deployment cache key - "model_name": "opus", - } - request_kwargs = { - "messages": [{"role": "user", "content": "hello world"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - - result = await async_io_token_pre_call_check(dual_cache, deployment) - - assert result is deployment - # No reservation is stashed, so nothing lands in a shared None:None bucket. - assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] - - @pytest.mark.asyncio - async def test_reconcile_uses_reservation_minute_key(self): - dual_cache = DualCache() - # Reservation was made on a fixed minute key; a call that finishes in a - # later minute must reconcile against that same key, never a key built - # from the response-time minute. - itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), - ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 4 - - @pytest.mark.asyncio - async def test_reconcile_missing_usage_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 8, - OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - ) - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 - assert ITPM_RESERVED_KEY not in kwargs["metadata"] - - @pytest.mark.asyncio - async def test_reconcile_total_tokens_only_keeps_reservation(self): - """ - A response usage object with only total_tokens (no prompt/completion - breakdown) can't be split into input/output, so it must be treated the - same as missing usage: keep the reservation instead of resolving to - (0, 0) and refunding it in full. - """ - dual_cache = DualCache() - itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 8, - OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = {"type": "message", "usage": {"total_tokens": 13}} - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 - - def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": {"total_tokens": 4}, - } - response = {"type": "message", "role": "assistant", "content": []} - - io_token_reconcile_success(dual_cache, kwargs, response) - - assert dual_cache.get_cache(key=itpm_key) == 10 - - @pytest.mark.asyncio - async def test_reconcile_falls_back_to_standard_logging_object(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "prompt_tokens": 4, - "completion_tokens": 0, - "total_tokens": 4, - }, - } - response = {"type": "message", "role": "assistant", "content": []} - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 4 - - def test_sync_reconcile_anthropic_dict_usage(self): - dual_cache = DualCache() - itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=6, ttl=60) - dual_cache.set_cache(key=otpm_key, value=4, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 6, - OTPM_RESERVED_KEY: 4, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = { - "type": "message", - "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, - } - - io_token_reconcile_success(dual_cache, kwargs, response) - - assert dual_cache.get_cache(key=itpm_key) == 2 - assert dual_cache.get_cache(key=otpm_key) == 2 - - @pytest.mark.asyncio - async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): - import logging - - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-mixed-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - deployment = { - "litellm_params": { - "model": deployment_name, - "itpm": 100, - "rpm": 1, - }, - "model_info": {"id": model_id}, - "model_name": "opus", - } - - minute = get_utc_datetime().strftime("%H-%M") - rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) - - request_kwargs = { - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - with pytest.raises(litellm.RateLimitError): - await check.async_pre_call_check(deployment) - - assert await dual_cache.async_get_cache(key=rpm_key) == 6 - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 - assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] - assert any("both limit types are enforced" in record.message for record in caplog.records) - - @pytest.mark.asyncio - async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): - """ - A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on - success, otherwise the tpm_key the pre-call check reads is never written - and the tpm_limit can never be enforced. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-tpm-mixed-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "model_id": model_id, - "total_tokens": 7, - "hidden_params": {"litellm_model_name": deployment_name}, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - await check.async_log_success_event(kwargs, response, None, None) - - # ITPM reconciled down from the 5-token reservation to actual usage (3). - assert await dual_cache.async_get_cache(key=itpm_key) == 3 - # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. - assert await dual_cache.async_get_cache(key=tpm_key) == 7 - - def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-tpm-mixed-sync-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - dual_cache.set_cache(key=itpm_key, value=5, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "model_id": model_id, - "total_tokens": 7, - "hidden_params": {"litellm_model_name": deployment_name}, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - check.log_success_event(kwargs, response, None, None) - - assert dual_cache.get_cache(key=itpm_key) == 3 - assert dual_cache.get_cache(key=tpm_key) == 7 - - @pytest.mark.asyncio - async def test_failure_refunds_itpm_reservation(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) - - reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} - kwargs = { - "standard_logging_object": { - "model_id": "io-refund-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": dict(reservation), - }, - "metadata": dict(reservation), - } - await check.async_log_failure_event(kwargs, None, None, None) - - current = await dual_cache.async_get_cache(key=itpm_key) - assert current == 0 - - -class TestRouterIOTokenIntegration: - @pytest.mark.asyncio - async def test_model_group_info_aggregates_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "opus", - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 20, - }, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - info = router.get_model_group_info("opus") - assert info is not None - assert info.itpm == 100 - assert info.otpm == 20 - - -class TestContextSlotRetention: - def test_setter_stores_kwargs_only_for_io_limited_deployments(self): - """ - The context slot pins the entire request kwargs (messages included) - for the lifetime of the surrounding asyncio context, and pooled - resources created mid-request (e.g. redis connections) capture that - context, extending the pin far past the request. Only ITPM/OTPM - pre-call checks read the slot, so the setter must store None for - deployments without io token limits and still clear reservation - sentinels from kwargs either way. - """ - kwargs = { - "messages": [{"role": "user", "content": "x" * 1000}], - "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, - } - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) - assert get_io_token_rate_limit_request_kwargs() is None - assert ITPM_RESERVED_KEY not in kwargs["metadata"] - assert ITPM_CACHE_KEY not in kwargs["metadata"] - - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) - assert get_io_token_rate_limit_request_kwargs() is kwargs - - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_does_not_pin_kwargs_without_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "plain", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, - } - ] - ) - set_io_token_rate_limit_request_kwargs(None) - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("plain") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_pins_kwargs_for_io_limited_deployment(self): - router = Router( - model_list=[ - { - "model_name": "limited", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - set_io_token_rate_limit_request_kwargs(None) - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("limited") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is kwargs - - -@pytest.mark.asyncio -async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): - from litellm.utils import get_utc_datetime - from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( - assert_loop_stayed_free, - timed_with_loop_lags, - warm_tokenizer, - ) - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - warm_tokenizer("anthropic/claude-fable-5") - deployment = { - "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, - "model_info": {"id": "io-loop-id"}, - "model_name": "claude", - } - set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) - - _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) - - minute = get_utc_datetime().strftime("%H-%M") - reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") - assert reserved > 100_000 - assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/types/llms/test_types_llms_bedrock.py b/tests/test_litellm/types/llms/test_types_llms_bedrock.py deleted file mode 100644 index a5ad882e775..00000000000 --- a/tests/test_litellm/types/llms/test_types_llms_bedrock.py +++ /dev/null @@ -1,46 +0,0 @@ -import pytest -from pydantic import ValidationError - -from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams - - -def test_model_validate_keeps_auth_params_and_ignores_request_params(): - auth_params = AwsAuthParams.model_validate( - { - "aws_role_name": "arn:aws:iam::999999999999:role/litellm-role", - "aws_session_name": "litellm-session", - "aws_external_id": "litellm-external-id", - "aws_region_name": "us-west-2", - "aws_bedrock_runtime_endpoint": "https://bedrock.example.com", - "model": "anthropic.claude-haiku-4-5-20251001-v1:0", - "temperature": 0.1, - "messages": [{"role": "user", "content": "hi"}], - } - ) - - assert auth_params.aws_role_name == "arn:aws:iam::999999999999:role/litellm-role" - assert auth_params.aws_session_name == "litellm-session" - assert auth_params.aws_external_id == "litellm-external-id" - assert auth_params.aws_access_key_id is None - assert set(auth_params.model_dump()) == set(AWS_AUTH_PARAM_KEYS) - assert not set(AWS_AUTH_PARAM_KEYS) & {"aws_region_name", "aws_bedrock_runtime_endpoint", "model", "temperature"} - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("aws_role_name", 1234), - ("aws_session_name", ["litellm-session"]), - ("aws_external_id", {"id": "x"}), - ], -) -def test_model_validate_rejects_non_string_credentials(field, value): - with pytest.raises(ValidationError): - AwsAuthParams.model_validate({field: value}) - - -def test_frozen_struct_rejects_field_assignment(): - auth_params = AwsAuthParams(aws_role_name="arn:aws:iam::999999999999:role/litellm-role") - - with pytest.raises(ValidationError): - auth_params.aws_role_name = "arn:aws:iam::999999999999:role/other-role" diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py deleted file mode 100644 index 64ec09838e8..00000000000 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ /dev/null @@ -1,591 +0,0 @@ -import asyncio -from typing import Optional -from unittest.mock import AsyncMock, patch - -import pytest - -import json - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent - - -@pytest.mark.parametrize("stream", (False, True)) -def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: - from typing import Final - - from litellm.types.llms.openai import ( - ChatCompletionReasoningItem, - ChatCompletionReasoningSummaryTextBlock, - ) - from litellm.types.utils import ( - Choices, - Delta, - Message, - ModelResponse, - ModelResponseStream, - StreamingChoices, - ) - - reasoning_item: Final = ChatCompletionReasoningItem( - type="reasoning", - id="rs_123", - encrypted_content="encrypted", - summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], - ) - response: Final = ( - ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) - if stream - else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) - ) - message_key: Final = "delta" if stream else "message" - assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] - - restored: Final = type(response).model_validate_json(response.model_dump_json()) - assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] - - -def test_generic_event(): - from litellm.types.llms.openai import GenericEvent - - event = {"type": "test", "test": "test"} - event = GenericEvent(**event) - assert event.type == "test" - assert event.test == "test" - - -def test_output_item_added_event(): - from litellm.types.llms.openai import OutputItemAddedEvent - - event = { - "type": "response.output_item.added", - "sequence_number": 4, - "output_index": 1, - "item": None, - } - event = OutputItemAddedEvent(**event) - assert event.type == "response.output_item.added" - assert event.sequence_number == 4 - assert event.output_index == 1 - assert event.item is None - - -class TestResponsesAPIResponseOutputText: - """Tests for the output_text property on ResponsesAPIResponse""" - - def test_output_text_with_single_message(self): - """Test output_text with a single message containing text output""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_123", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "Hello, world!", - } - ], - } - ], - ) - - assert response.output_text == "Hello, world!" - - def test_output_text_with_multiple_messages(self): - """Test output_text with multiple messages aggregates all text""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "First part. ", - } - ], - }, - { - "type": "message", - "id": "msg_2", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "Second part.", - } - ], - }, - ], - ) - - assert response.output_text == "First part. Second part." - - def test_output_text_with_no_text_content(self): - """Test output_text returns empty string when no output_text content exists""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "function_call", - "id": "call_123", - "status": "completed", - "name": "get_weather", - "arguments": "{}", - } - ], - ) - - assert response.output_text == "" - - def test_output_text_with_mixed_content(self): - """Test output_text only aggregates output_text type content""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "The weather is sunny. ", - }, - { - "type": "refusal", - "refusal": "I cannot do that.", - }, - ], - }, - { - "type": "function_call", - "id": "call_123", - "status": "completed", - "name": "get_weather", - "arguments": "{}", - }, - ], - ) - - assert response.output_text == "The weather is sunny. " - - def test_output_text_with_empty_output(self): - """Test output_text returns empty string with empty output list""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[], - ) - - assert response.output_text == "" - - -class TestAssistantMessageImageUrlContent: - """ - Regression tests for image_url blocks in assistant message content. - - Bug: ChatCompletionAssistantMessage.content did not include - ChatCompletionImageObject in its union, so Pydantic v2 silently dropped - image_url blocks (content → []) when serialising via AllMessageValues. - This affects users who store conversation history as JSON (e.g. in a DB) - and read it back typed as list[AllMessageValues]. - """ - - ASSISTANT_MESSAGE_WITH_IMAGE = { - "role": "assistant", - "content": [ - {"type": "text", "text": "Here is the image you requested:"}, - { - "type": "image_url", - "image_url": { - "url": ( - "data:image/png;base64," - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" - "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - ) - }, - }, - ], - } - - def test_assistant_message_image_url_preserved_single(self): - """ - TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive - validate_python → dump_python without being dropped or raising an error. - """ - from typing import List - - from pydantic import TypeAdapter - - from litellm.types.llms.openai import ChatCompletionAssistantMessage - - adapter = TypeAdapter(ChatCompletionAssistantMessage) - validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) - dumped = adapter.dump_python(validated) - - raw_content = dumped.get("content") - # Pydantic may return a lazy SerializationIterator for Iterable fields; - # convert to list to consume it — this must not raise ValidationError. - content_blocks = list(raw_content) if raw_content is not None else [] - - assert ( - len(content_blocks) == 2 - ), f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" - types = [b.get("type") for b in content_blocks if isinstance(b, dict)] - assert ( - "image_url" in types - ), f"image_url block was silently dropped; blocks: {content_blocks}" - - def test_assistant_message_image_url_preserved_in_all_message_values(self): - """ - TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an - assistant message must not be silently dropped during dump_python(mode='json'). - - This is the primary failing path: conversation history stored as JSON in a - database and read back typed as list[AllMessageValues]. - """ - from typing import List - - from pydantic import TypeAdapter - - from litellm.types.llms.openai import AllMessageValues - - conversation = [ - { - "role": "user", - "content": "Generate an image of a banana wearing a LiteLLM costume", - }, - self.ASSISTANT_MESSAGE_WITH_IMAGE, - ] - - adapter = TypeAdapter(List[AllMessageValues]) - validated = adapter.validate_python(conversation) - dumped = adapter.dump_python(validated, mode="json") - - assistant = next((m for m in dumped if m.get("role") == "assistant"), None) - assert assistant is not None, "Assistant message missing after serialisation" - - content = assistant.get("content", []) - assert isinstance( - content, list - ), f"content should be a list, got {type(content)}" - assert ( - len(content) == 2 - ), f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" - types = [b.get("type") for b in content if isinstance(b, dict)] - assert ( - "image_url" in types - ), f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" - - -class TestResponsesAPIReasoningNullFields: - """ - Tests for issue #16824: reasoning output items should not include null - status/content/encrypted_content fields. - - When a provider returns reasoning items without these fields, LiteLLM's - Pydantic parsing adds them as Optional defaults (None). Serializing them - as null breaks downstream SDKs (e.g., the OpenAI C# SDK crashes on - status=null). - - The fix uses a field_serializer on ResponsesAPIResponse.output that - mirrors the request-side filtering in - OpenAIResponsesAPIConfig._handle_reasoning_item(). - """ - - def _make_response(self, output): - from litellm.types.llms.openai import ResponsesAPIResponse - - return ResponsesAPIResponse( - id="resp_test", - created_at=1741476542, - model="gpt-5-mini", - object="response", - status="completed", - output=output, - ) - - def test_reasoning_item_null_fields_removed_model_dump(self): - """Null status/content/encrypted_content should be absent from model_dump.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert "encrypted_content" not in reasoning - - def test_reasoning_item_null_fields_removed_model_dump_json(self): - """Null fields should also be absent from model_dump_json.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - parsed = json.loads(response.model_dump_json()) - reasoning = parsed["output"][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert "encrypted_content" not in reasoning - - def test_reasoning_item_non_null_values_preserved(self): - """Non-null values on reasoning items should be kept.""" - response = self._make_response( - output=[ - { - "id": "rs_abc", - "type": "reasoning", - "summary": [], - "status": "completed", - "encrypted_content": "gAAAA...", - } - ] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert reasoning["status"] == "completed" - assert reasoning["encrypted_content"] == "gAAAA..." - - def test_message_item_not_affected(self): - """Non-reasoning output items should keep all their fields.""" - response = self._make_response( - output=[ - { - "id": "msg_abc", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "Hello!", - "annotations": [], - } - ], - } - ] - ) - dumped = response.model_dump() - message = dumped["output"][0] - assert message["status"] == "completed" - assert message["type"] == "message" - assert len(message["content"]) == 1 - - def test_mixed_output_reasoning_and_message(self): - """Reasoning items cleaned, message items untouched in same response.""" - response = self._make_response( - output=[ - {"id": "rs_abc", "type": "reasoning", "summary": []}, - { - "id": "msg_abc", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "Answer", - "annotations": [], - } - ], - }, - ] - ) - dumped = response.model_dump() - reasoning = [ - o - for o in dumped["output"] - if isinstance(o, dict) and o.get("type") == "reasoning" - ][0] - message = [ - o - for o in dumped["output"] - if isinstance(o, dict) and o.get("type") == "message" - ][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert message["status"] == "completed" - assert len(message["content"]) == 1 - - def test_reasoning_core_fields_preserved(self): - """id, type, summary should always be present on reasoning items.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": ["thinking..."]}] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert reasoning["id"] == "rs_abc" - assert reasoning["type"] == "reasoning" - assert reasoning["summary"] == ["thinking..."] - - def test_top_level_null_fields_unaffected(self): - """Top-level response fields with None should not be affected.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - dumped = response.model_dump() - assert "error" in dumped - assert dumped["error"] is None - assert "instructions" in dumped - assert dumped["instructions"] is None - - -def test_normalize_fine_tuning_job_dict_maps_azure_pending(): - from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict - - out = _normalize_fine_tuning_job_dict( - {"organization_id": None, "result_files": None, "status": "pending"}, - is_azure=True, - ) - assert out["organization_id"] == "" - assert out["result_files"] == [] - assert out["status"] == "queued" - - -def test_normalize_fine_tuning_job_dict_openai_unchanged(): - from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict - - data = {"organization_id": None, "result_files": None, "status": "pending"} - out = _normalize_fine_tuning_job_dict(data, is_azure=False) - assert out is data - - -def test_openai_file_object_accepts_pending_status(): - from litellm.types.llms.openai import OpenAIFileObject - - file_obj = OpenAIFileObject( - id="file-123", - bytes=1024, - created_at=1677610602, - filename="train.jsonl", - object="file", - purpose="fine-tune", - status="pending", - ) - assert file_obj.status == "pending" - - -class TestOpenAIFileObjectBatchGuardrailSerialization: - """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" - - @staticmethod - def _file_object(**overrides): - from litellm.types.llms.openai import OpenAIFileObject - - return OpenAIFileObject( - id="file-123", - object="file", - bytes=1024, - created_at=1677610602, - filename="batch.jsonl", - purpose="batch", - status="uploaded", - **overrides, - ) - - @staticmethod - def _report(): - from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport - - return BatchGuardrailReport( - submitted_records=3, - modified_records=(BatchGuardrailRecord(line=2, custom_id="dirty", action="redacted"),), - ) - - @pytest.mark.parametrize("mode", ["python", "json"]) - def test_key_absent_when_unset(self, mode): - assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode=mode) - - @pytest.mark.parametrize("mode", ["python", "json"]) - def test_key_present_when_set(self, mode): - dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode=mode) - assert dumped["litellm_batch_guardrail"]["submitted_records"] == 3 - - def test_nested_nulls_of_a_set_report_survive(self): - """`exclude_none=True` was rejected as the fix because it would strip these.""" - dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode="json") - assert dumped["litellm_batch_guardrail"]["modified_records"] == [ - {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} - ] - - def test_by_alias_dump_also_omits_the_key(self): - """Tripwire: the serializer filters a literal key name, which an added alias would bypass.""" - assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode="json", by_alias=True) - - def test_other_optional_fields_still_serialize_as_null(self): - dumped = self._file_object().model_dump(mode="json") - assert dumped["expires_at"] is None - assert dumped["status_details"] is None - - def test_round_trip_of_a_set_report_is_lossless(self): - from litellm.types.llms.openai import OpenAIFileObject - - original = self._file_object(litellm_batch_guardrail=self._report()) - assert OpenAIFileObject(**original.model_dump()) == original - - def test_serialization_json_schema_still_describes_the_model(self): - """A return annotation on the wrap serializer would collapse this to a bare object.""" - from litellm.types.llms.openai import OpenAIFileObject - - schema = OpenAIFileObject.model_json_schema(mode="serialization") - assert "litellm_batch_guardrail" in schema["properties"] - - def test_key_omitted_inside_a_file_list_page(self): - from litellm.types.llms.openai import FileListPage - - page = FileListPage(object="list", data=[self._file_object()], has_more=False) - assert "litellm_batch_guardrail" not in page.model_dump(mode="json")["data"][0] - - -def _binary_content(payload: bytes) -> HttpxBinaryResponseContent: - import httpx - - return HttpxBinaryResponseContent(httpx.Response(200, content=payload)) - - -def test_httpx_binary_response_content_hidden_params_are_per_instance(): - first = _binary_content(b"first") - second = _binary_content(b"second") - - first._hidden_params["response_cost"] = 0.5 - - assert second._hidden_params == {} - - -def test_set_response_cost_none_leaves_hidden_params_empty(): - binary_response = _binary_content(b"audio") - - binary_response.set_response_cost(None) - - assert "response_cost" not in binary_response._hidden_params - - binary_response.set_response_cost(0.25) - - assert binary_response._hidden_params["response_cost"] == 0.25 - - binary_response.set_response_cost(None) - - assert "response_cost" not in binary_response._hidden_params diff --git a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py deleted file mode 100644 index 2e5986d3ef8..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py +++ /dev/null @@ -1,168 +0,0 @@ -""" -Tests for pipeline type definitions. -""" - -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.pipeline_types import ( - GuardrailPipeline, - PipelineExecutionResult, - PipelineStep, - PipelineStepResult, -) -from litellm.types.proxy.policy_engine.policy_types import ( - Policy, - PolicyGuardrails, -) - - -def test_pipeline_step_defaults(): - step = PipelineStep(guardrail="my-guard") - assert step.on_fail == "block" - assert step.on_pass == "allow" - assert step.on_error is None - assert step.pass_data is False - assert step.modify_response_message is None - - -def test_pipeline_step_valid_actions(): - step = PipelineStep(guardrail="my-guard", on_fail="next", on_pass="next") - assert step.on_fail == "next" - assert step.on_pass == "next" - - -def test_pipeline_step_all_action_types(): - for action in ("allow", "block", "next", "modify_response"): - step = PipelineStep( - guardrail="g", on_fail=action, on_pass=action, on_error=action - ) - assert step.on_fail == action - assert step.on_pass == action - assert step.on_error == action - - -def test_pipeline_step_invalid_action_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_fail="invalid_action") - - -def test_pipeline_step_invalid_on_pass_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_pass="skip") - - -def test_pipeline_step_on_error_valid(): - step = PipelineStep( - guardrail="g", on_error="next", on_fail="block", on_pass="allow" - ) - assert step.on_error == "next" - - -def test_pipeline_step_invalid_on_error_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_error="invalid") - - -def test_pipeline_requires_at_least_one_step(): - with pytest.raises(ValidationError): - GuardrailPipeline(mode="pre_call", steps=[]) - - -def test_pipeline_invalid_mode_rejected(): - with pytest.raises(ValidationError): - GuardrailPipeline( - mode="during_call", - steps=[PipelineStep(guardrail="g")], - ) - - -def test_pipeline_valid_modes(): - for mode in ("pre_call", "post_call"): - pipeline = GuardrailPipeline( - mode=mode, - steps=[PipelineStep(guardrail="g")], - ) - assert pipeline.mode == mode - - -def test_pipeline_with_multiple_steps(): - pipeline = GuardrailPipeline( - mode="pre_call", - steps=[ - PipelineStep(guardrail="g1", on_fail="next", on_pass="allow"), - PipelineStep(guardrail="g2", on_fail="block", on_pass="allow"), - ], - ) - assert len(pipeline.steps) == 2 - assert pipeline.steps[0].guardrail == "g1" - assert pipeline.steps[1].guardrail == "g2" - - -def test_policy_with_pipeline_parses(): - policy = Policy( - guardrails=PolicyGuardrails(add=["g1", "g2"]), - pipeline=GuardrailPipeline( - mode="pre_call", - steps=[ - PipelineStep(guardrail="g1", on_fail="next"), - PipelineStep(guardrail="g2"), - ], - ), - ) - assert policy.pipeline is not None - assert len(policy.pipeline.steps) == 2 - - -def test_policy_without_pipeline(): - policy = Policy( - guardrails=PolicyGuardrails(add=["g1"]), - ) - assert policy.pipeline is None - - -def test_pipeline_step_result(): - result = PipelineStepResult( - guardrail_name="g1", - outcome="fail", - action_taken="next", - error_detail="Content policy violation", - duration_seconds=0.05, - ) - assert result.outcome == "fail" - assert result.action_taken == "next" - - -def test_pipeline_execution_result(): - result = PipelineExecutionResult( - terminal_action="block", - step_results=[ - PipelineStepResult( - guardrail_name="g1", - outcome="fail", - action_taken="next", - ), - PipelineStepResult( - guardrail_name="g2", - outcome="fail", - action_taken="block", - ), - ], - error_message="Content blocked", - ) - assert result.terminal_action == "block" - assert len(result.step_results) == 2 - - -def test_pipeline_step_extra_fields_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="g", unknown_field="value") - - -def test_pipeline_extra_fields_rejected(): - with pytest.raises(ValidationError): - GuardrailPipeline( - mode="pre_call", - steps=[PipelineStep(guardrail="g")], - unknown="value", - ) diff --git a/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py b/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py deleted file mode 100644 index bcd6d39aa4d..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py +++ /dev/null @@ -1,15 +0,0 @@ -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment - - -@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) -def test_policy_attachment_accepts_int32_priority(priority: int): - assert PolicyAttachment(policy="p", priority=priority).priority == priority - - -@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) -def test_policy_attachment_rejects_priority_outside_int32(priority: int): - with pytest.raises(ValidationError): - PolicyAttachment(policy="p", priority=priority) diff --git a/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py b/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py deleted file mode 100644 index f31b9d7e873..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Tests for pipeline field on policy CRUD types (resolver_types.py). -""" - -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.resolver_types import ( - PolicyAttachmentCreateRequest, - PolicyCreateRequest, - PolicyDBResponse, - PolicyUpdateRequest, -) - - -def test_policy_create_request_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, - {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, - ], - } - req = PolicyCreateRequest( - policy_name="test-policy", - guardrails_add=["g1", "g2"], - pipeline=pipeline_data, - ) - assert req.pipeline is not None - assert req.pipeline["mode"] == "pre_call" - assert len(req.pipeline["steps"]) == 2 - - -def test_policy_create_request_without_pipeline(): - req = PolicyCreateRequest( - policy_name="test-policy", - guardrails_add=["g1"], - ) - assert req.pipeline is None - - -def test_policy_update_request_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "block", "on_pass": "allow"}, - ], - } - req = PolicyUpdateRequest(pipeline=pipeline_data) - assert req.pipeline is not None - assert req.pipeline["steps"][0]["guardrail"] == "g1" - - -def test_policy_db_response_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, - {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, - ], - } - resp = PolicyDBResponse( - policy_id="test-id", - policy_name="test-policy", - guardrails_add=["g1", "g2"], - pipeline=pipeline_data, - ) - assert resp.pipeline is not None - assert resp.pipeline["mode"] == "pre_call" - dumped = resp.model_dump() - assert dumped["pipeline"]["steps"][0]["guardrail"] == "g1" - - -def test_policy_db_response_without_pipeline(): - resp = PolicyDBResponse( - policy_id="test-id", - policy_name="test-policy", - ) - assert resp.pipeline is None - dumped = resp.model_dump() - assert dumped["pipeline"] is None - - -def test_policy_create_request_roundtrip(): - pipeline_data = { - "mode": "post_call", - "steps": [ - { - "guardrail": "g1", - "on_fail": "modify_response", - "on_pass": "next", - "pass_data": True, - "modify_response_message": "custom msg", - }, - ], - } - req = PolicyCreateRequest( - policy_name="roundtrip-test", - guardrails_add=["g1"], - pipeline=pipeline_data, - ) - dumped = req.model_dump() - restored = PolicyCreateRequest(**dumped) - assert restored.pipeline == pipeline_data - - -@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) -def test_policy_attachment_create_request_accepts_int32_priority(priority: int): - assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority - - -@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) -def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int): - with pytest.raises(ValidationError): - PolicyAttachmentCreateRequest(policy_name="p", priority=priority) diff --git a/tests/test_litellm/videos/test_main.py b/tests/test_litellm/videos/test_main.py deleted file mode 100644 index 22e1e5c05eb..00000000000 --- a/tests/test_litellm/videos/test_main.py +++ /dev/null @@ -1,455 +0,0 @@ -""" -Dispatch-contract tests for litellm/videos/main.py - -Each public video operation is a pair: a sync `video_*` worker (decorated with -@client) that resolves the provider, fetches the provider config, logs, and then -forwards to exactly one `base_llm_http_handler.video_*_handler`; and an async -`avideo_*` wrapper that delegates to the sync worker in an executor. - -This file locks the contract of that layer so a regression fails loudly: - - 1. DISPATCH - the one correct handler fired and every sibling video handler - asserted NOT called. A copy-paste that calls the wrong handler - (e.g. remix -> edit) flips this. - 2. RESULT - the handler's return value is propagated by identity. - 3. PROVIDER - custom_llm_provider is decoded from an encoded video id when not - passed (status/content/remix/edit/extension), or defaults to - "openai" (list/create_character/get_character). This is the exact - surface of the historical "content defaulted to openai" bug. - 4. PAYLOAD - the provider config object and the operation's identifying args - (video_id/prompt/name/...) reach the handler; _is_async is False - on the sync path. - 5. SHORT-CIRCUIT - mock_response returns a typed object without any handler call. - 6. UNSUPPORTED - a None provider config raises before any handler fires. - 7. DELEGATION - avideo_* returns the sync worker's result untouched, sets - async_call=True, and pre-resolves the provider where it must. - -Seams mocked: the http handler (network), the provider-config registry lookup, -get_llm_provider, and the video-generation optional-param builders. The id decode -helper runs for real against genuinely-encoded ids, so the provider assertions -reflect production. -""" - -from contextlib import ExitStack -from dataclasses import dataclass -from typing import Any, Dict -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.types.videos.main import CharacterObject, VideoObject -from litellm.types.videos.utils import encode_video_id_with_provider -from litellm.videos import main as videos_main - -# A real model-encoded video id: decodes (for real) to provider "azure". Used to -# prove the sync workers derive custom_llm_provider from the id, not a hardcode. -AZURE_VIDEO_ID = encode_video_id_with_provider("video_raw", "azure", "deployment-1") - -# The nine sync handlers on base_llm_http_handler. Dispatch tests assert exactly -# one fired and the other eight did not. -SYNC_HANDLERS = ( - "video_generation_handler", - "video_content_handler", - "video_remix_handler", - "video_create_character_handler", - "video_get_character_handler", - "video_edit_handler", - "video_extension_handler", - "video_list_handler", - "video_status_handler", -) - -GEN_OPTIONAL_PARAMS = {"seconds": "8", "size": "720x1280"} - - -@dataclass -class Seams: - handler: MagicMock - get_config: MagicMock - config: MagicMock - - def kwargs_of(self, handler_name: str) -> Dict[str, Any]: - method = getattr(self.handler, handler_name) - assert method.call_count == 1 - return dict(method.call_args.kwargs) - - def assert_only(self, handler_name: str) -> None: - for name in SYNC_HANDLERS: - method = getattr(self.handler, name) - if name == handler_name: - method.assert_called_once() - else: - method.assert_not_called() - - -@pytest.fixture -def seams(): - handler = MagicMock(spec=BaseLLMHTTPHandler) - config = MagicMock(name="provider_video_config") - get_config = MagicMock(return_value=config) - - with ExitStack() as stack: - stack.enter_context(patch.object(videos_main, "base_llm_http_handler", handler)) - stack.enter_context( - patch.object( - videos_main.ProviderConfigManager, - "get_provider_video_config", - get_config, - ) - ) - # video_generation resolves model+provider through get_llm_provider and - # builds optional params; mock those so the dispatch payload is deterministic. - stack.enter_context( - patch.object( - videos_main, - "get_llm_provider", - MagicMock(return_value=("sora-2", "openai", None, None)), - ) - ) - stack.enter_context( - patch.object( - videos_main.VideoGenerationRequestUtils, - "get_requested_video_generation_optional_param", - MagicMock(return_value={"seconds": "8"}), - ) - ) - stack.enter_context( - patch.object( - videos_main.VideoGenerationRequestUtils, - "get_optional_params_video_generation", - MagicMock(return_value=dict(GEN_OPTIONAL_PARAMS)), - ) - ) - yield Seams(handler=handler, get_config=get_config, config=config) - - -# =========================================================================== # -# Dispatch contract - one rich test per sync worker. -# =========================================================================== # - - -def test_video_generation__dispatch(seams): - result = videos_main.video_generation(prompt="a sunset", model="sora-2") - - seams.assert_only("video_generation_handler") - assert result is seams.handler.video_generation_handler.return_value - kw = seams.kwargs_of("video_generation_handler") - assert kw["model"] == "sora-2" - assert kw["prompt"] == "a sunset" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_generation_provider_config"] is seams.config - assert kw["video_generation_optional_request_params"] == GEN_OPTIONAL_PARAMS - assert kw["_is_async"] is False - - -def test_video_status__dispatch_and_provider_from_id(seams): - result = videos_main.video_status(video_id=AZURE_VIDEO_ID) - - seams.assert_only("video_status_handler") - assert result is seams.handler.video_status_handler.return_value - kw = seams.kwargs_of("video_status_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["custom_llm_provider"] == "azure" # decoded from the id, not openai - assert kw["video_status_provider_config"] is seams.config - assert kw["_is_async"] is False - # provider config requested for the decoded provider, not a hardcode. - assert seams.get_config.call_args.kwargs["provider"] == litellm.LlmProviders.AZURE - - -def test_video_content__dispatch_and_provider_from_id(seams): - result = videos_main.video_content(video_id=AZURE_VIDEO_ID, variant="thumbnail") - - seams.assert_only("video_content_handler") - assert result is seams.handler.video_content_handler.return_value - kw = seams.kwargs_of("video_content_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["custom_llm_provider"] == "azure" - assert kw["variant"] == "thumbnail" - assert kw["video_content_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_content__plain_id_defaults_to_openai(seams): - videos_main.video_content(video_id="video_plain") - - assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai" - - -def test_video_remix__dispatch_and_provider_from_id(seams): - result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors") - - seams.assert_only("video_remix_handler") - assert result is seams.handler.video_remix_handler.return_value - kw = seams.kwargs_of("video_remix_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "new colors" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_remix_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_edit__dispatch_and_provider_from_id(seams): - result = videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="brighter") - - seams.assert_only("video_edit_handler") - assert result is seams.handler.video_edit_handler.return_value - kw = seams.kwargs_of("video_edit_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "brighter" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_extension__dispatch_and_provider_from_id(seams): - result = videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" - ) - - seams.assert_only("video_extension_handler") - assert result is seams.handler.video_extension_handler.return_value - kw = seams.kwargs_of("video_extension_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "continue" - assert kw["seconds"] == "5" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_list__dispatch_defaults_to_openai(seams): - result = videos_main.video_list(after="cur", limit=5, order="desc") - - seams.assert_only("video_list_handler") - assert result is seams.handler.video_list_handler.return_value - kw = seams.kwargs_of("video_list_handler") - assert kw["after"] == "cur" - assert kw["limit"] == 5 - assert kw["order"] == "desc" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_list_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_create_character__dispatch_defaults_to_openai(seams): - video = MagicMock(name="video_upload") - result = videos_main.video_create_character(name="hero", video=video) - - seams.assert_only("video_create_character_handler") - assert result is seams.handler.video_create_character_handler.return_value - kw = seams.kwargs_of("video_create_character_handler") - assert kw["name"] == "hero" - assert kw["video"] is video - assert kw["custom_llm_provider"] == "openai" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_get_character__dispatch_defaults_to_openai(seams): - result = videos_main.video_get_character(character_id="char_1") - - seams.assert_only("video_get_character_handler") - assert result is seams.handler.video_get_character_handler.return_value - kw = seams.kwargs_of("video_get_character_handler") - assert kw["character_id"] == "char_1" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_explicit_provider_beats_decoded_id(seams): - """An explicit custom_llm_provider wins over the one encoded in the id.""" - videos_main.video_status(video_id=AZURE_VIDEO_ID, custom_llm_provider="vertex_ai") - - assert seams.kwargs_of("video_status_handler")["custom_llm_provider"] == "vertex_ai" - - -# =========================================================================== # -# mock_response short-circuit - returns a typed object, no handler call. -# =========================================================================== # - - -def test_generation__mock_response_short_circuits(seams): - resp = videos_main.video_generation( - prompt="x", - model="sora-2", - mock_response={"id": "v1", "object": "video", "status": "queued"}, - ) - - assert isinstance(resp, VideoObject) - assert resp.id == "v1" - seams.handler.video_generation_handler.assert_not_called() - - -def test_list__mock_response_short_circuits(seams): - resp = videos_main.video_list( - mock_response=[{"id": "v1", "object": "video", "status": "completed"}] - ) - - assert isinstance(resp, list) - assert resp[0].id == "v1" - seams.handler.video_list_handler.assert_not_called() - - -def test_get_character__mock_response_short_circuits(seams): - resp = videos_main.video_get_character( - character_id="char_1", - mock_response={ - "id": "char_1", - "object": "character", - "created_at": 1, - "name": "hero", - }, - ) - - assert isinstance(resp, CharacterObject) - assert resp.id == "char_1" - seams.handler.video_get_character_handler.assert_not_called() - - -# =========================================================================== # -# Unsupported provider - a None provider config raises before any dispatch. -# =========================================================================== # - - -def test_unsupported_provider_raises_without_dispatch(seams): - seams.get_config.return_value = None - - with pytest.raises(litellm.APIConnectionError): - videos_main.video_status(video_id=AZURE_VIDEO_ID) - - seams.handler.video_status_handler.assert_not_called() - - -# =========================================================================== # -# Async-wrapper delegation - representative coverage. -# =========================================================================== # - - -@pytest.mark.asyncio -async def test_avideo_generation__delegates_with_async_flag(): - sentinel = VideoObject(id="v-async", object="video", status="queued") - with ( - patch.object( - videos_main, "video_generation", MagicMock(return_value=sentinel) - ) as sync, - patch.object( - litellm, - "get_llm_provider", - MagicMock(return_value=("sora-2", "openai", None, None)), - ), - ): - result = await videos_main.avideo_generation(prompt="x", model="sora-2") - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["custom_llm_provider"] == "openai" - - -@pytest.mark.asyncio -async def test_avideo_status__delegates_untouched(): - sentinel = VideoObject(id="v-async", object="video", status="queued") - with patch.object( - videos_main, "video_status", MagicMock(return_value=sentinel) - ) as sync: - result = await videos_main.avideo_status(video_id="video_plain") - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["video_id"] == "video_plain" - - -@pytest.mark.asyncio -async def test_avideo_content__pre_decodes_provider_before_delegating(): - """avideo_content resolves the provider from the encoded id itself before - handing off, so the sync worker receives the decoded provider, not None.""" - sentinel = b"mp4-bytes" - with patch.object( - videos_main, "video_content", MagicMock(return_value=sentinel) - ) as sync: - result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["custom_llm_provider"] == "azure" - - -# =========================================================================== # -# Credential passthrough - DB/YAML model-config credentials the router injects -# via kwargs must reach the provider call for EVERY video handler, carried in -# litellm_params. Distinct per-field values catch a cross-wired field. -# =========================================================================== # - -DB_YAML_CREDS = { - "api_key": "sk-db-credential", - "api_base": "https://db-resource.test", - "api_version": "2024-12-31", - "vertex_project": "db-project-xyz", -} - -CREDENTIAL_OPERATIONS = [ - ( - "video_generation_handler", - lambda: videos_main.video_generation( - prompt="p", model="sora-2", **DB_YAML_CREDS - ), - ), - ( - "video_status_handler", - lambda: videos_main.video_status(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), - ), - ( - "video_content_handler", - lambda: videos_main.video_content(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), - ), - ( - "video_remix_handler", - lambda: videos_main.video_remix( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), - ), - ( - "video_edit_handler", - lambda: videos_main.video_edit( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), - ), - ( - "video_extension_handler", - lambda: videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS - ), - ), - ( - "video_list_handler", - lambda: videos_main.video_list(**DB_YAML_CREDS), - ), - ( - "video_create_character_handler", - lambda: videos_main.video_create_character( - name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS - ), - ), - ( - "video_get_character_handler", - lambda: videos_main.video_get_character(character_id="char_1", **DB_YAML_CREDS), - ), -] - - -@pytest.mark.parametrize( - "handler_name,invoke", - CREDENTIAL_OPERATIONS, - ids=[op[0] for op in CREDENTIAL_OPERATIONS], -) -def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): - invoke() - - litellm_params = seams.kwargs_of(handler_name)["litellm_params"] - assert litellm_params.get("api_key") == DB_YAML_CREDS["api_key"] - assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] - assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] - assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] diff --git a/tests/test_litellm/videos/test_utils.py b/tests/test_litellm/videos/test_utils.py deleted file mode 100644 index 57fb549c23d..00000000000 --- a/tests/test_litellm/videos/test_utils.py +++ /dev/null @@ -1,193 +0,0 @@ -""" -Pure-logic contract tests for litellm/videos/main.py's request utils -(litellm/videos/utils.py: VideoGenerationRequestUtils). - -These lock the exact param-shaping behavior so a mutation that drops a filter, -flips a precedence, or stops removing a key fails loudly. The only seam is the -provider config's map_openai_params (a provider boundary); filter_out_litellm_params -runs for real, so the "litellm-internal params get stripped" assertions reflect -production. Every test asserts the exact resulting dict, never "ran without error". -""" - -from unittest.mock import MagicMock - - - -import litellm -from litellm.videos.utils import VideoGenerationRequestUtils - -get_requested = ( - VideoGenerationRequestUtils.get_requested_video_generation_optional_param -) -get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation - - -# =========================================================================== # -# get_requested_video_generation_optional_param -# -# Receives the caller's full local_vars; must return only the API-bound optional -# params. filter_out_litellm_params strips known internal keys for real; the -# values used below were chosen against the live set: seconds/size/user/foo_param/ -# vertex_project/extra/a/b survive, api_key/metadata/litellm_* are stripped. -# =========================================================================== # - - -def test_requested__drops_none_and_excluded_keys(): - result = get_requested( - { - "seconds": "8", - "size": None, # None -> dropped - "prompt": "a sunset", # excluded - "model": "sora-2", # excluded - "user": "u1", - } - ) - assert result == {"seconds": "8", "user": "u1"} - - -def test_requested__strips_litellm_internal_params(): - result = get_requested( - { - "seconds": "8", - "api_key": "sk-secret", - "metadata": {"x": 1}, - "litellm_call_id": "id-123", - } - ) - assert result == {"seconds": "8"} - - -def test_requested__timeout_always_removed(): - # timeout is NOT a litellm-internal param, so only the explicit pop removes it. - result = get_requested({"seconds": "8", "timeout": 30}) - assert result == {"seconds": "8"} - - -def test_requested__nested_kwargs_merge_and_override_base(): - result = get_requested( - {"seconds": "8", "kwargs": {"size": "720x1280", "seconds": "override"}} - ) - # nested kwargs win over the top-level base params on collision. - assert result == {"seconds": "override", "size": "720x1280"} - - -def test_requested__non_dict_kwargs_treated_as_empty(): - result = get_requested({"seconds": "8", "kwargs": "not-a-dict"}) - assert result == {"seconds": "8"} - - -def test_requested__none_input_returns_empty(): - assert get_requested(None) == {} - - -def test_requested__top_level_extra_body_spread_and_preserved(): - result = get_requested( - {"seconds": "8", "extra_body": {"vertex_project": "proj", "foo_param": "bar"}} - ) - # extra_body keys are both spread at top level AND kept under "extra_body". - assert result == { - "seconds": "8", - "vertex_project": "proj", - "foo_param": "bar", - "extra_body": {"vertex_project": "proj", "foo_param": "bar"}, - } - - -def test_requested__extra_body_kwargs_overrides_top_level(): - result = get_requested( - { - "extra_body": {"a": "top", "b": "top_b"}, - "kwargs": {"extra_body": {"a": "kw"}}, - } - ) - # kwargs' extra_body wins over the top-level extra_body on collision; the - # non-colliding top-level key survives. - assert result == { - "a": "kw", - "b": "top_b", - "extra_body": {"a": "kw", "b": "top_b"}, - } - - -def test_requested__extra_body_strips_litellm_internal_params(): - result = get_requested({"extra_body": {"api_key": "sk", "foo_param": "bar"}}) - # api_key filtered out of extra_body; only foo_param remains (and is spread). - assert result == {"foo_param": "bar", "extra_body": {"foo_param": "bar"}} - - -def test_requested__empty_extra_body_not_added(): - result = get_requested({"seconds": "8", "extra_body": {}}) - assert result == {"seconds": "8"} - assert "extra_body" not in result - - -# =========================================================================== # -# get_optional_params_video_generation -# -# Delegates mapping to the provider config (the seam) then folds extra_body in. -# =========================================================================== # - - -def _config(map_return): - config = MagicMock() - config.map_openai_params.return_value = map_return - return config - - -def test_optional__delegates_to_map_openai_params_with_drop_params(): - config = _config({"seconds": "8"}) - optional_params = {"seconds": "8"} - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params=optional_params, - ) - - assert result == {"seconds": "8"} - config.map_openai_params.assert_called_once_with( - video_create_optional_params=optional_params, - model="sora-2", - drop_params=litellm.drop_params, - ) - - -def test_optional__extra_body_overrides_mapped_and_is_removed(): - # mapped output carries a leftover extra_body that must be popped; the input - # extra_body overrides a colliding mapped key and is spread in. - config = _config({"seconds": "8", "size": "mapped", "extra_body": {"leftover": 1}}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={ - "extra_body": {"size": "override", "extra": "x"} - }, - ) - - assert result == {"seconds": "8", "size": "override", "extra": "x"} - assert "extra_body" not in result - - -def test_optional__no_extra_body_returns_mapped_unchanged(): - config = _config({"seconds": "8"}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={"seconds": "8"}, - ) - - assert result == {"seconds": "8"} - - -def test_optional__non_dict_extra_body_ignored(): - config = _config({"seconds": "8"}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={"seconds": "8", "extra_body": None}, - ) - - assert result == {"seconds": "8"} From 72abd11b4ce43fb3762b008d97e60c95bc76e313 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:01:32 +0000 Subject: [PATCH 063/146] test: assert the responses bridge forwards aws_region_name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 tests/unit/completion_extras/test_responses_bridge_provider_propagation.py diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py new file mode 100644 index 00000000000..09ef1889818 --- /dev/null +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -0,0 +1,59 @@ +from datetime import datetime +from unittest.mock import patch + +import pytest + +from litellm.completion_extras.litellm_responses_transformation.handler import ( + ResponsesToCompletionBridgeHandler, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.utils import ModelResponse + +MODEL = "openai.gpt-5.5" +REGION = "us-east-2" + + +def _bedrock_mantle_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="acompletion", + model=MODEL, + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": MODEL, + "custom_llm_provider": "bedrock_mantle", + "messages": messages, + "optional_params": {}, + "litellm_params": { + "aws_region_name": REGION, + "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", + "custom_llm_provider": "bedrock_mantle", + }, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +@pytest.mark.asyncio +async def test_acompletion_forwards_aws_region_name_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model=MODEL) + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_bedrock_mantle_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["aws_region_name"] == REGION + assert _fake_aresponses.kwargs["custom_llm_provider"] == "bedrock_mantle" From a69ca90ea5380bde22ad36ac9276604bae28652f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:01:40 +0000 Subject: [PATCH 064/146] test: migrate phase 14 wave 2 provider tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...sonx_audio_transcription_transformation.py | 337 ---------- .../test_litellm/llms/watsonx/test_watsonx.py | 577 ------------------ .../llms/xai/xai_responses/__init__.py | 1 - .../xai/xai_responses/test_transformation.py | 105 ---- .../count_tokens/__init__.py | 0 .../test_count_tokens_location.py | 2 + .../test_count_tokens_no_vertexai_sdk.py | 0 ...ai_partner_models_llama3_transformation.py | 0 ...i_partner_models_mistral_transformation.py | 15 + .../vertex_ai/vertex_gemma_models/__init__.py | 0 .../test_vertex_gemma_transformation.py | 0 tests/unit/llms/vertex_ai/videos/__init__.py | 0 .../test_vertex_video_transformation.py | 0 ...est_volcengine_responses_transformation.py | 24 - tests/unit/llms/voyage/rerank/__init__.py | 0 .../test_voyage_rerank_transformation.py | 0 .../test_voyage_contextual_embedding.py | 0 .../test_voyage_multimodal_embedding.py | 0 tests/unit/llms/watsonx/__init__.py | 0 .../watsonx/audio_transcription/__init__.py | 0 ...sonx_audio_transcription_transformation.py | 85 +++ .../test_watsonx_embedding_transformation.py | 0 ...test_watsonx_passthrough_transformation.py | 0 tests/unit/llms/watsonx/rerank/__init__.py | 0 .../watsonx/rerank/test_watsonx_rerank.py | 0 tests/unit/llms/watsonx/test_watsonx.py | 74 +++ .../llms/watsonx/test_watsonx_common_utils.py | 0 .../test_xai_responses_transformation.py | 0 tests/unit/llms/you_com/__init__.py | 0 .../llms/you_com/test_you_com_search.py | 0 .../llms/zai/test_zai_provider.py | 0 31 files changed, 176 insertions(+), 1044 deletions(-) delete mode 100644 tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py delete mode 100644 tests/test_litellm/llms/watsonx/test_watsonx.py delete mode 100644 tests/test_litellm/llms/xai/xai_responses/__init__.py delete mode 100644 tests/test_litellm/llms/xai/xai_responses/test_transformation.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py (98%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py (60%) create mode 100644 tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/videos/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/videos/test_vertex_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/volcengine/responses/test_volcengine_responses_transformation.py (94%) create mode 100644 tests/unit/llms/voyage/rerank/__init__.py rename tests/{test_litellm => unit}/llms/voyage/rerank/test_voyage_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/voyage/test_voyage_contextual_embedding.py (100%) rename tests/{test_litellm => unit}/llms/voyage/test_voyage_multimodal_embedding.py (100%) create mode 100644 tests/unit/llms/watsonx/__init__.py create mode 100644 tests/unit/llms/watsonx/audio_transcription/__init__.py create mode 100644 tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py rename tests/{test_litellm => unit}/llms/watsonx/embed/test_watsonx_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py (100%) create mode 100644 tests/unit/llms/watsonx/rerank/__init__.py rename tests/{test_litellm => unit}/llms/watsonx/rerank/test_watsonx_rerank.py (100%) create mode 100644 tests/unit/llms/watsonx/test_watsonx.py rename tests/{test_litellm => unit}/llms/watsonx/test_watsonx_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/xai/responses/test_xai_responses_transformation.py (100%) create mode 100644 tests/unit/llms/you_com/__init__.py rename tests/{test_litellm => unit}/llms/you_com/test_you_com_search.py (100%) rename tests/{test_litellm => unit}/llms/zai/test_zai_provider.py (100%) diff --git a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py deleted file mode 100644 index e269e782061..00000000000 --- a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py +++ /dev/null @@ -1,337 +0,0 @@ -""" -Tests for IBM WatsonX Audio Transcription. - -Validates that litellm.transcription transforms requests correctly for WatsonX. -""" - -import json -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - - -import litellm -from litellm.llms.watsonx.audio_transcription.transformation import ( - IBMWatsonXAudioTranscriptionConfig, -) -from litellm.types.utils import TranscriptionResponse - - -class TestWatsonXAudioTranscription: - """Tests for WatsonX audio transcription via litellm.transcription.""" - - @pytest.mark.asyncio - async def test_watsonx_transcription_url_and_headers(self): - """ - Test that litellm.transcription sends request to correct WatsonX URL with proper headers. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["url"] = str(kwargs.get("url", args[0] if args else None)) - captured_request["headers"] = kwargs.get("headers", {}) - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - # Validate URL contains WatsonX audio transcription endpoint - assert "/ml/v1/audio/transcriptions" in captured_request["url"] - assert "version=" in captured_request["url"] - # project_id should NOT be in URL (it should be in form data instead) - assert "project_id=test-project-123" not in captured_request["url"] - - # Validate headers contain WatsonX auth - assert "Authorization" in captured_request["headers"] - assert ( - "Bearer test-bearer-token" in captured_request["headers"]["Authorization"] - ) - - # Validate Content-Type is NOT set (httpx sets multipart/form-data automatically) - assert "Content-Type" not in captured_request["headers"] - - # Validate project_id is in form data, not URL - assert captured_request["data"].get("project_id") == "test-project-123" - - # Validate file is in files dict - assert "file" in captured_request["files"] - - @pytest.mark.asyncio - async def test_watsonx_transcription_request_body(self): - """ - Test that litellm.transcription sends correct request body for WatsonX. - - Validates that: - - Request uses multipart/form-data (data + files) - - Model name has watsonx/ prefix removed - - project_id is in form data, not URL - - Audio file is in files dict - - OpenAI params are included in form data - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - language="en", - temperature=0.5, - ) - except Exception: - pass # We just want to capture the request - - # Validate form data contains expected fields - data = captured_request.get("data", {}) - - print("JSON DUMPS captured_request:") - print(json.dumps(captured_request, indent=4, default=str)) - - # Model name should NOT have watsonx/ prefix - assert data.get("model") == "whisper-large-v3-turbo" - - # project_id should be in form data - assert data.get("project_id") == "test-project-123" - - # OpenAI params should be in form data - assert data.get("language") == "en" - assert data.get("temperature") == 0.5 - # response_format should NOT be set by default - only send what user specifies - assert "response_format" not in data - - # Validate file is in files dict (multipart/form-data) - files = captured_request.get("files", {}) - assert "file" in files - assert isinstance( - files["file"], tuple - ) # Should be (filename, content, content_type) - - @pytest.mark.asyncio - async def test_watsonx_transcription_only_user_params_sent_with_project_id(self): - """ - Test that only user-specified params are sent in request body to WatsonX. - - LiteLLM should NOT add extra params like response_format if user didn't specify them. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - # Minimal request - only required params - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - data = captured_request.get("data", {}) - - # These are the ONLY keys that should be in data - expected_keys = {"model", "project_id"} - actual_keys = set(data.keys()) - - assert actual_keys == expected_keys, ( - f"Request body should only contain {expected_keys}, " - f"but got {actual_keys}. " - f"Extra keys: {actual_keys - expected_keys}" - ) - - # Specifically verify response_format is NOT added - assert ( - "response_format" not in data - ), "response_format should NOT be added by default" - - # Verify file is sent separately - files = captured_request.get("files", {}) - assert "file" in files - - @pytest.mark.asyncio - async def test_watsonx_transcription_only_user_params_sent_with_space_id(self): - """ - Test that only user-specified params are sent in request body to WatsonX. - - LiteLLM should NOT add extra params like response_format if user didn't specify them. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - # Minimal request - only required params - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - space_id="test-space_id-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - data = captured_request.get("data", {}) - - # These are the ONLY keys that should be in data - expected_keys = {"model", "space_id"} - actual_keys = set(data.keys()) - - assert actual_keys == expected_keys, ( - f"Request body should only contain {expected_keys}, " - f"but got {actual_keys}. " - f"Extra keys: {actual_keys - expected_keys}" - ) - - # Specifically verify response_format is NOT added - assert ( - "response_format" not in data - ), "response_format should NOT be added by default" - - # Verify file is sent separately - files = captured_request.get("files", {}) - assert "file" in files - - def test_transform_audio_transcription_response_removes_model_field(self): - """ - Test that transform_audio_transcription_response removes the 'model' field - from WatsonX response before creating TranscriptionResponse. - - This test ensures that when WatsonX returns a response with a 'model' field, - it is removed before creating the TranscriptionResponse object, since - TranscriptionResponse doesn't accept a 'model' parameter. - """ - handler = IBMWatsonXAudioTranscriptionConfig() - - # Mock response with 'model' field (as WatsonX may return) - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello, this is a test transcription.", - "model": "whisper-large-v3-turbo", # This field should be removed - "duration": 5.5, - } - mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' - - # This should not raise a TypeError - model field should be removed - result = handler.transform_audio_transcription_response(mock_response) - - # Verify the result is a TranscriptionResponse - assert isinstance(result, TranscriptionResponse) - - # Verify the text is correct - assert result.text == "Hello, this is a test transcription." - - # Verify duration is set via dictionary assignment - assert result["duration"] == 5.5 - - # Verify the model field is NOT in the serialized result - # Check via model_dump() or dict() to ensure it's not in the output - try: - result_dict = result.model_dump() - except AttributeError: - # Fallback for pydantic v1 - result_dict = result.dict() - - # The 'model' field should not be in the result - assert "model" not in result_dict, "Model field should be removed from response" - - def test_transform_audio_transcription_response_without_model_field(self): - """ - Test that transform_audio_transcription_response works correctly - when WatsonX response doesn't include a 'model' field. - """ - handler = IBMWatsonXAudioTranscriptionConfig() - - # Mock response without 'model' field - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello, this is a test transcription.", - "duration": 5.5, - } - mock_response.text = ( - '{"text": "Hello, this is a test transcription.", "duration": 5.5}' - ) - - result = handler.transform_audio_transcription_response(mock_response) - - # Verify the result is a TranscriptionResponse - assert isinstance(result, TranscriptionResponse) - - # Verify the text is correct - assert result.text == "Hello, this is a test transcription." - - # Verify duration is set via dictionary assignment - assert result["duration"] == 5.5 diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py deleted file mode 100644 index 285afffefc0..00000000000 --- a/tests/test_litellm/llms/watsonx/test_watsonx.py +++ /dev/null @@ -1,577 +0,0 @@ -import json - -from typing import Optional -from unittest.mock import Mock, patch - -import pytest - -import litellm -from litellm import completion -from litellm.llms.custom_httpx.http_handler import HTTPHandler - - -@pytest.fixture -def watsonx_chat_completion_call(): - def _call( - model="watsonx/my-test-model", - messages=None, - api_key="test_api_key", - space_id: Optional[str] = None, - headers=None, - client=None, - patch_token_call=True, - ): - if messages is None: - messages = [{"role": "user", "content": "Hello, how are you?"}] - if client is None: - client = HTTPHandler() - - if patch_token_call: - mock_response = Mock() - mock_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_response.raise_for_status = Mock() # No-op to simulate no exception - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_response - ) as mock_get, - ): - try: - completion( - model=model, - messages=messages, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - - return mock_post, mock_get - else: - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - return mock_post, None - - return _call - - -def test_watsonx_deployment_model_id_not_in_payload( - monkeypatch, watsonx_chat_completion_call -): - """Test that deployment models do not include 'model_id' in the request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx/deployment/test-deployment-id" - messages = [{"role": "user", "content": "Test message"}] - - mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is not in the payload for deployment models - assert "model_id" not in json_data or json_data["model_id"] is None - # Ensure project_id is also not in the payload for deployment models - assert "project_id" not in json_data or json_data["project_id"] is None - - -def test_watsonx_regular_model_includes_model_id( - monkeypatch, watsonx_chat_completion_call -): - """Test that regular models include 'model_id' in the request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx/regular-model" - messages = [{"role": "user", "content": "Test message"}] - - mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is included in the payload for regular models - assert "model_id" in json_data - assert json_data["model_id"] == "regular-model" # Provider prefix is stripped - # Ensure project_id is also included for regular models - assert "project_id" in json_data - - -@pytest.fixture -def watsonx_completion_call(): - def _call( - model="watsonx_text/my-test-model", - prompt="Hello, how are you?", - api_key="test_api_key", - space_id: Optional[str] = None, - headers=None, - client=None, - patch_token_call=True, - ): - if client is None: - client = HTTPHandler() - - if patch_token_call: - mock_response = Mock() - mock_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_response.raise_for_status = Mock() - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_response - ) as mock_get, - ): - try: - litellm.text_completion( - model=model, - prompt=prompt, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - - return mock_post, mock_get - else: - with patch.object(client, "post") as mock_post: - try: - litellm.text_completion( - model=model, - prompt=prompt, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - return mock_post, None - - return _call - - -def test_watsonx_completion_deployment_model_id_not_in_payload( - monkeypatch, watsonx_completion_call -): - """Test that deployment models do not include 'model_id' in completion request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx_text/deployment/test-deployment-id" - prompt = "Test prompt" - - mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is not in the payload for deployment models - assert "model_id" not in json_data - # Ensure project_id is also not in the payload for deployment models - assert "project_id" not in json_data - - -def test_watsonx_completion_regular_model_includes_model_id( - monkeypatch, watsonx_completion_call -): - """Test that regular models include 'model_id' in completion request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx_text/regular-model" - prompt = "Test prompt" - - mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is included in the payload for regular models - assert "model_id" in json_data - assert json_data["model_id"] == "regular-model" # Provider prefix is stripped - # Ensure project_id is also included for regular models - assert "project_id" in json_data - - -def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): - """ - Test that gpt-oss-120b model transforms messages to proper format instead of simple concatenation. - - This test calls litellm.completion (sync) and verifies what gets sent in the final POST request body. - Input messages should be transformed using the HuggingFace chat template from openai/gpt-oss-120b, - not just concatenated as "You are chatgpt Hi there". - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - # Test with gpt-oss model using watsonx_text provider (text generation endpoint) - model = "watsonx_text/openai/gpt-oss-120b" - - # Input messages - messages = [ - {"role": "system", "content": "You are chatgpt"}, - {"role": "user", "content": "Hi there"}, - ] - - client = HTTPHandler() - - # Mock HuggingFace template fetch to make test deterministic and avoid network flakiness. - # The test verifies that prompt transformation occurs (not simple concatenation), not the exact - # HuggingFace template format. Using a mock template that produces the correct format is sufficient. - # - # Mock template that produces gpt-oss-120b-like format. - # Note: This is a simplified version of the actual template. The real template is more complex - # (adds metadata, handles tools, thinking messages, etc.), but this captures the key aspects: - # - Converts system role to developer (matching real template behavior) - # - Uses the same tag structure (<|start|>, <|message|>, <|end|>) - # - Preserves message content - mock_tokenizer_config = { - "status": "success", - "tokenizer": { - "chat_template": "{% for message in messages %}{% if message['role'] == 'system' %}<|start|>developer<|message|>{% else %}<|start|>{{ message['role'] }}<|message|>{% endif %}{{ message['content'] }}<|end|>{% endfor %}", - "bos_token": None, - "eos_token": None, - }, - } - - # Isolate known_tokenizer_config so parallel tests don't interfere. - # monkeypatch.setitem restores the original value on teardown. - hf_model = "openai/gpt-oss-120b" - monkeypatch.setitem(litellm.known_tokenizer_config, hf_model, mock_tokenizer_config) - - # Mock IAM token generation to avoid real HTTP calls. - mock_token_response = Mock() - mock_token_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_token_response.raise_for_status = Mock() - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_token_response - ), - ): - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the POST was called - assert ( - mock_post.call_count == 1 - ), f"POST should have been called exactly once, got {mock_post.call_count}" - - # Get the request body - call_args = mock_post.call_args - assert "data" in call_args.kwargs, "call_args.kwargs should contain 'data'" - json_data = json.loads(call_args.kwargs["data"]) - - # Verify the transformed input is in the request - assert "input" in json_data, "Request should have 'input' field" - transformed_prompt = json_data["input"] - - # Verify it's NOT simple concatenation - simple_concat = "You are chatgpt Hi there" - assert transformed_prompt != simple_concat, ( - f"Prompt should not be simple concatenation.\n" - f"Expected: Chat template with <|start|> tags\n" - f"Got: {transformed_prompt}" - ) - - # Verify it contains proper chat template formatting - assert "<|start|>" in transformed_prompt, "Prompt should contain <|start|> tag" - assert "<|message|>" in transformed_prompt, "Prompt should contain <|message|> tag" - assert "<|end|>" in transformed_prompt, "Prompt should contain <|end|> tag" - assert ( - "You are chatgpt" in transformed_prompt - ), "Prompt should contain system message content" - assert ( - "Hi there" in transformed_prompt - ), "Prompt should contain user message content" - - -@pytest.mark.asyncio -@pytest.mark.xdist_group("watsonx_heavy") -async def test_watsonx_gpt_oss_uses_async_http_handler(): - """ - Test that verifies async HTTP client is used when fetching HuggingFace templates. - """ - from unittest.mock import AsyncMock, MagicMock, patch - - from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import ( - _aget_chat_template_file, - ) - - # Mock the async HTTP client - mock_async_client = MagicMock() - mock_get = AsyncMock() - mock_async_client.get = mock_get - - # Create mock response for chat template file - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.content = b"test template content" - mock_get.return_value = mock_response - - # Test the async function directly - with patch( - "litellm.litellm_core_utils.prompt_templates.huggingface_template_handler.get_async_httpx_client", - return_value=mock_async_client, - ): - result = await _aget_chat_template_file(hf_model_name="test/model") - - # Verify async HTTP client was called - assert mock_get.called, "Async HTTP client's get method should be called" - assert mock_get.await_count > 0, "Async HTTP client's get should be awaited" - - # Verify it was called with HuggingFace URL - call_args = mock_get.call_args - assert call_args is not None, "get should have been called with arguments" - called_url = call_args.kwargs.get("url", "") - assert ( - "huggingface.co/test/model" in called_url - ), f"Should call HuggingFace API for test/model, got: {called_url}" - assert result["status"] == "success", "Should return success status" - - -@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) -async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( - monkeypatch, tokenizer_config_cached -): - import httpx - - from litellm._uuid import uuid - from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - - hf_model = f"openai/gpt-oss-{uuid.uuid4()}" - chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" - if tokenizer_config_cached: - cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} - monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) - expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" - else: - monkeypatch.setattr(litellm, "known_tokenizer_config", {}) - expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" - hf_fetched = [] - captured = {} - - def forbid_sync_client(): - raise AssertionError("sync HuggingFace fetch ran on the request path") - - async def serve_hf_file(url, **kwargs): - hf_fetched.append(url) - if url.endswith(".jinja"): - return httpx.Response(200, content=chat_template.encode()) - return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) - - monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) - monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) - - def handle(request): - captured["body"] = json.loads(request.content) - return httpx.Response( - 200, - json={ - "model_id": hf_model, - "results": [ - { - "generated_text": "Hi", - "generated_token_count": 1, - "input_token_count": 1, - "stop_reason": "eos_token", - } - ], - }, - ) - - client = AsyncHTTPHandler() - client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) - - response = await litellm.acompletion( - model=f"watsonx_text/{hf_model}", - messages=[{"role": "user", "content": "Hi there"}], - api_base="https://test-api.watsonx.ai", - project_id="test-project-id", - token="test-token", - client=client, - ) - - assert response.choices[0].message.content == "Hi" - assert hf_fetched == [expected_fetch] - assert captured["body"]["input"] == "<|user|>Hi there" - - -def test_watsonx_chat_completion_with_reasoning_effort(monkeypatch): - """ - Test that 'reasoning_effort' is correctly passed through to the WatsonX API payload. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - model = "watsonx/openai/gpt-oss-120b" - messages = [{"role": "user", "content": "Test message"}] - - client = HTTPHandler() - - # Mock the token generation call - mock_token_response = Mock() - mock_token_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_token_response.raise_for_status = Mock() - - # Call litellm.completion with the new parameter - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_token_response - ), - ): - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - reasoning_effort="low", - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the parameter is in the final request payload - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the JSON data sent in the POST request - request_kwargs = mock_post.call_args.kwargs - json_data = json.loads(request_kwargs["data"]) - - print("\nRequest payload sent to WatsonX API:") - print(json.dumps(json_data, indent=2)) - - # Check for the parameter at the top level of the payload - assert ( - "reasoning_effort" in json_data - ), "'reasoning_effort' should be at the top level of the payload." - assert ( - json_data["reasoning_effort"] == "low" - ), "The value of 'reasoning_effort' should be 'low'." - - -def test_watsonx_zen_api_key_from_client(monkeypatch, watsonx_chat_completion_call): - """ - Test that zen_api_key can be passed from client code and is used in Authorization header. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - model = "watsonx/ibm/granite-3-3-8b-instruct" - messages = [{"role": "user", "content": "What is your favorite color?"}] - - client = HTTPHandler() - - zen_api_key = "U1ZDLWQo=" - - # No need to patch token call since zen_api_key should skip token generation - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - zen_api_key=zen_api_key, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the request was made - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the headers sent in the POST request - request_kwargs = mock_post.call_args.kwargs - headers = request_kwargs["headers"] - - print("\nHeaders sent to WatsonX API:") - print(json.dumps(dict(headers), indent=2)) - - # Verify Authorization header uses ZenApiKey format - assert "Authorization" in headers, "Authorization header should be present." - assert headers["Authorization"] == f"ZenApiKey {zen_api_key}", ( - f"Authorization header should use ZenApiKey format. " - f"Expected: 'ZenApiKey {zen_api_key}', Got: '{headers['Authorization']}'" - ) - - -def test_watsonx_zen_api_key_from_env(monkeypatch, watsonx_chat_completion_call): - """ - Test that zen_api_key from environment variable is used in Authorization header. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - zen_api_key = "U1ZDLWxpdG--===" - monkeypatch.setenv("WATSONX_ZENAPIKEY", zen_api_key) - - model = "watsonx/ibm/granite-3-3-8b-instruct" - messages = [{"role": "user", "content": "What is your favorite color?"}] - - client = HTTPHandler() - - # No need to patch token call since zen_api_key should skip token generation - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the request was made - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the headers sent in the POST request - request_kwargs = mock_post.call_args.kwargs - headers = request_kwargs["headers"] - - print("\nHeaders sent to WatsonX API:") - print(json.dumps(dict(headers), indent=2)) - - # Verify Authorization header uses ZenApiKey format - assert "Authorization" in headers, "Authorization header should be present." - assert headers["Authorization"] == f"ZenApiKey {zen_api_key}", ( - f"Authorization header should use ZenApiKey format. " - f"Expected: 'ZenApiKey {zen_api_key}', Got: '{headers['Authorization']}'" - ) diff --git a/tests/test_litellm/llms/xai/xai_responses/__init__.py b/tests/test_litellm/llms/xai/xai_responses/__init__.py deleted file mode 100644 index 330e9f5a560..00000000000 --- a/tests/test_litellm/llms/xai/xai_responses/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# XAI Responses API tests diff --git a/tests/test_litellm/llms/xai/xai_responses/test_transformation.py b/tests/test_litellm/llms/xai/xai_responses/test_transformation.py deleted file mode 100644 index 3ea3fe631bd..00000000000 --- a/tests/test_litellm/llms/xai/xai_responses/test_transformation.py +++ /dev/null @@ -1,105 +0,0 @@ -""" -Tests for XAI Responses API transformation - -Tests the XAIResponsesAPIConfig class that handles XAI-specific -transformations for the Responses API. - -Source: litellm/llms/xai/responses/transformation.py -""" - - - -import pytest -from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager -from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams - - -class TestXAIResponsesAPITransformation: - """Test XAI Responses API configuration and transformations""" - - def test_xai_provider_config_registration(self): - """Test that XAI provider returns XAIResponsesAPIConfig""" - config = ProviderConfigManager.get_provider_responses_api_config( - model="xai/grok-4-fast", - provider=LlmProviders.XAI, - ) - - assert config is not None, "Config should not be None for XAI provider" - assert isinstance( - config, XAIResponsesAPIConfig - ), f"Expected XAIResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.XAI - ), "custom_llm_provider should be XAI" - - def test_code_interpreter_container_field_removed(self): - """Test that container field is removed from code_interpreter tools""" - config = XAIResponsesAPIConfig() - - params = ResponsesAPIOptionalRequestParams( - tools=[{"type": "code_interpreter", "container": {"type": "auto"}}] - ) - - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) - - assert "tools" in result - assert len(result["tools"]) == 1 - assert result["tools"][0]["type"] == "code_interpreter" - assert ( - "container" not in result["tools"][0] - ), "Container field should be removed" - - def test_instructions_parameter_forwarded(self): - """xAI supports 'instructions' on /v1/responses, so it must survive param mapping""" - config = XAIResponsesAPIConfig() - - params = ResponsesAPIOptionalRequestParams( - instructions="You are a helpful assistant.", temperature=0.7 - ) - - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) - - assert result.get("instructions") == "You are a helpful assistant." - assert result.get("temperature") == 0.7, "Other params should be preserved" - - def test_supported_params_includes_instructions(self): - """A system message bridged to 'instructions' must not be rejected for xAI""" - config = XAIResponsesAPIConfig() - supported = config.get_supported_openai_params("grok-4-fast") - - assert "instructions" in supported, "instructions should be supported" - assert "tools" in supported, "tools should be supported" - assert "temperature" in supported, "temperature should be supported" - assert "model" in supported, "model should be supported" - - def test_xai_responses_endpoint_url(self): - """Test that get_complete_url returns correct XAI endpoint""" - config = XAIResponsesAPIConfig() - - # Test with default XAI API base - url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.x.ai/v1/responses" - ), f"Expected XAI responses endpoint, got {url}" - - # Test with custom api_base - custom_url = config.get_complete_url( - api_base="https://custom.x.ai/v1", litellm_params={} - ) - assert ( - custom_url == "https://custom.x.ai/v1/responses" - ), f"Expected custom endpoint, got {custom_url}" - - # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.x.ai/v1/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.x.ai/v1/responses" - ), "Should handle trailing slash" diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py similarity index 98% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py index 4b710175a48..e2fb81bc240 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py @@ -98,6 +98,8 @@ class TestCountTokensLocationResolution: self, counter, monkeypatch ): """Claude models without any location should default to us-east5.""" + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) captured = {} async def fake_ensure_access_token( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py similarity index 60% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index f7df4507651..15df8e47af3 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -1,6 +1,21 @@ +import pytest + import litellm +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_reasoning_effort_stays_unsupported_on_vertex_partner_models(local_model_cost_map): assert "reasoning_effort" in litellm.get_supported_openai_params( model="mistral-medium-3", custom_llm_provider="mistral" diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py rename to tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py diff --git a/tests/unit/llms/vertex_ai/videos/__init__.py b/tests/unit/llms/vertex_ai/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py rename to tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py similarity index 94% rename from tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py rename to tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py index d42bf7b7a1c..5c8d67ecc70 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -137,30 +137,6 @@ class TestVolcengineResponsesAPITransformation: with pytest.raises(ValueError, match='Volcengine API key is required\\. Set ARK_API_KEY /'): config.validate_environment(headers={}, model="volcengine/demo", litellm_params={}) - def test_unsupported_params_are_dropped_with_extra_body(self): - """Unknown fields (including extra_body) should be dropped before send.""" - config = VolcEngineResponsesAPIConfig() - - request = config.transform_responses_api_request( - model="volcengine/demo-model", - input="hi", - response_api_optional_request_params={ - "unsupported_custom_param": 0.1, - "temperature": 0.2, - "metadata": {"k": "v"}, - "extra_body": {"unsupported_custom_param": 1, "temperature": 0.3}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - - assert "unsupported_custom_param" not in request - assert "metadata" not in request - assert request["temperature"] == 0.2 - assert "extra_body" in request - assert "unsupported_custom_param" not in request["extra_body"] - assert request["extra_body"]["temperature"] == 0.3 - def test_valid_thinking_caching_and_expire_at_pass(self): """Documented params should pass through without validation errors.""" config = VolcEngineResponsesAPIConfig() diff --git a/tests/unit/llms/voyage/rerank/__init__.py b/tests/unit/llms/voyage/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py rename to tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py b/tests/unit/llms/voyage/test_voyage_contextual_embedding.py similarity index 100% rename from tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py rename to tests/unit/llms/voyage/test_voyage_contextual_embedding.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py b/tests/unit/llms/voyage/test_voyage_multimodal_embedding.py similarity index 100% rename from tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py rename to tests/unit/llms/voyage/test_voyage_multimodal_embedding.py diff --git a/tests/unit/llms/watsonx/__init__.py b/tests/unit/llms/watsonx/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/audio_transcription/__init__.py b/tests/unit/llms/watsonx/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py new file mode 100644 index 00000000000..efe592f515e --- /dev/null +++ b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py @@ -0,0 +1,85 @@ +""" +Tests for IBM WatsonX Audio Transcription. + +Validates the WatsonX transcription response transformation. +""" + +from unittest.mock import MagicMock + +from litellm.llms.watsonx.audio_transcription.transformation import ( + IBMWatsonXAudioTranscriptionConfig, +) +from litellm.types.utils import TranscriptionResponse + + +class TestWatsonXAudioTranscription: + def test_transform_audio_transcription_response_removes_model_field(self): + """ + Test that transform_audio_transcription_response removes the 'model' field + from WatsonX response before creating TranscriptionResponse. + + This test ensures that when WatsonX returns a response with a 'model' field, + it is removed before creating the TranscriptionResponse object, since + TranscriptionResponse doesn't accept a 'model' parameter. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response with 'model' field (as WatsonX may return) + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "model": "whisper-large-v3-turbo", # This field should be removed + "duration": 5.5, + } + mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' + + # This should not raise a TypeError - model field should be removed + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 + + # Verify the model field is NOT in the serialized result + # Check via model_dump() or dict() to ensure it's not in the output + try: + result_dict = result.model_dump() + except AttributeError: + # Fallback for pydantic v1 + result_dict = result.dict() + + # The 'model' field should not be in the result + assert "model" not in result_dict, "Model field should be removed from response" + + def test_transform_audio_transcription_response_without_model_field(self): + """ + Test that transform_audio_transcription_response works correctly + when WatsonX response doesn't include a 'model' field. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response without 'model' field + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "duration": 5.5, + } + mock_response.text = ( + '{"text": "Hello, this is a test transcription.", "duration": 5.5}' + ) + + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 diff --git a/tests/test_litellm/llms/watsonx/embed/test_watsonx_embedding_transformation.py b/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/watsonx/embed/test_watsonx_embedding_transformation.py rename to tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py diff --git a/tests/test_litellm/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py b/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py rename to tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py diff --git a/tests/unit/llms/watsonx/rerank/__init__.py b/tests/unit/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py similarity index 100% rename from tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py rename to tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py diff --git a/tests/unit/llms/watsonx/test_watsonx.py b/tests/unit/llms/watsonx/test_watsonx.py new file mode 100644 index 00000000000..077539c9acd --- /dev/null +++ b/tests/unit/llms/watsonx/test_watsonx.py @@ -0,0 +1,74 @@ +import json +from unittest.mock import Mock + +import pytest + +import litellm + + +@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) +async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( + monkeypatch, tokenizer_config_cached +): + import httpx + + from litellm._uuid import uuid + from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + hf_model = f"openai/gpt-oss-{uuid.uuid4()}" + chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" + if tokenizer_config_cached: + cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} + monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" + else: + monkeypatch.setattr(litellm, "known_tokenizer_config", {}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" + hf_fetched = [] + captured = {} + + def forbid_sync_client(): + raise AssertionError("sync HuggingFace fetch ran on the request path") + + async def serve_hf_file(url, **kwargs): + hf_fetched.append(url) + if url.endswith(".jinja"): + return httpx.Response(200, content=chat_template.encode()) + return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) + + monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) + monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "model_id": hf_model, + "results": [ + { + "generated_text": "Hi", + "generated_token_count": 1, + "input_token_count": 1, + "stop_reason": "eos_token", + } + ], + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model=f"watsonx_text/{hf_model}", + messages=[{"role": "user", "content": "Hi there"}], + api_base="https://test-api.watsonx.ai", + project_id="test-project-id", + token="test-token", + client=client, + ) + + assert response.choices[0].message.content == "Hi" + assert hf_fetched == [expected_fetch] + assert captured["body"]["input"] == "<|user|>Hi there" diff --git a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py b/tests/unit/llms/watsonx/test_watsonx_common_utils.py similarity index 100% rename from tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py rename to tests/unit/llms/watsonx/test_watsonx_common_utils.py diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/unit/llms/xai/responses/test_xai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py rename to tests/unit/llms/xai/responses/test_xai_responses_transformation.py diff --git a/tests/unit/llms/you_com/__init__.py b/tests/unit/llms/you_com/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/you_com/test_you_com_search.py b/tests/unit/llms/you_com/test_you_com_search.py similarity index 100% rename from tests/test_litellm/llms/you_com/test_you_com_search.py rename to tests/unit/llms/you_com/test_you_com_search.py diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/unit/llms/zai/test_zai_provider.py similarity index 100% rename from tests/test_litellm/llms/zai/test_zai_provider.py rename to tests/unit/llms/zai/test_zai_provider.py From ef6237b9f2bcafba3e7e535a283dd542d67adc70 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:02:53 +0000 Subject: [PATCH 065/146] test: migrate phase 12 legacy llm provider tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++++++++ ...est_openrouter_responses_transformation.py | 0 .../parallel_ai/test_parallel_ai_search.py | 0 .../test_parallel_ai_search_gateway.py | 0 .../llms/parasail/test_parasail.py | 0 .../test_perplexity_chat_transformation.py | 0 ...est_perplexity_embedding_transformation.py | 0 ...est_perplexity_responses_transformation.py | 0 .../test_publicai_chat_transformation.py | 0 .../chat/test_ragflow_chat_transformation.py | 0 .../test_recraft_image_edit_transformation.py | 0 .../test_recraft_image_gen_transformation.py | 0 .../test_text_to_speech_transformation.py | 0 .../test_runway_video_transformation.py | 0 .../test_s3_vectors_transformation.py | 4 ---- .../llms/sap/test_sap_fetch_creds.py | 0 ...eway_audio_transcription_transformation.py | 0 .../test_snowflake_native_endpoints.py | 0 .../test_soniox_provider_registration.py | 0 .../test_stability_image_generation.py | 19 +------------------ .../chat/test_tencent_chat_transformation.py | 0 21 files changed, 15 insertions(+), 22 deletions(-) rename tests/{test_litellm => unit}/llms/openrouter/responses/test_openrouter_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/parallel_ai/test_parallel_ai_search.py (100%) rename tests/{test_litellm => unit}/llms/parallel_ai/test_parallel_ai_search_gateway.py (100%) rename tests/{test_litellm => unit}/llms/parasail/test_parasail.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/chat/test_perplexity_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/embedding/test_perplexity_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/responses/test_perplexity_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/publicai/test_publicai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/ragflow/chat/test_ragflow_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/recraft/image_edit/test_recraft_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/recraft/image_generation/test_recraft_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/runwayml/test_text_to_speech_transformation.py (100%) rename tests/{test_litellm => unit}/llms/runwayml/videos/test_runway_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py (99%) rename tests/{test_litellm => unit}/llms/sap/test_sap_fetch_creds.py (100%) rename tests/{test_litellm => unit}/llms/scaleway/test_scaleway_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/snowflake/test_snowflake_native_endpoints.py (100%) rename tests/{test_litellm => unit}/llms/soniox/test_soniox_provider_registration.py (100%) rename tests/{test_litellm => unit}/llms/stability/image_generation/test_stability_image_generation.py (93%) rename tests/{test_litellm => unit}/llms/tencent/chat/test_tencent_chat_transformation.py (100%) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 017e63ed1b8..fa91437964f 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,6 +1,7 @@ import os from typing import Final +import litellm import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -9,6 +10,19 @@ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) diff --git a/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py rename to tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py rename to tests/unit/llms/parallel_ai/test_parallel_ai_search.py diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search_gateway.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py rename to tests/unit/llms/parallel_ai/test_parallel_ai_search_gateway.py diff --git a/tests/test_litellm/llms/parasail/test_parasail.py b/tests/unit/llms/parasail/test_parasail.py similarity index 100% rename from tests/test_litellm/llms/parasail/test_parasail.py rename to tests/unit/llms/parasail/test_parasail.py diff --git a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py rename to tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py rename to tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py diff --git a/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py rename to tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py diff --git a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py b/tests/unit/llms/publicai/test_publicai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py rename to tests/unit/llms/publicai/test_publicai_chat_transformation.py diff --git a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py b/tests/unit/llms/ragflow/chat/test_ragflow_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py rename to tests/unit/llms/ragflow/chat/test_ragflow_chat_transformation.py diff --git a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py rename to tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py diff --git a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py rename to tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py diff --git a/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py b/tests/unit/llms/runwayml/test_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py rename to tests/unit/llms/runwayml/test_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py rename to tests/unit/llms/runwayml/videos/test_runway_video_transformation.py diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py similarity index 99% rename from tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py rename to tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 781e92ea7d9..c39887d86ce 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -55,10 +55,6 @@ def _search_kwargs(**overrides): class TestS3VectorsVectorStoreConfig: - def test_init(self): - config = S3VectorsVectorStoreConfig() - assert config is not None - def test_get_supported_openai_params(self): config = S3VectorsVectorStoreConfig() params = config.get_supported_openai_params("test-model") diff --git a/tests/test_litellm/llms/sap/test_sap_fetch_creds.py b/tests/unit/llms/sap/test_sap_fetch_creds.py similarity index 100% rename from tests/test_litellm/llms/sap/test_sap_fetch_creds.py rename to tests/unit/llms/sap/test_sap_fetch_creds.py diff --git a/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py b/tests/unit/llms/scaleway/test_scaleway_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py rename to tests/unit/llms/scaleway/test_scaleway_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py similarity index 100% rename from tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py rename to tests/unit/llms/snowflake/test_snowflake_native_endpoints.py diff --git a/tests/test_litellm/llms/soniox/test_soniox_provider_registration.py b/tests/unit/llms/soniox/test_soniox_provider_registration.py similarity index 100% rename from tests/test_litellm/llms/soniox/test_soniox_provider_registration.py rename to tests/unit/llms/soniox/test_soniox_provider_registration.py diff --git a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py similarity index 93% rename from tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py rename to tests/unit/llms/stability/image_generation/test_stability_image_generation.py index c5b3c8fbdc5..c5a78603f9c 100644 --- a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py +++ b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py @@ -10,10 +10,7 @@ from unittest.mock import MagicMock import httpx import pytest -from litellm.llms.stability.image_generation import ( - StabilityImageGenerationConfig, - get_stability_image_generation_config, -) +from litellm.llms.stability.image_generation import StabilityImageGenerationConfig from litellm.types.llms.stability import ( OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, STABILITY_GENERATION_MODELS, @@ -266,20 +263,6 @@ class TestStabilityImageGenerationConfig: assert "filtered" in str(exc_info.value).lower() -class TestFactoryFunction: - """Test the factory function""" - - def test_get_stability_image_generation_config(self): - """Test that factory returns correct config type""" - config = get_stability_image_generation_config("stability/sd3") - assert isinstance(config, StabilityImageGenerationConfig) - - def test_factory_returns_config_for_any_model(self): - """Test that factory works for any model name""" - config = get_stability_image_generation_config("stability/custom-model") - assert isinstance(config, StabilityImageGenerationConfig) - - class TestOpenAISizeMapping: """Test the size to aspect ratio mapping""" diff --git a/tests/test_litellm/llms/tencent/chat/test_tencent_chat_transformation.py b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/tencent/chat/test_tencent_chat_transformation.py rename to tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py From a3dd47ea11b420018b5082e431b1f0ea611a6fad Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:05:03 +0000 Subject: [PATCH 066/146] test(unit): clear ambient Azure credentials in entra token tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++++++++ .../azure/realtime/test_azure_realtime_handler.py | 2 +- .../test_azure_ai_image_edit_transformation.py | 4 ++-- .../test_mai_image_edit_transformation.py | 2 +- .../test_azure_ai_passthrough_transformation.py | 4 ++-- .../rerank/test_azure_ai_rerank_transformation.py | 2 +- 6 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 6b6f7c43760..b3bb19a8b8a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -12,6 +12,14 @@ import litellm.router as litellm_router_module # noqa: E402 # same import-time import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", +) def _allow_loopback_only() -> None: @@ -53,5 +61,11 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +@pytest.fixture +def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + for name in AMBIENT_AZURE_CREDENTIAL_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + def pytest_sessionfinish() -> None: enable_socket() diff --git a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index c1ba286f8c0..73f43ec8d8a 100644 --- a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -707,7 +707,7 @@ async def test_realtime_health_check_uses_bearer_token_when_no_api_key(monkeypat @pytest.mark.asyncio -async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch): +async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch, no_ambient_azure_credentials): """ The router binds a deployment's `azure_ad_token` to `_arealtime`'s named parameter rather than **kwargs, so it must still reach the handler. diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 39001c1795b..51ba2c34cd7 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -41,7 +41,7 @@ def test_azure_ai_url_generation(): assert complete_url == expected_url -def test_azure_ai_validate_environment_with_entra_token(monkeypatch): +def test_azure_ai_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) config = AzureFoundryFluxImageEditConfig() @@ -55,7 +55,7 @@ def test_azure_ai_validate_environment_with_entra_token(monkeypatch): assert headers == {"Authorization": "Bearer entra-token"} -def test_flux2_validate_environment_with_entra_token(monkeypatch): +def test_flux2_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) config = AzureFoundryFlux2ImageEditConfig() diff --git a/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 75e046825a3..2d6f0083194 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -174,7 +174,7 @@ class TestAzureMAIImageEdit: assert image_response.usage.total_tokens == 1024 -def test_mai_validate_environment_with_entra_token(monkeypatch): +def test_mai_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) diff --git a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index f00698a6624..f9fd9681db8 100644 --- a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -256,13 +256,13 @@ def test_serverless_host_gets_a_bearer_token(): assert "api-key" not in headers -def test_entra_token_is_used_when_the_deployment_has_no_api_key(): +def test_entra_token_is_used_when_the_deployment_has_no_api_key(no_ambient_azure_credentials): headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) assert headers["Authorization"] == "Bearer entra-token" -def test_no_credentials_at_all_raises(): +def test_no_credentials_at_all_raises(no_ambient_azure_credentials): with pytest.raises(ValueError, match="Missing Azure AI credentials"): _auth_headers(api_key=None, api_base=FOUNDRY_BASE) diff --git a/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index 91bf665f18d..3de27199e2e 100644 --- a/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -105,7 +105,7 @@ class TestAzureAIRerankConfigValidateEnvironment: assert headers["Authorization"] == "Bearer my-key" - def test_falls_back_to_entra_token(self, monkeypatch): + def test_falls_back_to_entra_token(self, monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "azure_key", None) From e0b92b2b257bd26650a5887343cfc0de3127129d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:07:38 +0000 Subject: [PATCH 067/146] refactor(types): keep a2a response_dict annotation as it was Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/a2a_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index faf3e98a3a7..834c16ba6dc 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -880,7 +880,7 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None _enqueue_fn() - response_dict: Final[dict[str, object]] = ( + response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response From 6ba4c2e3399bc16a9996c379c020d1b726f8247a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:40:11 +0000 Subject: [PATCH 068/146] refactor(types): keep guardrail metadata helper accepting dicts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/base_llm/guardrail_translation/base_translation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ace5af8124f..89ad67f0485 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -81,7 +81,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Optional["UserAPIKeyAuth"], + user_api_key_dict: Any | None, ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. From e4a58ef91acdebe2a25c66f7dbdf7d4bdc338fd7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:50:59 +0000 Subject: [PATCH 069/146] test(unit): make every tests/unit directory a package so pytest collection is unique Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/integrations/__init__.py | 0 tests/unit/integrations/levo/__init__.py | 0 tests/unit/integrations/litellm_agent/__init__.py | 0 tests/unit/integrations/mavvrik_focus/__init__.py | 0 tests/unit/integrations/opik/__init__.py | 0 tests/unit/integrations/pointfive/__init__.py | 0 .../vector_store_integrations/__init__.py | 0 tests/unit/litellm_core_utils/__init__.py | 0 .../unit/litellm_core_utils/audio_utils/__init__.py | 0 .../llm_response_utils/__init__.py | 0 tests/unit/llms/__init__.py | 0 tests/unit/llms/a2a/__init__.py | 0 tests/unit/llms/a2a/chat/__init__.py | 0 .../llms/a2a/chat/guardrail_translation/__init__.py | 0 tests/unit/llms/anthropic/__init__.py | 0 tests/unit/llms/anthropic/batches/__init__.py | 0 tests/unit/llms/base_llm/__init__.py | 0 tests/unit/llms/base_llm/batches/__init__.py | 0 tests/unit/llms/base_llm/realtime/__init__.py | 0 tests/unit/llms/baseten/__init__.py | 0 tests/unit/llms/baseten/chat/__init__.py | 0 tests/unit/llms/bedrock/__init__.py | 0 tests/unit/llms/bedrock/chat/__init__.py | 0 tests/unit/llms/bedrock/chat/agentcore/__init__.py | 0 .../bedrock/chat/invoke_transformations/__init__.py | 0 tests/unit/llms/bedrock/chat/mantle/__init__.py | 0 tests/unit/llms/bedrock/count_tokens/__init__.py | 0 tests/unit/llms/bedrock/files/__init__.py | 0 tests/unit/llms/bedrock/image/__init__.py | 0 tests/unit/llms/bedrock/image_edit/__init__.py | 0 tests/unit/llms/bedrock/invoke_agent/__init__.py | 0 tests/unit/llms/bedrock/passthrough/__init__.py | 0 .../passthrough/guardrail_translation/__init__.py | 0 tests/unit/llms/bedrock/realtime/__init__.py | 0 tests/unit/llms/bedrock/rerank/__init__.py | 0 tests/unit/llms/bedrock/vector_stores/__init__.py | 0 tests/unit/llms/bedrock_mantle/__init__.py | 0 .../unit/llms/bedrock_mantle/passthrough/__init__.py | 0 tests/unit/llms/black_forest_labs/__init__.py | 0 .../llms/black_forest_labs/image_edit/__init__.py | 0 .../black_forest_labs/image_generation/__init__.py | 0 tests/unit/llms/bytez/__init__.py | 0 tests/unit/llms/bytez/chat/__init__.py | 0 tests/unit/llms/cerebras/__init__.py | 0 tests/unit/llms/chat/__init__.py | 0 tests/unit/llms/chatgpt/__init__.py | 0 tests/unit/llms/chatgpt/chat/__init__.py | 0 tests/unit/llms/chatgpt/responses/__init__.py | 0 tests/unit/llms/cloudflare/__init__.py | 0 tests/unit/llms/cohere/__init__.py | 0 tests/unit/llms/cohere/chat/__init__.py | 0 tests/unit/llms/cohere/embed/__init__.py | 0 tests/unit/llms/cohere/ocr/__init__.py | 0 tests/unit/llms/cohere/rerank/__init__.py | 0 tests/unit/llms/crusoe/__init__.py | 0 tests/unit/llms/databricks/__init__.py | 0 tests/unit/llms/databricks/chat/__init__.py | 0 tests/unit/llms/databricks/responses/__init__.py | 0 tests/unit/llms/datarobot/__init__.py | 0 tests/unit/llms/datarobot/chat/__init__.py | 0 tests/unit/llms/deepseek/__init__.py | 0 tests/unit/llms/deepseek/chat/__init__.py | 0 tests/unit/llms/deepseek/messages/__init__.py | 0 tests/unit/llms/docker_model_runner/__init__.py | 0 tests/unit/llms/elevenlabs/__init__.py | 0 tests/unit/llms/fastcrw/__init__.py | 0 tests/unit/llms/fastcrw/search/__init__.py | 0 tests/unit/llms/fireworks_ai/__init__.py | 0 tests/unit/llms/fireworks_ai/chat/__init__.py | 0 tests/unit/llms/fireworks_ai/rerank/__init__.py | 0 tests/unit/llms/fireworks_ai/responses/__init__.py | 0 tests/unit/llms/gemini/__init__.py | 0 .../unit/llms/gemini/audio_transcription/__init__.py | 0 tests/unit/llms/gemini/files/__init__.py | 0 tests/unit/llms/gemini/google_genai/__init__.py | 0 .../google_genai/guardrail_translation/__init__.py | 0 tests/unit/llms/gemini/image_edit/__init__.py | 0 tests/unit/llms/gemini/realtime/__init__.py | 0 tests/unit/llms/gemini/videos/__init__.py | 0 tests/unit/llms/gigachat/__init__.py | 0 tests/unit/llms/gigachat/chat/__init__.py | 0 tests/unit/llms/gigachat/embedding/__init__.py | 0 tests/unit/llms/gigachat/passthrough/__init__.py | 0 tests/unit/llms/github_copilot/__init__.py | 0 tests/unit/llms/github_copilot/embedding/__init__.py | 0 tests/unit/llms/github_copilot/messages/__init__.py | 0 tests/unit/llms/github_copilot/responses/__init__.py | 0 tests/unit/llms/gradient_ai/__init__.py | 0 tests/unit/llms/gradient_ai/chat/__init__.py | 0 tests/unit/llms/groq/__init__.py | 0 tests/unit/llms/groq/chat/__init__.py | 0 tests/unit/llms/hosted_vllm/__init__.py | 0 tests/unit/llms/hosted_vllm/chat/__init__.py | 0 tests/unit/llms/hosted_vllm/embedding/__init__.py | 0 tests/unit/llms/hosted_vllm/image_edit/__init__.py | 0 tests/unit/llms/hosted_vllm/responses/__init__.py | 0 tests/unit/llms/hosted_vllm/videos/__init__.py | 0 tests/unit/llms/huggingface/__init__.py | 0 tests/unit/llms/huggingface/rerank/__init__.py | 0 tests/unit/llms/inception/__init__.py | 0 tests/unit/llms/jina_ai/__init__.py | 0 tests/unit/llms/jina_ai/embedding/__init__.py | 0 tests/unit/llms/langflow/__init__.py | 0 tests/unit/llms/langflow/chat/__init__.py | 0 tests/unit/llms/litellm_proxy/__init__.py | 0 tests/unit/llms/litellm_proxy/chat/__init__.py | 0 tests/unit/llms/litellm_proxy/skills/__init__.py | 0 tests/unit/llms/llamafile/__init__.py | 0 tests/unit/llms/llamafile/chat/__init__.py | 0 tests/unit/llms/meta/__init__.py | 0 tests/unit/llms/meta/realtime/__init__.py | 0 tests/unit/llms/meta_llama/__init__.py | 0 tests/unit/llms/mistral/audio_speech/__init__.py | 0 tests/unit/llms/modelscope/__init__.py | 0 .../llms/modelscope/image_generation/__init__.py | 0 tests/unit/llms/mongodb/__init__.py | 0 tests/unit/llms/mongodb/vector_stores/__init__.py | 0 tests/unit/llms/moonshot/__init__.py | 0 tests/unit/llms/neosantara/__init__.py | 0 tests/unit/llms/nimble/__init__.py | 0 tests/unit/llms/nimble/search/__init__.py | 0 tests/unit/llms/novita/__init__.py | 0 tests/unit/llms/novita/chat/__init__.py | 0 tests/unit/llms/nscale/__init__.py | 0 tests/unit/llms/nscale/chat/__init__.py | 0 tests/unit/llms/nvidia_nim/__init__.py | 0 tests/unit/llms/nvidia_nim/passthrough/__init__.py | 0 tests/unit/llms/nvidia_nim/rerank/__init__.py | 0 tests/unit/llms/nvidia_riva/__init__.py | 0 .../llms/nvidia_riva/audio_transcription/__init__.py | 0 tests/unit/llms/oci/__init__.py | 0 tests/unit/llms/oci/chat/__init__.py | 0 tests/unit/llms/oci/embed/__init__.py | 0 tests/unit/llms/ocr/__init__.py | 0 .../unit/llms/ocr/guardrail_translation/__init__.py | 0 tests/unit/llms/oobabooga/__init__.py | 0 tests/unit/llms/oobabooga/chat/__init__.py | 0 tests/unit/llms/openai/__init__.py | 0 tests/unit/llms/openai/chat/__init__.py | 0 .../openai/chat/guardrail_translation/__init__.py | 0 tests/unit/llms/openai/completion/__init__.py | 0 tests/unit/test_package_layout.py | 12 ++++++++++++ 143 files changed, 12 insertions(+) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/integrations/__init__.py create mode 100644 tests/unit/integrations/levo/__init__.py create mode 100644 tests/unit/integrations/litellm_agent/__init__.py create mode 100644 tests/unit/integrations/mavvrik_focus/__init__.py create mode 100644 tests/unit/integrations/opik/__init__.py create mode 100644 tests/unit/integrations/pointfive/__init__.py create mode 100644 tests/unit/integrations/vector_store_integrations/__init__.py create mode 100644 tests/unit/litellm_core_utils/__init__.py create mode 100644 tests/unit/litellm_core_utils/audio_utils/__init__.py create mode 100644 tests/unit/litellm_core_utils/llm_response_utils/__init__.py create mode 100644 tests/unit/llms/__init__.py create mode 100644 tests/unit/llms/a2a/__init__.py create mode 100644 tests/unit/llms/a2a/chat/__init__.py create mode 100644 tests/unit/llms/a2a/chat/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/anthropic/__init__.py create mode 100644 tests/unit/llms/anthropic/batches/__init__.py create mode 100644 tests/unit/llms/base_llm/__init__.py create mode 100644 tests/unit/llms/base_llm/batches/__init__.py create mode 100644 tests/unit/llms/base_llm/realtime/__init__.py create mode 100644 tests/unit/llms/baseten/__init__.py create mode 100644 tests/unit/llms/baseten/chat/__init__.py create mode 100644 tests/unit/llms/bedrock/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/agentcore/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/mantle/__init__.py create mode 100644 tests/unit/llms/bedrock/count_tokens/__init__.py create mode 100644 tests/unit/llms/bedrock/files/__init__.py create mode 100644 tests/unit/llms/bedrock/image/__init__.py create mode 100644 tests/unit/llms/bedrock/image_edit/__init__.py create mode 100644 tests/unit/llms/bedrock/invoke_agent/__init__.py create mode 100644 tests/unit/llms/bedrock/passthrough/__init__.py create mode 100644 tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/bedrock/realtime/__init__.py create mode 100644 tests/unit/llms/bedrock/rerank/__init__.py create mode 100644 tests/unit/llms/bedrock/vector_stores/__init__.py create mode 100644 tests/unit/llms/bedrock_mantle/__init__.py create mode 100644 tests/unit/llms/bedrock_mantle/passthrough/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/image_edit/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/image_generation/__init__.py create mode 100644 tests/unit/llms/bytez/__init__.py create mode 100644 tests/unit/llms/bytez/chat/__init__.py create mode 100644 tests/unit/llms/cerebras/__init__.py create mode 100644 tests/unit/llms/chat/__init__.py create mode 100644 tests/unit/llms/chatgpt/__init__.py create mode 100644 tests/unit/llms/chatgpt/chat/__init__.py create mode 100644 tests/unit/llms/chatgpt/responses/__init__.py create mode 100644 tests/unit/llms/cloudflare/__init__.py create mode 100644 tests/unit/llms/cohere/__init__.py create mode 100644 tests/unit/llms/cohere/chat/__init__.py create mode 100644 tests/unit/llms/cohere/embed/__init__.py create mode 100644 tests/unit/llms/cohere/ocr/__init__.py create mode 100644 tests/unit/llms/cohere/rerank/__init__.py create mode 100644 tests/unit/llms/crusoe/__init__.py create mode 100644 tests/unit/llms/databricks/__init__.py create mode 100644 tests/unit/llms/databricks/chat/__init__.py create mode 100644 tests/unit/llms/databricks/responses/__init__.py create mode 100644 tests/unit/llms/datarobot/__init__.py create mode 100644 tests/unit/llms/datarobot/chat/__init__.py create mode 100644 tests/unit/llms/deepseek/__init__.py create mode 100644 tests/unit/llms/deepseek/chat/__init__.py create mode 100644 tests/unit/llms/deepseek/messages/__init__.py create mode 100644 tests/unit/llms/docker_model_runner/__init__.py create mode 100644 tests/unit/llms/elevenlabs/__init__.py create mode 100644 tests/unit/llms/fastcrw/__init__.py create mode 100644 tests/unit/llms/fastcrw/search/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/chat/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/rerank/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/responses/__init__.py create mode 100644 tests/unit/llms/gemini/__init__.py create mode 100644 tests/unit/llms/gemini/audio_transcription/__init__.py create mode 100644 tests/unit/llms/gemini/files/__init__.py create mode 100644 tests/unit/llms/gemini/google_genai/__init__.py create mode 100644 tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/gemini/image_edit/__init__.py create mode 100644 tests/unit/llms/gemini/realtime/__init__.py create mode 100644 tests/unit/llms/gemini/videos/__init__.py create mode 100644 tests/unit/llms/gigachat/__init__.py create mode 100644 tests/unit/llms/gigachat/chat/__init__.py create mode 100644 tests/unit/llms/gigachat/embedding/__init__.py create mode 100644 tests/unit/llms/gigachat/passthrough/__init__.py create mode 100644 tests/unit/llms/github_copilot/__init__.py create mode 100644 tests/unit/llms/github_copilot/embedding/__init__.py create mode 100644 tests/unit/llms/github_copilot/messages/__init__.py create mode 100644 tests/unit/llms/github_copilot/responses/__init__.py create mode 100644 tests/unit/llms/gradient_ai/__init__.py create mode 100644 tests/unit/llms/gradient_ai/chat/__init__.py create mode 100644 tests/unit/llms/groq/__init__.py create mode 100644 tests/unit/llms/groq/chat/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/chat/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/embedding/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/image_edit/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/responses/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/videos/__init__.py create mode 100644 tests/unit/llms/huggingface/__init__.py create mode 100644 tests/unit/llms/huggingface/rerank/__init__.py create mode 100644 tests/unit/llms/inception/__init__.py create mode 100644 tests/unit/llms/jina_ai/__init__.py create mode 100644 tests/unit/llms/jina_ai/embedding/__init__.py create mode 100644 tests/unit/llms/langflow/__init__.py create mode 100644 tests/unit/llms/langflow/chat/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/chat/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/skills/__init__.py create mode 100644 tests/unit/llms/llamafile/__init__.py create mode 100644 tests/unit/llms/llamafile/chat/__init__.py create mode 100644 tests/unit/llms/meta/__init__.py create mode 100644 tests/unit/llms/meta/realtime/__init__.py create mode 100644 tests/unit/llms/meta_llama/__init__.py create mode 100644 tests/unit/llms/mistral/audio_speech/__init__.py create mode 100644 tests/unit/llms/modelscope/__init__.py create mode 100644 tests/unit/llms/modelscope/image_generation/__init__.py create mode 100644 tests/unit/llms/mongodb/__init__.py create mode 100644 tests/unit/llms/mongodb/vector_stores/__init__.py create mode 100644 tests/unit/llms/moonshot/__init__.py create mode 100644 tests/unit/llms/neosantara/__init__.py create mode 100644 tests/unit/llms/nimble/__init__.py create mode 100644 tests/unit/llms/nimble/search/__init__.py create mode 100644 tests/unit/llms/novita/__init__.py create mode 100644 tests/unit/llms/novita/chat/__init__.py create mode 100644 tests/unit/llms/nscale/__init__.py create mode 100644 tests/unit/llms/nscale/chat/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/passthrough/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/rerank/__init__.py create mode 100644 tests/unit/llms/nvidia_riva/__init__.py create mode 100644 tests/unit/llms/nvidia_riva/audio_transcription/__init__.py create mode 100644 tests/unit/llms/oci/__init__.py create mode 100644 tests/unit/llms/oci/chat/__init__.py create mode 100644 tests/unit/llms/oci/embed/__init__.py create mode 100644 tests/unit/llms/ocr/__init__.py create mode 100644 tests/unit/llms/ocr/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/oobabooga/__init__.py create mode 100644 tests/unit/llms/oobabooga/chat/__init__.py create mode 100644 tests/unit/llms/openai/__init__.py create mode 100644 tests/unit/llms/openai/chat/__init__.py create mode 100644 tests/unit/llms/openai/chat/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/completion/__init__.py create mode 100644 tests/unit/test_package_layout.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/__init__.py b/tests/unit/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/levo/__init__.py b/tests/unit/integrations/levo/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/litellm_agent/__init__.py b/tests/unit/integrations/litellm_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/mavvrik_focus/__init__.py b/tests/unit/integrations/mavvrik_focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/opik/__init__.py b/tests/unit/integrations/opik/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/pointfive/__init__.py b/tests/unit/integrations/pointfive/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/vector_store_integrations/__init__.py b/tests/unit/integrations/vector_store_integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/__init__.py b/tests/unit/litellm_core_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/audio_utils/__init__.py b/tests/unit/litellm_core_utils/audio_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/llm_response_utils/__init__.py b/tests/unit/litellm_core_utils/llm_response_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/__init__.py b/tests/unit/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/__init__.py b/tests/unit/llms/a2a/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/chat/__init__.py b/tests/unit/llms/a2a/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/chat/guardrail_translation/__init__.py b/tests/unit/llms/a2a/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/__init__.py b/tests/unit/llms/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/batches/__init__.py b/tests/unit/llms/anthropic/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/__init__.py b/tests/unit/llms/base_llm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/batches/__init__.py b/tests/unit/llms/base_llm/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/realtime/__init__.py b/tests/unit/llms/base_llm/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/baseten/__init__.py b/tests/unit/llms/baseten/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/baseten/chat/__init__.py b/tests/unit/llms/baseten/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/__init__.py b/tests/unit/llms/bedrock/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/__init__.py b/tests/unit/llms/bedrock/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/agentcore/__init__.py b/tests/unit/llms/bedrock/chat/agentcore/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py b/tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/mantle/__init__.py b/tests/unit/llms/bedrock/chat/mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/count_tokens/__init__.py b/tests/unit/llms/bedrock/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/files/__init__.py b/tests/unit/llms/bedrock/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image/__init__.py b/tests/unit/llms/bedrock/image/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image_edit/__init__.py b/tests/unit/llms/bedrock/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/invoke_agent/__init__.py b/tests/unit/llms/bedrock/invoke_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/passthrough/__init__.py b/tests/unit/llms/bedrock/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/realtime/__init__.py b/tests/unit/llms/bedrock/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/rerank/__init__.py b/tests/unit/llms/bedrock/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/vector_stores/__init__.py b/tests/unit/llms/bedrock/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock_mantle/__init__.py b/tests/unit/llms/bedrock_mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock_mantle/passthrough/__init__.py b/tests/unit/llms/bedrock_mantle/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/__init__.py b/tests/unit/llms/black_forest_labs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/image_edit/__init__.py b/tests/unit/llms/black_forest_labs/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/image_generation/__init__.py b/tests/unit/llms/black_forest_labs/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bytez/__init__.py b/tests/unit/llms/bytez/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bytez/chat/__init__.py b/tests/unit/llms/bytez/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cerebras/__init__.py b/tests/unit/llms/cerebras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chat/__init__.py b/tests/unit/llms/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/__init__.py b/tests/unit/llms/chatgpt/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/chat/__init__.py b/tests/unit/llms/chatgpt/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/responses/__init__.py b/tests/unit/llms/chatgpt/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cloudflare/__init__.py b/tests/unit/llms/cloudflare/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/__init__.py b/tests/unit/llms/cohere/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/chat/__init__.py b/tests/unit/llms/cohere/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/embed/__init__.py b/tests/unit/llms/cohere/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/ocr/__init__.py b/tests/unit/llms/cohere/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/rerank/__init__.py b/tests/unit/llms/cohere/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/crusoe/__init__.py b/tests/unit/llms/crusoe/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/__init__.py b/tests/unit/llms/databricks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/chat/__init__.py b/tests/unit/llms/databricks/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/responses/__init__.py b/tests/unit/llms/databricks/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/datarobot/__init__.py b/tests/unit/llms/datarobot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/datarobot/chat/__init__.py b/tests/unit/llms/datarobot/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/__init__.py b/tests/unit/llms/deepseek/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/chat/__init__.py b/tests/unit/llms/deepseek/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/messages/__init__.py b/tests/unit/llms/deepseek/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/docker_model_runner/__init__.py b/tests/unit/llms/docker_model_runner/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/elevenlabs/__init__.py b/tests/unit/llms/elevenlabs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fastcrw/__init__.py b/tests/unit/llms/fastcrw/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fastcrw/search/__init__.py b/tests/unit/llms/fastcrw/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/__init__.py b/tests/unit/llms/fireworks_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/chat/__init__.py b/tests/unit/llms/fireworks_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/rerank/__init__.py b/tests/unit/llms/fireworks_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/responses/__init__.py b/tests/unit/llms/fireworks_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/__init__.py b/tests/unit/llms/gemini/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/audio_transcription/__init__.py b/tests/unit/llms/gemini/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/files/__init__.py b/tests/unit/llms/gemini/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/google_genai/__init__.py b/tests/unit/llms/gemini/google_genai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/image_edit/__init__.py b/tests/unit/llms/gemini/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/realtime/__init__.py b/tests/unit/llms/gemini/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/videos/__init__.py b/tests/unit/llms/gemini/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/__init__.py b/tests/unit/llms/gigachat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/chat/__init__.py b/tests/unit/llms/gigachat/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/embedding/__init__.py b/tests/unit/llms/gigachat/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/passthrough/__init__.py b/tests/unit/llms/gigachat/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/__init__.py b/tests/unit/llms/github_copilot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/embedding/__init__.py b/tests/unit/llms/github_copilot/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/messages/__init__.py b/tests/unit/llms/github_copilot/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/responses/__init__.py b/tests/unit/llms/github_copilot/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gradient_ai/__init__.py b/tests/unit/llms/gradient_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gradient_ai/chat/__init__.py b/tests/unit/llms/gradient_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/groq/__init__.py b/tests/unit/llms/groq/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/groq/chat/__init__.py b/tests/unit/llms/groq/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/__init__.py b/tests/unit/llms/hosted_vllm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/chat/__init__.py b/tests/unit/llms/hosted_vllm/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/embedding/__init__.py b/tests/unit/llms/hosted_vllm/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/image_edit/__init__.py b/tests/unit/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/responses/__init__.py b/tests/unit/llms/hosted_vllm/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/videos/__init__.py b/tests/unit/llms/hosted_vllm/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/huggingface/__init__.py b/tests/unit/llms/huggingface/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/huggingface/rerank/__init__.py b/tests/unit/llms/huggingface/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/inception/__init__.py b/tests/unit/llms/inception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/__init__.py b/tests/unit/llms/jina_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/embedding/__init__.py b/tests/unit/llms/jina_ai/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/langflow/__init__.py b/tests/unit/llms/langflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/langflow/chat/__init__.py b/tests/unit/llms/langflow/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/__init__.py b/tests/unit/llms/litellm_proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/chat/__init__.py b/tests/unit/llms/litellm_proxy/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/skills/__init__.py b/tests/unit/llms/litellm_proxy/skills/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/llamafile/__init__.py b/tests/unit/llms/llamafile/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/llamafile/chat/__init__.py b/tests/unit/llms/llamafile/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta/__init__.py b/tests/unit/llms/meta/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta/realtime/__init__.py b/tests/unit/llms/meta/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta_llama/__init__.py b/tests/unit/llms/meta_llama/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mistral/audio_speech/__init__.py b/tests/unit/llms/mistral/audio_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/modelscope/__init__.py b/tests/unit/llms/modelscope/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/modelscope/image_generation/__init__.py b/tests/unit/llms/modelscope/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mongodb/__init__.py b/tests/unit/llms/mongodb/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mongodb/vector_stores/__init__.py b/tests/unit/llms/mongodb/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/moonshot/__init__.py b/tests/unit/llms/moonshot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/neosantara/__init__.py b/tests/unit/llms/neosantara/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nimble/__init__.py b/tests/unit/llms/nimble/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nimble/search/__init__.py b/tests/unit/llms/nimble/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/novita/__init__.py b/tests/unit/llms/novita/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/novita/chat/__init__.py b/tests/unit/llms/novita/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nscale/__init__.py b/tests/unit/llms/nscale/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nscale/chat/__init__.py b/tests/unit/llms/nscale/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/__init__.py b/tests/unit/llms/nvidia_nim/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/passthrough/__init__.py b/tests/unit/llms/nvidia_nim/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/rerank/__init__.py b/tests/unit/llms/nvidia_nim/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_riva/__init__.py b/tests/unit/llms/nvidia_riva/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/__init__.py b/tests/unit/llms/nvidia_riva/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/__init__.py b/tests/unit/llms/oci/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/chat/__init__.py b/tests/unit/llms/oci/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/embed/__init__.py b/tests/unit/llms/oci/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ocr/__init__.py b/tests/unit/llms/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ocr/guardrail_translation/__init__.py b/tests/unit/llms/ocr/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oobabooga/__init__.py b/tests/unit/llms/oobabooga/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oobabooga/chat/__init__.py b/tests/unit/llms/oobabooga/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/__init__.py b/tests/unit/llms/openai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/chat/__init__.py b/tests/unit/llms/openai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/chat/guardrail_translation/__init__.py b/tests/unit/llms/openai/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/completion/__init__.py b/tests/unit/llms/openai/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/test_package_layout.py b/tests/unit/test_package_layout.py new file mode 100644 index 00000000000..4ea68fc06ba --- /dev/null +++ b/tests/unit/test_package_layout.py @@ -0,0 +1,12 @@ +import os + +TESTS_UNIT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def test_every_directory_under_tests_unit_is_a_package(): + missing = [] + for root, dirs, _files in os.walk(TESTS_UNIT_DIR): + dirs[:] = [d for d in dirs if d != "__pycache__"] + if not os.path.isfile(os.path.join(root, "__init__.py")): + missing.append(os.path.relpath(root, TESTS_UNIT_DIR)) + assert missing == [] From 5599c59923c5362dccc5c17012ba0c02a2d7254f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:07 +0000 Subject: [PATCH 070/146] test: add package initializers to migrated unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openai/embeddings/__init__.py | 0 .../unit/llms/openai/embeddings/guardrail_translation/__init__.py | 0 tests/unit/llms/openai/evals/__init__.py | 0 tests/unit/llms/openai/image_generation/__init__.py | 0 tests/unit/llms/openai/speech/__init__.py | 0 tests/unit/llms/openai/transcriptions/__init__.py | 0 tests/unit/llms/openai/vector_store_files/__init__.py | 0 tests/unit/llms/openai/vector_stores/__init__.py | 0 tests/unit/llms/openai/videos/__init__.py | 0 tests/unit/llms/openai_like/__init__.py | 0 tests/unit/llms/openai_like/chat/__init__.py | 0 tests/unit/llms/openai_like/embedding/__init__.py | 0 tests/unit/llms/openai_like/messages/__init__.py | 0 tests/unit/llms/openrouter/__init__.py | 0 tests/unit/llms/openrouter/chat/__init__.py | 0 tests/unit/llms/openrouter/image_edit/__init__.py | 0 tests/unit/llms/openrouter/image_generation/__init__.py | 0 17 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openai/embeddings/__init__.py create mode 100644 tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/evals/__init__.py create mode 100644 tests/unit/llms/openai/image_generation/__init__.py create mode 100644 tests/unit/llms/openai/speech/__init__.py create mode 100644 tests/unit/llms/openai/transcriptions/__init__.py create mode 100644 tests/unit/llms/openai/vector_store_files/__init__.py create mode 100644 tests/unit/llms/openai/vector_stores/__init__.py create mode 100644 tests/unit/llms/openai/videos/__init__.py create mode 100644 tests/unit/llms/openai_like/__init__.py create mode 100644 tests/unit/llms/openai_like/chat/__init__.py create mode 100644 tests/unit/llms/openai_like/embedding/__init__.py create mode 100644 tests/unit/llms/openai_like/messages/__init__.py create mode 100644 tests/unit/llms/openrouter/__init__.py create mode 100644 tests/unit/llms/openrouter/chat/__init__.py create mode 100644 tests/unit/llms/openrouter/image_edit/__init__.py create mode 100644 tests/unit/llms/openrouter/image_generation/__init__.py diff --git a/tests/unit/llms/openai/embeddings/__init__.py b/tests/unit/llms/openai/embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py b/tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/evals/__init__.py b/tests/unit/llms/openai/evals/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/image_generation/__init__.py b/tests/unit/llms/openai/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/speech/__init__.py b/tests/unit/llms/openai/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/transcriptions/__init__.py b/tests/unit/llms/openai/transcriptions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/vector_store_files/__init__.py b/tests/unit/llms/openai/vector_store_files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/vector_stores/__init__.py b/tests/unit/llms/openai/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/videos/__init__.py b/tests/unit/llms/openai/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/__init__.py b/tests/unit/llms/openai_like/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/chat/__init__.py b/tests/unit/llms/openai_like/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/embedding/__init__.py b/tests/unit/llms/openai_like/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/messages/__init__.py b/tests/unit/llms/openai_like/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/__init__.py b/tests/unit/llms/openrouter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/chat/__init__.py b/tests/unit/llms/openrouter/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/image_edit/__init__.py b/tests/unit/llms/openrouter/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/image_generation/__init__.py b/tests/unit/llms/openrouter/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 49cd32affe1cfd597c223cdb62049d52ed55e5b5 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:09 +0000 Subject: [PATCH 071/146] test: add __init__.py to every tests/unit directory this migration touches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/vertex_ai/__init__.py | 0 tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py | 0 .../llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py | 0 .../llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py | 0 tests/unit/llms/volcengine/__init__.py | 0 tests/unit/llms/volcengine/responses/__init__.py | 0 tests/unit/llms/voyage/__init__.py | 0 tests/unit/llms/watsonx/embed/__init__.py | 0 tests/unit/llms/watsonx/passthrough/__init__.py | 0 tests/unit/llms/xai/__init__.py | 0 tests/unit/llms/xai/responses/__init__.py | 0 tests/unit/llms/zai/__init__.py | 0 12 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py create mode 100644 tests/unit/llms/volcengine/__init__.py create mode 100644 tests/unit/llms/volcengine/responses/__init__.py create mode 100644 tests/unit/llms/voyage/__init__.py create mode 100644 tests/unit/llms/watsonx/embed/__init__.py create mode 100644 tests/unit/llms/watsonx/passthrough/__init__.py create mode 100644 tests/unit/llms/xai/__init__.py create mode 100644 tests/unit/llms/xai/responses/__init__.py create mode 100644 tests/unit/llms/zai/__init__.py diff --git a/tests/unit/llms/vertex_ai/__init__.py b/tests/unit/llms/vertex_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/volcengine/__init__.py b/tests/unit/llms/volcengine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/volcengine/responses/__init__.py b/tests/unit/llms/volcengine/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/voyage/__init__.py b/tests/unit/llms/voyage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/embed/__init__.py b/tests/unit/llms/watsonx/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/passthrough/__init__.py b/tests/unit/llms/watsonx/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/__init__.py b/tests/unit/llms/xai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/responses/__init__.py b/tests/unit/llms/xai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/zai/__init__.py b/tests/unit/llms/zai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 434af084e7572d2b2f23101f56d7c722189aa33d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:38 +0000 Subject: [PATCH 072/146] test: add __init__.py to every tests/unit directory phase 16 touches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/sandbox/__init__.py | 0 tests/unit/skills/__init__.py | 0 tests/unit/test_router/__init__.py | 0 tests/unit/types/__init__.py | 0 tests/unit/types/llms/__init__.py | 0 tests/unit/types/proxy/__init__.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/sandbox/__init__.py create mode 100644 tests/unit/skills/__init__.py create mode 100644 tests/unit/test_router/__init__.py create mode 100644 tests/unit/types/__init__.py create mode 100644 tests/unit/types/llms/__init__.py create mode 100644 tests/unit/types/proxy/__init__.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/sandbox/__init__.py b/tests/unit/sandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/skills/__init__.py b/tests/unit/skills/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/test_router/__init__.py b/tests/unit/test_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/__init__.py b/tests/unit/types/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/llms/__init__.py b/tests/unit/types/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/__init__.py b/tests/unit/types/proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From baf40ea5e896674dd8e75d6ac8b1937fae07e475 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:53:14 +0000 Subject: [PATCH 073/146] test(unit): add package markers to migrated unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/llms/__init__.py | 0 tests/unit/llms/anthropic/__init__.py | 0 tests/unit/llms/anthropic/experimental_pass_through/__init__.py | 0 tests/unit/llms/anthropic/files/__init__.py | 0 tests/unit/llms/anthropic/messages/__init__.py | 0 tests/unit/llms/apiserpent/__init__.py | 0 tests/unit/llms/azure/__init__.py | 0 tests/unit/llms/azure/image_edit/__init__.py | 0 tests/unit/llms/azure/image_generation/__init__.py | 0 tests/unit/llms/azure/passthrough/__init__.py | 0 tests/unit/llms/azure/realtime/__init__.py | 0 tests/unit/llms/azure/response/__init__.py | 0 tests/unit/llms/azure/search/__init__.py | 0 tests/unit/llms/azure/text_to_speech/__init__.py | 0 tests/unit/llms/azure/vector_stores/__init__.py | 0 tests/unit/llms/azure_ai/__init__.py | 0 tests/unit/llms/azure_ai/chat/__init__.py | 0 tests/unit/llms/azure_ai/embed/__init__.py | 0 tests/unit/llms/azure_ai/image_edit/__init__.py | 0 tests/unit/llms/azure_ai/ocr/__init__.py | 0 tests/unit/llms/azure_ai/passthrough/__init__.py | 0 tests/unit/llms/azure_ai/rerank/__init__.py | 0 tests/unit/llms/azure_ai/responses/__init__.py | 0 24 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/llms/__init__.py create mode 100644 tests/unit/llms/anthropic/__init__.py create mode 100644 tests/unit/llms/anthropic/experimental_pass_through/__init__.py create mode 100644 tests/unit/llms/anthropic/files/__init__.py create mode 100644 tests/unit/llms/anthropic/messages/__init__.py create mode 100644 tests/unit/llms/apiserpent/__init__.py create mode 100644 tests/unit/llms/azure/__init__.py create mode 100644 tests/unit/llms/azure/image_edit/__init__.py create mode 100644 tests/unit/llms/azure/image_generation/__init__.py create mode 100644 tests/unit/llms/azure/passthrough/__init__.py create mode 100644 tests/unit/llms/azure/realtime/__init__.py create mode 100644 tests/unit/llms/azure/response/__init__.py create mode 100644 tests/unit/llms/azure/search/__init__.py create mode 100644 tests/unit/llms/azure/text_to_speech/__init__.py create mode 100644 tests/unit/llms/azure/vector_stores/__init__.py create mode 100644 tests/unit/llms/azure_ai/__init__.py create mode 100644 tests/unit/llms/azure_ai/chat/__init__.py create mode 100644 tests/unit/llms/azure_ai/embed/__init__.py create mode 100644 tests/unit/llms/azure_ai/image_edit/__init__.py create mode 100644 tests/unit/llms/azure_ai/ocr/__init__.py create mode 100644 tests/unit/llms/azure_ai/passthrough/__init__.py create mode 100644 tests/unit/llms/azure_ai/rerank/__init__.py create mode 100644 tests/unit/llms/azure_ai/responses/__init__.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/__init__.py b/tests/unit/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/__init__.py b/tests/unit/llms/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/experimental_pass_through/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/files/__init__.py b/tests/unit/llms/anthropic/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/messages/__init__.py b/tests/unit/llms/anthropic/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/apiserpent/__init__.py b/tests/unit/llms/apiserpent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/__init__.py b/tests/unit/llms/azure/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/image_edit/__init__.py b/tests/unit/llms/azure/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/image_generation/__init__.py b/tests/unit/llms/azure/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/passthrough/__init__.py b/tests/unit/llms/azure/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/realtime/__init__.py b/tests/unit/llms/azure/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/response/__init__.py b/tests/unit/llms/azure/response/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/search/__init__.py b/tests/unit/llms/azure/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/text_to_speech/__init__.py b/tests/unit/llms/azure/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/vector_stores/__init__.py b/tests/unit/llms/azure/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/__init__.py b/tests/unit/llms/azure_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/chat/__init__.py b/tests/unit/llms/azure_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/embed/__init__.py b/tests/unit/llms/azure_ai/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/image_edit/__init__.py b/tests/unit/llms/azure_ai/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/ocr/__init__.py b/tests/unit/llms/azure_ai/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/passthrough/__init__.py b/tests/unit/llms/azure_ai/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/rerank/__init__.py b/tests/unit/llms/azure_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/responses/__init__.py b/tests/unit/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 30a422087c9bf28cc3fc94a6d68b241b703325a1 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:53:49 +0000 Subject: [PATCH 074/146] test: add __init__.py to migrated tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/a2a_protocol/__init__.py | 0 tests/unit/a2a_protocol/providers/__init__.py | 0 tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py | 0 tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py | 0 tests/unit/anthropic_interface/__init__.py | 0 tests/unit/anthropic_interface/exceptions/__init__.py | 0 tests/unit/batches/__init__.py | 0 tests/unit/chat_completions/__init__.py | 0 tests/unit/completion_extras/__init__.py | 0 tests/unit/compression/__init__.py | 0 tests/unit/endpoints/__init__.py | 0 tests/unit/endpoints/speech/__init__.py | 0 .../unit/endpoints/speech/speech_to_completion_bridge/__init__.py | 0 tests/unit/enterprise/__init__.py | 0 tests/unit/enterprise/enterprise_callbacks/__init__.py | 0 tests/unit/integrations/__init__.py | 0 tests/unit/integrations/compression_interception/__init__.py | 0 tests/unit/integrations/gcs_bucket/__init__.py | 0 tests/unit/integrations/gcs_pubsub/__init__.py | 0 tests/unit/integrations/helicone/__init__.py | 0 20 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/a2a_protocol/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py create mode 100644 tests/unit/anthropic_interface/__init__.py create mode 100644 tests/unit/anthropic_interface/exceptions/__init__.py create mode 100644 tests/unit/batches/__init__.py create mode 100644 tests/unit/chat_completions/__init__.py create mode 100644 tests/unit/completion_extras/__init__.py create mode 100644 tests/unit/compression/__init__.py create mode 100644 tests/unit/endpoints/__init__.py create mode 100644 tests/unit/endpoints/speech/__init__.py create mode 100644 tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py create mode 100644 tests/unit/enterprise/__init__.py create mode 100644 tests/unit/enterprise/enterprise_callbacks/__init__.py create mode 100644 tests/unit/integrations/__init__.py create mode 100644 tests/unit/integrations/compression_interception/__init__.py create mode 100644 tests/unit/integrations/gcs_bucket/__init__.py create mode 100644 tests/unit/integrations/gcs_pubsub/__init__.py create mode 100644 tests/unit/integrations/helicone/__init__.py diff --git a/tests/unit/a2a_protocol/__init__.py b/tests/unit/a2a_protocol/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/__init__.py b/tests/unit/a2a_protocol/providers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/anthropic_interface/__init__.py b/tests/unit/anthropic_interface/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/anthropic_interface/exceptions/__init__.py b/tests/unit/anthropic_interface/exceptions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/batches/__init__.py b/tests/unit/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/chat_completions/__init__.py b/tests/unit/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/completion_extras/__init__.py b/tests/unit/completion_extras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/compression/__init__.py b/tests/unit/compression/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/__init__.py b/tests/unit/endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/speech/__init__.py b/tests/unit/endpoints/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/__init__.py b/tests/unit/enterprise/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/enterprise_callbacks/__init__.py b/tests/unit/enterprise/enterprise_callbacks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/__init__.py b/tests/unit/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/compression_interception/__init__.py b/tests/unit/integrations/compression_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/gcs_bucket/__init__.py b/tests/unit/integrations/gcs_bucket/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/gcs_pubsub/__init__.py b/tests/unit/integrations/gcs_pubsub/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/helicone/__init__.py b/tests/unit/integrations/helicone/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 78a751c04972d0079ef01bb87fef58a6e472be65 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:55:43 +0000 Subject: [PATCH 075/146] test: migrate phase 15 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/messages/__init__.py | 0 .../messages/test_dispatch.py | 11 +- .../models/test_models.py | 46 ++--- tests/unit/ocr/__init__.py | 0 .../ocr/test_dispatch.py | 0 tests/{test_litellm => unit}/ocr/test_main.py | 0 .../ocr/test_ocr_file_input.py | 48 ++--- tests/unit/passthrough/__init__.py | 0 .../test_async_streaming_error_propagation.py | 27 ++- .../passthrough/test_passthrough_main.py | 77 ++------ ...test_streaming_interrupt_spend_tracking.py | 31 +--- tests/unit/rag/ingestion/__init__.py | 0 .../ingestion/test_s3_vectors_ingestion.py | 12 +- .../realtime_api/test_main.py | 11 +- .../repositories/test_repositories.py | 48 ++--- .../repositories/test_unit_of_work.py | 0 .../complexity_router/test_jev_classifier.py | 0 .../test_deployment_affinity_check.py | 70 ++++---- .../test_encrypted_content_affinity_check.py | 169 ++++++++++-------- .../test_prompt_caching_deployment_check.py | 47 +++-- .../test_responses_api_deployment_check.py | 26 +-- .../test_session_id_affinity.py | 70 ++++---- tests/unit/rust_bridge/__init__.py | 0 .../rust_bridge/chat_completions/__init__.py | 0 .../chat_completions/test_route_host.py | 0 tests/unit/rust_bridge/messages/__init__.py | 0 .../rust_bridge/messages/test_route_host.py | 0 27 files changed, 290 insertions(+), 403 deletions(-) create mode 100644 tests/unit/messages/__init__.py rename tests/{test_litellm => unit}/messages/test_dispatch.py (97%) rename tests/{test_litellm => unit}/models/test_models.py (93%) create mode 100644 tests/unit/ocr/__init__.py rename tests/{test_litellm => unit}/ocr/test_dispatch.py (100%) rename tests/{test_litellm => unit}/ocr/test_main.py (100%) rename tests/{test_litellm => unit}/ocr/test_ocr_file_input.py (92%) create mode 100644 tests/unit/passthrough/__init__.py rename tests/{test_litellm => unit}/passthrough/test_async_streaming_error_propagation.py (92%) rename tests/{test_litellm => unit}/passthrough/test_passthrough_main.py (94%) rename tests/{test_litellm => unit}/passthrough/test_streaming_interrupt_spend_tracking.py (91%) create mode 100644 tests/unit/rag/ingestion/__init__.py rename tests/{test_litellm => unit}/rag/ingestion/test_s3_vectors_ingestion.py (94%) rename tests/{test_litellm => unit}/realtime_api/test_main.py (98%) rename tests/{test_litellm => unit}/repositories/test_repositories.py (98%) rename tests/{test_litellm => unit}/repositories/test_unit_of_work.py (100%) rename tests/{test_litellm => unit}/router_strategy/complexity_router/test_jev_classifier.py (100%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_deployment_affinity_check.py (95%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py (96%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py (94%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_responses_api_deployment_check.py (95%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_session_id_affinity.py (95%) create mode 100644 tests/unit/rust_bridge/__init__.py create mode 100644 tests/unit/rust_bridge/chat_completions/__init__.py rename tests/{test_litellm => unit}/rust_bridge/chat_completions/test_route_host.py (100%) create mode 100644 tests/unit/rust_bridge/messages/__init__.py rename tests/{test_litellm => unit}/rust_bridge/messages/test_route_host.py (100%) diff --git a/tests/unit/messages/__init__.py b/tests/unit/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py similarity index 97% rename from tests/test_litellm/messages/test_dispatch.py rename to tests/unit/messages/test_dispatch.py index 2eaf4cd9a50..586b77d9a25 100644 --- a/tests/test_litellm/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -29,9 +29,7 @@ RUST_RULES: Final[Rules] = (Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),) def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]: - binding: Final[NativeBinding[NativeMessages]] = NativeBinding( - "anthropic_messages_handler", validate=lambda _: None - ) + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("anthropic_messages_handler", validate=lambda _: None) binding.override(native) return binding @@ -99,7 +97,8 @@ async def test_async_python_route_forwards_original_call_shape() -> None: expected: Final = response() async def python( - *call_args: object, **call_kwargs: object # kwargs-ok: records call shape + *call_args: object, + **call_kwargs: object, # kwargs-ok: records call shape ) -> AnthropicMessagesResponse: captured.append((call_args, call_kwargs)) return expected @@ -217,7 +216,9 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] expected: Final = response() - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call + def python( + *call_args: object, **call_kwargs: object + ) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call captured.append((call_args, call_kwargs)) return expected diff --git a/tests/test_litellm/models/test_models.py b/tests/unit/models/test_models.py similarity index 93% rename from tests/test_litellm/models/test_models.py rename to tests/unit/models/test_models.py index 777b4a265ac..b8bf55f1b4a 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/unit/models/test_models.py @@ -5,7 +5,7 @@ Tests for backend domain models. from datetime import datetime, timezone import pytest -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.models.autorouter_session import LiteLLM_AutoRouterSession @@ -19,7 +19,6 @@ from litellm.models.credentials import CreateCredentialItem, CredentialItem from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.managed_files import ( LiteLLM_ManagedFileTable, - LiteLLM_ManagedObjectTable, LiteLLM_ManagedVectorStoresTable, ) from litellm.models.mcp_server import LiteLLM_MCPServerTable @@ -41,7 +40,6 @@ from litellm.models.verification_token import ( LiteLLM_DeletedVerificationToken, LiteLLM_VerificationToken, ) -from pydantic import ValidationError class TestBudget: @@ -121,9 +119,7 @@ class TestCredentials: assert item.credential_values is None def test_create_credential_item_requires_values_or_model_id(self): - with pytest.raises( - ValueError, match="Either credential_values or model_id must be set" - ): + with pytest.raises(ValueError, match="Either credential_values or model_id must be set"): CreateCredentialItem(credential_name="bad", credential_info={}) @@ -141,12 +137,8 @@ class TestModel: assert model.team_public_model_name == "my-gpt4" def test_is_blocked(self): - model_blocked = LiteLLM_ProxyModelTable( - model_id="m1", model_name="test", litellm_params={}, blocked=True - ) - model_unblocked = LiteLLM_ProxyModelTable( - model_id="m2", model_name="test", litellm_params={}, blocked=False - ) + model_blocked = LiteLLM_ProxyModelTable(model_id="m1", model_name="test", litellm_params={}, blocked=True) + model_unblocked = LiteLLM_ProxyModelTable(model_id="m2", model_name="test", litellm_params={}, blocked=False) assert model_blocked.is_blocked assert not model_unblocked.is_blocked @@ -188,9 +180,7 @@ class TestModel: assert model.blocked is True def test_team_helpers_none_when_no_model_info(self): - model = LiteLLM_ProxyModelTable( - model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None - ) + model = LiteLLM_ProxyModelTable(model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None) assert model.team_id is None assert model.team_public_model_name is None @@ -292,9 +282,7 @@ class TestTeam: assert team.model_max_budget == {"gpt-4": 5.0} def test_cached_team(self): - cached = LiteLLM_TeamTableCachedObj( - team_id="t1", last_refreshed_at=1234567890.0 - ) + cached = LiteLLM_TeamTableCachedObj(team_id="t1", last_refreshed_at=1234567890.0) assert cached.last_refreshed_at == 1234567890.0 def test_deleted_team(self): @@ -345,9 +333,7 @@ class TestUser: assert "password" not in user.model_dump() assert "password" not in user.model_dump_json() - with_keys = LiteLLM_UserTableWithKeyCount( - user_id="u1", user_email="a@b.c", password=secret, key_count=2 - ) + with_keys = LiteLLM_UserTableWithKeyCount(user_id="u1", user_email="a@b.c", password=secret, key_count=2) assert with_keys.password == secret assert "password" not in with_keys.model_dump() assert "password" not in with_keys.model_dump_json() @@ -479,9 +465,7 @@ class TestEndUserTable: class TestBudgetTableFull: def test_full_adds_server_managed_fields(self): now = datetime.now() - budget = LiteLLM_BudgetTableFull( - budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now - ) + budget = LiteLLM_BudgetTableFull(budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now) assert budget.created_at == now assert budget.budget_reset_at == now assert budget.max_budget == 10.0 @@ -493,9 +477,7 @@ class TestBudgetTableFull: class TestTeamMemberTable: def test_tracks_user_within_team(self): - member = LiteLLM_TeamMemberTable( - user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0 - ) + member = LiteLLM_TeamMemberTable(user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0) assert member.user_id == "u1" assert member.team_id == "t1" assert member.spend == 3.0 @@ -585,9 +567,7 @@ class TestSpendLogs: assert log.updated_at == updated_at def test_error_logs_creation(self): - log = LiteLLM_ErrorLogs( - request_id="r1", startTime=None, endTime=None, status_code="500" - ) + log = LiteLLM_ErrorLogs(request_id="r1", startTime=None, endTime=None, status_code="500") assert log.request_id == "r1" assert log.status_code == "500" @@ -603,12 +583,6 @@ class TestManagedTables: assert table.model_mappings == {"gpt-4": "file-abc"} assert table.flat_model_file_ids == ["file-abc"] - def test_managed_object_table_requires_purpose(self): - with pytest.raises(ValidationError): - LiteLLM_ManagedObjectTable( - unified_object_id="o1", model_object_id="m1", file_object={} - ) - def test_managed_vector_stores_table(self): table = LiteLLM_ManagedVectorStoresTable( vector_store_id="vs1", diff --git a/tests/unit/ocr/__init__.py b/tests/unit/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/ocr/test_dispatch.py b/tests/unit/ocr/test_dispatch.py similarity index 100% rename from tests/test_litellm/ocr/test_dispatch.py rename to tests/unit/ocr/test_dispatch.py diff --git a/tests/test_litellm/ocr/test_main.py b/tests/unit/ocr/test_main.py similarity index 100% rename from tests/test_litellm/ocr/test_main.py rename to tests/unit/ocr/test_main.py diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/unit/ocr/test_ocr_file_input.py similarity index 92% rename from tests/test_litellm/ocr/test_ocr_file_input.py rename to tests/unit/ocr/test_ocr_file_input.py index 4ac27d286e1..d67f5280195 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/unit/ocr/test_ocr_file_input.py @@ -73,9 +73,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "document_url" assert result["document_url"].startswith("data:application/pdf;base64,") @@ -95,9 +93,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -112,9 +108,7 @@ class TestConvertFileDocumentToUrlDocument: request handler the value is attacker-controlled, and opening it as a path is an arbitrary local file read on the proxy host.""" with pytest.raises(ValueError, match="does not accept bare str values"): - convert_file_document_to_url_document( - {"type": "file", "file": "/etc/passwd"} - ) + convert_file_document_to_url_document({"type": "file", "file": "/etc/passwd"}) def test_should_convert_pathlib_path(self): """pathlib.Path objects should work the same as string paths.""" @@ -126,9 +120,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "document_url" assert result["document_url"].startswith("data:application/pdf;base64,") @@ -139,9 +131,7 @@ class TestConvertFileDocumentToUrlDocument: """Raw bytes should be converted using a fallback MIME type.""" content = b"raw bytes content" - result = convert_file_document_to_url_document( - {"type": "file", "file": content} - ) + result = convert_file_document_to_url_document({"type": "file", "file": content}) assert result["type"] == "document_url" assert "base64," in result["document_url"] @@ -164,9 +154,7 @@ class TestConvertFileDocumentToUrlDocument: """Raw bytes with an image MIME type should produce type=image_url.""" content = b"raw image content" - result = convert_file_document_to_url_document( - {"type": "file", "file": content, "mime_type": "image/jpeg"} - ) + result = convert_file_document_to_url_document({"type": "file", "file": content, "mime_type": "image/jpeg"}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/jpeg;base64,") @@ -176,9 +164,7 @@ class TestConvertFileDocumentToUrlDocument: content = b"file-like content" file_obj = BytesIO(content) - result = convert_file_document_to_url_document( - {"type": "file", "file": file_obj} - ) + result = convert_file_document_to_url_document({"type": "file", "file": file_obj}) assert result["type"] == "document_url" assert "base64," in result["document_url"] @@ -189,9 +175,7 @@ class TestConvertFileDocumentToUrlDocument: file_obj = BytesIO(content) file_obj.name = "test_image.png" - result = convert_file_document_to_url_document( - {"type": "file", "file": file_obj} - ) + result = convert_file_document_to_url_document({"type": "file", "file": file_obj}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -204,9 +188,7 @@ class TestConvertFileDocumentToUrlDocument: def test_should_raise_error_for_nonexistent_pathlib_path(self): """Non-existent pathlib.Path should raise FileNotFoundError.""" with pytest.raises(FileNotFoundError, match="File not found"): - convert_file_document_to_url_document( - {"type": "file", "file": Path("/nonexistent/path/to/file.pdf")} - ) + convert_file_document_to_url_document({"type": "file", "file": Path("/nonexistent/path/to/file.pdf")}) def test_should_raise_error_for_empty_file(self): """Empty file should raise ValueError.""" @@ -215,9 +197,7 @@ class TestConvertFileDocumentToUrlDocument: try: with pytest.raises(ValueError, match="File is empty"): - convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + convert_file_document_to_url_document({"type": "file", "file": tmp_path}) finally: os.unlink(str(tmp_path)) @@ -248,9 +228,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path, "mime_type": "image/png"} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path, "mime_type": "image/png"}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -477,9 +455,7 @@ class TestProxySecurityGuard: result = await self._parse_multipart(mock_request) assert result["document"]["type"] == "document_url" - assert result["document"]["document_url"].startswith( - "data:application/pdf;base64," - ) + assert result["document"]["document_url"].startswith("data:application/pdf;base64,") assert result["model"] == "mistral/mistral-ocr-latest" diff --git a/tests/unit/passthrough/__init__.py b/tests/unit/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/unit/passthrough/test_async_streaming_error_propagation.py similarity index 92% rename from tests/test_litellm/passthrough/test_async_streaming_error_propagation.py rename to tests/unit/passthrough/test_async_streaming_error_propagation.py index 9f2b436d2d8..cb93183957c 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/unit/passthrough/test_async_streaming_error_propagation.py @@ -21,9 +21,7 @@ def _make_mock_response(status_code: int, body: bytes, headers: dict = None): # def _raise_for_status(): if status_code >= 400: - request = httpx.Request( - "POST", "https://azure.example.com/openai/responses" - ) + request = httpx.Request("POST", "https://azure.example.com/openai/responses") real_response = httpx.Response( status_code=status_code, content=body, @@ -55,16 +53,15 @@ def _make_mock_logging_obj(): async def test_async_streaming_429_raises(): """429 from upstream should raise HTTPStatusError, not yield error bytes.""" from litellm.passthrough.main import AsyncPassthroughStreamingResponse - - error_body = json.dumps( - {"error": {"code": "429", "message": "Rate limit exceeded."}} - ).encode() + + error_body = json.dumps({"error": {"code": "429", "message": "Rate limit exceeded."}}).encode() mock_response = _make_mock_response(429, error_body) - + async def response_coro(): return mock_response - + chunks = [] + async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -84,15 +81,13 @@ async def test_async_streaming_429_raises(): async def test_async_streaming_500_raises(): """500 from upstream should also raise, not yield error bytes.""" from litellm.passthrough.main import AsyncPassthroughStreamingResponse - - error_body = json.dumps( - {"error": {"code": "500", "message": "Internal server error"}} - ).encode() + + error_body = json.dumps({"error": {"code": "500", "message": "Internal server error"}}).encode() mock_response = _make_mock_response(500, error_body) - + async def response_coro(): return mock_response - + with pytest.raises(httpx.HTTPStatusError) as exc_info: async for _ in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -100,7 +95,7 @@ async def test_async_streaming_500_raises(): provider_config=MagicMock(), ): pass - + assert exc_info.value.response.status_code == 500 diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/unit/passthrough/test_passthrough_main.py similarity index 94% rename from tests/test_litellm/passthrough/test_passthrough_main.py rename to tests/unit/passthrough/test_passthrough_main.py index 3f2c434cc00..82825ec2802 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/unit/passthrough/test_passthrough_main.py @@ -3,14 +3,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from fastapi.testclient import TestClient - -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - - - import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.passthrough.main import allm_passthrough_route, llm_passthrough_route @@ -37,10 +32,7 @@ def test_llm_passthrough_route(): client=client, ) - assert ( - mock_post.call_args.kwargs["request"].url - == "http://localhost:8090/v1/chat/completions" - ) + assert mock_post.call_args.kwargs["request"].url == "http://localhost:8090/v1/chat/completions" assert response.status_code == 200 assert response.json == {"message": "Hello, world!"} @@ -74,12 +66,9 @@ def test_bedrock_application_inference_profile_url_encoding(): "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base"), ), - patch.object( - client.client, "send", return_value=MagicMock(status_code=200) - ) as mock_send, + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -132,12 +121,9 @@ def test_bedrock_non_application_inference_profile_no_encoding(): "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base"), ), - patch.object( - client.client, "send", return_value=MagicMock(status_code=200) - ) as mock_send, + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -202,7 +188,6 @@ def test_update_stream_param_based_on_request_body(): @pytest.fixture def mock_request(): """Create a mock request with headers""" - from typing import Optional class QueryParams: def __init__(self): @@ -215,9 +200,7 @@ def mock_request(): return self._dict.items() class MockRequest: - def __init__( - self, headers=None, method="POST", request_body: Optional[dict] = None - ): + def __init__(self, headers=None, method="POST", request_body: dict | None = None): self.headers = headers or {} self.query_params = QueryParams() self.method = method @@ -245,9 +228,7 @@ def mock_user_api_key_dict(): @pytest.mark.asyncio -async def test_pass_through_request_stream_param_override( - mock_request, mock_user_api_key_dict -): +async def test_pass_through_request_stream_param_override(mock_request, mock_user_api_key_dict): """ Test that when stream=None is passed as parameter but stream=True is in request body, the request body value takes precedence and @@ -346,9 +327,7 @@ async def test_pass_through_request_stream_param_override( @pytest.mark.asyncio -async def test_pass_through_request_stream_param_no_override( - mock_request, mock_user_api_key_dict -): +async def test_pass_through_request_stream_param_no_override(mock_request, mock_user_api_key_dict): """ Test that when stream=False is passed as parameter and no stream is in request body, the function parameter is used and @@ -448,15 +427,11 @@ def test_azure_with_custom_api_base_and_key(): # Mock the provider config and its methods mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01" - ), + httpx.URL("https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01"), "https://my-custom-base", ) mock_provider_config.get_api_key.return_value = "my-custom-key" - mock_provider_config.validate_environment.return_value = { - "api-key": "my-custom-key" - } + mock_provider_config.validate_environment.return_value = {"api-key": "my-custom-key"} mock_provider_config.sign_request.return_value = ( {"api-key": "my-custom-key"}, None, @@ -484,13 +459,10 @@ def test_azure_with_custom_api_base_and_key(): patch.object( client.client, "send", - return_value=MagicMock( - status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []} - ), - ) as mock_send, + return_value=MagicMock(status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []}), + ), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -541,9 +513,7 @@ def test_content_param_forwarded_to_build_request(): mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-azure.openai.azure.com/openai/deployments/gpt-4/chat/completions" - ), + httpx.URL("https://my-azure.openai.azure.com/openai/deployments/gpt-4/chat/completions"), "https://my-azure.openai.azure.com", ) mock_provider_config.get_api_key.return_value = "test-key" @@ -575,7 +545,6 @@ def test_content_param_forwarded_to_build_request(): patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -656,15 +625,11 @@ async def test_allm_passthrough_route_429_streaming_raises(): """ mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-azure.openai.azure.com/openai/deployments/gpt-4/responses" - ), + httpx.URL("https://my-azure.openai.azure.com/openai/deployments/gpt-4/responses"), "https://my-azure.openai.azure.com", ) mock_provider_config.get_api_key.return_value = "fake-azure-key" - mock_provider_config.validate_environment.return_value = { - "api-key": "fake-azure-key" - } + mock_provider_config.validate_environment.return_value = {"api-key": "fake-azure-key"} mock_provider_config.sign_request.return_value = ( {"api-key": "fake-azure-key"}, None, @@ -752,9 +717,7 @@ def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): headers={"content-type": "application/json"}, ) - sync_client = HTTPHandler( - client=httpx.Client(transport=httpx.MockTransport(_handler)) - ) + sync_client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_handler))) mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( @@ -762,18 +725,14 @@ def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): "https://gigachat.devices.sberbank.ru/api/v1", ) mock_provider_config.get_api_key.return_value = "fake-key" - mock_provider_config.validate_environment.return_value = { - "Authorization": "Bearer fake-key" - } + mock_provider_config.validate_environment.return_value = {"Authorization": "Bearer fake-key"} mock_provider_config.sign_request.return_value = ( {"Authorization": "Bearer fake-key"}, None, ) mock_provider_config.is_streaming_request.return_value = True - mock_provider_config.get_error_class.side_effect = ( - lambda error_message, status_code, headers: BaseLLMException( - status_code=status_code, message=error_message, headers=headers - ) + mock_provider_config.get_error_class.side_effect = lambda error_message, status_code, headers: BaseLLMException( + status_code=status_code, message=error_message, headers=headers ) mock_logging_obj = MagicMock() diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py similarity index 91% rename from tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py rename to tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py index 5e13db9439b..922643f9834 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py @@ -68,9 +68,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] mock_response = _make_streaming_response(chunks) - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def response_coro(): return mock_response @@ -88,7 +86,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): received.append(chunk) assert received == chunks - + assert received_response.headers["content-type"] == "application/octet-stream" assert received_response.headers["x-request-id"] == "req-123" @@ -107,9 +105,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): b'{"chunk": 3, "outputTokens": 8}', ] mock_response = _make_streaming_response(chunks) - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def response_coro(): return mock_response @@ -138,17 +134,13 @@ async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx(): err_response = MagicMock(spec=httpx.Response) err_response.status_code = 429 - err_response.headers = httpx.Headers( - {"content-type": "application/octet-stream"} - ) + err_response.headers = httpx.Headers({"content-type": "application/octet-stream"}) def _raise(): raise httpx.HTTPStatusError( "429", request=httpx.Request("POST", "https://example.com"), - response=httpx.Response( - 429, request=httpx.Request("POST", "https://example.com") - ), + response=httpx.Response(429, request=httpx.Request("POST", "https://example.com")), ) err_response.raise_for_status = _raise @@ -180,9 +172,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w mock_response.status_code = 200 mock_response.raise_for_status = MagicMock(return_value=None) mock_response.aclose = AsyncMock() - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def _aiter_bytes_then_raise(): for c in partial_chunks: @@ -197,6 +187,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w mock_logging_obj = _make_logging_obj() received = [] + async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -222,9 +213,7 @@ def test_passthroughstreamingresponse_flushes_on_normal_completion(): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) def _iter_bytes(): yield from chunks @@ -258,9 +247,7 @@ def test_passthroughstreamingresponse_flushes_on_early_close(): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) def _iter_bytes(): yield from chunks diff --git a/tests/unit/rag/ingestion/__init__.py b/tests/unit/rag/ingestion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py b/tests/unit/rag/ingestion/test_s3_vectors_ingestion.py similarity index 94% rename from tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py rename to tests/unit/rag/ingestion/test_s3_vectors_ingestion.py index 07fd2b765f3..1e5b62456b6 100644 --- a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py +++ b/tests/unit/rag/ingestion/test_s3_vectors_ingestion.py @@ -21,10 +21,14 @@ class _RecordingRouter: def _ingestion(embedding=REQUEST_EMBEDDING, router=None, **vector_store): vector_store_options = {"custom_llm_provider": "s3_vectors", "aws_region_name": "us-west-2", **vector_store} - ingest_options = {"vector_store": vector_store_options} if embedding is None else { - "embedding": embedding, - "vector_store": vector_store_options, - } + ingest_options = ( + {"vector_store": vector_store_options} + if embedding is None + else { + "embedding": embedding, + "vector_store": vector_store_options, + } + ) return S3VectorsRAGIngestion(ingest_options=ingest_options, router=router) diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py similarity index 98% rename from tests/test_litellm/realtime_api/test_main.py rename to tests/unit/realtime_api/test_main.py index 86b25b2f9c8..5d3276dfae1 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -12,6 +12,15 @@ from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import _with_resolved_session_model +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class FakeLogging: def update_from_kwargs(self, **kwargs): pass @@ -502,8 +511,8 @@ async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypat async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location: str | None): - from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig captured: dict[str, object] = {} diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py similarity index 98% rename from tests/test_litellm/repositories/test_repositories.py rename to tests/unit/repositories/test_repositories.py index 63fde9b2b8f..87cf2fc4268 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -78,17 +78,11 @@ class MockTable: record_data = dict(data) if self._pk_field and self._pk_field not in record_data: record_data[self._pk_field] = f"{self._pk_field}-{len(self._records)}" - key = ( - record_data.get(self._pk_field) - if self._pk_field - else record_data.get("id", str(len(self._records))) - ) + key = record_data.get(self._pk_field) if self._pk_field else record_data.get("id", str(len(self._records))) self._records[key] = record_data return MockRecord(record_data) - async def update( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Optional[MockRecord]: + async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Optional[MockRecord]: key_field = list(where.keys())[0] key_value = where[key_field] if key_value in self._records: @@ -140,9 +134,7 @@ class MockPrismaClient: self.db.litellm_config = MockTable() self.db.litellm_organizationtable = MockTable() self.db.litellm_projecttable = MockTable(pk_field="project_id") - self.db.litellm_objectpermissiontable = MockTable( - pk_field="object_permission_id" - ) + self.db.litellm_objectpermissiontable = MockTable(pk_field="object_permission_id") self.db.litellm_credentialstable = MockTable() @@ -200,9 +192,7 @@ class TestBaseRepository: prisma_client.db.litellm_budgettable._records = { "b1": {"budget_id": "b1", "max_budget": 100.0}, } - budgets = await repo.find_many( - where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"} - ) + budgets = await repo.find_many(where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"}) assert len(budgets) == 1 def test_record_to_dict_branches(self): @@ -1518,9 +1508,7 @@ class TestVerificationTokenRepositoryExtended: class MockTx: def __init__(self, client): - self.litellm_deletedverificationtoken = ( - client.db.litellm_deletedverificationtoken - ) + self.litellm_deletedverificationtoken = client.db.litellm_deletedverificationtoken self.litellm_verificationtoken = client.db.litellm_verificationtoken async def __aenter__(self): @@ -1563,9 +1551,7 @@ class TestVerificationTokenRepositoryExtended: class MockTx: def __init__(self, client): - self.litellm_deletedverificationtoken = ( - client.db.litellm_deletedverificationtoken - ) + self.litellm_deletedverificationtoken = client.db.litellm_deletedverificationtoken self.litellm_verificationtoken = client.db.litellm_verificationtoken async def __aenter__(self): @@ -1578,9 +1564,7 @@ class TestVerificationTokenRepositoryExtended: await repo.delete_token("sk-arch", deleted_by="admin") - archived = list( - repo._prisma_client.db.litellm_deletedverificationtoken._records.values() - )[0] + archived = list(repo._prisma_client.db.litellm_deletedverificationtoken._records.values())[0] assert isinstance(archived["aliases"], str) assert json.loads(archived["aliases"]) == {"a": "b"} @@ -1599,9 +1583,7 @@ class TestVerificationTokenRepositoryExtended: ): assert relation_field not in archived - assert ( - "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records - ) + assert "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records @pytest.mark.asyncio async def test_find_by_id_maps_org_and_budget_columns(self, repo): @@ -1977,9 +1959,7 @@ class TestDomainModelExtended: DomainModel.from_db_record(None) def test_from_db_record_dict(self): - model = _SampleDomainModel.from_db_record( - {"budget_id": "b1", "max_budget": 100.0} - ) + model = _SampleDomainModel.from_db_record({"budget_id": "b1", "max_budget": 100.0}) assert model.budget_id == "b1" def test_from_db_record_model_dump(self): @@ -2174,9 +2154,7 @@ class TestPrismaTableRepository: assert self.CONFIG_SYNCED_TABLE_NAMES <= seen -def _json_path_equals( - metadata: Optional[Dict[str, Any]], path: List[str], expected: Any -) -> bool: +def _json_path_equals(metadata: Optional[Dict[str, Any]], path: List[str], expected: Any) -> bool: """Reproduce Postgres jsonb path-equals semantics: a missing path yields SQL NULL, which never matches `equals`.""" value: Any = metadata @@ -2201,11 +2179,7 @@ class _ScimAwareUserTable: json_filter = where["metadata"] path = json_filter["path"] expected = getattr(json_filter["equals"], "data", json_filter["equals"]) - return sum( - 1 - for metadata in self._metadatas - if _json_path_equals(metadata, path, expected) - ) + return sum(1 for metadata in self._metadatas if _json_path_equals(metadata, path, expected)) class TestCountBillableUsers: diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py similarity index 100% rename from tests/test_litellm/repositories/test_unit_of_work.py rename to tests/unit/repositories/test_unit_of_work.py diff --git a/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py similarity index 100% rename from tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py rename to tests/unit/router_strategy/complexity_router/test_jev_classifier.py diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py rename to tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index b5651062098..d8bc4c45ab8 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -14,6 +14,30 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) +@pytest.fixture(autouse=True) +def isolate_litellm_router_state(): + saved = { + name: getattr(litellm, name).copy() + if isinstance(getattr(litellm, name, None), list) + else getattr(litellm, name, None) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + "model_fallbacks", + "cache", + ) + if hasattr(litellm, name) + } + yield + for name, value in saved.items(): + setattr(litellm, name, value) + + class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data @@ -43,9 +67,7 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -348,9 +370,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): model_group=model_group, user_key=user_api_key_hash, ) - await router.cache.async_set_cache( - affinity_cache_key, {"model_id": other_model_id}, ttl=3600 - ) + await router.cache.async_set_cache(affinity_cache_key, {"model_id": other_model_id}, ttl=3600) # Even though user-key affinity points elsewhere, previous_response_id should pin # to the deployment that created the original response. @@ -519,9 +539,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -542,9 +560,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s model="some-router-model-group", healthy_deployments=healthy_deployments, messages=None, - request_kwargs={ - "metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"} - }, + request_kwargs={"metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"}}, parent_otel_span=None, ) @@ -580,9 +596,7 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -618,9 +632,7 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -660,9 +672,7 @@ async def test_async_user_key_affinity_ttl_expiry_allows_reroute(): }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -706,9 +716,7 @@ def test_cache_key_does_not_double_hash_user_api_key_hash(): The affinity cache key should not hash it again. """ - user_api_key_hash = ( - "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1" - ) + user_api_key_hash = "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1" key = DeploymentAffinityCheck.get_affinity_cache_key( model_group="any-model-group", user_key=user_api_key_hash, @@ -746,9 +754,7 @@ def test_get_effective_flags_returns_per_group_config(): assert session_id is True # unconfigured-model: falls back to global flags - user_key, responses_api, session_id = callback._get_effective_flags( - "unconfigured-model" - ) + user_key, responses_api, session_id = callback._get_effective_flags("unconfigured-model") assert user_key is True assert responses_api is True assert session_id is False @@ -980,12 +986,8 @@ async def test_model_group_affinity_config_overrides_global(): ] # Set up user-key affinity cache for claude-3 - cache_key = DeploymentAffinityCheck.get_affinity_cache_key( - model_group=stable_model_map_key, user_key=user_key - ) - await callback.cache.async_set_cache( - cache_key, {"model_id": "deployment-1"}, ttl=60 - ) + cache_key = DeploymentAffinityCheck.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key) + await callback.cache.async_set_cache(cache_key, {"model_id": "deployment-1"}, ttl=60) # claude-3 has per-group config (session_affinity only), so user-key affinity # should NOT apply even though it's globally enabled @@ -1050,7 +1052,7 @@ async def test_async_jwt_user_affinity_routes_to_same_deployment(): return seq[0] return seq[1] if len(seq) > 1 else seq[0] - with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy + with patch( "litellm.router_strategy.simple_shuffle.random.choice", side_effect=deterministic_choice, ): diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py similarity index 96% rename from tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py rename to tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index b93b8c1cdfc..aa34fbd6bf7 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -25,6 +25,31 @@ from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse + +@pytest.fixture(autouse=True) +def isolate_litellm_router_state(): + saved = { + name: getattr(litellm, name).copy() + if isinstance(getattr(litellm, name, None), list) + else getattr(litellm, name, None) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + "model_fallbacks", + "cache", + ) + if hasattr(litellm, name) + } + yield + for name, value in saved.items(): + setattr(litellm, name, value) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -1088,21 +1113,19 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) @@ -1114,21 +1137,19 @@ def test_boundary_key_matches_named_credential_precedence(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { @@ -1146,21 +1167,19 @@ def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { @@ -1178,37 +1197,35 @@ def test_boundary_fallback_matches_deployments_with_same_named_credential_values EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ), - CredentialItem( - credential_name="account-a-peer", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ), - CredentialItem( - credential_name="account-b", - credential_values={ - "api_base": "https://account-b.example.com", - "api_key": "credential-key-b", - }, - credential_info={}, - ), - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-a-peer", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-b", + credential_values={ + "api_base": "https://account-b.example.com", + "api_key": "credential-key-b", + }, + credential_info={}, + ), + ], ): router = litellm.Router( model_list=[ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py similarity index 94% rename from tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py rename to tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 333e7b2ff31..849edc8c537 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1,10 +1,9 @@ import asyncio import copy -from typing import List, cast +from typing import cast import pytest - import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT @@ -22,6 +21,15 @@ MODEL_GROUP_ALIAS = "my-claude-group" OPUS_4_6_MIN_TOKENS = 4096 +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + @pytest.fixture(autouse=True) def _local_model_cost_map_autouse(local_model_cost_map): """Every test here reads `prompt_cache_min_tokens`, which only the in-repo map @@ -30,8 +38,7 @@ def _local_model_cost_map_autouse(local_model_cost_map): yield - -def _deployments(*models: str) -> List[dict]: +def _deployments(*models: str) -> list[dict]: return [ { "model_name": MODEL_GROUP_ALIAS, @@ -42,9 +49,9 @@ def _deployments(*models: str) -> List[dict]: ] -def _messages(word_count: int) -> List[AllMessageValues]: +def _messages(word_count: int) -> list[AllMessageValues]: return cast( - List[AllMessageValues], + list[AllMessageValues], [ { "role": "user", @@ -84,7 +91,9 @@ def test_write_gate_is_what_prevents_a_pin_below_the_model_minimum(): """ messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True + ) assert 1024 < token_count < 4096 assert is_prompt_caching_valid_prompt(model="anthropic/claude-opus-4-5", messages=messages) is False @@ -110,7 +119,9 @@ async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minim deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT < token_count < OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -136,7 +147,9 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=5000) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert token_count > OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -197,10 +210,10 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" -def _auto_caching_messages() -> List[AllMessageValues]: +def _auto_caching_messages() -> list[AllMessageValues]: """A prompt over the model minimum that carries no client cache_control.""" return cast( - List[AllMessageValues], + list[AllMessageValues], [ {"role": "system", "content": "word " * 3000}, {"role": "user", "content": "hello"}, @@ -208,7 +221,7 @@ def _auto_caching_messages() -> List[AllMessageValues]: ) -def _affinity_messages(messages: List[AllMessageValues]) -> List[AllMessageValues]: +def _affinity_messages(messages: list[AllMessageValues]) -> list[AllMessageValues]: """The messages the check keys deployment affinity on, for a group of `AUTO_CACHING_MODEL`.""" return AnthropicCacheControlHook.messages_with_default_injections( messages=messages, @@ -218,7 +231,7 @@ def _affinity_messages(messages: List[AllMessageValues]) -> List[AllMessageValue class _SentMessagesCapture(CustomLogger): def __init__(self): - self.messages: List[AllMessageValues] | None = None + self.messages: list[AllMessageValues] | None = None async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): standard_logging_object = kwargs.get("standard_logging_object") @@ -338,7 +351,7 @@ async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_aff cache = DualCache() check = PromptCachingDeploymentCheck(cache=cache) deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) - messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) + messages = cast(list[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) request_kwargs = { "system": [ { @@ -441,7 +454,7 @@ def test_client_supplied_cache_control_keeps_its_own_prefix_boundary(monkeypatch """ monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) messages = cast( - List[AllMessageValues], + list[AllMessageValues], [ { "role": "system", @@ -491,7 +504,7 @@ async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): warm_tokenizer("anthropic/claude-fable-5") check = PromptCachingDeploymentCheck(cache=DualCache()) deployments = _deployments("anthropic/claude-fable-5") - messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + messages = cast(list[AllMessageValues], [{"role": "user", "content": text * 100}]) result, took, lags = await timed_with_loop_lags( lambda: check.async_filter_deployments( @@ -516,7 +529,7 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): cache = DualCache() check = PromptCachingDeploymentCheck(cache=cache) messages = cast( - List[AllMessageValues], + list[AllMessageValues], [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], ) standard_logging_object = { diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py rename to tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py index ee7fab7d19f..78cafbec70a 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py @@ -1,21 +1,10 @@ import asyncio -from typing import Optional +import json from unittest.mock import AsyncMock, patch import pytest -import json - import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.llms.openai import ( - IncompleteDetails, - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIResponse, -) -from litellm.types.utils import StandardLoggingPayload @pytest.mark.asyncio @@ -119,14 +108,11 @@ async def test_async_responses_api_routing_with_previous_response_id(): input="Hello, how are you?", truncation="auto", ) - print("RESPONSE", response) # Store the model_id from the response expected_model_id = response._hidden_params["model_id"] response_id = response.id - print("Response ID=", response_id, "came from model_id=", expected_model_id) - # Make 10 other requests with previous_response_id, assert that they are sent to the same model_id for i in range(10): # Reset the mock for the next call @@ -137,7 +123,7 @@ async def test_async_responses_api_routing_with_previous_response_id(): response = await router.aresponses( model=MODEL, - input=f"Follow-up question {i+1}", + input=f"Follow-up question {i + 1}", truncation="auto", previous_response_id=response_id, ) @@ -163,9 +149,7 @@ async def test_async_routing_without_previous_response_id(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -266,9 +250,7 @@ async def test_async_routing_without_previous_response_id(): used_model_ids.add(response._hidden_params["model_id"]) # We should have used more than one model_id if load balancing is working - assert ( - len(used_model_ids) > 1 - ), "Load balancing isn't working, only one deployment was used" + assert len(used_model_ids) > 1, "Load balancing isn't working, only one deployment was used" @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py rename to tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 780300bf9e1..9bbaed0ae1b 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,12 +1,10 @@ import asyncio +import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest - -import json - import litellm from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache @@ -46,9 +44,7 @@ async def test_async_session_id_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -164,9 +160,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): ) await callback.cache.async_set_cache( - DeploymentAffinityCheck.get_session_affinity_cache_key( - "model_group", "session1", user_key="user1" - ), + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "session1", user_key="user1"), {"model_id": "deployment-2"}, ) @@ -175,9 +169,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): model="model_group", healthy_deployments=healthy_deployments, messages=[], - request_kwargs={ - "metadata": {"user_api_key_hash": "user1", "session_id": "session1"} - }, + request_kwargs={"metadata": {"user_api_key_hash": "user1", "session_id": "session1"}}, ) assert len(filtered) == 1 @@ -575,16 +567,17 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): (None, {"model": "second"}), ], ) -async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( - stored: object, expected: object -) -> None: +async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10) clock.return_value = 105.0 winner: Final = await claim_affinity_pin( - cache, "tier-pin", {"model": "second"}, 30, + cache, + "tier-pin", + {"model": "second"}, + 30, eligible_values=({"model": "first"}, {"model": "second"}), ) @@ -600,13 +593,18 @@ async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( async def test_concurrent_eligible_claims_return_one_winner() -> None: cache: Final = DualCache() candidates: Final = ({"model": "first"}, {"model": "second"}) - winners: Final = await asyncio.gather(*( - claim_affinity_pin( - cache, "tier-pin", candidates[index % 2], 30, - eligible_values=candidates, + winners: Final = await asyncio.gather( + *( + claim_affinity_pin( + cache, + "tier-pin", + candidates[index % 2], + 30, + eligible_values=candidates, + ) + for index in range(20) ) - for index in range(20) - )) + ) assert winners == [{"model": "first"}] * 20 assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"} @@ -628,23 +626,19 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) callback: Final = DeploymentAffinityCheck( - cache=cache, ttl_seconds=30, - enable_user_key_affinity=False, enable_responses_api_affinity=False, + cache=cache, + ttl_seconds=30, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, ) cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) clock.return_value = 105.0 - winner: Final = await callback._claim_pin( - "deployment-pin", {"model_id": "7"}, 30 - ) + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30) assert winner == expected - assert cache.in_memory_cache.ttl_dict["deployment-pin"] == ( - 135.0 if refresh else 110.0 - ) - assert cache.in_memory_cache.get_cache("deployment-pin") == ( - {"model_id": "7"} if refresh else stored - ) + assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (135.0 if refresh else 110.0) + assert cache.in_memory_cache.get_cache("deployment-pin") == ({"model_id": "7"} if refresh else stored) @pytest.mark.asyncio @@ -668,13 +662,13 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding( redis.async_register_script.return_value = AsyncMock(return_value=raw) cache: Final = DualCache(redis_cache=redis) callback: Final = DeploymentAffinityCheck( - cache=cache, ttl_seconds=30, - enable_user_key_affinity=False, enable_responses_api_affinity=False, + cache=cache, + ttl_seconds=30, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, ) - winner: Final = await callback._claim_pin( - "deployment-pin", {"model_id": "candidate"}, 30 - ) + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) assert winner == expected assert cache.in_memory_cache.get_cache("deployment-pin") == stored diff --git a/tests/unit/rust_bridge/__init__.py b/tests/unit/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/chat_completions/__init__.py b/tests/unit/rust_bridge/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/chat_completions/test_route_host.py b/tests/unit/rust_bridge/chat_completions/test_route_host.py similarity index 100% rename from tests/test_litellm/rust_bridge/chat_completions/test_route_host.py rename to tests/unit/rust_bridge/chat_completions/test_route_host.py diff --git a/tests/unit/rust_bridge/messages/__init__.py b/tests/unit/rust_bridge/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py similarity index 100% rename from tests/test_litellm/rust_bridge/messages/test_route_host.py rename to tests/unit/rust_bridge/messages/test_route_host.py From 73a35abeb303dff68f232719bac0893cddf69366 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:57:53 +0000 Subject: [PATCH 076/146] refactor(types): keep object permission parsing as it was Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_helpers/object_permission_utils.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 61e432daa16..daab38d3662 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Any, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -156,13 +156,12 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - raw_object_permission: Final[dict | str | None] = data_json.pop("object_permission", None) - if raw_object_permission is None: + new_object_permission: dict | str | None = data_json.pop("object_permission", None) + if new_object_permission is None: return None - new_object_permission: Final[object] = ( - json.loads(raw_object_permission) if isinstance(raw_object_permission, str) else raw_object_permission - ) + if isinstance(new_object_permission, str): + new_object_permission = json.loads(new_object_permission) upsert: Final = await prepare_object_permission_upsert( new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, @@ -231,7 +230,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), From eea00175655063cd398f7200e9663cc8be2aa6c5 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:36:13 +0000 Subject: [PATCH 077/146] test: add __init__.py to new tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/models/__init__.py | 0 tests/unit/realtime_api/__init__.py | 0 tests/unit/repositories/__init__.py | 0 tests/unit/router_strategy/complexity_router/__init__.py | 0 tests/unit/router_utils/pre_call_checks/__init__.py | 0 5 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/models/__init__.py create mode 100644 tests/unit/realtime_api/__init__.py create mode 100644 tests/unit/repositories/__init__.py create mode 100644 tests/unit/router_strategy/complexity_router/__init__.py create mode 100644 tests/unit/router_utils/pre_call_checks/__init__.py diff --git a/tests/unit/models/__init__.py b/tests/unit/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/realtime_api/__init__.py b/tests/unit/realtime_api/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/repositories/__init__.py b/tests/unit/repositories/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/complexity_router/__init__.py b/tests/unit/router_strategy/complexity_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_utils/pre_call_checks/__init__.py b/tests/unit/router_utils/pre_call_checks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 42dd6a130016083166339c3ede9539a2c94a162c Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:36:15 +0000 Subject: [PATCH 078/146] test: add __init__.py to phase 12 unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openrouter/responses/__init__.py | 0 tests/unit/llms/parallel_ai/__init__.py | 0 tests/unit/llms/parasail/__init__.py | 0 tests/unit/llms/perplexity/chat/__init__.py | 0 tests/unit/llms/perplexity/embedding/__init__.py | 0 tests/unit/llms/perplexity/responses/__init__.py | 0 tests/unit/llms/publicai/__init__.py | 0 tests/unit/llms/ragflow/chat/__init__.py | 0 tests/unit/llms/recraft/image_edit/__init__.py | 0 tests/unit/llms/recraft/image_generation/__init__.py | 0 tests/unit/llms/runwayml/__init__.py | 0 tests/unit/llms/runwayml/videos/__init__.py | 0 tests/unit/llms/s3_vectors/vector_stores/__init__.py | 0 tests/unit/llms/sap/__init__.py | 0 tests/unit/llms/scaleway/__init__.py | 0 tests/unit/llms/snowflake/__init__.py | 0 tests/unit/llms/soniox/__init__.py | 0 tests/unit/llms/stability/image_generation/__init__.py | 0 tests/unit/llms/tencent/chat/__init__.py | 0 19 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openrouter/responses/__init__.py create mode 100644 tests/unit/llms/parallel_ai/__init__.py create mode 100644 tests/unit/llms/parasail/__init__.py create mode 100644 tests/unit/llms/perplexity/chat/__init__.py create mode 100644 tests/unit/llms/perplexity/embedding/__init__.py create mode 100644 tests/unit/llms/perplexity/responses/__init__.py create mode 100644 tests/unit/llms/publicai/__init__.py create mode 100644 tests/unit/llms/ragflow/chat/__init__.py create mode 100644 tests/unit/llms/recraft/image_edit/__init__.py create mode 100644 tests/unit/llms/recraft/image_generation/__init__.py create mode 100644 tests/unit/llms/runwayml/__init__.py create mode 100644 tests/unit/llms/runwayml/videos/__init__.py create mode 100644 tests/unit/llms/s3_vectors/vector_stores/__init__.py create mode 100644 tests/unit/llms/sap/__init__.py create mode 100644 tests/unit/llms/scaleway/__init__.py create mode 100644 tests/unit/llms/snowflake/__init__.py create mode 100644 tests/unit/llms/soniox/__init__.py create mode 100644 tests/unit/llms/stability/image_generation/__init__.py create mode 100644 tests/unit/llms/tencent/chat/__init__.py diff --git a/tests/unit/llms/openrouter/responses/__init__.py b/tests/unit/llms/openrouter/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parallel_ai/__init__.py b/tests/unit/llms/parallel_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parasail/__init__.py b/tests/unit/llms/parasail/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/chat/__init__.py b/tests/unit/llms/perplexity/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/embedding/__init__.py b/tests/unit/llms/perplexity/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/responses/__init__.py b/tests/unit/llms/perplexity/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/publicai/__init__.py b/tests/unit/llms/publicai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ragflow/chat/__init__.py b/tests/unit/llms/ragflow/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/image_edit/__init__.py b/tests/unit/llms/recraft/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/image_generation/__init__.py b/tests/unit/llms/recraft/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/__init__.py b/tests/unit/llms/runwayml/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/videos/__init__.py b/tests/unit/llms/runwayml/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/s3_vectors/vector_stores/__init__.py b/tests/unit/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/sap/__init__.py b/tests/unit/llms/sap/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/scaleway/__init__.py b/tests/unit/llms/scaleway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/snowflake/__init__.py b/tests/unit/llms/snowflake/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/soniox/__init__.py b/tests/unit/llms/soniox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/stability/image_generation/__init__.py b/tests/unit/llms/stability/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/chat/__init__.py b/tests/unit/llms/tencent/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From cd07acbdea3bc1d579e88e612c870b1b92b8ea28 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:37:21 +0000 Subject: [PATCH 079/146] test: add __init__.py to intermediate phase 12 unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openrouter/__init__.py | 0 tests/unit/llms/perplexity/__init__.py | 0 tests/unit/llms/ragflow/__init__.py | 0 tests/unit/llms/recraft/__init__.py | 0 tests/unit/llms/s3_vectors/__init__.py | 0 tests/unit/llms/stability/__init__.py | 0 tests/unit/llms/tencent/__init__.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openrouter/__init__.py create mode 100644 tests/unit/llms/perplexity/__init__.py create mode 100644 tests/unit/llms/ragflow/__init__.py create mode 100644 tests/unit/llms/recraft/__init__.py create mode 100644 tests/unit/llms/s3_vectors/__init__.py create mode 100644 tests/unit/llms/stability/__init__.py create mode 100644 tests/unit/llms/tencent/__init__.py diff --git a/tests/unit/llms/openrouter/__init__.py b/tests/unit/llms/openrouter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/__init__.py b/tests/unit/llms/perplexity/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ragflow/__init__.py b/tests/unit/llms/ragflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/__init__.py b/tests/unit/llms/recraft/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/s3_vectors/__init__.py b/tests/unit/llms/s3_vectors/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/stability/__init__.py b/tests/unit/llms/stability/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/__init__.py b/tests/unit/llms/tencent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 4fb1a3e61043a8308de3159417f1fca3a329d52e Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:55:27 +0000 Subject: [PATCH 080/146] test: add __init__.py to intermediate tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/rag/__init__.py | 0 tests/unit/router_strategy/__init__.py | 0 tests/unit/router_utils/__init__.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/rag/__init__.py create mode 100644 tests/unit/router_strategy/__init__.py create mode 100644 tests/unit/router_utils/__init__.py diff --git a/tests/unit/rag/__init__.py b/tests/unit/rag/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/__init__.py b/tests/unit/router_strategy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_utils/__init__.py b/tests/unit/router_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 30f3075010964a024480df032cc2bf818ac43607 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 13:44:30 +0000 Subject: [PATCH 081/146] test: use main's local_model_cost_map fixture in tests/unit/conftest.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 73a555e4455..b3bb19a8b8a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,7 +2,6 @@ import os from collections.abc import Iterator from typing import Final -import litellm import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -23,19 +22,6 @@ AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( ) -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - - def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) From d086f8574dcaac12c494fa9d37d60cab3b5a5ab7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 13:46:52 +0000 Subject: [PATCH 082/146] test: add sync and async custom_llm_provider bridge propagation tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py index 09ef1889818..f2e36137a19 100644 --- a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -40,6 +40,58 @@ def _bedrock_mantle_kwargs() -> dict: } +def _openai_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="completion", + model="gpt-5.5", + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": "gpt-5.5", + "custom_llm_provider": "openai", + "messages": messages, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +def test_completion_forwards_custom_llm_provider_to_responses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + with patch("litellm.responses", return_value=cached) as fake_responses: + result = bridge.completion(**_openai_kwargs()) + + assert result is cached + assert fake_responses.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_custom_llm_provider_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_openai_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["custom_llm_provider"] == "openai" + + @pytest.mark.asyncio async def test_acompletion_forwards_aws_region_name_to_aresponses(): bridge = ResponsesToCompletionBridgeHandler() From e1556ce32bc87fa17ccac09fd1abf27cc3fb3979 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 14:18:49 +0000 Subject: [PATCH 083/146] test: migrate legacy provider tests to tests/unit (wave 2, phase 13) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../files/test_vertex_ai_files_integration.py | 214 ------------------ tests/unit/llms/tencent/__init__.py | 0 tests/unit/llms/tencent/messages/__init__.py | 0 ...ncent_anthropic_messages_transformation.py | 0 tests/unit/llms/together_ai/__init__.py | 0 tests/unit/llms/together_ai/chat/__init__.py | 0 .../test_together_ai_chat_transformation.py | 0 tests/unit/llms/valkey/__init__.py | 0 .../llms/valkey/vector_stores/__init__.py | 0 .../test_valkey_transformation.py | 10 +- tests/unit/llms/vercel_ai_gateway/__init__.py | 0 .../llms/vercel_ai_gateway/chat/__init__.py | 0 .../test_vercel_ai_gateway_transformation.py | 0 .../vercel_ai_gateway/embedding/__init__.py | 0 .../test_vercel_ai_gateway_embedding.py | 0 tests/unit/llms/vertex_ai/__init__.py | 0 .../llms/vertex_ai/agent_engine/__init__.py | 0 .../agent_engine/test_transformation.py | 0 .../vertex_ai/context_caching/__init__.py | 0 .../test_context_caching_ttl.py | 0 .../test_vertex_ai_context_caching.py | 94 -------- tests/unit/llms/vertex_ai/files/__init__.py | 0 .../test_file_retrieve_provider_routing.py | 0 .../test_vertex_ai_binary_file_upload.py | 68 ------ .../files/test_vertex_ai_files_handler.py | 104 --------- .../files/test_vertex_ai_files_integration.py | 93 ++++++++ .../files/test_vertex_ai_files_streaming.py | 0 .../test_vertex_ai_files_transformation.py | 88 ------- .../vertex_ai/gemini_embeddings/__init__.py | 0 ...test_batch_embed_content_transformation.py | 0 .../llms/vertex_ai/image_edit/__init__.py | 0 ...est_vertex_ai_image_edit_transformation.py | 0 .../llms/vertex_ai/interactions/__init__.py | 0 ...t_vertex_ai_interactions_transformation.py | 0 .../multimodal_embeddings/__init__.py | 0 ..._ai_multimodal_embedding_transformation.py | 0 .../unit/llms/vertex_ai/realtime/__init__.py | 0 .../test_vertex_ai_realtime_transformation.py | 44 ---- .../llms/vertex_ai/text_to_speech/__init__.py | 0 .../text_to_speech/test_transformation.py | 0 40 files changed, 94 insertions(+), 621 deletions(-) delete mode 100644 tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py create mode 100644 tests/unit/llms/tencent/__init__.py create mode 100644 tests/unit/llms/tencent/messages/__init__.py rename tests/{test_litellm => unit}/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py (100%) create mode 100644 tests/unit/llms/together_ai/__init__.py create mode 100644 tests/unit/llms/together_ai/chat/__init__.py rename tests/{test_litellm => unit}/llms/together_ai/chat/test_together_ai_chat_transformation.py (100%) create mode 100644 tests/unit/llms/valkey/__init__.py create mode 100644 tests/unit/llms/valkey/vector_stores/__init__.py rename tests/{test_litellm => unit}/llms/valkey/vector_stores/test_valkey_transformation.py (97%) create mode 100644 tests/unit/llms/vercel_ai_gateway/__init__.py create mode 100644 tests/unit/llms/vercel_ai_gateway/chat/__init__.py rename tests/{test_litellm => unit}/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py (100%) create mode 100644 tests/unit/llms/vercel_ai_gateway/embedding/__init__.py rename tests/{test_litellm => unit}/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py (100%) create mode 100644 tests/unit/llms/vertex_ai/__init__.py create mode 100644 tests/unit/llms/vertex_ai/agent_engine/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/agent_engine/test_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/context_caching/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/context_caching/test_context_caching_ttl.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py (95%) create mode 100644 tests/unit/llms/vertex_ai/files/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_file_retrieve_provider_routing.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py (71%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_files_handler.py (75%) create mode 100644 tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_files_streaming.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_files_transformation.py (95%) create mode 100644 tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/image_edit/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/interactions/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/multimodal_embeddings/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/realtime/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py (92%) create mode 100644 tests/unit/llms/vertex_ai/text_to_speech/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/text_to_speech/test_transformation.py (100%) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py deleted file mode 100644 index 8f9acafa49d..00000000000 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ /dev/null @@ -1,214 +0,0 @@ -""" -Test Vertex AI files integration with main files API -""" - -import pytest -from unittest.mock import AsyncMock, MagicMock, patch - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent - - -class TestVertexAIFilesIntegration: - """Test integration of Vertex AI files with main litellm API""" - - @pytest.mark.asyncio - async def test_litellm_afile_content_vertex_ai_provider(self): - """Test litellm.afile_content with vertex_ai provider""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content since the code - # now routes through ProviderConfigManager -> base_llm_http_handler - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - # Make it return a coroutine for async path - mock_retrieve.return_value = mock_result - - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="vertex_ai", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called - mock_retrieve.assert_called_once() - - def test_litellm_file_content_vertex_ai_provider(self): - """Test litellm.file_content with vertex_ai provider (sync)""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - return_value=mock_result, - ) as mock_retrieve: - result = litellm.file_content( - file_id=file_id, - custom_llm_provider="vertex_ai", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called - mock_retrieve.assert_called_once() - - def test_litellm_file_content_vertex_ai_with_model_provider_detection(self): - """Test litellm.file_content with model parameter for provider detection""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - return_value=mock_result, - ): - # Mock get_llm_provider to return vertex_ai - with patch("litellm.files.main.get_llm_provider") as mock_get_provider: - mock_get_provider.return_value = ( - "vertex_ai/gemini-pro", - "vertex_ai", - None, - None, - ) - - # Call litellm.file_content with model to trigger provider detection - result = litellm.file_content( - file_id=file_id, - model="vertex_ai/gemini-pro", - vertex_project="test-project", - vertex_location="us-central1", - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - - # Verify provider detection was called - mock_get_provider.assert_called_once() - - def test_litellm_file_content_vertex_ai_error_cases(self): - """Test error handling in vertex_ai file_content""" - # Test missing file_id - the VertexAI provider config's - # transform_file_content_request should handle empty file_id. - # Since the code now goes through base_llm_http_handler, we mock - # ProviderConfigManager to return None so it falls through to the - # old vertex_ai code path that validates file_id. - with patch( - "litellm.files.main.ProviderConfigManager.get_provider_files_config", - return_value=None, - ): - with pytest.raises(ValueError, match="file_id is required"): - litellm.file_content( - file_id="", # Empty file_id should cause error - custom_llm_provider="vertex_ai", - vertex_project="test-project", - ) - - def test_vertex_ai_provider_in_supported_providers_list(self): - """Test that vertex_ai is included in supported providers for file_content""" - # This test ensures the type annotations and error messages include vertex_ai - - # Test that calling with unsupported provider raises appropriate error - with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info: - litellm.file_content( - file_id="test-file-id", - custom_llm_provider="unsupported_provider", # This should fail - ) - - # The error message should mention supported providers including vertex_ai - error_message = str(exc_info.value) - assert "vertex_ai" in error_message or "supported" in error_message.lower() - - @pytest.mark.asyncio - async def test_vertex_ai_file_content_with_timeout_and_retries(self): - """Test vertex_ai file_content with timeout and retry configuration""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - mock_retrieve.return_value = mock_result - - # Call with custom timeout and max_retries - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="vertex_ai", - vertex_project="test-project", - vertex_location="us-central1", - timeout=120, - max_retries=5, - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - - # Verify the mock was called - mock_retrieve.assert_called_once() - # Verify the timeout was passed through - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["timeout"] == 120 diff --git a/tests/unit/llms/tencent/__init__.py b/tests/unit/llms/tencent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/messages/__init__.py b/tests/unit/llms/tencent/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py b/tests/unit/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py rename to tests/unit/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py diff --git a/tests/unit/llms/together_ai/__init__.py b/tests/unit/llms/together_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/together_ai/chat/__init__.py b/tests/unit/llms/together_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py rename to tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py diff --git a/tests/unit/llms/valkey/__init__.py b/tests/unit/llms/valkey/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/valkey/vector_stores/__init__.py b/tests/unit/llms/valkey/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py similarity index 97% rename from tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py rename to tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py index aa114f128c5..64e15008536 100644 --- a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py +++ b/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py @@ -1,8 +1,7 @@ import struct -import sys from types import SimpleNamespace from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock from urllib.parse import unquote, urlsplit import httpx @@ -355,13 +354,6 @@ def test_search_treats_an_explicit_null_max_num_results_as_the_default(): assert client.index.searched_query.query_string() == "*=>[KNN 10 @embedding $vec AS vector_distance]" -def test_missing_redis_dependency_raises_actionable_error(): - config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=FakeEmbeddingFn([1.0])) - blocked = {name: None for name in list(sys.modules) if name == "redis" or name.startswith("redis.")} - - with patch.dict(sys.modules, blocked): - with pytest.raises(ValueError, match="pip install redis"): - _search(config) @pytest.mark.asyncio diff --git a/tests/unit/llms/vercel_ai_gateway/__init__.py b/tests/unit/llms/vercel_ai_gateway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vercel_ai_gateway/chat/__init__.py b/tests/unit/llms/vercel_ai_gateway/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py b/tests/unit/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py rename to tests/unit/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py diff --git a/tests/unit/llms/vercel_ai_gateway/embedding/__init__.py b/tests/unit/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py b/tests/unit/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py rename to tests/unit/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py diff --git a/tests/unit/llms/vertex_ai/__init__.py b/tests/unit/llms/vertex_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/agent_engine/__init__.py b/tests/unit/llms/vertex_ai/agent_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/agent_engine/test_transformation.py rename to tests/unit/llms/vertex_ai/agent_engine/test_transformation.py diff --git a/tests/unit/llms/vertex_ai/context_caching/__init__.py b/tests/unit/llms/vertex_ai/context_caching/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/unit/llms/vertex_ai/context_caching/test_context_caching_ttl.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py rename to tests/unit/llms/vertex_ai/context_caching/test_context_caching_ttl.py diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py similarity index 95% rename from tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py rename to tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 34c00e84d2e..4ba2d69a38b 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -9,7 +9,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( - MAX_PAGINATION_PAGES, ContextCachingEndpoints, ) @@ -1892,100 +1891,7 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): - """Test that pagination stops after MAX_PAGINATION_PAGES iterations""" - # Setup - mock_get_token_url.return_value = ("token", "https://test-url.com") - cache_key_to_find = "nonexistent_cache_key" - # Create mock response that always has nextPageToken (infinite pagination scenario) - def create_page_response(page_num): - response = MagicMock() - response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], - "nextPageToken": f"token_page_{page_num + 1}", - } - return response - - # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken - self.mock_client.get.side_effect = [ - create_page_response(i) for i in range(MAX_PAGINATION_PAGES) - ] - - # Execute - result = self.context_caching.check_cache( - cache_key=cache_key_to_find, - client=self.mock_client, - headers={"Authorization": "Bearer token"}, - api_key="test_key", - api_base=None, - logging_obj=self.mock_logging, - custom_llm_provider=custom_llm_provider, - vertex_project="test_project", - vertex_location="us-central1", - vertex_auth_header="Bearer test-token", - ) - - # Assert - should return None after exhausting all pages without finding match - assert result is None - # Verify exactly MAX_PAGINATION_PAGES API calls were made (not more) - assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): - """Test that async pagination stops after MAX_PAGINATION_PAGES iterations""" - # Setup - mock_get_token_url.return_value = ("token", "https://test-url.com") - cache_key_to_find = "nonexistent_cache_key" - - # Create mock response that always has nextPageToken (infinite pagination scenario) - def create_page_response(page_num): - response = MagicMock() - response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], - "nextPageToken": f"token_page_{page_num + 1}", - } - return response - - # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken - self.mock_async_client.get = AsyncMock( - side_effect=[create_page_response(i) for i in range(MAX_PAGINATION_PAGES)] - ) - - # Execute - result = await self.context_caching.async_check_cache( - cache_key=cache_key_to_find, - client=self.mock_async_client, - headers={"Authorization": "Bearer token"}, - api_key="test_key", - api_base=None, - logging_obj=self.mock_logging, - custom_llm_provider=custom_llm_provider, - vertex_project="test_project", - vertex_location="us-central1", - vertex_auth_header="Bearer test-token", - ) - - # Assert - should return None after exhausting all pages without finding match - assert result is None - # Verify exactly MAX_PAGINATION_PAGES async API calls were made (not more) - assert self.mock_async_client.get.call_count == MAX_PAGINATION_PAGES class TestVertexAIGlobalLocation: diff --git a/tests/unit/llms/vertex_ai/files/__init__.py b/tests/unit/llms/vertex_ai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py b/tests/unit/llms/vertex_ai/files/test_file_retrieve_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py rename to tests/unit/llms/vertex_ai/files/test_file_retrieve_provider_routing.py diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py similarity index 71% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index d2ee9d7d659..f4aee11c140 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -11,8 +11,6 @@ import io import json import pytest -import httpx - from litellm.llms.custom_httpx.llm_http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig from litellm.types.llms.openai import CreateFileRequest @@ -96,39 +94,6 @@ class TestVertexAIBinaryFileUpload: assert isinstance(transformed_request, bytes) assert transformed_request == mock_png_content - @pytest.mark.asyncio - async def test_http_handler_accepts_bytes_without_decoding(self): - """ - Test that httpx correctly accepts binary data without decoding. - - This test verifies that bytes can be passed to httpx's post/put methods - without needing UTF-8 decoding, which is the core of our fix. - """ - # Create mock binary data with non-UTF-8 bytes - mock_binary_data = b"\x00\x01\x02\x03\xff\xfe\xfd\xc4\xe5\xf2" - - # Test that httpx accepts bytes in the data parameter - # We're testing the behavior, not making an actual request - - # Verify that attempting to decode would fail (proving it's binary) - with pytest.raises(UnicodeDecodeError): - mock_binary_data.decode("utf-8") - - # Verify that httpx Request accepts bytes - try: - request = httpx.Request( - method="POST", - url="https://example.com/upload", - data=mock_binary_data, - headers={"Content-Type": "application/octet-stream"}, - ) - # If we get here, httpx accepts bytes - which is what we need - assert request.content == mock_binary_data - except Exception as e: - pytest.fail(f"httpx should accept bytes in data parameter: {e}") - - # Document the expected behavior - assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes" @pytest.mark.asyncio async def test_jsonl_file_upload_returns_streaming_body(self): @@ -224,36 +189,3 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) assert isinstance(result3, bytes) - - def test_bytes_type_preservation_documentation(self): - """ - Documentation test: Verify that bytes are the correct type for binary uploads. - - This test documents the expected behavior: - - Binary files (PDF, images, etc.) should remain as bytes - - Text files (JSONL) should be strings - - httpx accepts both bytes and strings in the 'data' parameter - - bytes should NEVER be decoded to UTF-8 for binary files - """ - # This is a documentation test - it always passes - # but serves as a reference for the expected behavior - - expected_behavior = { - "binary_files": { - "input_type": "bytes", - "output_type": "bytes", - "examples": ["PDF", "PNG", "JPEG", "binary data"], - "http_method": "POST or PUT", - "encoding": "none - preserve raw bytes", - }, - "text_files": { - "input_type": "str or bytes", - "output_type": "bytes", - "examples": ["JSONL", "CSV", "TXT"], - "http_method": "POST", - "encoding": "UTF-8", - }, - } - - assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" - assert expected_behavior["text_files"]["encoding"] == "UTF-8" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py similarity index 75% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 0a44f0a9a74..e0f0b7e5c0b 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -2,14 +2,11 @@ Test Vertex AI files handler functionality """ -import asyncio import re from types import MappingProxyType import pytest from unittest.mock import AsyncMock, patch -import httpx - from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent @@ -312,104 +309,3 @@ class TestVertexAIFilesHandler: assert isinstance(result, HttpxBinaryResponseContent) dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"] assert dynamic_params["gcs_bucket_name"] == "my-model-bucket" - - def test_file_content_sync_success(self): - """Test successful sync file content retrieval""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None) - - # Create expected response - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - expected_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock asyncio.run to return our expected result - with patch("asyncio.run") as mock_run: - mock_run.return_value = expected_result - - result = self.handler.file_content( - _is_async=False, - file_content_request=file_content_request, - api_base="", - vertex_credentials=None, - vertex_project="test-project", - vertex_location="us-central1", - timeout=60.0, - max_retries=3, - ) - - # Verify the result - assert result == expected_result - - # Verify asyncio.run was called (indicating sync execution) - mock_run.assert_called_once() - - @pytest.mark.asyncio - async def test_file_content_async_mode(self): - """Test async file content retrieval when _is_async=True""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None) - - # Mock the afile_content method - with patch.object(self.handler, "afile_content", new_callable=AsyncMock) as mock_afile_content: - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_afile_content.return_value = HttpxBinaryResponseContent(response=mock_response) - - # Call the method with _is_async=True - result = self.handler.file_content( - _is_async=True, - file_content_request=file_content_request, - api_base="", - vertex_credentials=None, - vertex_project="test-project", - vertex_location="us-central1", - timeout=60.0, - max_retries=3, - ) - - # Should return a coroutine since _is_async=True - assert asyncio.iscoroutine(result) - - # Await the result - final_result = await result - assert isinstance(final_result, HttpxBinaryResponseContent) - assert final_result.response.content == expected_content - - def test_httpx_response_compatibility(self): - """Test that the created HttpxBinaryResponseContent is compatible with expected interface""" - # Test the mock response creation logic - expected_content = b"test file content" - decoded_path = "gs://test-bucket/test-file.txt" - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=decoded_path), - ) - - result = HttpxBinaryResponseContent(response=mock_response) - - # Verify the response properties - assert result.response.status_code == 200 - assert result.response.content == expected_content - assert result.response.headers["content-type"] == "application/octet-stream" - - # Verify it has the expected interface (matching OpenAI file content response) - assert hasattr(result, "response") - assert hasattr(result.response, "content") - assert hasattr(result.response, "status_code") - assert hasattr(result.response, "headers") diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py new file mode 100644 index 00000000000..402c463d134 --- /dev/null +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -0,0 +1,93 @@ +""" +Test Vertex AI files integration with main files API +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class TestVertexAIFilesIntegration: + """Test integration of Vertex AI files with main litellm API""" + + + + + def test_litellm_file_content_vertex_ai_error_cases(self): + """Test error handling in vertex_ai file_content""" + # Test missing file_id - the VertexAI provider config's + # transform_file_content_request should handle empty file_id. + # Since the code now goes through base_llm_http_handler, we mock + # ProviderConfigManager to return None so it falls through to the + # old vertex_ai code path that validates file_id. + with patch( + "litellm.files.main.ProviderConfigManager.get_provider_files_config", + return_value=None, + ): + with pytest.raises(ValueError, match="file_id is required"): + litellm.file_content( + file_id="", # Empty file_id should cause error + custom_llm_provider="vertex_ai", + vertex_project="test-project", + ) + + def test_vertex_ai_provider_in_supported_providers_list(self): + """Test that vertex_ai is included in supported providers for file_content""" + # This test ensures the type annotations and error messages include vertex_ai + + # Test that calling with unsupported provider raises appropriate error + with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info: + litellm.file_content( + file_id="test-file-id", + custom_llm_provider="unsupported_provider", # This should fail + ) + + # The error message should mention supported providers including vertex_ai + error_message = str(exc_info.value) + assert "vertex_ai" in error_message or "supported" in error_message.lower() + + @pytest.mark.asyncio + async def test_vertex_ai_file_content_with_timeout_and_retries(self): + """Test vertex_ai file_content with timeout and retry configuration""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + + # Create a mock HttpxBinaryResponseContent response + import httpx + + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), + ) + mock_result = HttpxBinaryResponseContent(response=mock_response) + + # Mock the base_llm_http_handler.retrieve_file_content + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file_content", + new_callable=MagicMock, + ) as mock_retrieve: + mock_retrieve.return_value = mock_result + + # Call with custom timeout and max_retries + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="us-central1", + timeout=120, + max_retries=5, + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + + # Verify the mock was called + mock_retrieve.assert_called_once() + # Verify the timeout was passed through + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs["timeout"] == 120 diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py similarity index 95% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8a249820cbd..48464e79876 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -860,94 +860,6 @@ class TestVertexBatchOutputTransformation: binary = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\n" + b"\x00\x01\x02\xff\xfe" * 64 assert config._try_transform_vertex_batch_output_to_openai(binary) == binary - def test_streaming_transform_peaks_below_list_pipeline(self, config): - """The output transform must stream row-by-row, not build a list of every - parsed row and a second list of transformed rows. This guards against a - regression to the list pipeline, which peaks at several full copies and - OOMs on large result files. The relative comparison cancels shared noise - (per-row transform cost, GC timing) and only the list overhead differs. - """ - import gc - import tracemalloc - - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - def vertex_row(index: int) -> dict: - return { - "status": "", - "processed_time": "2024-11-01T18:13:16.826+00:00", - "request": { - "contents": [{"role": "user", "parts": [{"text": "hi"}]}], - "labels": {"litellm_custom_id": f"r-{index}"}, - }, - "response": { - "candidates": [ - { - "content": { - "parts": [{"text": "hello " * 20}], - "role": "model", - }, - "finishReason": "STOP", - } - ], - "modelVersion": "gemini-2.0-flash-001", - "usageMetadata": { - "promptTokenCount": 10, - "candidatesTokenCount": 20, - "totalTokenCount": 30, - }, - }, - } - - content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode("utf-8") - - def list_pipeline() -> bytes: - gemini_config = VertexGeminiConfig() - logging_obj = Logging( - model="", - messages=[], - stream=False, - call_type="batch_transform", - start_time=0.1, - litellm_call_id="", - function_id="", - ) - logging_obj.optional_params = {} - mock_response = httpx.Response( - status_code=200, - headers={"content-type": "application/json"}, - request=httpx.Request("POST", "https://example.com"), - ) - rows = content.decode("utf-8").strip().split("\n") - transformed = [ - json.dumps( - config._transform_single_vertex_batch_output_to_openai( - json.loads(row), gemini_config, logging_obj, mock_response - ) - ) - for row in rows - ] - return "\n".join(transformed).encode("utf-8") - - def peak_of(fn) -> int: - gc.collect() - tracemalloc.start() - try: - fn() - return tracemalloc.get_traced_memory()[1] - finally: - tracemalloc.stop() - - streaming_peak = peak_of(lambda: config._try_transform_vertex_batch_output_to_openai(content)) - list_peak = peak_of(list_pipeline) - - assert streaming_peak < list_peak * 0.75, ( - f"streaming peak {streaming_peak} is not a clear win over the list " - f"pipeline {list_peak} (ratio {streaming_peak / list_peak:.2f})" - ) class TestTryTransformDoesNotMutateCallerLoggingObj: diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py b/tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py rename to tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py diff --git a/tests/unit/llms/vertex_ai/image_edit/__init__.py b/tests/unit/llms/vertex_ai/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py rename to tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py diff --git a/tests/unit/llms/vertex_ai/interactions/__init__.py b/tests/unit/llms/vertex_ai/interactions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py rename to tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/__init__.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py rename to tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py diff --git a/tests/unit/llms/vertex_ai/realtime/__init__.py b/tests/unit/llms/vertex_ai/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py similarity index 92% rename from tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py rename to tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index d4cf58bc0b4..a71ef21d997 100644 --- a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -387,50 +387,6 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) ) -@pytest.mark.asyncio -async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( - monkeypatch, -): - """Regression: forwarding client ?model=/?intent= to the Vertex Live WSS URL causes 1007 errors. - - Exercises ``async_realtime`` end-to-end so that re-adding ``_append_query_params`` - (the reverted bug) would push ``model=``/``intent=`` onto the backend URL and fail here. - """ - import websockets - - from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler - - cfg = VertexAIRealtimeConfig( - access_token="tok", project="my-proj", location="us-central1" - ) - - captured: dict = {} - - def fake_connect(url, *args, **kwargs): - captured["url"] = url - raise RuntimeError("stop before establishing the backend connection") - - monkeypatch.setattr(websockets, "connect", fake_connect) - - try: - await BaseLLMHTTPHandler().async_realtime( - model="gemini-live-2.5-flash-preview-native-audio-09-2025", - websocket=AsyncMock(), - logging_obj=MagicMock(), - provider_config=cfg, - headers={}, - query_params={ - "model": "gemini-live-2.5-flash-preview-native-audio-09-2025", - "intent": "chat", - }, - ) - except (RuntimeError, Exception): - pass - - assert "url" in captured, "websockets.connect was never called" - assert "?" not in captured["url"] - assert "model=" not in captured["url"] - assert "intent=" not in captured["url"] def test_vertex_function_call_output_omits_id(): diff --git a/tests/unit/llms/vertex_ai/text_to_speech/__init__.py b/tests/unit/llms/vertex_ai/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py rename to tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py From a5033fff6e243dcd58dfa108bb83f3a638287283 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 14:20:04 +0000 Subject: [PATCH 084/146] test: define local_model_cost_map fixture for migrated context caching tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_context_caching.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 4ba2d69a38b..a438b62c94c 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -13,6 +14,25 @@ from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( ) +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled in-repo cost map so capability and pricing assertions do not + depend on the network-fetched ``main`` copy, which lags this branch until merge. + + ``get_model_info`` is lru_cached, so swapping ``model_cost`` is not enough on its + own; clear on the way in and out so entries warmed against either map never leak + across tests.""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + class TestContextCachingEndpoints: """Test class for ContextCachingEndpoints methods""" From 5ddf6ff3963d0700bebbb08213f303bcc0adbfdc Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 14:53:26 +0000 Subject: [PATCH 085/146] test: restore pagination limit tests and realtime query param regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_context_caching.py | 81 +++++++++++++++++++ .../test_vertex_ai_realtime_transformation.py | 37 +++++++++ 2 files changed, 118 insertions(+) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index a438b62c94c..7913700c8a7 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1911,6 +1911,87 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + self.mock_client.get.side_effect = [ + create_page_response(i) for i in range(100) + ] + + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + assert result is None + assert self.mock_client.get.call_count == 100 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + self.mock_async_client.get = AsyncMock( + side_effect=[create_page_response(i) for i in range(100)] + ) + + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + assert result is None + assert self.mock_async_client.get.call_count == 100 + diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index a71ef21d997..14b3bdb48a1 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -387,6 +387,43 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) ) +@pytest.mark.asyncio +async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( + monkeypatch, +): + import websockets + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + captured = {} + + def fake_connect(url, *args, **kwargs): + captured["url"] = url + raise RuntimeError("stop before establishing the backend connection") + + monkeypatch.setattr(websockets, "connect", fake_connect) + + await BaseLLMHTTPHandler().async_realtime( + model="gemini-live-2.5-flash-preview-native-audio-09-2025", + websocket=AsyncMock(), + logging_obj=MagicMock(), + provider_config=cfg, + headers={}, + query_params={ + "model": "gemini-live-2.5-flash-preview-native-audio-09-2025", + "intent": "chat", + }, + ) + + assert "?" not in captured["url"] + assert "model=" not in captured["url"] + assert "intent=" not in captured["url"] + + def test_vertex_function_call_output_omits_id(): From d93cc8defda03a9c2cb7fb86903190cce1dcdc07 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sun, 20 Sep 2026 08:39:51 -0700 Subject: [PATCH 086/146] test(mcp): run public client regressions in the MCP shard --- .github/workflows/test-unit.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index aa82a0bf3ee..49e6d7040d4 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -51,7 +51,7 @@ jobs: include: - shard: mcp-integration artifact-name: mcp-integration - test-path: "tests/mcp_tests" + test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" workers: 2 reruns: 0 timeout-minutes: 20 @@ -113,7 +113,6 @@ jobs: tests/test_litellm/compression tests/test_litellm/containers tests/test_litellm/endpoints - tests/test_litellm/experimental_mcp_client tests/test_litellm/models tests/test_litellm/repositories tests/test_litellm/images From e3911c71f74a46eef73d789c433cfa72ed72c3fd Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 16:24:36 +0000 Subject: [PATCH 087/146] test: make path-sourced streaming peak test differential Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../files/test_vertex_ai_files_streaming.py | 53 +++++++++---------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index b94ea1ea269..a824a1b75d9 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -99,6 +99,17 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st ) +def _measure_peak(fn) -> int: + gc.collect() + tracemalloc.start() + try: + fn() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + class TestStreamingOutputParity: def test_transform_create_file_request_returns_streaming_body_parity(self): cfg = VertexAIFilesConfig() @@ -258,16 +269,6 @@ class TestStreamingPeakMemory: measurement removes any garbage the previous run left behind. """ - def _measure(self, fn): - gc.collect() - tracemalloc.start() - try: - fn() - _, peak = tracemalloc.get_traced_memory() - finally: - tracemalloc.stop() - return peak - def test_streaming_peak_well_below_list_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(8000) @@ -279,8 +280,8 @@ class TestStreamingPeakMemory: for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes(): pass - streaming_peak = self._measure(drain_stream) - list_peak = self._measure(lambda: _reference_vertex_jsonl_string(cfg, content_str)) + streaming_peak = _measure_peak(drain_stream) + list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str)) # Core guard: the lazily consumed streaming body peaks well under a list # pipeline that materializes every transformed row. Building full @@ -298,7 +299,7 @@ class TestStreamingPeakMemory: # The payload bytes already exist before measurement starts, so a lazy # first-row parse should allocate only a small fraction of the payload; # parsing every row would blow past this bound. - peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch")) + peak = _measure_peak(lambda: cfg.get_object_name(file_data, purpose="batch")) assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload" @@ -344,12 +345,12 @@ class TestPathSourcedStreaming: first_labels = json.loads(lines[0])["request"]["labels"] assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" - def test_path_source_peak_stays_below_payload(self, tmp_path): + def test_path_source_peak_stays_below_list_pipeline(self, tmp_path): cfg = VertexAIFilesConfig() - path, raw = self._write_jsonl(tmp_path, 8000) + path, _ = self._write_jsonl(tmp_path, 8000) data = self._batch_request(path) - def run(): + def drain_stream(): cfg.get_complete_file_url( api_base=None, api_key=None, @@ -362,19 +363,17 @@ class TestPathSourcedStreaming: model="", create_file_data=data, optional_params={}, litellm_params={} ) for _ in _upload_stream(out).iter_bytes(): - pass # drain without accumulating + pass - gc.collect() - tracemalloc.start() - try: - run() - _, peak = tracemalloc.get_traced_memory() - finally: - tracemalloc.stop() + streaming_peak = _measure_peak(drain_stream) + list_peak = _measure_peak( + lambda: _reference_vertex_jsonl_string(cfg, path.read_bytes().decode("utf-8")) + ) - # Streaming from disk must not materialize the payload. Reading the whole - # file into bytes (the pre-fix path) would push peak past the file size. - assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (ratio {peak / len(raw):.2f})" + assert streaming_peak < list_peak * 0.3, ( + f"path-sourced streaming peak {streaming_peak} not a clear win over list pipeline " + f"{list_peak} (ratio {streaming_peak / list_peak:.2f})" + ) def test_path_source_stream_is_reiterable(self, tmp_path): cfg = VertexAIFilesConfig() From 8554d1f8275d9b6e3eebaa0c7def5d4b24354c02 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 17:08:16 +0000 Subject: [PATCH 088/146] test: prepare reference input before tracing path-sourced peak Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/test_vertex_ai_files_streaming.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index a824a1b75d9..29eaf38b427 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -347,8 +347,9 @@ class TestPathSourcedStreaming: def test_path_source_peak_stays_below_list_pipeline(self, tmp_path): cfg = VertexAIFilesConfig() - path, _ = self._write_jsonl(tmp_path, 8000) + path, raw = self._write_jsonl(tmp_path, 8000) data = self._batch_request(path) + content_str = raw.decode("utf-8") def drain_stream(): cfg.get_complete_file_url( @@ -366,9 +367,7 @@ class TestPathSourcedStreaming: pass streaming_peak = _measure_peak(drain_stream) - list_peak = _measure_peak( - lambda: _reference_vertex_jsonl_string(cfg, path.read_bytes().decode("utf-8")) - ) + list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str)) assert streaming_peak < list_peak * 0.3, ( f"path-sourced streaming peak {streaming_peak} not a clear win over list pipeline " From 3ebd5add3230f0acb01fb134434c55bbda4e4104 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 17:20:16 +0000 Subject: [PATCH 089/146] fix(anthropic): forward safeguards and anthropic-beta unchanged on native /v1/messages Native Anthropic Messages requests derived their allowlist from AnthropicMessagesRequestOptionalParams, which lacked safeguards, and the shared beta-header filter dropped betas unknown to the provider mapping even when the upstream is api.anthropic.com itself. Claude Code auto mode then saw no safeguard_results and fell back to billed classifier calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/transformation.py | 4 + litellm/types/llms/anthropic.py | 1 + .../anthropic_messages/anthropic_response.py | 1 + ...erimental_pass_through_messages_handler.py | 108 ++++++++++++++++++ 4 files changed, 114 insertions(+) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 5fa686b7560..eed30c2698c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -79,10 +79,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): "speed", "output_config", "reasoning_effort", + "safeguards", # TODO: Add Anthropic `metadata` support # "metadata", ] + def should_filter_anthropic_beta_headers(self) -> bool: + return self._resolved_provider != "anthropic" + def _remove_scope_from_cache_control(self, anthropic_messages_request: dict) -> None: """ Remove `scope` field from cache_control blocks. diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index bcd24695f25..f4fe7a0bf14 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -411,6 +411,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): output_config: AnthropicOutputConfig | None # Configuration for Claude's output behavior cache_control: dict[str, Any] | None # Automatic prompt caching reasoning_effort: str | None + safeguards: ReadOnly[dict[str, object] | None] class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False): diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 038a23a3ca2..41060e96d85 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -97,3 +97,4 @@ class AnthropicMessagesResponse(TypedDict, total=False): type: Literal["message"] | None usage: AnthropicUsage | None context_management: NotRequired[ContextManagementResponse] + safeguard_results: NotRequired[ReadOnly[dict[str, object]]] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 997a97c6fd3..4246e70bbbf 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1438,3 +1438,111 @@ async def test_anthropic_messages_leaves_non_provider_failures_unmapped(): ) assert "Traceback" not in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthropic(): + """Regression test for LIT-8232. Claude Code auto mode sends a `safeguards` body + field paired with a beta value the gateway has never seen. Both must reach + api.anthropic.com unchanged or the session falls back to billed classifier calls.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + safeguards = {"auto_mode": {"enabled": True, "version": "2026-09-01"}} + client_betas = "safeguards-2026-09-01,interleaved-thinking-2025-05-14" + captured: dict[str, object] = {} + + def upstream_records_the_request(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + captured["anthropic-beta"] = request.headers.get("anthropic-beta") + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + "safeguard_results": {"verdict": "allow"}, + }, + request=request, + ) + + upstream = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream_records_the_request)) + + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="anthropic/claude-haiku-4-5", + custom_llm_provider="anthropic", + api_key="sk-test", + client=upstream, + safeguards=safeguards, + extra_headers={"anthropic-beta": client_betas}, + ) + + assert captured["body"]["safeguards"] == safeguards + assert set(captured["anthropic-beta"].split(",")) == set(client_betas.split(",")) + assert response["safeguard_results"] == {"verdict": "allow"} + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_forwards_safeguards_and_keeps_safeguard_results(): + """Streaming sibling of the LIT-8232 regression: the request must still carry + `safeguards` and the `safeguard_results` Anthropic emits on `message_start` and + `message_delta` must reach the client byte for byte.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + safeguards = {"auto_mode": {"enabled": True, "version": "2026-09-01"}} + safeguard_results = {"verdict": "allow", "checks": ["shell_command"]} + captured: dict[str, object] = {} + message_start = { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + "safeguard_results": safeguard_results, + }, + } + message_delta = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 1}, + } + sse = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" + for event in (message_start, message_delta, {"type": "message_stop"}) + ) + + def upstream_streams_safeguard_results(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse.encode(), request=request) + + upstream = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream_streams_safeguard_results)) + + stream = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="anthropic/claude-haiku-4-5", + custom_llm_provider="anthropic", + api_key="sk-test", + client=upstream, + stream=True, + safeguards=safeguards, + ) + raw = b"".join([chunk async for chunk in stream]).decode() + events = [json.loads(line[len("data: ") :]) for line in raw.splitlines() if line.startswith("data: ")] + + assert captured["body"]["safeguards"] == safeguards + assert events[0]["message"]["safeguard_results"] == safeguard_results + assert [e for e in events if e["type"] == "message_delta"][0]["delta"]["safeguard_results"] == safeguard_results From a6361dc08cc1d9c50b072ecf38d92bfa27df8049 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 17:30:49 +0000 Subject: [PATCH 090/146] chore(prices): sync OpenRouter prices: 11 models openrouter/~deepseek/deepseek-flash-latest: off_peak_pricing, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-flash-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/meta/muse-glimmer-30b: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token openrouter/mistralai/mistral-small-3.1-24b-instruct: supports_tool_choice, supports_function_calling openrouter/qwen/qwen-2.5-coder-32b-instruct: supports_response_schema openrouter/qwen/qwen3.8-27b: input_cost_per_token, output_cost_per_token openrouter/z-ai/glm-5.3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 65 ++++++++++--------- model_prices_and_context_window.json | 65 ++++++++++--------- 2 files changed, 66 insertions(+), 64 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bfb4729ef32..e92712fd6d9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,21 +41579,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.3724e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost": 1.7908e-08, "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, @@ -42097,12 +42097,12 @@ "max_tokens": 102400, "mode": "chat", "output_cost_per_token": 5.55e-07, - "supports_tool_choice": false, + "supports_tool_choice": true, "max_input_tokens": 128000, "max_output_tokens": 102400, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, - "supports_function_calling": false, + "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, @@ -42775,7 +42775,7 @@ "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_response_schema": false, + "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": false, "supports_web_search": false @@ -66570,9 +66570,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.96e-07, - "output_cost_per_token": 2.816e-06, - "cache_read_input_token_cost": 1.664e-07, + "input_cost_per_token": 9.1e-07, + "output_cost_per_token": 2.86e-06, + "cache_read_input_token_cost": 1.69e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -66591,8 +66591,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2.55e-06, "cache_read_input_token_cost": 8.5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.668e-08, - "output_cost_per_token": 7.336e-08, - "cache_read_input_token_cost": 7.336e-09, + "input_cost_per_token": 3.556e-08, + "output_cost_per_token": 7.112e-08, + "cache_read_input_token_cost": 7.112e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71296,14 +71296,15 @@ "supports_web_search": true }, "openrouter/~deepseek/deepseek-flash-latest": { - "cache_read_input_token_cost": 2.6e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 5.2e-07, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71316,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.9228e-08, - "input_cost_per_token": 5.7684e-07, + "cache_read_input_token_cost": 1.7908e-08, + "input_cost_per_token": 5.3724e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.73052e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71568,14 +71569,14 @@ "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 7.5e-08, + "cache_read_input_token_cost": 1.8e-08, + "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71588,14 +71589,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.5678e-07, - "input_cost_per_token": 8.442e-07, + "cache_read_input_token_cost": 1.69e-07, + "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.6532e-06, + "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -73022,13 +73023,13 @@ }, "openrouter/meta/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, - "input_cost_per_token": 3.5e-07, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 117964, - "max_tokens": 117964, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bfb4729ef32..e92712fd6d9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,21 +41579,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.3724e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost": 1.7908e-08, "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, @@ -42097,12 +42097,12 @@ "max_tokens": 102400, "mode": "chat", "output_cost_per_token": 5.55e-07, - "supports_tool_choice": false, + "supports_tool_choice": true, "max_input_tokens": 128000, "max_output_tokens": 102400, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, - "supports_function_calling": false, + "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, @@ -42775,7 +42775,7 @@ "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_response_schema": false, + "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": false, "supports_web_search": false @@ -66570,9 +66570,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.96e-07, - "output_cost_per_token": 2.816e-06, - "cache_read_input_token_cost": 1.664e-07, + "input_cost_per_token": 9.1e-07, + "output_cost_per_token": 2.86e-06, + "cache_read_input_token_cost": 1.69e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -66591,8 +66591,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2.55e-06, "cache_read_input_token_cost": 8.5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.668e-08, - "output_cost_per_token": 7.336e-08, - "cache_read_input_token_cost": 7.336e-09, + "input_cost_per_token": 3.556e-08, + "output_cost_per_token": 7.112e-08, + "cache_read_input_token_cost": 7.112e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71296,14 +71296,15 @@ "supports_web_search": true }, "openrouter/~deepseek/deepseek-flash-latest": { - "cache_read_input_token_cost": 2.6e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 5.2e-07, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71316,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.9228e-08, - "input_cost_per_token": 5.7684e-07, + "cache_read_input_token_cost": 1.7908e-08, + "input_cost_per_token": 5.3724e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.73052e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71568,14 +71569,14 @@ "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 7.5e-08, + "cache_read_input_token_cost": 1.8e-08, + "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71588,14 +71589,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.5678e-07, - "input_cost_per_token": 8.442e-07, + "cache_read_input_token_cost": 1.69e-07, + "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.6532e-06, + "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -73022,13 +73023,13 @@ }, "openrouter/meta/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, - "input_cost_per_token": 3.5e-07, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 117964, - "max_tokens": 117964, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From b59028d525352454de672621c2b223c5d75e57b1 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 17:37:56 +0000 Subject: [PATCH 091/146] fix(anthropic): strip safeguards on the adapter path and type it on streaming chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../experimental_pass_through/adapters/handler.py | 2 +- litellm/types/llms/anthropic.py | 2 ++ .../test_handler_output_config_passthrough.py | 13 +++++++++++++ ...ic_experimental_pass_through_messages_handler.py | 6 ------ 4 files changed, 16 insertions(+), 7 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 87a29ca50ba..54d10837d74 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -35,7 +35,7 @@ if TYPE_CHECKING: from litellm.router import Router # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. -ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"}) +ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config", "safeguards"}) _AnthropicMessages: TypeAlias = "list[dict[str, object]]" _AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None" diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index f4fe7a0bf14..f57591d0262 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -531,6 +531,7 @@ class AnthropicStopDetails(TypedDict, total=False): class MessageDelta(TypedDict, total=False): stop_reason: str | None stop_details: ReadOnly[AnthropicStopDetails] + safeguard_results: ReadOnly[dict[str, object]] class ServerToolUsage(TypedDict, total=False): @@ -601,6 +602,7 @@ class MessageChunk(TypedDict, total=False): stop_reason: str | None stop_sequence: str | None usage: UsageDelta + safeguard_results: ReadOnly[dict[str, object]] class MessageStartBlock(TypedDict): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py index a944afc6152..d6de6372e0b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py @@ -110,6 +110,19 @@ class TestOutputConfigStrippedFromCompletionKwargs: "reject it with 400 'Extra inputs are not permitted'" ) + def test_safeguards_is_stripped_for_non_anthropic_target(self): + extra_kwargs = { + "custom_llm_provider": "azure", + "safeguards": {"auto_mode": {"enabled": True, "version": "2026-09-01"}}, + } + + result = _call_prepare(extra_kwargs=extra_kwargs) + + completion_kwargs = result[0] if isinstance(result, tuple) else result + assert "safeguards" not in completion_kwargs, ( + "safeguards is an Anthropic-only field; OpenAI-format backends reject it with 400" + ) + def test_output_config_format_translated_to_response_format(self): """When ``output_config`` carries structured-output ``format``, the translator now maps it to OpenAI's ``response_format`` so non-Anthropic diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 4246e70bbbf..0acb9d634a3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1442,9 +1442,6 @@ async def test_anthropic_messages_leaves_non_provider_failures_unmapped(): @pytest.mark.asyncio async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthropic(): - """Regression test for LIT-8232. Claude Code auto mode sends a `safeguards` body - field paired with a beta value the gateway has never seen. Both must reach - api.anthropic.com unchanged or the session falls back to billed classifier calls.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler safeguards = {"auto_mode": {"enabled": True, "version": "2026-09-01"}} @@ -1491,9 +1488,6 @@ async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthro @pytest.mark.asyncio async def test_anthropic_messages_streaming_forwards_safeguards_and_keeps_safeguard_results(): - """Streaming sibling of the LIT-8232 regression: the request must still carry - `safeguards` and the `safeguard_results` Anthropic emits on `message_start` and - `message_delta` must reach the client byte for byte.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler safeguards = {"auto_mode": {"enabled": True, "version": "2026-09-01"}} From 229d5cbfcbeac05a4e30d5dc24e22d0713016900 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 18:00:52 +0000 Subject: [PATCH 092/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e92712fd6d9..771657cd4d9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3724e-07, + "input_cost_per_token": 5.346e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7908e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.782e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7908e-08, - "input_cost_per_token": 5.3724e-07, + "cache_read_input_token_cost": 1.782e-08, + "input_cost_per_token": 5.346e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e92712fd6d9..771657cd4d9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3724e-07, + "input_cost_per_token": 5.346e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7908e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.782e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7908e-08, - "input_cost_per_token": 5.3724e-07, + "cache_read_input_token_cost": 1.782e-08, + "input_cost_per_token": 5.346e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 010a5a5b1c71dd5b8e548051236ed6212dffbd80 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 18:30:53 +0000 Subject: [PATCH 093/146] chore(prices): sync OpenRouter prices: 5 models openrouter/~anthropic/claude-fable-latest: supports_tool_choice openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/anthropic/claude-fable-5.1: supports_tool_choice openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/qwen/qwen3.6-35b-a3b: input_cost_per_token, output_cost_per_token --- ...odel_prices_and_context_window_backup.json | 30 +++++++++---------- model_prices_and_context_window.json | 30 +++++++++---------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 771657cd4d9..bf65675f3c5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.346e-07, + "input_cost_per_token": 5.3328e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.782e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, + "cache_read_input_token_cost": 1.6968e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65764,7 +65764,7 @@ "thinking_always_on": true, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, @@ -67155,8 +67155,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.6-35b-a3b": { - "input_cost_per_token": 1e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1e-06, "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -71225,7 +71225,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true }, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.782e-08, - "input_cost_per_token": 5.346e-07, + "cache_read_input_token_cost": 1.6968e-08, + "input_cost_per_token": 5.3328e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 771657cd4d9..bf65675f3c5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.346e-07, + "input_cost_per_token": 5.3328e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.782e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, + "cache_read_input_token_cost": 1.6968e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65764,7 +65764,7 @@ "thinking_always_on": true, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, @@ -67155,8 +67155,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.6-35b-a3b": { - "input_cost_per_token": 1e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1e-06, "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -71225,7 +71225,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true }, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.782e-08, - "input_cost_per_token": 5.346e-07, + "cache_read_input_token_cost": 1.6968e-08, + "input_cost_per_token": 5.3328e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 6a8b131fbdc0b041ffe8112924919ea4bb2dd2df Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:00:55 +0000 Subject: [PATCH 094/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 22 +++++++++---------- model_prices_and_context_window.json | 22 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bf65675f3c5..d325f46017a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3328e-07, + "input_cost_per_token": 5.2932e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6968e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7644e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6968e-08, - "input_cost_per_token": 5.3328e-07, + "cache_read_input_token_cost": 1.7644e-08, + "input_cost_per_token": 5.2932e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bf65675f3c5..d325f46017a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3328e-07, + "input_cost_per_token": 5.2932e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6968e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7644e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6968e-08, - "input_cost_per_token": 5.3328e-07, + "cache_read_input_token_cost": 1.7644e-08, + "input_cost_per_token": 5.2932e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From a220fb7d30115785768590fbe90974ef0b2f5bb2 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 19:23:52 +0000 Subject: [PATCH 095/146] test(a2a): migrate a2a_protocol legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../providers/bedrock_agentcore/__init__.py | 0 .../test_bedrock_agentcore_a2a.py | 130 ++---------------- .../test_a2a_exception_mapping_utils.py | 5 +- .../test_a2a_streaming_iterator.py | 23 +--- .../a2a_protocol/test_card_resolver.py | 8 +- .../test_completion_bridge_streaming.py | 6 +- .../a2a_protocol/test_cost_calculator.py | 16 +-- .../a2a_protocol/test_main.py | 17 +-- .../test_send_message_response.py | 29 +--- .../a2a_protocol/test_utils.py | 0 10 files changed, 33 insertions(+), 201 deletions(-) create mode 100644 tests/unit/a2a_protocol/providers/bedrock_agentcore/__init__.py rename tests/{test_litellm => unit}/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py (82%) rename tests/{test_litellm => unit}/a2a_protocol/test_a2a_exception_mapping_utils.py (98%) rename tests/{test_litellm => unit}/a2a_protocol/test_a2a_streaming_iterator.py (89%) rename tests/{test_litellm => unit}/a2a_protocol/test_card_resolver.py (97%) rename tests/{test_litellm => unit}/a2a_protocol/test_completion_bridge_streaming.py (98%) rename tests/{test_litellm => unit}/a2a_protocol/test_cost_calculator.py (96%) rename tests/{test_litellm => unit}/a2a_protocol/test_main.py (97%) rename tests/{test_litellm => unit}/a2a_protocol/test_send_message_response.py (77%) rename tests/{test_litellm => unit}/a2a_protocol/test_utils.py (100%) diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/__init__.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py similarity index 82% rename from tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py rename to tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index a8fe464ec32..dcea462d7a1 100644 --- a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -42,13 +42,11 @@ class TestTransformation: BedrockAgentCoreA2ATransformation, ) - url, headers, body = ( - BedrockAgentCoreA2ATransformation.get_url_and_signed_request( - request_id="req-001", - params=SAMPLE_PARAMS, - litellm_params=SAMPLE_LITELLM_PARAMS, - method="message/send", - ) + url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id="req-001", + params=SAMPLE_PARAMS, + litellm_params=SAMPLE_LITELLM_PARAMS, + method="message/send", ) body_dict = json.loads(body) assert body_dict["jsonrpc"] == "2.0" @@ -201,10 +199,7 @@ class TestTransformation: # Runtime user id is the value set from litellm_params, NOT the spoof. assert normalized["x-amzn-bedrock-agentcore-runtime-user-id"] == "legit-user" # Session id is the auto-generated one, not the spoofed value. - assert ( - normalized["x-amzn-bedrock-agentcore-runtime-session-id"] - != "spoofed-session" - ) + assert normalized["x-amzn-bedrock-agentcore-runtime-session-id"] != "spoofed-session" # Authorization is the JWT bearer set by the signer, not the spoof. assert normalized["authorization"] == "Bearer test-jwt-token" # Host / x-amz-* must not have been carried over from the client. @@ -259,43 +254,6 @@ class TestTransformation: # Non-reserved header still makes it into the signed dict. assert captured.get("x-mcp-token") == "mcp-abc" - def test_sigv4_auth_when_no_api_key(self): - """When no api_key, falls through to SigV4 signing.""" - from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( - BedrockAgentCoreA2ATransformation, - ) - - litellm_params_no_key = { - "model": SAMPLE_MODEL, - "custom_llm_provider": "bedrock", - "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", - "aws_region_name": "us-west-2", - } - - # Mock _sign_request to avoid hitting real botocore credential resolution - fake_sigv4_headers = { - "Authorization": "AWS4-HMAC-SHA256 Credential=AKIA.../bedrock-agentcore/aws4_request", - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - } - fake_body = b'{"jsonrpc":"2.0"}' - - with patch( - "litellm.llms.bedrock.chat.agentcore.transformation.AmazonAgentCoreConfig._sign_request", - return_value=(fake_sigv4_headers, fake_body), - ): - _, headers, _ = ( - BedrockAgentCoreA2ATransformation.get_url_and_signed_request( - request_id="req-001", - params=SAMPLE_PARAMS, - litellm_params=litellm_params_no_key, - ) - ) - # SigV4 produces an Authorization header starting with "AWS4-HMAC-SHA256" - assert "Authorization" in headers - assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") - SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" CONTEXT_ID = "conversation-alpha-0001-0000000000000000" @@ -571,39 +529,6 @@ class TestNonStreaming: sent_headers = mock_client.post.call_args.kwargs["headers"] assert sent_headers.get("x-mcp-token") == "mcp-abc" - @pytest.mark.asyncio - async def test_a2a_error_response_passthrough(self): - """JSON-RPC error responses from the agent are returned as-is.""" - from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( - BedrockAgentCoreA2AConfig, - ) - - error_response = { - "jsonrpc": "2.0", - "id": "req-001", - "error": {"code": -32600, "message": "Bad request"}, - } - mock_response = MagicMock() - mock_response.json.return_value = error_response - mock_response.raise_for_status = MagicMock() - - with patch( - "litellm.a2a_protocol.providers.bedrock_agentcore.handler.get_async_httpx_client" - ) as mock_get_client: - mock_client = AsyncMock() - mock_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_client - - config = BedrockAgentCoreA2AConfig() - result = await config.handle_non_streaming( - request_id="req-001", - params=SAMPLE_PARAMS, - litellm_params=SAMPLE_LITELLM_PARAMS, - ) - - assert result["error"]["code"] == -32600 - assert result["error"]["message"] == "Bad request" - class TestConfigManager: """Test that config manager routes 'bedrock' correctly.""" @@ -616,9 +541,7 @@ class TestConfigManager: A2AProviderConfigManager, ) - config = A2AProviderConfigManager.get_provider_config( - "bedrock", model=SAMPLE_MODEL - ) + config = A2AProviderConfigManager.get_provider_config("bedrock", model=SAMPLE_MODEL) assert config is not None assert isinstance(config, BedrockAgentCoreA2AConfig) @@ -628,9 +551,7 @@ class TestConfigManager: A2AProviderConfigManager, ) - config = A2AProviderConfigManager.get_provider_config( - "bedrock", model="bedrock/anthropic.claude-3-sonnet" - ) + config = A2AProviderConfigManager.get_provider_config("bedrock", model="bedrock/anthropic.claude-3-sonnet") assert config is None def test_unknown_provider_returns_none(self): @@ -644,37 +565,6 @@ class TestConfigManager: class TestHandlerIntegration: """Test handler.py changes — litellm_params passed through, api_base not required.""" - @pytest.mark.asyncio - async def test_provider_config_receives_litellm_params(self): - """Verify handler passes litellm_params to provider config via kwargs.""" - from litellm.a2a_protocol.litellm_completion_bridge.handler import ( - A2ACompletionBridgeHandler, - ) - - mock_config = AsyncMock() - mock_config.handle_non_streaming = AsyncMock( - return_value={"jsonrpc": "2.0", "id": "req-001", "result": {}} - ) - - with patch( - "litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config", - return_value=mock_config, - ): - await A2ACompletionBridgeHandler.handle_non_streaming( - request_id="req-001", - params=SAMPLE_PARAMS, - litellm_params=SAMPLE_LITELLM_PARAMS, - api_base=None, - ) - - mock_config.handle_non_streaming.assert_called_once_with( - request_id="req-001", - params=SAMPLE_PARAMS, - api_base=None, - litellm_params=SAMPLE_LITELLM_PARAMS, - agent_extra_headers=None, - ) - @pytest.mark.asyncio async def test_api_base_none_allowed_with_provider_config(self): """api_base=None no longer raises when a provider config is registered.""" @@ -683,9 +573,7 @@ class TestHandlerIntegration: ) mock_config = AsyncMock() - mock_config.handle_non_streaming = AsyncMock( - return_value={"jsonrpc": "2.0", "id": "req-001", "result": {}} - ) + mock_config.handle_non_streaming = AsyncMock(return_value={"jsonrpc": "2.0", "id": "req-001", "result": {}}) with patch( "litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config", diff --git a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py b/tests/unit/a2a_protocol/test_a2a_exception_mapping_utils.py similarity index 98% rename from tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py rename to tests/unit/a2a_protocol/test_a2a_exception_mapping_utils.py index c31d50960b1..5f097570bc2 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py +++ b/tests/unit/a2a_protocol/test_a2a_exception_mapping_utils.py @@ -38,9 +38,7 @@ async def test_localhost_retry_reuses_stashed_httpx_client(): patch.object(emu, "A2A_SDK_AVAILABLE", True), patch.object(emu, "set_agent_card_url") as mock_set_url, patch.object(emu, "ClientConfig", side_effect=fake_client_config), - patch.object( - emu, "create_client", new=AsyncMock(return_value=new_client) - ) as mock_create, + patch.object(emu, "create_client", new=AsyncMock(return_value=new_client)) as mock_create, ): result = await emu.handle_a2a_localhost_retry( error=_localhost_error(), @@ -171,6 +169,7 @@ async def test_stream_with_retry_raises_after_localhost_retries_exhausted(): api_base="https://agent.example", agent_name="test-agent", ) + async def _drain(): async for _chunk in stream: pytest.fail("expected retry exhaustion to raise before yielding") diff --git a/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py similarity index 89% rename from tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py rename to tests/unit/a2a_protocol/test_a2a_streaming_iterator.py index 2603d135dce..abf6a6dda31 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py @@ -43,25 +43,6 @@ class RecordingExecutor: return [fn for fn in self.submits if getattr(fn, "__self__", None) is logging_obj] -@pytest.fixture(autouse=True) -def _isolate_callbacks(): - saved = ( - litellm.callbacks, - litellm.success_callback, - litellm._async_success_callback, - litellm.failure_callback, - litellm._async_failure_callback, - ) - yield - ( - litellm.callbacks, - litellm.success_callback, - litellm._async_success_callback, - litellm.failure_callback, - litellm._async_failure_callback, - ) = saved - - @pytest.mark.asyncio async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch): recording_executor = RecordingExecutor(thread_pool_executor_module.executor) @@ -69,8 +50,8 @@ async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch monkeypatch.setattr(a2a_streaming_iterator_module, "executor", recording_executor, raising=False) recorder = RecordingCustomLogger() - litellm.success_callback = [recorder] - litellm._async_success_callback = [recorder] + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) logging_obj = LitellmLogging( model="a2a/test-agent", diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/unit/a2a_protocol/test_card_resolver.py similarity index 97% rename from tests/test_litellm/a2a_protocol/test_card_resolver.py rename to tests/unit/a2a_protocol/test_card_resolver.py index 88dc835df0e..fdfb51987a3 100644 --- a/tests/test_litellm/a2a_protocol/test_card_resolver.py +++ b/tests/unit/a2a_protocol/test_card_resolver.py @@ -36,9 +36,7 @@ async def test_card_resolver_fallback_from_new_to_old_path(): paths_called = [] # Create a mock for the parent's get_agent_card method - async def mock_parent_get_agent_card( - self, relative_card_path=None, http_kwargs=None - ): + async def mock_parent_get_agent_card(self, relative_card_path=None, http_kwargs=None): paths_called.append(relative_card_path) if relative_card_path == "/.well-known/agent-card.json": # First call (new path) fails @@ -57,9 +55,7 @@ async def test_card_resolver_fallback_from_new_to_old_path(): "get_agent_card", mock_parent_get_agent_card, ): - resolver = LiteLLMA2ACardResolver( - httpx_client=mock_httpx_client, base_url="http://test-agent:8000" - ) + resolver = LiteLLMA2ACardResolver(httpx_client=mock_httpx_client, base_url="http://test-agent:8000") result = await resolver.get_agent_card() # Verify both paths were tried in correct order diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/unit/a2a_protocol/test_completion_bridge_streaming.py similarity index 98% rename from tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py rename to tests/unit/a2a_protocol/test_completion_bridge_streaming.py index 8fd35369cf2..913c917bd2d 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/unit/a2a_protocol/test_completion_bridge_streaming.py @@ -344,11 +344,7 @@ async def test_handle_streaming_keeps_agent_card_path_out_of_the_completion_call chunk.choices[0].delta.content = "Hello" yield chunk - with ( - patch( # test-quality-ok: the bridge calls litellm.acompletion directly; the sibling tests capture its kwargs through the same seam - "litellm.acompletion", new_callable=AsyncMock - ) as mock_acompletion - ): + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: mock_acompletion.return_value = mock_streaming_response() events = [ diff --git a/tests/test_litellm/a2a_protocol/test_cost_calculator.py b/tests/unit/a2a_protocol/test_cost_calculator.py similarity index 96% rename from tests/test_litellm/a2a_protocol/test_cost_calculator.py rename to tests/unit/a2a_protocol/test_cost_calculator.py index a29f012170f..56d3d57c89e 100644 --- a/tests/test_litellm/a2a_protocol/test_cost_calculator.py +++ b/tests/unit/a2a_protocol/test_cost_calculator.py @@ -122,7 +122,7 @@ class CostLogger(CustomLogger): @pytest.mark.asyncio -async def test_asend_message_uses_cost_per_query(): +async def test_asend_message_uses_cost_per_query(monkeypatch): """ Test that asend_message uses cost_per_query param for response_cost. """ @@ -131,7 +131,7 @@ async def test_asend_message_uses_cost_per_query(): # Setup logger litellm.logging_callback_manager._reset_all_callbacks() cost_logger = CostLogger() - litellm.callbacks = [cost_logger] + monkeypatch.setattr(litellm, "callbacks", [cost_logger]) # Mock A2A client mock_client = MagicMock() @@ -157,7 +157,7 @@ async def test_asend_message_uses_cost_per_query(): @pytest.mark.asyncio -async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(): +async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(monkeypatch): """ Proxy passes agent pricing as the litellm_params dict param (not top-level kwargs). Regression for cost_per_query landing at $0 on the native path. @@ -166,7 +166,7 @@ async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(): litellm.logging_callback_manager._reset_all_callbacks() cost_logger = CostLogger() - litellm.callbacks = [cost_logger] + monkeypatch.setattr(litellm, "callbacks", [cost_logger]) mock_client = MagicMock() mock_client._litellm_agent_card = MagicMock() @@ -217,7 +217,7 @@ class TokenAndCostLogger(CustomLogger): @pytest.mark.asyncio -async def test_asend_message_uses_input_output_cost_per_token(): +async def test_asend_message_uses_input_output_cost_per_token(monkeypatch): """ Test that asend_message calculates cost using input_cost_per_token and output_cost_per_token. Validates exact cost calculation: cost = (prompt_tokens * input_cost) + (completion_tokens * output_cost) @@ -227,7 +227,7 @@ async def test_asend_message_uses_input_output_cost_per_token(): # Setup logger litellm.logging_callback_manager._reset_all_callbacks() token_cost_logger = TokenAndCostLogger() - litellm.callbacks = [token_cost_logger] + monkeypatch.setattr(litellm, "callbacks", [token_cost_logger]) # Mock A2A client mock_client = MagicMock() @@ -292,7 +292,7 @@ class AgentIdLogger(CustomLogger): @pytest.mark.asyncio -async def test_asend_message_passes_agent_id_to_callback(): +async def test_asend_message_passes_agent_id_to_callback(monkeypatch): """ Test that asend_message passes agent_id to callbacks via kwargs. """ @@ -301,7 +301,7 @@ async def test_asend_message_passes_agent_id_to_callback(): # Setup logger litellm.logging_callback_manager._reset_all_callbacks() agent_id_logger = AgentIdLogger() - litellm.callbacks = [agent_id_logger] + monkeypatch.setattr(litellm, "callbacks", [agent_id_logger]) # Mock A2A client mock_client = MagicMock() diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py similarity index 97% rename from tests/test_litellm/a2a_protocol/test_main.py rename to tests/unit/a2a_protocol/test_main.py index f00ac16f7b3..c65d171246d 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/unit/a2a_protocol/test_main.py @@ -115,9 +115,7 @@ async def test_streaming_trace_id_prefers_logging_trace_id(): captured["extra_headers"] = extra_headers raise RuntimeError("stop") - with patch.object( - a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture) - ): + with patch.object(a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture)): with pytest.raises(RuntimeError, match="stop"): async for _ in a2a_main.asend_message_streaming( request=request, @@ -229,9 +227,7 @@ _LOWERCASE_BINDING_CARD = { "defaultInputModes": ["text/plain"], "defaultOutputModes": ["text/plain"], "skills": [], - "supportedInterfaces": [ - {"url": "http://127.0.0.1:9/", "protocolBinding": "jsonrpc", "protocolVersion": "1.0"} - ], + "supportedInterfaces": [{"url": "http://127.0.0.1:9/", "protocolBinding": "jsonrpc", "protocolVersion": "1.0"}], } @@ -289,11 +285,10 @@ async def _seed_shared_a2a_client( @pytest.fixture -def isolated_client_cache(): - previous = getattr(litellm, "in_memory_llm_clients_cache", None) - litellm.in_memory_llm_clients_cache = LLMClientCache() - yield litellm.in_memory_llm_clients_cache - litellm.in_memory_llm_clients_cache = previous +def isolated_client_cache(monkeypatch): + cache = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", cache) + return cache def _send_request(request_id): diff --git a/tests/test_litellm/a2a_protocol/test_send_message_response.py b/tests/unit/a2a_protocol/test_send_message_response.py similarity index 77% rename from tests/test_litellm/a2a_protocol/test_send_message_response.py rename to tests/unit/a2a_protocol/test_send_message_response.py index ade7c72fc2e..599e97e4923 100644 --- a/tests/test_litellm/a2a_protocol/test_send_message_response.py +++ b/tests/unit/a2a_protocol/test_send_message_response.py @@ -9,9 +9,7 @@ def test_from_dict_backfills_id_on_agent_error_response(): "error": {"code": -32054, "message": "Session not found"}, } - response = LiteLLMSendMessageResponse.from_dict( - agent_error, request_id="r1" - ) + response = LiteLLMSendMessageResponse.from_dict(agent_error, request_id="r1") assert response.id == "r1" assert response.error == {"code": -32054, "message": "Session not found"} @@ -25,9 +23,7 @@ def test_from_dict_preserves_existing_id(): "error": {"code": -32001, "message": "Task not found"}, } - response = LiteLLMSendMessageResponse.from_dict( - payload, request_id="r1" - ) + response = LiteLLMSendMessageResponse.from_dict(payload, request_id="r1") assert response.id == "upstream-id" @@ -82,9 +78,7 @@ def test_from_dict_accepts_null_id_when_the_error_cannot_be_correlated(): """JSON-RPC 2.0 section 5 requires ``id`` to be null on an error that cannot be matched to a request, which is exactly the case where the caller supplied no id for the backfill to use. Rejecting it turned an agent's error into a proxy 500.""" - response = LiteLLMSendMessageResponse.from_dict( - {"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}} - ) + response = LiteLLMSendMessageResponse.from_dict({"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}}) assert response.id is None assert response.error == {"code": -32054, "message": "x"} @@ -100,23 +94,6 @@ def test_from_dict_accepts_null_id_echoed_by_upstream(): assert response.id is None -def test_id_accepts_every_member_of_the_json_rpc_union_and_nothing_else(): - """One test pinning the whole ``string | integer | null`` union the spec defines, - so widening the annotation cannot silently become "accept anything".""" - for accepted in ("s1", 42, 0, None): - assert LiteLLMSendMessageResponse(id=accepted).id == accepted - - # ``True``/``False`` are in here because bool subclasses int: a non-strict integer - # half would accept them and relay them as 1/0. Direct construction bypasses - # normalization, so the model has to hold this line on its own. - for rejected in (True, False, 1.5, ["a"], {"a": 1}): - try: - LiteLLMSendMessageResponse(id=rejected) - except Exception: - continue - raise AssertionError(f"id={rejected!r} is outside the JSON-RPC union and must be rejected") - - def test_boolean_id_is_never_relayed_as_an_integer(): """``bool`` subclasses ``int``, so widening the annotation to accept integers also made pydantic coerce a boolean id to 1 or 0. That is worse than rejecting it: an id diff --git a/tests/test_litellm/a2a_protocol/test_utils.py b/tests/unit/a2a_protocol/test_utils.py similarity index 100% rename from tests/test_litellm/a2a_protocol/test_utils.py rename to tests/unit/a2a_protocol/test_utils.py From 9e84b8b048d2cec984bd4094c7243a48f6b95b3d Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:30:56 +0000 Subject: [PATCH 096/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d325f46017a..a5419899773 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2932e-07, + "input_cost_per_token": 5.2668e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7644e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7556e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7644e-08, - "input_cost_per_token": 5.2932e-07, + "cache_read_input_token_cost": 1.7556e-08, + "input_cost_per_token": 5.2668e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d325f46017a..a5419899773 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2932e-07, + "input_cost_per_token": 5.2668e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7644e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7556e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7644e-08, - "input_cost_per_token": 5.2932e-07, + "cache_read_input_token_cost": 1.7556e-08, + "input_cost_per_token": 5.2668e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 83e1f0834ebd3ea17738af52733cc01adaa19194 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 20:30:54 +0000 Subject: [PATCH 097/146] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-v4-flash-latest: output_cost_per_token openrouter/deepseek/deepseek-v4-flash-0731: output_cost_per_token openrouter/deepseek/deepseek-v4-flash-vision-exp: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a5419899773..a59ba1082da 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66550,13 +66550,13 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-vision-exp": { - "input_cost_per_token": 2.156e-07, - "output_cost_per_token": 6.468e-07, - "cache_read_input_token_cost": 6.86e-09, + "input_cost_per_token": 2.2e-07, + "output_cost_per_token": 6.6e-07, + "cache_read_input_token_cost": 7e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 262144, - "max_tokens": 262144, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a5419899773..a59ba1082da 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66550,13 +66550,13 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-vision-exp": { - "input_cost_per_token": 2.156e-07, - "output_cost_per_token": 6.468e-07, - "cache_read_input_token_cost": 6.86e-09, + "input_cost_per_token": 2.2e-07, + "output_cost_per_token": 6.6e-07, + "cache_read_input_token_cost": 7e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 262144, - "max_tokens": 262144, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 8a8a15af932ffb5ef3c5cf25a875d0c41b05fe5a Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:00:50 +0000 Subject: [PATCH 098/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-v4-flash-latest: output_cost_per_token openrouter/deepseek/deepseek-v4-flash-0731: output_cost_per_token --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a59ba1082da..d25cf0c78df 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a59ba1082da..d25cf0c78df 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 84c26978a3dc144e44c5e5e55b41e6a3d129093e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 14:10:31 -0700 Subject: [PATCH 099/146] refactor(rust): split token counter backends --- litellm-rust/Cargo.lock | 83 ++++- litellm-rust/Cargo.toml | 4 + litellm-rust/crates/python-bridge/Cargo.toml | 7 +- .../crates/python-bridge/src/token_counter.rs | 72 ++++- .../crates/token-counter-fast/Cargo.toml | 19 ++ .../src/byte_level.rs | 20 +- .../src/cl100k.rs | 0 .../crates/token-counter-fast/src/error.rs | 13 + .../crates/token-counter-fast/src/lib.rs | 70 +++++ .../src/o200k.rs | 0 .../src/scanner.rs | 0 .../crates/token-counter-fast/src/tiktoken.rs | 222 ++++++++++++++ .../src/unicode_classes.rs | 0 .../tests/fixtures/cl100k/requests.jsonl | 0 .../tests/fixtures/cl100k/texts.jsonl | 0 .../tests/fixtures/generate.py | 0 .../tests/fixtures/o200k/requests.jsonl | 0 .../tests/fixtures/o200k/texts.jsonl | 0 .../token-counter-huggingface/Cargo.toml | 10 + .../token-counter-huggingface/src/lib.rs | 29 ++ .../crates/token-counter-tiktoken/Cargo.toml | 10 + .../crates/token-counter-tiktoken/src/lib.rs | 54 ++++ litellm-rust/crates/token-counter/Cargo.toml | 15 +- litellm-rust/crates/token-counter/README.md | 23 ++ .../token-counter/benches/allocations.rs | 2 +- .../token-counter/benches/token_counter.rs | 2 +- .../crates/token-counter/src/counter.rs | 65 +--- .../crates/token-counter/src/error.rs | 6 +- litellm-rust/crates/token-counter/src/fast.rs | 41 +++ .../crates/token-counter/src/huggingface.rs | 27 ++ litellm-rust/crates/token-counter/src/lib.rs | 15 +- .../crates/token-counter/src/tiktoken.rs | 226 +------------- .../crates/token-counter/src/tokenizer.rs | 31 ++ .../token-counter/tests/token_counter.rs | 288 +++++++++--------- litellm/rust_bridge/_native.pyi | 4 +- litellm/rust_bridge/token_counter.py | 40 ++- .../spend_tracking/test_budget_reservation.py | 14 +- .../rust_bridge/test_token_counter.py | 33 +- 38 files changed, 944 insertions(+), 501 deletions(-) create mode 100644 litellm-rust/crates/token-counter-fast/Cargo.toml rename litellm-rust/crates/{token-counter => token-counter-fast}/src/byte_level.rs (97%) rename litellm-rust/crates/{token-counter => token-counter-fast}/src/cl100k.rs (100%) create mode 100644 litellm-rust/crates/token-counter-fast/src/error.rs create mode 100644 litellm-rust/crates/token-counter-fast/src/lib.rs rename litellm-rust/crates/{token-counter => token-counter-fast}/src/o200k.rs (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/src/scanner.rs (100%) create mode 100644 litellm-rust/crates/token-counter-fast/src/tiktoken.rs rename litellm-rust/crates/{token-counter => token-counter-fast}/src/unicode_classes.rs (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/cl100k/requests.jsonl (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/cl100k/texts.jsonl (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/generate.py (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/o200k/requests.jsonl (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/o200k/texts.jsonl (100%) create mode 100644 litellm-rust/crates/token-counter-huggingface/Cargo.toml create mode 100644 litellm-rust/crates/token-counter-huggingface/src/lib.rs create mode 100644 litellm-rust/crates/token-counter-tiktoken/Cargo.toml create mode 100644 litellm-rust/crates/token-counter-tiktoken/src/lib.rs create mode 100644 litellm-rust/crates/token-counter/README.md create mode 100644 litellm-rust/crates/token-counter/src/fast.rs create mode 100644 litellm-rust/crates/token-counter/src/huggingface.rs create mode 100644 litellm-rust/crates/token-counter/src/tokenizer.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2720cf01f2e..62bde5806a4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -61,6 +61,12 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + [[package]] name = "arc-swap" version = "1.9.2" @@ -598,6 +604,17 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", + "regex-automata", + "serde_core", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -1181,6 +1198,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fancy-regex" version = "0.19.2" @@ -1926,6 +1954,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -2083,7 +2117,7 @@ dependencies = [ name = "litellm-core-utils" version = "0.1.0" dependencies = [ - "fancy-regex", + "fancy-regex 0.19.2", "litellm-types", "rstest", "serde", @@ -2212,10 +2246,24 @@ dependencies = [ name = "litellm-token-counter" version = "0.1.0" dependencies = [ - "base64 0.22.1", "criterion", "indexmap 2.14.0", "itoa", + "litellm-token-counter-fast", + "litellm-token-counter-huggingface", + "litellm-token-counter-tiktoken", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "litellm-token-counter-fast" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", "rand 0.8.7", "rstest", "rustc-hash", @@ -2226,6 +2274,22 @@ dependencies = [ "unicode-normalization-alignments", ] +[[package]] +name = "litellm-token-counter-huggingface" +version = "0.1.0" +dependencies = [ + "thiserror 2.0.19", + "tokenizers", +] + +[[package]] +name = "litellm-token-counter-tiktoken" +version = "0.1.0" +dependencies = [ + "thiserror 2.0.19", + "tiktoken-rs", +] + [[package]] name = "litellm-types" version = "0.1.0" @@ -3794,6 +3858,21 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "tiktoken-rs" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex 0.17.0", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "time" version = "0.3.53" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a6185632871..109a2edaf4d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -24,6 +24,9 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-token-counter = { path = "crates/token-counter" } +litellm-token-counter-fast = { path = "crates/token-counter-fast" } +litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } +litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } bytes = "1" @@ -45,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" +tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 3a4a579efa3..2c868ebc769 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,10 +10,13 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3"] +default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] +token-counter-fast = ["litellm-token-counter/fast"] +token-counter-huggingface = ["litellm-token-counter/huggingface"] +token-counter-tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true @@ -26,7 +29,7 @@ litellm-http.workspace = true litellm-llms.workspace = true litellm-types.workspace = true litellm-host-python.workspace = true -litellm-token-counter.workspace = true +litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 7dc86b78ad6..a7c6bce4d0b 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -1,6 +1,19 @@ -use std::{num::NonZero, sync::Arc, thread::available_parallelism}; +use std::sync::Arc; -use litellm_host_python::{release_gil, run_async}; +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] +use std::{num::NonZero, thread::available_parallelism}; + +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] +use litellm_host_python::release_gil; +use litellm_host_python::run_async; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, }; @@ -28,17 +41,47 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) + #[cfg(feature = "token-counter-huggingface")] + { + Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) + } + #[cfg(not(feature = "token-counter-huggingface"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-huggingface feature", + )) + } } #[staticmethod] - fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file)) + fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { + #[cfg(feature = "token-counter-fast")] + { + Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + } + #[cfg(not(feature = "token-counter-fast"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-fast feature", + )) + } } #[staticmethod] - fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) + fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + #[cfg(feature = "token-counter-tiktoken")] + { + Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) + } + #[cfg(not(feature = "token-counter-tiktoken"))] + { + let _ = (py, encoding); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-tiktoken feature", + )) + } } fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { @@ -62,6 +105,11 @@ impl TokenCounter { } impl TokenCounter { + #[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" + ))] fn load( py: Python<'_>, load: impl FnOnce() -> Result + Send, @@ -74,6 +122,11 @@ impl TokenCounter { } } +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] fn encode_parallelism() -> usize { available_parallelism().map_or(1, NonZero::get) } @@ -86,7 +139,10 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result PyErr { let message = error.to_string(); match error { - Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message), + Error::Load(_) + | Error::Ranks(_) + | Error::UnicodeClasses + | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), Error::RequestParse(_) | Error::MissingInput | Error::FloatText diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml new file mode 100644 index 00000000000..4950277f5b0 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-token-counter-fast" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +base64.workspace = true +rustc-hash = "2.1.3" +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +unicode-normalization-alignments = "0.1.12" + +[dev-dependencies] +rand.workspace = true +rstest.workspace = true +serde.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/token-counter/src/byte_level.rs b/litellm-rust/crates/token-counter-fast/src/byte_level.rs similarity index 97% rename from litellm-rust/crates/token-counter/src/byte_level.rs rename to litellm-rust/crates/token-counter-fast/src/byte_level.rs index ec6134a252e..6fc7d9146b9 100644 --- a/litellm-rust/crates/token-counter/src/byte_level.rs +++ b/litellm-rust/crates/token-counter-fast/src/byte_level.rs @@ -473,13 +473,13 @@ mod tests { _ => unreachable!(), } assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none()); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", " stop"] { assert_eq!( - counter.count_text(text).expect("count"), + counter.count_tokens(text).expect("count"), reference_count(&anthropic_tokenizer, text) ); } @@ -545,7 +545,7 @@ mod tests { .rstrip(rstrip)]) .expect("add token"); let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); @@ -557,7 +557,7 @@ mod tests { ] { assert_eq!(fast.count(&anthropic_tokenizer, text), None); assert_eq!( - counter.count_text(text).expect("count"), + counter.count_tokens(text).expect("count"), reference_count(&anthropic_tokenizer, text) ); } @@ -571,17 +571,17 @@ mod tests { assert_eq!(fast.count(&tokenizer, "hello"), None); assert!(tokenizer.encode_fast("hello", true).is_err()); let counter = - crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize")) + crate::FastTokenizer::from_json(&tokenizer.to_string(false).expect("serialize")) .expect("load"); assert!(matches!( - counter.count_text("hello"), + counter.count_tokens("hello"), Err(crate::Error::Encode(_)) )); } #[rstest] fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) { - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); @@ -598,7 +598,7 @@ mod tests { scope.spawn(move || { for _ in 0..100 { for (text, count) in inputs.iter().zip(expected) { - assert_eq!(counter.count_text(text).expect("count"), count); + assert_eq!(counter.count_tokens(text).expect("count"), count); } } }); @@ -614,10 +614,10 @@ mod tests { let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1); assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); - assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1); + assert_eq!(counter.count_tokens("ABCD EFGH").expect("count"), 1); } } diff --git a/litellm-rust/crates/token-counter/src/cl100k.rs b/litellm-rust/crates/token-counter-fast/src/cl100k.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/cl100k.rs rename to litellm-rust/crates/token-counter-fast/src/cl100k.rs diff --git a/litellm-rust/crates/token-counter-fast/src/error.rs b/litellm-rust/crates/token-counter-fast/src/error.rs new file mode 100644 index 00000000000..e63ccf3ad39 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/error.rs @@ -0,0 +1,13 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("failed to load tokenizer: tiktoken rank file: {0}")] + Ranks(String), + #[error("failed to load tokenizer: Unicode character classes are unavailable")] + UnicodeClasses, + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} diff --git a/litellm-rust/crates/token-counter-fast/src/lib.rs b/litellm-rust/crates/token-counter-fast/src/lib.rs new file mode 100644 index 00000000000..ce91af642ea --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/lib.rs @@ -0,0 +1,70 @@ +#![forbid(unsafe_code)] + +mod byte_level; +mod cl100k; +mod error; +mod o200k; +mod scanner; +mod tiktoken; +mod unicode_classes; + +use byte_level::ByteLevelCounter; +use scanner::{SplitPattern, TiktokenCounter}; + +pub use error::Error; + +enum Encoder { + HuggingFace { + tokenizer: Box, + byte_level: Option, + }, + Tiktoken(TiktokenCounter), +} + +pub struct FastTokenizer(Encoder); + +impl FastTokenizer { + pub fn from_json(json: &str) -> Result { + let tokenizer = json.parse::().map_err(Error::Load)?; + let byte_level = ByteLevelCounter::detect(&tokenizer); + Ok(Self(Encoder::HuggingFace { + tokenizer: Box::new(tokenizer), + byte_level, + })) + } + + pub fn from_cl100k_ranks(ranks: &str) -> Result { + Self::from_ranks(SplitPattern::Cl100k, ranks) + } + + pub fn from_o200k_ranks(ranks: &str) -> Result { + Self::from_ranks(SplitPattern::O200k, ranks) + } + + fn from_ranks(split: SplitPattern, ranks: &str) -> Result { + TiktokenCounter::from_ranks(split, ranks) + .map(Encoder::Tiktoken) + .map(Self) + } + + pub fn count_tokens(&self, text: &str) -> Result { + match &self.0 { + Encoder::Tiktoken(counter) => Ok(counter.count(text)), + Encoder::HuggingFace { + tokenizer, + byte_level, + } => { + if let Some(count) = byte_level + .as_ref() + .and_then(|counter| counter.count(tokenizer, text)) + { + return Ok(count); + } + tokenizer + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } + } + } +} diff --git a/litellm-rust/crates/token-counter/src/o200k.rs b/litellm-rust/crates/token-counter-fast/src/o200k.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/o200k.rs rename to litellm-rust/crates/token-counter-fast/src/o200k.rs diff --git a/litellm-rust/crates/token-counter/src/scanner.rs b/litellm-rust/crates/token-counter-fast/src/scanner.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/scanner.rs rename to litellm-rust/crates/token-counter-fast/src/scanner.rs diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs new file mode 100644 index 00000000000..7a9e71ed587 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -0,0 +1,222 @@ +//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and +//! the merge loop that turns one regex piece into tokens. The merge order is +//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is +//! identical, but pairs are tracked in a heap so a long piece costs +//! `O(n log n)` instead of tiktoken's `O(n^2)`. + +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +use base64::Engine; +use base64::engine::general_purpose::STANDARD; +use rustc_hash::FxHashMap; + +use crate::Error; + +type Rank = u32; + +const NO_RANK: Rank = Rank::MAX; +const END: usize = usize::MAX; + +pub(super) struct MergeRanks(FxHashMap, Rank>); + +impl MergeRanks { + pub(super) fn parse(text: &str) -> Result { + let ranks = text + .lines() + .filter(|line| !line.is_empty()) + .map(parse_line) + .collect::, _>>()?; + if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { + return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); + } + Ok(Self(ranks)) + } + + fn rank(&self, bytes: &[u8]) -> Rank { + self.0.get(bytes).copied().unwrap_or(NO_RANK) + } + + /// Token count of one regex piece, as `encode_ordinary` would produce. + pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize { + if piece.len() < 2 || self.0.contains_key(piece) { + return 1; + } + scratch.reset(piece.len()); + for start in 0..piece.len() - 1 { + scratch.set_rank(start, self.rank(&piece[start..start + 2])); + } + let mut parts = piece.len(); + while let Some(Reverse((rank, start))) = scratch.heap.pop() { + if scratch.next[start] == END || scratch.rank[start] != rank { + continue; + } + let merged = scratch.next[start]; + let after = scratch.next[merged]; + scratch.next[merged] = END; + scratch.next[start] = after; + parts -= 1; + if after < piece.len() { + scratch.prev[after] = start; + scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)])); + } else { + scratch.rank[start] = NO_RANK; + } + let before = scratch.prev[start]; + if before != END { + scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)])); + } + } + parts + } +} + +fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { + let (token, rank) = line + .split_once(' ') + .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?; + let bytes = STANDARD + .decode(token) + .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?; + let rank = rank + .parse() + .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; + Ok((bytes.into_boxed_slice(), rank)) +} + +/// Buffers reused across the pieces of one text. Parts are addressed by the +/// byte offset they start at, which also gives the leftmost-pair tie break. +#[derive(Default)] +pub(super) struct MergeScratch { + next: Vec, + prev: Vec, + rank: Vec, + heap: BinaryHeap>, +} + +impl MergeScratch { + fn reset(&mut self, len: usize) { + self.next.clear(); + self.next.extend(1..=len); + self.prev.clear(); + self.prev.push(END); + self.prev.extend(0..len - 1); + self.rank.clear(); + self.rank.resize(len, NO_RANK); + self.heap.clear(); + } + + fn end(&self, start: usize) -> usize { + self.next[start] + } + + fn set_rank(&mut self, start: usize, rank: Rank) { + self.rank[start] = rank; + if rank != NO_RANK { + self.heap.push(Reverse((rank, start))); + } + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + use super::*; + + fn ranks() -> MergeRanks { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" + ); + MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo")) + .expect("rank file parses") + } + + /// tiktoken's `_byte_pair_merge`, transcribed, as the reference. + fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize { + if piece.len() < 2 || ranks.0.contains_key(piece) { + return 1; + } + let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1) + .map(|index| (index, ranks.rank(&piece[index..index + 2]))) + .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)]) + .collect(); + let get_rank = |parts: &[(usize, Rank)], index: usize| { + if index + 3 < parts.len() { + ranks.rank(&piece[parts[index].0..parts[index + 3].0]) + } else { + NO_RANK + } + }; + loop { + let Some(index) = parts[..parts.len() - 1] + .iter() + .enumerate() + .filter(|(_, (_, rank))| *rank != NO_RANK) + .min_by_key(|(index, (_, rank))| (*rank, *index)) + .map(|(index, _)| index) + else { + return parts.len() - 1; + }; + if index > 0 { + parts[index - 1].1 = get_rank(&parts, index - 1); + } + parts[index].1 = get_rank(&parts, index); + parts.remove(index + 1); + } + } + + #[test] + fn every_byte_is_a_token() { + let ranks = ranks(); + assert_eq!(ranks.0.len(), 100_256); + assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK)); + } + + #[test] + fn heap_merge_matches_tiktokens_merge_loop() { + let ranks = ranks(); + let mut scratch = MergeScratch::default(); + let mut rng = StdRng::seed_from_u64(99); + let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123"; + for _ in 0..20_000 { + let piece: Vec = (0..rng.gen_range(1..24)) + .map(|_| alphabet[rng.gen_range(0..alphabet.len())]) + .collect(); + assert_eq!( + ranks.count_piece(&piece, &mut scratch), + reference_count(&ranks, &piece), + "piece {:?}", + String::from_utf8_lossy(&piece) + ); + } + } + + #[test] + fn long_repeated_runs_cost_close_to_linear() { + let ranks = ranks(); + let mut scratch = MergeScratch::default(); + let mut time = |len: usize| { + let piece = vec![b' '; len]; + let started = std::time::Instant::now(); + assert!(ranks.count_piece(&piece, &mut scratch) > 0); + started.elapsed() + }; + let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); + let large = time(1 << 18); + assert!( + large < small * 64, + "{small:?} for 2^14 bytes, {large:?} for 2^18" + ); + } + + #[test] + fn malformed_rank_files_are_rejected() { + assert!(MergeRanks::parse("IQ==").is_err()); + assert!(MergeRanks::parse("IQ== x").is_err()); + assert!(MergeRanks::parse("!!! 1").is_err()); + assert!(MergeRanks::parse("IQ== 1").is_err()); + } +} diff --git a/litellm-rust/crates/token-counter/src/unicode_classes.rs b/litellm-rust/crates/token-counter-fast/src/unicode_classes.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/unicode_classes.rs rename to litellm-rust/crates/token-counter-fast/src/unicode_classes.rs diff --git a/litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/requests.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/requests.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/texts.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/texts.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/generate.py b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/generate.py rename to litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py diff --git a/litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/requests.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/requests.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/texts.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/texts.jsonl diff --git a/litellm-rust/crates/token-counter-huggingface/Cargo.toml b/litellm-rust/crates/token-counter-huggingface/Cargo.toml new file mode 100644 index 00000000000..ad15d9f19d0 --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "litellm-token-counter-huggingface" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs new file mode 100644 index 00000000000..184956b5f39 --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -0,0 +1,29 @@ +#![forbid(unsafe_code)] + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} + +pub struct HuggingFaceTokenizer(Box); + +impl HuggingFaceTokenizer { + pub fn from_json(json: &str) -> Result { + json.parse::() + .map(Box::new) + .map(Self) + .map_err(Error::Load) + } + + pub fn count_tokens(&self, text: &str) -> Result { + self.0 + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } +} diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml new file mode 100644 index 00000000000..494a9233e69 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "litellm-token-counter-tiktoken" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +thiserror.workspace = true +tiktoken-rs.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs new file mode 100644 index 00000000000..b35c4876d00 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -0,0 +1,54 @@ +#![forbid(unsafe_code)] + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +#[error("unsupported tokenizer: {0}")] +pub struct UnsupportedTokenizer(pub String); + +pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); + +impl TiktokenTokenizer { + pub fn from_name(name: &str) -> Result { + let tokenizer = match name { + "cl100k_base" => tiktoken_rs::cl100k_base_singleton(), + "o200k_base" => tiktoken_rs::o200k_base_singleton(), + "o200k_harmony" => tiktoken_rs::o200k_harmony_singleton(), + "p50k_base" => tiktoken_rs::p50k_base_singleton(), + "p50k_edit" => tiktoken_rs::p50k_edit_singleton(), + "r50k_base" | "gpt2" => tiktoken_rs::r50k_base_singleton(), + _ => return Err(UnsupportedTokenizer(name.to_owned())), + }; + Ok(Self(tokenizer)) + } + + pub fn count_tokens(&self, text: &str) -> usize { + self.0.count_ordinary(text) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn special_tokens_are_counted_as_ordinary_text() { + let counter = TiktokenTokenizer::from_name("cl100k_base").unwrap(); + assert!(counter.count_tokens("<|endoftext|>") > 1); + } + + #[test] + fn all_python_tiktoken_encodings_are_available() { + for name in [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ] { + assert!(TiktokenTokenizer::from_name(name).is_ok(), "{name}"); + } + } +} diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index d0369631682..59e4f9a6a1d 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -5,16 +5,21 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +default = ["huggingface", "tiktoken"] +fast = ["dep:litellm-token-counter-fast"] +huggingface = ["dep:litellm-token-counter-huggingface"] +tiktoken = ["dep:litellm-token-counter-tiktoken"] + [dependencies] -base64.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } itoa = "1.0" -rustc-hash = "2.1.3" +litellm-token-counter-fast = { workspace = true, optional = true } +litellm-token-counter-huggingface = { workspace = true, optional = true } +litellm-token-counter-tiktoken = { workspace = true, optional = true } serde.workspace = true serde_json.workspace = true thiserror.workspace = true -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } -unicode-normalization-alignments = "0.1.12" [dev-dependencies] criterion.workspace = true @@ -24,7 +29,9 @@ rstest.workspace = true [[bench]] name = "token_counter" harness = false +required-features = ["fast"] [[bench]] name = "allocations" harness = false +required-features = ["fast"] diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md new file mode 100644 index 00000000000..c8ed0737a7e --- /dev/null +++ b/litellm-rust/crates/token-counter/README.md @@ -0,0 +1,23 @@ +# Token counting + +`Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface + +The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast`, `TokenCounter::from_cl100k_ranks`, and `TokenCounter::from_o200k_ranks` use this implementation + +The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through the upstream `tokenizers` library. `TokenCounter::from_json` uses this implementation + +The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` + +The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend + +Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits + +Run the feature matrix with: + +```sh +cargo test -p litellm-token-counter +cargo test -p litellm-token-counter --no-default-features +cargo test -p litellm-token-counter --no-default-features --features fast +cargo test -p litellm-token-counter --no-default-features --features huggingface +cargo test -p litellm-token-counter --no-default-features --features tiktoken +``` diff --git a/litellm-rust/crates/token-counter/benches/allocations.rs b/litellm-rust/crates/token-counter/benches/allocations.rs index 343f815749d..7aebceb6e9c 100644 --- a/litellm-rust/crates/token-counter/benches/allocations.rs +++ b/litellm-rust/crates/token-counter/benches/allocations.rs @@ -87,7 +87,7 @@ fn main() { }, ); - let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads"); + let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("tokenizer loads"); let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses"); counter .count_request(&object) diff --git a/litellm-rust/crates/token-counter/benches/token_counter.rs b/litellm-rust/crates/token-counter/benches/token_counter.rs index a7c3177b88a..5e30cf6f77e 100644 --- a/litellm-rust/crates/token-counter/benches/token_counter.rs +++ b/litellm-rust/crates/token-counter/benches/token_counter.rs @@ -42,7 +42,7 @@ fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> { } fn token_counter(c: &mut Criterion) { - let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load"); + let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("token counter should load"); let tokenizer = TOKENIZER_JSON .parse::() .expect("reference tokenizer should load"); diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs index 7eedc449dd1..ce08e225be4 100644 --- a/litellm-rust/crates/token-counter/src/counter.rs +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -1,9 +1,7 @@ use serde::Serialize; use crate::Error; -use crate::byte_level::ByteLevelCounter; use crate::python_json; -use crate::scanner::{SplitPattern, TiktokenCounter}; use crate::tools::format_function_definitions; use crate::types::{ ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice, @@ -24,73 +22,22 @@ pub struct InputTokenCount { pub input_tokens: usize, } -enum Encoder { - HuggingFace { - tokenizer: Box, - byte_level: Option, - }, - Tiktoken(TiktokenCounter), -} - /// A loaded tokenizer plus the message accounting Python applies on top of /// it. Encoding is CPU-bound and synchronous; hosts run it off their event /// loop. pub struct TokenCounter { - encoder: Encoder, + encoder: Box, } impl TokenCounter { - /// Load a HuggingFace `tokenizer.json` document. The host reads the file. - pub fn from_json(tokenizer_json: &str) -> Result { - let tokenizer = tokenizer_json - .parse::() - .map_err(Error::Load)?; - let byte_level = ByteLevelCounter::detect(&tokenizer); - Ok(Self { - encoder: Encoder::HuggingFace { - tokenizer: Box::new(tokenizer), - byte_level, - }, - }) - } - - /// Load tiktoken's `cl100k_base` rank file (`base64(token) rank` lines). - /// The host reads the file. - pub fn from_cl100k_ranks(rank_file: &str) -> Result { - Self::from_tiktoken_ranks(SplitPattern::Cl100k, rank_file) - } - - /// Load tiktoken's `o200k_base` rank file (`base64(token) rank` lines). - /// The host reads the file. - pub fn from_o200k_ranks(rank_file: &str) -> Result { - Self::from_tiktoken_ranks(SplitPattern::O200k, rank_file) - } - - fn from_tiktoken_ranks(split: SplitPattern, rank_file: &str) -> Result { - Ok(Self { - encoder: Encoder::Tiktoken(TiktokenCounter::from_ranks(split, rank_file)?), - }) + pub fn new(tokenizer: impl crate::Tokenizer + 'static) -> Self { + Self { + encoder: Box::new(tokenizer), + } } pub fn count_text(&self, text: &str) -> Result { - match &self.encoder { - Encoder::Tiktoken(counter) => Ok(counter.count(text)), - Encoder::HuggingFace { - tokenizer, - byte_level, - } => { - if let Some(count) = byte_level - .as_ref() - .and_then(|counter| counter.count(tokenizer, text)) - { - return Ok(count); - } - tokenizer - .encode_fast(text, true) - .map(|encoding| encoding.len()) - .map_err(Error::Encode) - } - } + self.encoder.count_tokens(text) } /// Mirrors the host's key precedence: `messages`, then `prompt`, then diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index 6b8668fe182..b05ce007e46 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -4,8 +4,10 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { + #[error("unsupported tokenizer: {0}")] + UnsupportedTokenizer(String), #[error("failed to load tokenizer: {0}")] - Load(#[source] tokenizers::Error), + Load(#[source] Box), #[error("failed to load tokenizer: tiktoken rank file: {0}")] Ranks(String), #[error("failed to load tokenizer: Unicode character classes are unavailable")] @@ -29,7 +31,7 @@ pub enum Error { #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")] JsonUtf8(#[source] FromUtf8Error), #[error("tokenization failed: {0}")] - Encode(#[source] tokenizers::Error), + Encode(#[source] Box), #[error("token counting task failed: {0}")] Task(String), } diff --git a/litellm-rust/crates/token-counter/src/fast.rs b/litellm-rust/crates/token-counter/src/fast.rs new file mode 100644 index 00000000000..de3f86abd68 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/fast.rs @@ -0,0 +1,41 @@ +use litellm_token_counter_fast::Error as BackendError; +pub use litellm_token_counter_fast::FastTokenizer; + +use crate::{Error, TokenCounter, Tokenizer}; + +impl TokenCounter { + pub fn from_json_fast(tokenizer_json: &str) -> Result { + FastTokenizer::from_json(tokenizer_json) + .map(Self::new) + .map_err(Error::from) + } + + pub fn from_cl100k_ranks(rank_file: &str) -> Result { + FastTokenizer::from_cl100k_ranks(rank_file) + .map(Self::new) + .map_err(Error::from) + } + + pub fn from_o200k_ranks(rank_file: &str) -> Result { + FastTokenizer::from_o200k_ranks(rank_file) + .map(Self::new) + .map_err(Error::from) + } +} + +impl Tokenizer for FastTokenizer { + fn count_tokens(&self, text: &str) -> Result { + FastTokenizer::count_tokens(self, text).map_err(Error::from) + } +} + +impl From for Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Load(source) => Self::Load(source), + BackendError::Ranks(message) => Self::Ranks(message), + BackendError::UnicodeClasses => Self::UnicodeClasses, + BackendError::Encode(source) => Self::Encode(source), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs new file mode 100644 index 00000000000..fb7683b373e --- /dev/null +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -0,0 +1,27 @@ +use litellm_token_counter_huggingface::Error as BackendError; +pub use litellm_token_counter_huggingface::HuggingFaceTokenizer; + +use crate::{Error, TokenCounter, Tokenizer}; + +impl TokenCounter { + pub fn from_json(tokenizer_json: &str) -> Result { + HuggingFaceTokenizer::from_json(tokenizer_json) + .map(Self::new) + .map_err(Error::from) + } +} + +impl Tokenizer for HuggingFaceTokenizer { + fn count_tokens(&self, text: &str) -> Result { + HuggingFaceTokenizer::count_tokens(self, text).map_err(Error::from) + } +} + +impl From for Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Load(source) => Self::Load(source), + BackendError::Encode(source) => Self::Encode(source), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index fa0014e2bad..446c91049de 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -4,18 +4,21 @@ #![forbid(unsafe_code)] -mod byte_level; -mod cl100k; mod counter; mod error; -mod o200k; mod python_json; -mod scanner; -mod tiktoken; +mod tokenizer; mod tools; mod types; -mod unicode_classes; + +#[cfg(feature = "fast")] +pub mod fast; +#[cfg(feature = "huggingface")] +pub mod huggingface; +#[cfg(feature = "tiktoken")] +pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; pub use error::Error; +pub use tokenizer::Tokenizer; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index 7a9e71ed587..07c1c9f5b73 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,222 +1,24 @@ -//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and -//! the merge loop that turns one regex piece into tokens. The merge order is -//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is -//! identical, but pairs are tracked in a heap so a long piece costs -//! `O(n log n)` instead of tiktoken's `O(n^2)`. +pub use litellm_token_counter_tiktoken::TiktokenTokenizer; +use litellm_token_counter_tiktoken::UnsupportedTokenizer; -use std::cmp::Reverse; -use std::collections::BinaryHeap; +use crate::{Error, TokenCounter, Tokenizer}; -use base64::Engine; -use base64::engine::general_purpose::STANDARD; -use rustc_hash::FxHashMap; - -use crate::Error; - -type Rank = u32; - -const NO_RANK: Rank = Rank::MAX; -const END: usize = usize::MAX; - -pub(super) struct MergeRanks(FxHashMap, Rank>); - -impl MergeRanks { - pub(super) fn parse(text: &str) -> Result { - let ranks = text - .lines() - .filter(|line| !line.is_empty()) - .map(parse_line) - .collect::, _>>()?; - if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { - return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); - } - Ok(Self(ranks)) - } - - fn rank(&self, bytes: &[u8]) -> Rank { - self.0.get(bytes).copied().unwrap_or(NO_RANK) - } - - /// Token count of one regex piece, as `encode_ordinary` would produce. - pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize { - if piece.len() < 2 || self.0.contains_key(piece) { - return 1; - } - scratch.reset(piece.len()); - for start in 0..piece.len() - 1 { - scratch.set_rank(start, self.rank(&piece[start..start + 2])); - } - let mut parts = piece.len(); - while let Some(Reverse((rank, start))) = scratch.heap.pop() { - if scratch.next[start] == END || scratch.rank[start] != rank { - continue; - } - let merged = scratch.next[start]; - let after = scratch.next[merged]; - scratch.next[merged] = END; - scratch.next[start] = after; - parts -= 1; - if after < piece.len() { - scratch.prev[after] = start; - scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)])); - } else { - scratch.rank[start] = NO_RANK; - } - let before = scratch.prev[start]; - if before != END { - scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)])); - } - } - parts +impl TokenCounter { + pub fn from_tiktoken(encoding: &str) -> Result { + TiktokenTokenizer::from_name(encoding) + .map(Self::new) + .map_err(Error::from) } } -fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { - let (token, rank) = line - .split_once(' ') - .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?; - let bytes = STANDARD - .decode(token) - .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?; - let rank = rank - .parse() - .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; - Ok((bytes.into_boxed_slice(), rank)) -} - -/// Buffers reused across the pieces of one text. Parts are addressed by the -/// byte offset they start at, which also gives the leftmost-pair tie break. -#[derive(Default)] -pub(super) struct MergeScratch { - next: Vec, - prev: Vec, - rank: Vec, - heap: BinaryHeap>, -} - -impl MergeScratch { - fn reset(&mut self, len: usize) { - self.next.clear(); - self.next.extend(1..=len); - self.prev.clear(); - self.prev.push(END); - self.prev.extend(0..len - 1); - self.rank.clear(); - self.rank.resize(len, NO_RANK); - self.heap.clear(); - } - - fn end(&self, start: usize) -> usize { - self.next[start] - } - - fn set_rank(&mut self, start: usize, rank: Rank) { - self.rank[start] = rank; - if rank != NO_RANK { - self.heap.push(Reverse((rank, start))); - } +impl Tokenizer for TiktokenTokenizer { + fn count_tokens(&self, text: &str) -> Result { + Ok(TiktokenTokenizer::count_tokens(self, text)) } } -#[cfg(test)] -mod tests { - use rand::rngs::StdRng; - use rand::{Rng, SeedableRng}; - - use super::*; - - fn ranks() -> MergeRanks { - let path = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" - ); - MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo")) - .expect("rank file parses") - } - - /// tiktoken's `_byte_pair_merge`, transcribed, as the reference. - fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize { - if piece.len() < 2 || ranks.0.contains_key(piece) { - return 1; - } - let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1) - .map(|index| (index, ranks.rank(&piece[index..index + 2]))) - .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)]) - .collect(); - let get_rank = |parts: &[(usize, Rank)], index: usize| { - if index + 3 < parts.len() { - ranks.rank(&piece[parts[index].0..parts[index + 3].0]) - } else { - NO_RANK - } - }; - loop { - let Some(index) = parts[..parts.len() - 1] - .iter() - .enumerate() - .filter(|(_, (_, rank))| *rank != NO_RANK) - .min_by_key(|(index, (_, rank))| (*rank, *index)) - .map(|(index, _)| index) - else { - return parts.len() - 1; - }; - if index > 0 { - parts[index - 1].1 = get_rank(&parts, index - 1); - } - parts[index].1 = get_rank(&parts, index); - parts.remove(index + 1); - } - } - - #[test] - fn every_byte_is_a_token() { - let ranks = ranks(); - assert_eq!(ranks.0.len(), 100_256); - assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK)); - } - - #[test] - fn heap_merge_matches_tiktokens_merge_loop() { - let ranks = ranks(); - let mut scratch = MergeScratch::default(); - let mut rng = StdRng::seed_from_u64(99); - let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123"; - for _ in 0..20_000 { - let piece: Vec = (0..rng.gen_range(1..24)) - .map(|_| alphabet[rng.gen_range(0..alphabet.len())]) - .collect(); - assert_eq!( - ranks.count_piece(&piece, &mut scratch), - reference_count(&ranks, &piece), - "piece {:?}", - String::from_utf8_lossy(&piece) - ); - } - } - - #[test] - fn long_repeated_runs_cost_close_to_linear() { - let ranks = ranks(); - let mut scratch = MergeScratch::default(); - let mut time = |len: usize| { - let piece = vec![b' '; len]; - let started = std::time::Instant::now(); - assert!(ranks.count_piece(&piece, &mut scratch) > 0); - started.elapsed() - }; - let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); - let large = time(1 << 18); - assert!( - large < small * 64, - "{small:?} for 2^14 bytes, {large:?} for 2^18" - ); - } - - #[test] - fn malformed_rank_files_are_rejected() { - assert!(MergeRanks::parse("IQ==").is_err()); - assert!(MergeRanks::parse("IQ== x").is_err()); - assert!(MergeRanks::parse("!!! 1").is_err()); - assert!(MergeRanks::parse("IQ== 1").is_err()); +impl From for Error { + fn from(error: UnsupportedTokenizer) -> Self { + Self::UnsupportedTokenizer(error.0) } } diff --git a/litellm-rust/crates/token-counter/src/tokenizer.rs b/litellm-rust/crates/token-counter/src/tokenizer.rs new file mode 100644 index 00000000000..88c29c672a7 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/tokenizer.rs @@ -0,0 +1,31 @@ +use crate::Error; + +pub trait Tokenizer: Send + Sync { + fn count_tokens(&self, text: &str) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{CountableRequest, TokenCounter}; + + struct Characters; + + impl Tokenizer for Characters { + fn count_tokens(&self, text: &str) -> Result { + Ok(text.chars().count()) + } + } + + #[test] + fn request_accounting_works_with_an_injected_backend() { + let counter = TokenCounter::new(Characters); + let request = + CountableRequest::parse(br#"{"messages":[{"role":"user","content":"hello"}]}"#) + .unwrap(); + assert_eq!( + counter.count_request(&request).unwrap().input_tokens, + 3 + 4 + 5 + 3 + ); + } +} diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 12c59768952..378a16d5d72 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,5 +1,6 @@ +#![cfg(feature = "huggingface")] + use rstest::rstest; -use serde::Deserialize; use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; @@ -176,146 +177,155 @@ fn loading_a_bad_tokenizer_is_a_load_error() { assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); } -/// A tiktoken encoding: its fixture directory, the vendored rank file Python -/// loads, the constructor, and the model `generate.py` counted the requests for. -#[derive(Clone, Copy)] -struct TiktokenEncoding { - fixtures: &'static str, - rank_file: &'static str, - load: fn(&str) -> Result, - model: &'static str, -} +#[cfg(feature = "fast")] +mod fast { + use super::*; + use serde::Deserialize; -const CL100K: TiktokenEncoding = TiktokenEncoding { - fixtures: "cl100k", - rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", - load: TokenCounter::from_cl100k_ranks, - model: "gpt-4", -}; + /// A tiktoken encoding: its fixture directory, the vendored rank file Python + /// loads, the constructor, and the model `generate.py` counted the requests for. + #[derive(Clone, Copy)] + struct TiktokenEncoding { + fixtures: &'static str, + rank_file: &'static str, + load: fn(&str) -> Result, + model: &'static str, + } -const O200K: TiktokenEncoding = TiktokenEncoding { - fixtures: "o200k", - rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", - load: TokenCounter::from_o200k_ranks, - model: "gpt-4o", -}; - -fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { - let path = format!( - "{}/../../../litellm/litellm_core_utils/tokenizers/{}", - env!("CARGO_MANIFEST_DIR"), - encoding.rank_file - ); - let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); - (encoding.load)(&ranks).expect("ranks load") -} - -fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { - let path = format!( - "{}/tests/fixtures/{}/{name}", - env!("CARGO_MANIFEST_DIR"), - encoding.fixtures - ); - std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") -} - -#[derive(Deserialize)] -struct TextFixture { - text: String, - tokens: usize, -} - -#[derive(Deserialize)] -struct RequestFixture { - body: String, - input_tokens: usize, -} - -/// Reference counts come from `tiktoken.get_encoding(name)`; see -/// `tests/fixtures/generate.py`. -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { - let counter = tiktoken_counter(encoding); - let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") - .lines() - .map(|line| serde_json::from_str(line).expect("fixture line is json")) - .collect(); - assert!(fixtures.len() > 3000); - let mismatches: Vec<_> = fixtures - .iter() - .filter_map(|fixture| { - let count = counter.count_text(&fixture.text).expect("text counts"); - (count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count)) - }) - .collect(); - assert!( - mismatches.is_empty(), - "(text, tiktoken, rust): {mismatches:?}" - ); -} - -/// Reference counts come from the proxy's admission counter -/// (`_count_input_tokens(body, model)`), so this pins the shared message, -/// tool and reply-priming accounting on the tiktoken paths as well. -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { - let counter = tiktoken_counter(encoding); - let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") - .lines() - .map(|line| serde_json::from_str(line).expect("fixture line is json")) - .collect(); - let counts: Vec = fixtures - .iter() - .map(|fixture| { - let request = CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses"); - let count = counter.count_request(&request).expect("fixture counts"); - assert_eq!(count.model.as_deref(), Some(encoding.model)); - assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body); - count.input_tokens - }) - .collect(); - assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000)); -} - -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( - #[case] encoding: TiktokenEncoding, -) { - let counter = tiktoken_counter(encoding); - let count = |body: &str| { - counter - .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) - .expect("counts") - .input_tokens + const CL100K: TiktokenEncoding = TiktokenEncoding { + fixtures: "cl100k", + rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", + load: TokenCounter::from_cl100k_ranks, + model: "gpt-4", }; - let text = |text: &str| counter.count_text(text).expect("counts"); - let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); - assert_eq!(base, 3 + text("user") + text("hi") + 3); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#), - base + text("al") + 1 - ); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), - base + 1 - ); -} -#[rstest] -#[case::empty("")] -#[case::not_base64("!!!! 0")] -#[case::missing_rank("YQ==")] -#[case::rank_not_a_number("YQ== x")] -#[case::single_byte_tokens_missing("YWI= 0")] -fn loading_a_bad_rank_file_is_a_load_error( - #[case] rank_file: &str, - #[values(CL100K, O200K)] encoding: TiktokenEncoding, -) { - assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + const O200K: TiktokenEncoding = TiktokenEncoding { + fixtures: "o200k", + rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", + load: TokenCounter::from_o200k_ranks, + model: "gpt-4o", + }; + + fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { + let path = format!( + "{}/../../../litellm/litellm_core_utils/tokenizers/{}", + env!("CARGO_MANIFEST_DIR"), + encoding.rank_file + ); + let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); + (encoding.load)(&ranks).expect("ranks load") + } + + fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { + let path = format!( + "{}/../token-counter-fast/tests/fixtures/{}/{name}", + env!("CARGO_MANIFEST_DIR"), + encoding.fixtures + ); + std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") + } + + #[derive(Deserialize)] + struct TextFixture { + text: String, + tokens: usize, + } + + #[derive(Deserialize)] + struct RequestFixture { + body: String, + input_tokens: usize, + } + + /// Reference counts come from `tiktoken.get_encoding(name)`; see + /// `tests/fixtures/generate.py`. + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { + let counter = tiktoken_counter(encoding); + let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") + .lines() + .map(|line| serde_json::from_str(line).expect("fixture line is json")) + .collect(); + assert!(fixtures.len() > 3000); + let mismatches: Vec<_> = fixtures + .iter() + .filter_map(|fixture| { + let count = counter.count_text(&fixture.text).expect("text counts"); + (count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count)) + }) + .collect(); + assert!( + mismatches.is_empty(), + "(text, tiktoken, rust): {mismatches:?}" + ); + } + + /// Reference counts come from the proxy's admission counter + /// (`_count_input_tokens(body, model)`), so this pins the shared message, + /// tool and reply-priming accounting on the tiktoken paths as well. + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { + let counter = tiktoken_counter(encoding); + let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") + .lines() + .map(|line| serde_json::from_str(line).expect("fixture line is json")) + .collect(); + let counts: Vec = fixtures + .iter() + .map(|fixture| { + let request = + CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses"); + let count = counter.count_request(&request).expect("fixture counts"); + assert_eq!(count.model.as_deref(), Some(encoding.model)); + assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body); + count.input_tokens + }) + .collect(); + assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000)); + } + + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( + #[case] encoding: TiktokenEncoding, + ) { + let counter = tiktoken_counter(encoding); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let text = |text: &str| counter.count_text(text).expect("counts"); + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!(base, 3 + text("user") + text("hi") + 3); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#), + base + text("al") + 1 + ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# + ), + base + 1 + ); + } + + #[rstest] + #[case::empty("")] + #[case::not_base64("!!!! 0")] + #[case::missing_rank("YQ==")] + #[case::rank_not_a_number("YQ== x")] + #[case::single_byte_tokens_missing("YWI= 0")] + fn loading_a_bad_rank_file_is_a_load_error( + #[case] rank_file: &str, + #[values(CL100K, O200K)] encoding: TiktokenEncoding, + ) { + assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c0a06364261..7d20fc5a02f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -97,9 +97,9 @@ class ResponsesWebSocketConnection: class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_cl100k_ranks(rank_file: str) -> TokenCounter: ... + def from_json_fast(tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_o200k_ranks(rank_file: str) -> TokenCounter: ... + def from_tiktoken(encoding: str) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... def gil_stats() -> dict[str, int]: ... diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d36234f56c1..e816f1ab388 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -11,14 +11,21 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt from litellm.utils import claude_json_str, huggingface_tokenizer_kind -RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"] +RustTokenizer = Literal[ + "anthropic", + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", +] class RustTokenCounter(Protocol): @@ -30,10 +37,7 @@ class RustTokenCounterFactory(Protocol): def __call__(self, tokenizer_json: str) -> RustTokenCounter: raise NotImplementedError - def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter: - raise NotImplementedError - - def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter: + def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter: raise NotImplementedError @@ -63,9 +67,8 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: """The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count. Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace - downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust - prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in - Python.""" + downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every + message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python.""" if litellm.disable_token_counter is True: return None kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model) @@ -73,24 +76,19 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: return "anthropic" if kind is not None or uses_legacy_message_accounting(model): return None - match openai_tokenizer_encoding(model).name: - case "cl100k_base": - return "cl100k_base" - case "o200k_base": - return "o200k_base" - case _: - return None + encoding: Final = openai_tokenizer_encoding(model).name + if encoding in {"cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base"}: + return cast(RustTokenizer, encoding) + return None -@lru_cache(maxsize=4) +@lru_cache(maxsize=8) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: match tokenizer: case "anthropic": return factory(claude_json_str) - case "cl100k_base": - return factory.from_cl100k_ranks(cl100k_base_rank_file()) - case "o200k_base": - return factory.from_o200k_ranks(o200k_base_rank_file()) + case _: + return factory.from_tiktoken(tokenizer) async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 0e0025f5194..c785a7b657a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -256,7 +256,7 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`.""" + """Stands in for the native `TokenCounter` class.""" def __init__(self) -> None: self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = [] @@ -264,11 +264,8 @@ class _RecordingFactory: def __call__(self, tokenizer_json: str) -> _RecordingCounter: return _RecordingCounter(self, "anthropic") - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "o200k_base") + def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter: + return _RecordingCounter(self, encoding) class _DecliningCounter: @@ -280,10 +277,7 @@ class _DecliningFactory: def __call__(self, tokenizer_json: str) -> _DecliningCounter: return _DecliningCounter() - def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter: - return _DecliningCounter() - - def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter: + def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter: return _DecliningCounter() diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index 71aa79cc4bb..b95e6a62e9c 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -8,7 +8,6 @@ cases need the extension and are skipped when it is not built. from __future__ import annotations import json -from types import MappingProxyType from typing import Final import pytest @@ -27,7 +26,6 @@ MODEL: Final = "claude-sonnet-4-5-20250929" CL100K_MODEL: Final = "gpt-4" O200K_MODEL: Final = "gpt-4o" TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base") -RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998}) BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() @@ -55,24 +53,20 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files.""" + """Stands in for the native `TokenCounter` class.""" def __init__(self) -> None: self.counters: list[_RecordingCounter] = [] - self.rank_files: list[str] = [] + self.encodings: list[str] = [] def __call__(self, tokenizer_json: str) -> _RecordingCounter: counter = _RecordingCounter(tokenizer_json) self.counters.append(counter) return counter - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("o200k_base") + def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter: + self.encodings.append(encoding) + return self(encoding) class _RaisingCounter: @@ -92,10 +86,7 @@ class _RaisingFactory: def __call__(self, tokenizer_json: str) -> _RaisingCounter: return _RaisingCounter(self.error) - def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter: - return _RaisingCounter(self.error) - - def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter: + def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter: return _RaisingCounter(self.error) @@ -140,7 +131,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No @pytest.mark.asyncio @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base")) -async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None: +async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None: factory: Final = _RecordingFactory() litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) @@ -149,9 +140,7 @@ async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokeni second: Final = await bridge.count_input_tokens(BODY, tokenizer) assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42) - assert len(factory.rank_files) == 1 - assert factory.rank_files[0].startswith("IQ== 0\n") - assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer] + assert factory.encodings == [tokenizer] assert factory.counters[0].tokenizer_json == tokenizer assert factory.counters[0].bodies == [BODY, BODY] @@ -235,13 +224,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: ("model", "python_encoding"), (("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")), ) -def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have( - monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str +def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust( + monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer ) -> None: monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model}) assert openai_tokenizer_encoding(model).name == python_encoding - assert bridge.rust_tokenizer(model) is None + assert bridge.rust_tokenizer(model) == python_encoding def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None: From 9b2b3d0b90346c0da91a8abf4b52be7561e1dcfe Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:21:52 +0000 Subject: [PATCH 100/146] refactor(rust): align tokenizer docs and lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/token-counter/README.md | 2 +- litellm/rust_bridge/token_counter.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index c8ed0737a7e..effa5324822 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -2,7 +2,7 @@ `Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface -The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast`, `TokenCounter::from_cl100k_ranks`, and `TokenCounter::from_o200k_ranks` use this implementation +The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast` uses this implementation The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through the upstream `tokenizers` library. `TokenCounter::from_json` uses this implementation diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index e816f1ab388..d232db19a4a 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -77,7 +77,14 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: if kind is not None or uses_legacy_message_accounting(model): return None encoding: Final = openai_tokenizer_encoding(model).name - if encoding in {"cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base"}: + if encoding in ( + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + ): return cast(RustTokenizer, encoding) return None From a619b765fc06b9ba9c7124debcf75f0c79033f58 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:25:19 +0000 Subject: [PATCH 101/146] refactor(rust): simplify tokenizer bridge features Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/Cargo.toml | 8 ++-- .../crates/python-bridge/src/token_counter.rs | 42 ++++++------------- 2 files changed, 17 insertions(+), 33 deletions(-) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 2c868ebc769..67f662408c4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,13 +10,13 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"] +default = ["abi3", "huggingface", "tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] -token-counter-fast = ["litellm-token-counter/fast"] -token-counter-huggingface = ["litellm-token-counter/huggingface"] -token-counter-tiktoken = ["litellm-token-counter/tiktoken"] +fast = ["litellm-token-counter/fast"] +huggingface = ["litellm-token-counter/huggingface"] +tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index a7c6bce4d0b..12749486679 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -1,17 +1,9 @@ use std::sync::Arc; -#[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" -))] +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use std::{num::NonZero, thread::available_parallelism}; -#[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" -))] +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use litellm_host_python::release_gil; use litellm_host_python::run_async; use litellm_token_counter::{ @@ -41,45 +33,45 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "token-counter-huggingface")] + #[cfg(feature = "huggingface")] { Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) } - #[cfg(not(feature = "token-counter-huggingface"))] + #[cfg(not(feature = "huggingface"))] { let _ = (py, tokenizer_json); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the token-counter-huggingface feature", + "tokenizer backend requires the huggingface feature", )) } } #[staticmethod] fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "token-counter-fast")] + #[cfg(feature = "fast")] { Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) } - #[cfg(not(feature = "token-counter-fast"))] + #[cfg(not(feature = "fast"))] { let _ = (py, tokenizer_json); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the token-counter-fast feature", + "tokenizer backend requires the fast feature", )) } } #[staticmethod] fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { - #[cfg(feature = "token-counter-tiktoken")] + #[cfg(feature = "tiktoken")] { Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) } - #[cfg(not(feature = "token-counter-tiktoken"))] + #[cfg(not(feature = "tiktoken"))] { let _ = (py, encoding); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the token-counter-tiktoken feature", + "tokenizer backend requires the tiktoken feature", )) } } @@ -105,11 +97,7 @@ impl TokenCounter { } impl TokenCounter { - #[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" - ))] + #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] fn load( py: Python<'_>, load: impl FnOnce() -> Result + Send, @@ -122,11 +110,7 @@ impl TokenCounter { } } -#[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" -))] +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] fn encode_parallelism() -> usize { available_parallelism().map_or(1, NonZero::get) } From be7b7d2ade72d0b158b13c8d1fed426d6a24cd92 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:30:55 +0000 Subject: [PATCH 102/146] chore(prices): sync OpenRouter prices: 1 model openrouter/~z-ai/glm-latest: max_tokens, max_output_tokens --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d25cf0c78df..850e0e5ecb4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d25cf0c78df..850e0e5ecb4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 831810248ae95a407d8cd6c3f1c94be695d894e0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:49:39 +0000 Subject: [PATCH 103/146] refactor(rust): keep the Python token counter on the fast backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/Cargo.toml | 2 +- .../crates/python-bridge/src/token_counter.rs | 31 +++++++++--- litellm-rust/crates/token-counter/README.md | 2 +- litellm/rust_bridge/_native.pyi | 4 +- litellm/rust_bridge/token_counter.py | 47 +++++++++---------- .../spend_tracking/test_budget_reservation.py | 14 ++++-- .../rust_bridge/test_token_counter.py | 33 ++++++++----- 7 files changed, 83 insertions(+), 50 deletions(-) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 67f662408c4..a76b069935f 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,7 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3", "huggingface", "tiktoken"] +default = ["abi3", "fast"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 12749486679..244401e6696 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -33,28 +33,47 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "huggingface")] + #[cfg(feature = "fast")] + { + Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + } + #[cfg(all(not(feature = "fast"), feature = "huggingface"))] { Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) } - #[cfg(not(feature = "huggingface"))] + #[cfg(not(any(feature = "fast", feature = "huggingface")))] { let _ = (py, tokenizer_json); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the huggingface feature", + "tokenizer backend requires the fast or huggingface feature", )) } } #[staticmethod] - fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { + fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { #[cfg(feature = "fast")] { - Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file)) } #[cfg(not(feature = "fast"))] { - let _ = (py, tokenizer_json); + let _ = (py, rank_file); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the fast feature", + )) + } + } + + #[staticmethod] + fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { + #[cfg(feature = "fast")] + { + Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) + } + #[cfg(not(feature = "fast"))] + { + let _ = (py, rank_file); Err(RustBridgeDeclined::new_err( "tokenizer backend requires the fast feature", )) diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index effa5324822..3ad21a00697 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -8,7 +8,7 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` -The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend +The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7d20fc5a02f..05a6df6d5af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -97,7 +97,9 @@ class ResponsesWebSocketConnection: class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_json_fast(tokenizer_json: str) -> TokenCounter: ... + def from_cl100k_ranks(rank_file: str) -> TokenCounter: ... + @staticmethod + def from_o200k_ranks(rank_file: str) -> TokenCounter: ... @staticmethod def from_tiktoken(encoding: str) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d232db19a4a..d36234f56c1 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -11,21 +11,14 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt from litellm.utils import claude_json_str, huggingface_tokenizer_kind -RustTokenizer = Literal[ - "anthropic", - "cl100k_base", - "o200k_base", - "o200k_harmony", - "p50k_base", - "p50k_edit", - "r50k_base", -] +RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"] class RustTokenCounter(Protocol): @@ -37,7 +30,10 @@ class RustTokenCounterFactory(Protocol): def __call__(self, tokenizer_json: str) -> RustTokenCounter: raise NotImplementedError - def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter: + def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter: + raise NotImplementedError + + def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter: raise NotImplementedError @@ -67,8 +63,9 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: """The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count. Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace - downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every - message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python.""" + downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust + prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in + Python.""" if litellm.disable_token_counter is True: return None kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model) @@ -76,26 +73,24 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: return "anthropic" if kind is not None or uses_legacy_message_accounting(model): return None - encoding: Final = openai_tokenizer_encoding(model).name - if encoding in ( - "cl100k_base", - "o200k_base", - "o200k_harmony", - "p50k_base", - "p50k_edit", - "r50k_base", - ): - return cast(RustTokenizer, encoding) - return None + match openai_tokenizer_encoding(model).name: + case "cl100k_base": + return "cl100k_base" + case "o200k_base": + return "o200k_base" + case _: + return None -@lru_cache(maxsize=8) +@lru_cache(maxsize=4) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: match tokenizer: case "anthropic": return factory(claude_json_str) - case _: - return factory.from_tiktoken(tokenizer) + case "cl100k_base": + return factory.from_cl100k_ranks(cl100k_base_rank_file()) + case "o200k_base": + return factory.from_o200k_ranks(o200k_base_rank_file()) async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index c785a7b657a..0e0025f5194 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -256,7 +256,7 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class.""" + """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`.""" def __init__(self) -> None: self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = [] @@ -264,8 +264,11 @@ class _RecordingFactory: def __call__(self, tokenizer_json: str) -> _RecordingCounter: return _RecordingCounter(self, "anthropic") - def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter: - return _RecordingCounter(self, encoding) + def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: + return _RecordingCounter(self, "cl100k_base") + + def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: + return _RecordingCounter(self, "o200k_base") class _DecliningCounter: @@ -277,7 +280,10 @@ class _DecliningFactory: def __call__(self, tokenizer_json: str) -> _DecliningCounter: return _DecliningCounter() - def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter: + def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter: + return _DecliningCounter() + + def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter: return _DecliningCounter() diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index b95e6a62e9c..71aa79cc4bb 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -8,6 +8,7 @@ cases need the extension and are skipped when it is not built. from __future__ import annotations import json +from types import MappingProxyType from typing import Final import pytest @@ -26,6 +27,7 @@ MODEL: Final = "claude-sonnet-4-5-20250929" CL100K_MODEL: Final = "gpt-4" O200K_MODEL: Final = "gpt-4o" TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base") +RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998}) BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() @@ -53,20 +55,24 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class.""" + """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files.""" def __init__(self) -> None: self.counters: list[_RecordingCounter] = [] - self.encodings: list[str] = [] + self.rank_files: list[str] = [] def __call__(self, tokenizer_json: str) -> _RecordingCounter: counter = _RecordingCounter(tokenizer_json) self.counters.append(counter) return counter - def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter: - self.encodings.append(encoding) - return self(encoding) + def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: + self.rank_files.append(rank_file) + return self("cl100k_base") + + def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: + self.rank_files.append(rank_file) + return self("o200k_base") class _RaisingCounter: @@ -86,7 +92,10 @@ class _RaisingFactory: def __call__(self, tokenizer_json: str) -> _RaisingCounter: return _RaisingCounter(self.error) - def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter: + def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter: + return _RaisingCounter(self.error) + + def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter: return _RaisingCounter(self.error) @@ -131,7 +140,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No @pytest.mark.asyncio @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base")) -async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None: +async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None: factory: Final = _RecordingFactory() litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) @@ -140,7 +149,9 @@ async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridg second: Final = await bridge.count_input_tokens(BODY, tokenizer) assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42) - assert factory.encodings == [tokenizer] + assert len(factory.rank_files) == 1 + assert factory.rank_files[0].startswith("IQ== 0\n") + assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer] assert factory.counters[0].tokenizer_json == tokenizer assert factory.counters[0].bodies == [BODY, BODY] @@ -224,13 +235,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: ("model", "python_encoding"), (("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")), ) -def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust( - monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer +def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have( + monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str ) -> None: monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model}) assert openai_tokenizer_encoding(model).name == python_encoding - assert bridge.rust_tokenizer(model) == python_encoding + assert bridge.rust_tokenizer(model) is None def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None: From 56c5d31e73d7b7ecc5f6a91d2b49ed355fb8772e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:52:48 +0000 Subject: [PATCH 104/146] test(rust): run fast token counter parity tests by default Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + litellm-rust/Cargo.toml | 1 + .../tests/fixtures/generate.py | 4 +- litellm-rust/crates/token-counter/Cargo.toml | 3 +- litellm-rust/crates/token-counter/README.md | 2 +- .../token-counter/tests/token_counter.rs | 191 +++++++++++------- 6 files changed, 129 insertions(+), 73 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 62bde5806a4..af3a31ddbfa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2257,6 +2257,7 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.19", + "tokenizers", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 109a2edaf4d..4250fa81d26 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -48,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } diff --git a/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py index 1bfbdf00218..2ca853cbd62 100644 --- a/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py +++ b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py @@ -2,8 +2,8 @@ Run from the repository root with the project environment, once per encoding: - uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py cl100k_base - uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py o200k_base + uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py cl100k_base + uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py o200k_base `/texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens` counted with `tiktoken.get_encoding(name).encode(text, disallowed_special=())`, diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index 59e4f9a6a1d..67e5bd60537 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -default = ["huggingface", "tiktoken"] +default = ["fast", "huggingface", "tiktoken"] fast = ["dep:litellm-token-counter-fast"] huggingface = ["dep:litellm-token-counter-huggingface"] tiktoken = ["dep:litellm-token-counter-tiktoken"] @@ -25,6 +25,7 @@ thiserror.workspace = true criterion.workspace = true rand.workspace = true rstest.workspace = true +tokenizers.workspace = true [[bench]] name = "token_counter" diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index 3ad21a00697..a6b2b50aac0 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -8,7 +8,7 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` -The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend +All three backends are enabled by default. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 378a16d5d72..ebbe2f00ffc 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,4 +1,4 @@ -#![cfg(feature = "huggingface")] +#![cfg(any(feature = "fast", feature = "huggingface"))] use rstest::rstest; @@ -6,13 +6,15 @@ use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCount /// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` /// so this test also guards Python parity. -fn counter() -> TokenCounter { +type JsonLoader = fn(&str) -> Result; + +fn counter(load: JsonLoader) -> TokenCounter { let path = concat!( env!("CARGO_MANIFEST_DIR"), "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" ); let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); - TokenCounter::from_json(&json).expect("anthropic tokenizer loads") + load(&json).expect("anthropic tokenizer loads") } const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; @@ -63,21 +65,15 @@ const EMBEDDINGS_TOKEN_IDS: &str = const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; -/// Expected counts are pinned from -/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`. -#[rstest] -#[case::text_only(SIMPLE, 14)] -#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)] -#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)] -#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)] -#[case::completions_prompt(COMPLETIONS_PROMPT, 7)] -#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)] -#[case::responses_input_items(RESPONSES_INPUT, 62)] -#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)] -#[case::rerank_query_and_documents(RERANK, 41)] -fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) { +fn assert_count_request_matches_python_token_counter( + load: JsonLoader, + body: &str, + expected: usize, +) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter().count_request(&request).expect("fixture counts"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); assert_eq!( count, InputTokenCount { @@ -87,63 +83,24 @@ fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expect ); } -#[rstest] -#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)] -#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] -#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] -#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] -fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { +fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter().count_request(&request).expect("fixture counts"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); assert_eq!(count.input_tokens, expected); } -#[rstest] -#[case::not_json(b"not json" as &[u8])] -#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] -#[case::message_with_tool_calls( - br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# -)] -#[case::dict_content( - br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# -)] -#[case::float_enum( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# -)] -#[case::anthropic_tool_choice_without_function( - br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# -)] -fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { - assert!(matches!( - CountableRequest::parse(body), - Err(Error::RequestParse(_)) - )); -} - -#[rstest] -#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] -#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] -#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] -#[case::image_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# -)] -#[case::tool_result_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# -)] -#[case::array_without_items( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# -)] -fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { +fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { let request = CountableRequest::parse(body).expect("shape parses"); assert!(matches!( - counter().count_request(&request), + counter(load).count_request(&request), Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) )); } -#[test] -fn tool_choice_and_system_discount_change_the_count() { - let counter = counter(); +fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { + let counter = counter(load); let count = |body: &str| { counter .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) @@ -172,9 +129,104 @@ fn tool_choice_and_system_discount_change_the_count() { ); } -#[test] -fn loading_a_bad_tokenizer_is_a_load_error() { - assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); +fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { + assert!(matches!(load("{}"), Err(Error::Load(_)))); +} + +macro_rules! json_backend_tests { + ($loader:path) => { + #[rstest] + #[case::text_only(super::SIMPLE, 14)] + #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] + #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] + #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] + #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] + #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] + #[case::responses_input_items(super::RESPONSES_INPUT, 62)] + #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] + #[case::rerank_query_and_documents(super::RERANK, 41)] + fn count_request_matches_python_token_counter( + #[case] body: &str, + #[case] expected: usize, + ) { + super::assert_count_request_matches_python_token_counter($loader, body, expected); + } + + #[rstest] + #[case::null_messages_win_over_prompt( + r#"{"model":"m","messages":null,"prompt":"ignored"}"#, + 3 + )] + #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] + #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] + #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] + fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + super::assert_key_presence_follows_python($loader, body, expected); + } + + #[rstest] + #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] + #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] + #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] + #[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# + )] + #[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# + )] + #[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# + )] + fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); + } + + #[test] + fn tool_choice_and_system_discount_change_the_count() { + super::assert_tool_choice_and_system_discount_change_the_count($loader); + } + + #[test] + fn loading_a_bad_tokenizer_is_a_load_error() { + super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); + } + }; +} + +#[rstest] +#[case::not_json(b"not json" as &[u8])] +#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] +#[case::message_with_tool_calls( + br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# +)] +#[case::dict_content( + br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# +)] +#[case::float_enum( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# +)] +#[case::anthropic_tool_choice_without_function( + br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# +)] +fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { + assert!(matches!( + CountableRequest::parse(body), + Err(Error::RequestParse(_)) + )); +} + +#[cfg(feature = "fast")] +mod fast_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json_fast); +} + +#[cfg(feature = "huggingface")] +mod huggingface_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json); } #[cfg(feature = "fast")] @@ -222,7 +274,8 @@ mod fast { env!("CARGO_MANIFEST_DIR"), encoding.fixtures ); - std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") + std::fs::read_to_string(&path) + .expect("fixture generated by token-counter-fast/tests/fixtures/generate.py") } #[derive(Deserialize)] @@ -238,7 +291,7 @@ mod fast { } /// Reference counts come from `tiktoken.get_encoding(name)`; see - /// `tests/fixtures/generate.py`. + /// `token-counter-fast/tests/fixtures/generate.py`. #[rstest] #[case::cl100k(CL100K)] #[case::o200k(O200K)] From 10715b1b4607ff2260f5ad22d7bc94fbd43d48f3 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 22:00:46 +0000 Subject: [PATCH 105/146] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-latest: max_tokens, max_output_tokens openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 26 +++++++++---------- model_prices_and_context_window.json | 26 +++++++++---------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 850e0e5ecb4..dc570f14a2d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2668e-07, + "input_cost_per_token": 5.2536e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7556e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.6716e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7556e-08, - "input_cost_per_token": 5.2668e-07, + "cache_read_input_token_cost": 1.6716e-08, + "input_cost_per_token": 5.2536e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 850e0e5ecb4..dc570f14a2d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2668e-07, + "input_cost_per_token": 5.2536e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7556e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.6716e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7556e-08, - "input_cost_per_token": 5.2668e-07, + "cache_read_input_token_cost": 1.6716e-08, + "input_cost_per_token": 5.2536e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 661da87c910deebd487953d4f2c38c0211d88398 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 15:13:59 -0700 Subject: [PATCH 106/146] fix(rust): validate tokenizer ranks and cover backend features --- .github/workflows/test-rust.yml | 7 + .../crates/token-counter-fast/Cargo.toml | 2 +- .../crates/token-counter-fast/src/tiktoken.rs | 18 + .../token-counter-huggingface/Cargo.toml | 2 +- .../token-counter-huggingface/src/error.rs | 9 + .../token-counter-huggingface/src/lib.rs | 10 +- .../token-counter-tiktoken/src/error.rs | 5 + .../crates/token-counter-tiktoken/src/lib.rs | 54 ++- .../token-counter/tests/token_counter.rs | 423 ++++++++++-------- 9 files changed, 326 insertions(+), 204 deletions(-) create mode 100644 litellm-rust/crates/token-counter-huggingface/src/error.rs create mode 100644 litellm-rust/crates/token-counter-tiktoken/src/error.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 551f783d4f9..70d828b5c8b 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -120,6 +120,13 @@ jobs: - run: cargo test --workspace --doc --locked + - name: Test token counter feature combinations + run: | + for features in '' fast huggingface tiktoken fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken; do + cargo test -p litellm-token-counter --locked --no-default-features --features "$features" + cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}" + done + rust-wheel: runs-on: ubuntu-latest timeout-minutes: 30 diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml index 4950277f5b0..5127dc17f22 100644 --- a/litellm-rust/crates/token-counter-fast/Cargo.toml +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -9,7 +9,7 @@ repository.workspace = true base64.workspace = true rustc-hash = "2.1.3" thiserror.workspace = true -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +tokenizers.workspace = true unicode-normalization-alignments = "0.1.12" [dev-dependencies] diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs index 7a9e71ed587..16172b7a688 100644 --- a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -81,6 +81,9 @@ fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { let rank = rank .parse() .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; + if rank == NO_RANK { + return Err(Error::Ranks(format!("rank {rank} is reserved"))); + } Ok((bytes.into_boxed_slice(), rank)) } @@ -212,6 +215,21 @@ mod tests { ); } + #[test] + fn reserved_merge_rank_is_rejected() { + let bytes = (0..=u8::MAX) + .map(|byte| format!("{} {byte}\n", STANDARD.encode([byte]))) + .collect::(); + let rank_file = format!("{bytes}{} {NO_RANK}\n", STANDARD.encode(b"ab")); + assert!(matches!( + MergeRanks::parse(&rank_file), + Err(Error::Ranks(_)) + )); + let valid_rank_file = format!("{bytes}{} {}\n", STANDARD.encode(b"ab"), NO_RANK - 1); + let ranks = MergeRanks::parse(&valid_rank_file).unwrap(); + assert_eq!(ranks.count_piece(b"aab", &mut MergeScratch::default()), 2); + } + #[test] fn malformed_rank_files_are_rejected() { assert!(MergeRanks::parse("IQ==").is_err()); diff --git a/litellm-rust/crates/token-counter-huggingface/Cargo.toml b/litellm-rust/crates/token-counter-huggingface/Cargo.toml index ad15d9f19d0..6d8cb85e524 100644 --- a/litellm-rust/crates/token-counter-huggingface/Cargo.toml +++ b/litellm-rust/crates/token-counter-huggingface/Cargo.toml @@ -7,4 +7,4 @@ repository.workspace = true [dependencies] thiserror.workspace = true -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +tokenizers.workspace = true diff --git a/litellm-rust/crates/token-counter-huggingface/src/error.rs b/litellm-rust/crates/token-counter-huggingface/src/error.rs new file mode 100644 index 00000000000..adc4551886f --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/src/error.rs @@ -0,0 +1,9 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs index 184956b5f39..8e05c2cca46 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/lib.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -1,14 +1,8 @@ #![forbid(unsafe_code)] -use thiserror::Error as ThisError; +mod error; -#[derive(Debug, ThisError)] -pub enum Error { - #[error("failed to load tokenizer: {0}")] - Load(#[source] tokenizers::Error), - #[error("tokenization failed: {0}")] - Encode(#[source] tokenizers::Error), -} +pub use error::Error; pub struct HuggingFaceTokenizer(Box); diff --git a/litellm-rust/crates/token-counter-tiktoken/src/error.rs b/litellm-rust/crates/token-counter-tiktoken/src/error.rs new file mode 100644 index 00000000000..e28cbbfb620 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/error.rs @@ -0,0 +1,5 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +#[error("unsupported tokenizer: {0}")] +pub struct UnsupportedTokenizer(pub String); diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index b35c4876d00..ecdb3946eee 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -1,10 +1,8 @@ #![forbid(unsafe_code)] -use thiserror::Error as ThisError; +mod error; -#[derive(Debug, ThisError)] -#[error("unsupported tokenizer: {0}")] -pub struct UnsupportedTokenizer(pub String); +pub use error::UnsupportedTokenizer; pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); @@ -32,23 +30,41 @@ mod tests { use super::*; #[test] - fn special_tokens_are_counted_as_ordinary_text() { - let counter = TiktokenTokenizer::from_name("cl100k_base").unwrap(); - assert!(counter.count_tokens("<|endoftext|>") > 1); + fn named_encodings_match_their_reference_counts() { + let encodings = [ + ("cl100k_base", tiktoken_rs::cl100k_base_singleton()), + ("o200k_base", tiktoken_rs::o200k_base_singleton()), + ("o200k_harmony", tiktoken_rs::o200k_harmony_singleton()), + ("p50k_base", tiktoken_rs::p50k_base_singleton()), + ("p50k_edit", tiktoken_rs::p50k_edit_singleton()), + ("r50k_base", tiktoken_rs::r50k_base_singleton()), + ("gpt2", tiktoken_rs::r50k_base_singleton()), + ]; + let texts = [ + "", + "Hello, how are you today?", + "é e\u{301} 漢字 ع ३ 🙂 ＡfiⅣ", + " def function():\n return 123456789\r\n", + "<|endoftext|><|fim_prefix|><|start|>assistant<|message|>", + ]; + for (name, reference) in encodings { + let counter = TiktokenTokenizer::from_name(name).unwrap(); + for text in texts { + assert_eq!( + counter.count_tokens(text), + reference.encode_ordinary(text).len(), + "{name}: {text:?}", + ); + } + } } #[test] - fn all_python_tiktoken_encodings_are_available() { - for name in [ - "cl100k_base", - "o200k_base", - "o200k_harmony", - "p50k_base", - "p50k_edit", - "r50k_base", - "gpt2", - ] { - assert!(TiktokenTokenizer::from_name(name).is_ok(), "{name}"); - } + fn unsupported_encoding_preserves_its_name() { + let Err(UnsupportedTokenizer(name)) = TiktokenTokenizer::from_name("unknown-encoding") + else { + panic!("unknown encoding must be rejected"); + }; + assert_eq!(name, "unknown-encoding"); } } diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index ebbe2f00ffc..542bd4a1fc4 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,25 +1,30 @@ -#![cfg(any(feature = "fast", feature = "huggingface"))] - use rstest::rstest; -use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] +use litellm_token_counter::TokenCounter; +use litellm_token_counter::{CountableRequest, Error}; -/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` -/// so this test also guards Python parity. -type JsonLoader = fn(&str) -> Result; +#[cfg(any(feature = "fast", feature = "huggingface"))] +mod json { + use super::*; + use litellm_token_counter::InputTokenCount; -fn counter(load: JsonLoader) -> TokenCounter { - let path = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" - ); - let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); - load(&json).expect("anthropic tokenizer loads") -} + /// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` + /// so this test also guards Python parity. + type JsonLoader = fn(&str) -> Result; -const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; + fn counter(load: JsonLoader) -> TokenCounter { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + ); + let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); + load(&json).expect("anthropic tokenizer loads") + } -const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ + const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; + + const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ {"role":"system","content":"You are a terse assistant."}, {"role":"user","name":"alice","content":[ {"type":"text","text":"Summarise this paragraph about ships and harbours."}, @@ -28,7 +33,7 @@ const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ {"type":"tool_reference","tool_name":"get_weather"}]}, {"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#; -const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}], + const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}], "tools":[ {"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{ "type":"object", @@ -43,154 +48,188 @@ const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":" {"type":"function","function":{"name":"noop"}}], "tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#; -const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", + const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", "messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}], "tools":[{"name":"get_weather","description":"Get weather","input_schema":{ "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}], "tool_choice":"none"}"#; -const COMPLETIONS_PROMPT: &str = - r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; + const COMPLETIONS_PROMPT: &str = + r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; -const COMPLETIONS_PROMPT_LIST: &str = - r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; + const COMPLETIONS_PROMPT_LIST: &str = + r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; -const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ + const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]}, {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#; -const EMBEDDINGS_TOKEN_IDS: &str = - r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; + const EMBEDDINGS_TOKEN_IDS: &str = + r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; -const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", + const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; -fn assert_count_request_matches_python_token_counter( - load: JsonLoader, - body: &str, - expected: usize, -) { - let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter(load) - .count_request(&request) - .expect("fixture counts"); - assert_eq!( - count, - InputTokenCount { - model: Some("claude-sonnet-4-5".to_string()), - input_tokens: expected, - } - ); -} + fn assert_count_request_matches_python_token_counter( + load: JsonLoader, + body: &str, + expected: usize, + ) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); + assert_eq!( + count, + InputTokenCount { + model: Some("claude-sonnet-4-5".to_string()), + input_tokens: expected, + } + ); + } -fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { - let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter(load) - .count_request(&request) - .expect("fixture counts"); - assert_eq!(count.input_tokens, expected); -} + fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); + assert_eq!(count.input_tokens, expected); + } -fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { - let request = CountableRequest::parse(body).expect("shape parses"); - assert!(matches!( - counter(load).count_request(&request), - Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) - )); -} + fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { + let request = CountableRequest::parse(body).expect("shape parses"); + assert!(matches!( + counter(load).count_request(&request), + Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) + )); + } -fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { - let counter = counter(load); - let count = |body: &str| { - counter - .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) - .expect("counts") - .input_tokens - }; - let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), - base + 1 - ); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#), - base - ); - let with_tools = count( - r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, - ); - let with_tools_and_system = count( - r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, - ); - assert_eq!(with_tools - with_tools_and_system, 4); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#), - base - ); -} + fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { + let counter = counter(load); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# + ), + base + 1 + ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"# + ), + base + ); + let with_tools = count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + let with_tools_and_system = count( + r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + assert_eq!(with_tools - with_tools_and_system, 4); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#), + base + ); + } -fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { - assert!(matches!(load("{}"), Err(Error::Load(_)))); -} + fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { + assert!(matches!(load("{}"), Err(Error::Load(_)))); + } -macro_rules! json_backend_tests { - ($loader:path) => { - #[rstest] - #[case::text_only(super::SIMPLE, 14)] - #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] - #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] - #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] - #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] - #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] - #[case::responses_input_items(super::RESPONSES_INPUT, 62)] - #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] - #[case::rerank_query_and_documents(super::RERANK, 41)] - fn count_request_matches_python_token_counter( - #[case] body: &str, - #[case] expected: usize, - ) { - super::assert_count_request_matches_python_token_counter($loader, body, expected); - } + macro_rules! json_backend_tests { + ($loader:path) => { + #[rstest] + #[case::text_only(super::SIMPLE, 14)] + #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] + #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] + #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] + #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] + #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] + #[case::responses_input_items(super::RESPONSES_INPUT, 62)] + #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] + #[case::rerank_query_and_documents(super::RERANK, 41)] + fn count_request_matches_python_token_counter( + #[case] body: &str, + #[case] expected: usize, + ) { + super::assert_count_request_matches_python_token_counter($loader, body, expected); + } - #[rstest] - #[case::null_messages_win_over_prompt( - r#"{"model":"m","messages":null,"prompt":"ignored"}"#, - 3 - )] - #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] - #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] - #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] - fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { - super::assert_key_presence_follows_python($loader, body, expected); - } + #[rstest] + #[case::null_messages_win_over_prompt( + r#"{"model":"m","messages":null,"prompt":"ignored"}"#, + 3 + )] + #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] + #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] + #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] + fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + super::assert_key_presence_follows_python($loader, body, expected); + } - #[rstest] - #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] - #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] - #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] - #[case::image_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# - )] - #[case::tool_result_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# - )] - #[case::array_without_items( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# - )] - fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { - super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); - } + #[rstest] + #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] + #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] + #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] + #[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# + )] + #[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# + )] + #[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# + )] + fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); + } - #[test] - fn tool_choice_and_system_discount_change_the_count() { - super::assert_tool_choice_and_system_discount_change_the_count($loader); - } + #[test] + fn tool_choice_and_system_discount_change_the_count() { + super::assert_tool_choice_and_system_discount_change_the_count($loader); + } - #[test] - fn loading_a_bad_tokenizer_is_a_load_error() { - super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); - } - }; + #[test] + fn encoding_errors_preserve_the_backend_source() { + use std::error::Error as _; + + let tokenizer = tokenizers::Tokenizer::new( + tokenizers::models::wordpiece::WordPiece::default(), + ); + let expected = tokenizer.encode_fast("hello", true).unwrap_err(); + let counter = $loader(&tokenizer.to_string(false).unwrap()).unwrap(); + let request = CountableRequest::parse(br#"{"prompt":"hello"}"#).unwrap(); + let error = counter.count_request(&request).unwrap_err(); + assert!(matches!(error, Error::Encode(_))); + assert_eq!(error.source().unwrap().to_string(), expected.to_string()); + } + + #[test] + fn loading_a_bad_tokenizer_is_a_load_error() { + super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); + } + }; + } + + #[cfg(feature = "fast")] + mod fast_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json_fast); + } + + #[cfg(feature = "huggingface")] + mod huggingface_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json); + } } #[rstest] @@ -215,22 +254,8 @@ fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { )); } -#[cfg(feature = "fast")] -mod fast_json { - use super::*; - - json_backend_tests!(TokenCounter::from_json_fast); -} - -#[cfg(feature = "huggingface")] -mod huggingface_json { - use super::*; - - json_backend_tests!(TokenCounter::from_json); -} - -#[cfg(feature = "fast")] -mod fast { +#[cfg(any(feature = "fast", feature = "tiktoken"))] +mod tiktoken { use super::*; use serde::Deserialize; @@ -239,33 +264,65 @@ mod fast { #[derive(Clone, Copy)] struct TiktokenEncoding { fixtures: &'static str, - rank_file: &'static str, + source: TokenizerSource, load: fn(&str) -> Result, model: &'static str, } + #[cfg(feature = "fast")] const CL100K: TiktokenEncoding = TiktokenEncoding { fixtures: "cl100k", - rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", + source: TokenizerSource::RankFile("9b5ad71b2ce5302211f9c61530b329a4922fc6a4"), load: TokenCounter::from_cl100k_ranks, model: "gpt-4", }; + #[cfg(feature = "fast")] const O200K: TiktokenEncoding = TiktokenEncoding { fixtures: "o200k", - rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", + source: TokenizerSource::RankFile("fb374d419588a4632f3f557e76b4b70aebbca790"), load: TokenCounter::from_o200k_ranks, model: "gpt-4o", }; + #[cfg(feature = "tiktoken")] + const TIKTOKEN_CL100K: TiktokenEncoding = TiktokenEncoding { + fixtures: "cl100k", + source: TokenizerSource::Name("cl100k_base"), + load: TokenCounter::from_tiktoken, + model: "gpt-4", + }; + + #[cfg(feature = "tiktoken")] + const TIKTOKEN_O200K: TiktokenEncoding = TiktokenEncoding { + fixtures: "o200k", + source: TokenizerSource::Name("o200k_base"), + load: TokenCounter::from_tiktoken, + model: "gpt-4o", + }; + + #[derive(Clone, Copy)] + enum TokenizerSource { + #[cfg(feature = "fast")] + RankFile(&'static str), + #[cfg(feature = "tiktoken")] + Name(&'static str), + } + fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { - let path = format!( - "{}/../../../litellm/litellm_core_utils/tokenizers/{}", - env!("CARGO_MANIFEST_DIR"), - encoding.rank_file - ); - let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); - (encoding.load)(&ranks).expect("ranks load") + match encoding.source { + #[cfg(feature = "tiktoken")] + TokenizerSource::Name(name) => (encoding.load)(name).expect("encoding loads"), + #[cfg(feature = "fast")] + TokenizerSource::RankFile(file) => { + let path = format!( + "{}/../../../litellm/litellm_core_utils/tokenizers/{file}", + env!("CARGO_MANIFEST_DIR"), + ); + let ranks = std::fs::read_to_string(path).expect("rank file is in the repo"); + (encoding.load)(&ranks).expect("ranks load") + } + } } fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { @@ -293,8 +350,10 @@ mod fast { /// Reference counts come from `tiktoken.get_encoding(name)`; see /// `token-counter-fast/tests/fixtures/generate.py`. #[rstest] - #[case::cl100k(CL100K)] - #[case::o200k(O200K)] + #[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))] + #[cfg_attr(feature = "fast", case::fast_o200k(O200K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_O200K))] fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { let counter = tiktoken_counter(encoding); let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") @@ -319,8 +378,10 @@ mod fast { /// (`_count_input_tokens(body, model)`), so this pins the shared message, /// tool and reply-priming accounting on the tiktoken paths as well. #[rstest] - #[case::cl100k(CL100K)] - #[case::o200k(O200K)] + #[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))] + #[cfg_attr(feature = "fast", case::fast_o200k(O200K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_O200K))] fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { let counter = tiktoken_counter(encoding); let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") @@ -342,8 +403,10 @@ mod fast { } #[rstest] - #[case::cl100k(CL100K)] - #[case::o200k(O200K)] + #[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))] + #[cfg_attr(feature = "fast", case::fast_o200k(O200K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_O200K))] fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( #[case] encoding: TiktokenEncoding, ) { @@ -369,6 +432,7 @@ mod fast { ); } + #[cfg(feature = "fast")] #[rstest] #[case::empty("")] #[case::not_base64("!!!! 0")] @@ -382,3 +446,12 @@ mod fast { assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); } } + +#[cfg(feature = "tiktoken")] +#[test] +fn unsupported_encoding_reaches_the_counter_caller() { + assert!(matches!( + TokenCounter::from_tiktoken("unknown-encoding"), + Err(Error::UnsupportedTokenizer(name)) if name == "unknown-encoding" + )); +} From 668a89fbfd6273c9ba92da38fd08c34f287d2109 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 22:30:41 +0000 Subject: [PATCH 107/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 22 +++++++++---------- model_prices_and_context_window.json | 22 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dc570f14a2d..39c418de524 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2536e-07, + "input_cost_per_token": 5.2404e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6716e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 1.7468e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6716e-08, - "input_cost_per_token": 5.2536e-07, + "cache_read_input_token_cost": 1.7468e-08, + "input_cost_per_token": 5.2404e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dc570f14a2d..39c418de524 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2536e-07, + "input_cost_per_token": 5.2404e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6716e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 1.7468e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6716e-08, - "input_cost_per_token": 5.2536e-07, + "cache_read_input_token_cost": 1.7468e-08, + "input_cost_per_token": 5.2404e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 3fd2dd635b119ddd3afd1771c1cda217d71df33c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 15:45:38 -0700 Subject: [PATCH 108/146] refactor(rust): split auth facade from shared types --- litellm-rust/Cargo.lock | 26 +++++--- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-aws/Cargo.toml | 2 +- litellm-rust/crates/auth-aws/src/error.rs | 6 +- litellm-rust/crates/auth-azure/Cargo.toml | 2 +- .../src/credential_provider_cache.rs | 2 +- litellm-rust/crates/auth-azure/src/native.rs | 13 ++-- litellm-rust/crates/auth-azure/src/resolve.rs | 16 ++--- litellm-rust/crates/auth-azure/src/types.rs | 4 +- litellm-rust/crates/auth-gcp/Cargo.toml | 2 +- litellm-rust/crates/auth-gcp/src/lib.rs | 2 +- litellm-rust/crates/auth-types/Cargo.toml | 15 +++++ .../{auth => auth-types}/src/credential.rs | 4 +- .../crates/{auth => auth-types}/src/error.rs | 0 .../crates/{auth => auth-types}/src/http.rs | 3 - litellm-rust/crates/auth-types/src/lib.rs | 57 +++++++++++++++++ .../crates/{auth => auth-types}/src/policy.rs | 6 +- .../crates/{auth => auth-types}/src/secret.rs | 0 .../crates/{auth => auth-types}/src/token.rs | 4 +- litellm-rust/crates/auth/Cargo.toml | 17 +++--- litellm-rust/crates/auth/src/lib.rs | 61 +++---------------- litellm-rust/crates/auth/tests/facade.rs | 33 ++++++++++ 22 files changed, 172 insertions(+), 104 deletions(-) create mode 100644 litellm-rust/crates/auth-types/Cargo.toml rename litellm-rust/crates/{auth => auth-types}/src/credential.rs (98%) rename litellm-rust/crates/{auth => auth-types}/src/error.rs (100%) rename litellm-rust/crates/{auth => auth-types}/src/http.rs (92%) create mode 100644 litellm-rust/crates/auth-types/src/lib.rs rename litellm-rust/crates/{auth => auth-types}/src/policy.rs (96%) rename litellm-rust/crates/{auth => auth-types}/src/secret.rs (100%) rename litellm-rust/crates/{auth => auth-types}/src/token.rs (95%) create mode 100644 litellm-rust/crates/auth/tests/facade.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index af3a31ddbfa..5cf9116d752 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1976,11 +1976,10 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" name = "litellm-auth" version = "0.1.0" dependencies = [ - "serde", - "subtle", - "thiserror 2.0.19", - "tokio", - "veil", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", ] [[package]] @@ -1993,7 +1992,7 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "litellm-auth", + "litellm-auth-types", "litellm-http", "moka", "reqwest 0.12.28", @@ -2009,7 +2008,7 @@ version = "0.1.0" dependencies = [ "azure_core", "azure_identity", - "litellm-auth", + "litellm-auth-types", "moka", "rstest", "serde_json", @@ -2024,13 +2023,24 @@ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ "gcp_auth", - "litellm-auth", + "litellm-auth-types", "moka", "serde_json", "sha2 0.10.9", "tokio", ] +[[package]] +name = "litellm-auth-types" +version = "0.1.0" +dependencies = [ + "serde", + "subtle", + "thiserror 2.0.19", + "tokio", + "veil", +] + [[package]] name = "litellm-cache" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 4250fa81d26..b781d564805 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -14,6 +14,7 @@ litellm-host = { path = "crates/host" } litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } +litellm-auth-types = { path = "crates/auth-types" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 1f27c7bc990..1a35af48574 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true litellm-http.workspace = true moka = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index f80fbce456e..d4c6ae8cf6c 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -22,7 +22,7 @@ pub enum Error { AwsMissingWebIdentityCredentials, } -impl From for litellm_auth::Error { +impl From for litellm_auth_types::Error { fn from(error: Error) -> Self { Self::ProviderAuthentication(error.to_string()) } @@ -34,11 +34,11 @@ mod tests { #[test] fn converts_to_shared_auth_error_without_losing_context() { - let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into())); + let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into())); assert_eq!( error, - litellm_auth::Error::ProviderAuthentication( + litellm_auth_types::Error::ProviderAuthentication( "AWS profile credentials failed: profile not found".into() ) ); diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 8099506d2e5..1fd9d39d113 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true moka.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs index ab9ffc719df..cd16b27f66d 100644 --- a/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs +++ b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use azure_core::credentials::TokenCredential; use moka::future::Cache; -use litellm_auth::Error; +use litellm_auth_types::Error; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub(crate) struct AzureCredentialProviderCacheKey { diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 5f913a8ad01..d635e559641 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -12,8 +12,8 @@ use azure_identity::{ }; use sha2::{Digest, Sha256}; -use litellm_auth::Error; -use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::Error; +use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -484,7 +484,7 @@ mod tests { use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, SecretValue, Sourced}; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) @@ -649,7 +649,7 @@ mod tests { assert!(matches!( error, - litellm_auth::Error::MixedAzureCredentialSources + litellm_auth_types::Error::MixedAzureCredentialSources )); } @@ -679,7 +679,10 @@ mod tests { authority, )) .unwrap_err(); - assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority)); + assert!(matches!( + error, + litellm_auth_types::Error::InvalidAzureAuthority + )); } } } diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 4e18cbb89aa..4d564b6e68a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,5 +1,5 @@ -use litellm_auth::Error; -use litellm_auth::{ +use litellm_auth_types::Error; +use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, SecretValue, Sourced, TokenProviderHandle, }; @@ -451,9 +451,9 @@ mod tests { }; use crate::native::ValidatedAzureRequest; use crate::types::AzureAuthInputs; - use litellm_auth::Error; - use litellm_auth::ResolvedCredential; - use litellm_auth::{ + use litellm_auth_types::Error; + use litellm_auth_types::ResolvedCredential; + use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, }; @@ -661,8 +661,8 @@ mod tests { #[derive(Debug)] struct CallerToken(&'static str); - impl litellm_auth::TokenProvider for CallerToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + impl litellm_auth_types::TokenProvider for CallerToken { + fn acquire(&self) -> litellm_auth_types::TokenFuture<'_> { Box::pin(async move { Ok(ResolvedCredential::AccessToken { token: SecretValue::new(self.0), @@ -675,7 +675,7 @@ mod tests { fn caller_inputs(token: &'static str) -> AzureAuthInputs { let params = json!({"azure_ad_token": "static-token"}); AzureAuthInputs { - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + azure_ad_token_provider: Some(litellm_auth_types::TokenProviderHandle::new(Arc::new( CallerToken(token), ))), ..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap() diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index 87e883a6a54..d5a00f09751 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use litellm_auth::{ +use litellm_auth_types::{ CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, }; use serde_json::{Map, Value}; @@ -126,7 +126,7 @@ fn source_for(sources: &BTreeMap, name: &str) -> InputSourc mod tests { use std::collections::BTreeMap; - use litellm_auth::{InputSource, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; use serde_json::json; use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index f24582db13e..045320efc80 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true moka.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index bf619fee144..1ff9487e3f4 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc}; use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth::{ +use litellm_auth_types::{ CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential, }; use moka::future::Cache; diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml new file mode 100644 index 00000000000..cd65412127d --- /dev/null +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-auth-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +subtle.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +tokio.workspace = true diff --git a/litellm-rust/crates/auth/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs similarity index 98% rename from litellm-rust/crates/auth/src/credential.rs rename to litellm-rust/crates/auth-types/src/credential.rs index 8ed1867622a..a5d14a43b71 100644 --- a/litellm-rust/crates/auth/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,9 +5,7 @@ use std::sync::Arc; use veil::Redact; -use crate::Error; - -use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { diff --git a/litellm-rust/crates/auth/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs similarity index 100% rename from litellm-rust/crates/auth/src/error.rs rename to litellm-rust/crates/auth-types/src/error.rs diff --git a/litellm-rust/crates/auth/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs similarity index 92% rename from litellm-rust/crates/auth/src/http.rs rename to litellm-rust/crates/auth-types/src/http.rs index dd87d00e70f..0cb5839f965 100644 --- a/litellm-rust/crates/auth/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -40,9 +40,6 @@ pub fn apply_credential( ) } -/// How the upstream call is authenticated. API-key strategies become headers -/// in `prepare`; SigV4 covers the serialized body, so it is applied where the -/// outbound request is built. #[derive(Clone, Debug, PartialEq, Eq)] pub enum RequestAuth { Header { diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs new file mode 100644 index 00000000000..9d399249c05 --- /dev/null +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -0,0 +1,57 @@ +#![forbid(unsafe_code)] + +mod credential; +mod error; +pub mod http; +mod policy; +mod secret; +mod token; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSource { + Request, + #[default] + Deployment, + Environment, +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Sourced { + value: T, + source: InputSource, +} + +impl Sourced { + pub fn new(value: T, source: InputSource) -> Self { + Self { value, source } + } + + pub fn value(&self) -> &T { + &self.value + } + + pub fn source(&self) -> InputSource { + self.source + } + + pub fn into_value(self) -> T { + self.value + } + + pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { + Sourced::new(map(self.value), self.source) + } +} + +pub use credential::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, + CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, +}; +pub use error::Error; +pub use http::{CredentialPlacement, RequestAuth}; +pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; +pub use secret::SecretValue; +pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs similarity index 96% rename from litellm-rust/crates/auth/src/policy.rs rename to litellm-rust/crates/auth-types/src/policy.rs index 4a1f5eeecf9..4c5c0365f0b 100644 --- a/litellm-rust/crates/auth/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -1,7 +1,5 @@ -use crate::Error; - -use super::http::apply_credential; -use super::{CredentialPlacement, ResolvedCredential}; +use crate::http::apply_credential; +use crate::{CredentialPlacement, Error, ResolvedCredential}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CredentialPlanKind { diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth-types/src/secret.rs similarity index 100% rename from litellm-rust/crates/auth/src/secret.rs rename to litellm-rust/crates/auth-types/src/secret.rs diff --git a/litellm-rust/crates/auth/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs similarity index 95% rename from litellm-rust/crates/auth/src/token.rs rename to litellm-rust/crates/auth-types/src/token.rs index 94da5f259fb..4175641ce10 100644 --- a/litellm-rust/crates/auth/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,9 +5,7 @@ use std::time::SystemTime; use veil::Redact; -use crate::Error; - -use super::secret::SecretValue; +use crate::{Error, SecretValue}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index 128a05c1a25..ee4900ebcc2 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -5,11 +5,14 @@ edition.workspace = true license.workspace = true repository.workspace = true -[dependencies] -serde.workspace = true -subtle.workspace = true -thiserror.workspace = true -veil.workspace = true +[features] +default = [] +aws = ["dep:litellm-auth-aws"] +azure = ["dep:litellm-auth-azure"] +gcp = ["dep:litellm-auth-gcp"] -[dev-dependencies] -tokio.workspace = true +[dependencies] +litellm-auth-types.workspace = true +litellm-auth-aws = { workspace = true, optional = true } +litellm-auth-azure = { workspace = true, optional = true } +litellm-auth-gcp = { workspace = true, optional = true } diff --git a/litellm-rust/crates/auth/src/lib.rs b/litellm-rust/crates/auth/src/lib.rs index c8d73c239b0..622a5b2d58b 100644 --- a/litellm-rust/crates/auth/src/lib.rs +++ b/litellm-rust/crates/auth/src/lib.rs @@ -1,55 +1,10 @@ -mod credential; -mod error; -pub mod http; -mod policy; -mod secret; -mod token; +#![forbid(unsafe_code)] -use serde::{Deserialize, Serialize}; +pub use litellm_auth_types::*; -#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] -#[serde(rename_all = "snake_case")] -pub enum InputSource { - Request, - #[default] - Deployment, - Environment, -} - -#[derive(Clone, Debug, Default, Eq, PartialEq)] -pub struct Sourced { - value: T, - source: InputSource, -} - -impl Sourced { - pub fn new(value: T, source: InputSource) -> Self { - Self { value, source } - } - - pub fn value(&self) -> &T { - &self.value - } - - pub fn source(&self) -> InputSource { - self.source - } - - pub fn into_value(self) -> T { - self.value - } - - pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { - Sourced::new(map(self.value), self.source) - } -} - -pub use credential::{ - CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, - CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, -}; -pub use error::Error; -pub use http::{CredentialPlacement, RequestAuth}; -pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; -pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +#[cfg(feature = "aws")] +pub use litellm_auth_aws as aws; +#[cfg(feature = "azure")] +pub use litellm_auth_azure as azure; +#[cfg(feature = "gcp")] +pub use litellm_auth_gcp as gcp; diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs new file mode 100644 index 00000000000..f1092b15def --- /dev/null +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -0,0 +1,33 @@ +use litellm_auth::{ + CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, + ProviderAuthPolicy, ResolvedCredential, SecretValue, +}; + +const RULES: &[CredentialRule] = &[CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("x-api-key"), +}]; + +#[test] +fn facade_applies_shared_auth_policy() { + let policy = ProviderAuthPolicy { + rules: RULES, + accepted_existing_headers: &["x-api-key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: None, + audience: None, + }; + + let headers = policy + .apply( + Vec::new(), + CredentialPlanKind::Static, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .expect("facade policy applies"); + + assert_eq!( + headers, + vec![("x-api-key".to_string(), "secret".to_string())] + ); +} From 82bc67b1220389ffa7e396e285e2443e6933d48e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 16:09:20 -0700 Subject: [PATCH 109/146] feat(rust): add typed secret managers and shared auth adapters --- .github/workflows/test-rust.yml | 7 + litellm-rust/Cargo.lock | 1037 ++++++++++++++++- litellm-rust/Cargo.toml | 6 + litellm-rust/crates/auth-aws/src/constants.rs | 2 + litellm-rust/crates/auth-gcp/Cargo.toml | 5 + litellm-rust/crates/auth-gcp/src/lib.rs | 60 +- litellm-rust/crates/auth-gcp/src/sdk.rs | 106 ++ litellm-rust/crates/secrets-aws/Cargo.toml | 24 + litellm-rust/crates/secrets-aws/src/auth.rs | 79 ++ litellm-rust/crates/secrets-aws/src/error.rs | 29 + litellm-rust/crates/secrets-aws/src/kms.rs | 63 + litellm-rust/crates/secrets-aws/src/lib.rs | 10 + .../crates/secrets-aws/src/secret_manager.rs | 297 +++++ litellm-rust/crates/secrets-aws/tests/kms.rs | 59 + .../secrets-aws/tests/secret_manager.rs | 292 +++++ litellm-rust/crates/secrets-google/Cargo.toml | 28 + .../crates/secrets-google/src/auth.rs | 21 + .../crates/secrets-google/src/error.rs | 43 + litellm-rust/crates/secrets-google/src/kms.rs | 67 ++ litellm-rust/crates/secrets-google/src/lib.rs | 10 + .../secrets-google/src/secret_manager.rs | 169 +++ .../crates/secrets-google/tests/kms.rs | 49 + .../secrets-google/tests/secret_manager.rs | 176 +++ litellm-rust/crates/secrets-types/Cargo.toml | 17 + .../secrets-types/src/base_secret_manager.rs | 58 + .../crates/secrets-types/src/config.rs | 92 ++ .../crates/secrets-types/src/error.rs | 9 + litellm-rust/crates/secrets-types/src/lib.rs | 12 + .../crates/secrets-types/src/value.rs | 32 + .../crates/secrets-types/tests/config.rs | 60 + .../crates/secrets-types/tests/rotation.rs | 105 ++ litellm-rust/crates/secrets/Cargo.toml | 37 + litellm-rust/crates/secrets/src/error.rs | 39 + litellm-rust/crates/secrets/src/handler.rs | 118 ++ litellm-rust/crates/secrets/src/lib.rs | 21 + litellm-rust/crates/secrets/src/oidc.rs | 264 +++++ litellm-rust/crates/secrets/src/resolver.rs | 140 +++ litellm-rust/crates/secrets/src/state.rs | 105 ++ litellm-rust/crates/secrets/tests/handler.rs | 107 ++ litellm-rust/crates/secrets/tests/oidc.rs | 295 +++++ .../crates/secrets/tests/resolution.rs | 343 ++++++ 41 files changed, 4459 insertions(+), 34 deletions(-) create mode 100644 litellm-rust/crates/auth-gcp/src/sdk.rs create mode 100644 litellm-rust/crates/secrets-aws/Cargo.toml create mode 100644 litellm-rust/crates/secrets-aws/src/auth.rs create mode 100644 litellm-rust/crates/secrets-aws/src/error.rs create mode 100644 litellm-rust/crates/secrets-aws/src/kms.rs create mode 100644 litellm-rust/crates/secrets-aws/src/lib.rs create mode 100644 litellm-rust/crates/secrets-aws/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-aws/tests/kms.rs create mode 100644 litellm-rust/crates/secrets-aws/tests/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-google/Cargo.toml create mode 100644 litellm-rust/crates/secrets-google/src/auth.rs create mode 100644 litellm-rust/crates/secrets-google/src/error.rs create mode 100644 litellm-rust/crates/secrets-google/src/kms.rs create mode 100644 litellm-rust/crates/secrets-google/src/lib.rs create mode 100644 litellm-rust/crates/secrets-google/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-google/tests/kms.rs create mode 100644 litellm-rust/crates/secrets-google/tests/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-types/Cargo.toml create mode 100644 litellm-rust/crates/secrets-types/src/base_secret_manager.rs create mode 100644 litellm-rust/crates/secrets-types/src/config.rs create mode 100644 litellm-rust/crates/secrets-types/src/error.rs create mode 100644 litellm-rust/crates/secrets-types/src/lib.rs create mode 100644 litellm-rust/crates/secrets-types/src/value.rs create mode 100644 litellm-rust/crates/secrets-types/tests/config.rs create mode 100644 litellm-rust/crates/secrets-types/tests/rotation.rs create mode 100644 litellm-rust/crates/secrets/Cargo.toml create mode 100644 litellm-rust/crates/secrets/src/error.rs create mode 100644 litellm-rust/crates/secrets/src/handler.rs create mode 100644 litellm-rust/crates/secrets/src/lib.rs create mode 100644 litellm-rust/crates/secrets/src/oidc.rs create mode 100644 litellm-rust/crates/secrets/src/resolver.rs create mode 100644 litellm-rust/crates/secrets/src/state.rs create mode 100644 litellm-rust/crates/secrets/tests/handler.rs create mode 100644 litellm-rust/crates/secrets/tests/oidc.rs create mode 100644 litellm-rust/crates/secrets/tests/resolution.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 70d828b5c8b..278fa7c425f 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -127,6 +127,13 @@ jobs: cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}" done + - name: Test secret manager feature combinations + run: | + cargo test -p litellm-auth-gcp --locked --no-default-features + for features in '' aws google aws,google; do + cargo test -p litellm-secrets --locked --no-default-features --features "$features" + done + rust-wheel: runs-on: ubuntu-latest timeout-minutes: 30 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5cf9116d752..81fcabaf122 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -82,6 +82,16 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-compression" version = "0.4.46" @@ -191,9 +201,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.8.1" +version = "1.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7816e98ee912159f45d307e5ee6bfea4a335a55aee15f7f3e32f81a6f3000f1d" +checksum = "25b43ad47adc2517efe3d706559d94b97e50e80e0321b3cadbc9f77cee88adcd" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -214,6 +224,58 @@ dependencies = [ "uuid", ] +[[package]] +name = "aws-sdk-kms" +version = "1.120.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6b0fe38fee2ba5b6cd24d32d08365b314ae1adea123649e061a7eb6300b6f5b" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sdk-secretsmanager" +version = "1.117.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d32d781b34ab083e0dc54b4c68fc5e89ddf35e97d97bdcb9386d21325c14767" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sts" version = "1.108.0" @@ -243,9 +305,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.5.1" +version = "1.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "723c2234ad7511ceef63eab016b7ba6ff7c55590fefb96fa8467af014a07309f" +checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -308,9 +370,9 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.2.0" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e" +checksum = "7bd25384a4e437aa8d8f339afad4b69e786b936a7cb10db668a7aaf66717b1a8" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -338,9 +400,9 @@ dependencies = [ [[package]] name = "aws-smithy-json" -version = "0.63.0" +version = "0.63.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3dc65a121adb4b33729919fcfa14fa36fb33c1555a8f06bb0e2188dbfdc1d9ef" +checksum = "3385d469edbe8b60cc72002784652b5efca39178192aa9cc4b44c9875c6bdc18" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-schema", @@ -368,9 +430,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.12.0" +version = "1.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +checksum = "3296253d3a91b3f938a3f2bcce4daebadcb4aa4228153fcec90f9d23532b4484" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -394,9 +456,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.13.0" +version = "1.16.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +checksum = "6d881a7b7ad179fd6611680c9de89f716fb00ab40299a9a7b8c6913e8f7511a8" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", @@ -423,9 +485,9 @@ dependencies = [ [[package]] name = "aws-smithy-schema" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" +checksum = "e8f395d93304280b64b7632fea798d177e74897fe7f063416ce627cd6fa24829" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -434,9 +496,9 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.6.1" +version = "1.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8" +checksum = "0b791f3ac597193fe1d08b82366986eb1f5bc31f2ac6c194c0855276116c76cd" dependencies = [ "base64-simd", "bytes", @@ -472,9 +534,9 @@ dependencies = [ [[package]] name = "aws-types" -version = "1.4.0" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e957a6c6dbce82b7a91f44231c09273159703769f447cbe85e854dfe9cf67f86" +checksum = "209f3a6d82a6e9e5f94abbed94c7a26e1c052341002bf57a5fb5481f625896fc" dependencies = [ "aws-credential-types", "aws-smithy-async", @@ -580,6 +642,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.1" @@ -632,6 +700,9 @@ name = "bytes" version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" +dependencies = [ + "serde", +] [[package]] name = "bytes-utils" @@ -1065,6 +1136,55 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be1e0bca6c3637f992fc1cc7cbc52a78c1ef6db076dbf1059c4323d6a2048376" +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.19", +] + [[package]] name = "deranged" version = "0.5.8" @@ -1395,6 +1515,15 @@ dependencies = [ "version_check", ] +[[package]] +name = "getopts" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" +dependencies = [ + "unicode-width", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -1440,6 +1569,224 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "google-cloud-auth" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff461519b1a948200f163574be072753bcfb462a323f0eb426629d89872dd685" +dependencies = [ + "async-trait", + "base64 0.23.1", + "bytes", + "google-cloud-gax", + "hex", + "hmac", + "http 1.4.2", + "jiff", + "reqwest 0.13.5", + "rustc_version", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.19", + "time", + "tokio", + "url", +] + +[[package]] +name = "google-cloud-gax" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5615cff28ee59cfe52fbb4c11b8b1e77f650296e2ea4f4c2b7757ac6b19e752" +dependencies = [ + "bytes", + "futures", + "google-cloud-rpc", + "google-cloud-wkt", + "http 1.4.2", + "pin-project", + "rand 0.10.2", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", +] + +[[package]] +name = "google-cloud-gax-internal" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2766757d877a7a8ac23da9884cb0e3f10ed9b75a0ce59801ce6b19bf9d5819e" +dependencies = [ + "bytes", + "futures", + "google-cloud-auth", + "google-cloud-gax", + "google-cloud-rpc", + "google-cloud-wkt", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "lazy_static", + "opentelemetry", + "opentelemetry-semantic-conventions", + "opentelemetry_sdk", + "percent-encoding", + "pin-project", + "prost", + "prost-types", + "reqwest 0.13.5", + "rustc_version", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", + "tower", + "tracing", + "tracing-opentelemetry", +] + +[[package]] +name = "google-cloud-iam-v1" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f962b40234b1531e6ef73f7558871c96e117231e962086c98804323fb8d2c82" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-type", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-kms-v1" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d0f6eab19254d9abd98035cd54e3f2522d2c49abf9d93bf5425ee738d198c15" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-iam-v1", + "google-cloud-location", + "google-cloud-longrunning", + "google-cloud-lro", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-location" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "280d5acdba8fcb1232c0719ed788d85b7e362b82cbb425b7050d3ce46f075ede" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-longrunning" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c0363c5389ffda2b55cd8a86eef4b19a3481a48467c91dc5d77f626a9572766" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-lro" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47af3deef75c14a2983c430898d960c765bddbcc9f9188ca0563108e9227cfe7" +dependencies = [ + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-longrunning", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "tokio", + "tracing", +] + +[[package]] +name = "google-cloud-rpc" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2162c08a89118130979ba261080e960e44cdcb2d6e2ab8ca9b1da245285d353" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-type" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63acc3a92a85f96bab021c3a3e29b53bbacc97651e1b524d4c2991960a63eb82" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-wkt" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fccf98cfd5481a5f5a285181ab0c62123d7d47cd2bb7299448440649349e4e7" +dependencies = [ + "base64 0.22.1", + "bytes", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.19", + "time", + "url", +] + [[package]] name = "h2" version = "0.3.27" @@ -1507,6 +1854,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + [[package]] name = "hex" version = "0.4.3" @@ -1636,6 +1989,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -1675,6 +2029,19 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.10.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -1860,6 +2227,27 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "is-macro" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8267aa6001e25494f3015f9663bbd88a18240c74483afa5f0934a1b3e4c388e9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.13.0" @@ -1884,6 +2272,43 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jiff" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ab1baf72f08796de0260609515130699b890ac25f30e610ad894bc5856cafdb" +dependencies = [ + "defmt", + "jiff-core", + "jiff-static", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", +] + +[[package]] +name = "jiff-core" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e52fe76043ccecc9005d2305ebaadf7d7fc0cc89ca6baa10a94d6bc68c7128c" +dependencies = [ + "defmt", + "log", +] + +[[package]] +name = "jiff-static" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378268a1116ad67ae6228701118ac9f491d78fda38a40a1f1a9e1348de6f7212" +dependencies = [ + "jiff-core", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "jni" version = "0.22.4" @@ -1954,6 +2379,27 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "11.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e75fe14a82d81e5f5af639997db37d8b96045938a7ac6ab18cdbe1c7467e05e1" +dependencies = [ + "base64 0.22.1", + "getrandom 0.2.17", + "js-sys", + "serde", + "serde_json", + "signature", + "zeroize", +] + +[[package]] +name = "lalrpop-util" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553" + [[package]] name = "lazy_static" version = "1.5.0" @@ -2023,6 +2469,8 @@ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ "gcp_auth", + "google-cloud-auth", + "http 1.4.2", "litellm-auth-types", "moka", "serde_json", @@ -2252,6 +2700,90 @@ dependencies = [ "tokio-tungstenite", ] +[[package]] +name = "litellm-secrets" +version = "0.1.0" +dependencies = [ + "aws-sdk-kms", + "base64 0.22.1", + "google-cloud-auth", + "google-cloud-kms-v1", + "jsonwebtoken", + "litellm-core-utils", + "litellm-secrets-aws", + "litellm-secrets-google", + "litellm-secrets-types", + "moka", + "reqwest 0.12.28", + "rstest", + "rustpython-parser", + "serde", + "serde_json", + "strum", + "tempfile", + "thiserror 2.0.19", + "tokio", + "tracing", + "wiremock", +] + +[[package]] +name = "litellm-secrets-aws" +version = "0.1.0" +dependencies = [ + "aws-credential-types", + "aws-sdk-kms", + "aws-sdk-secretsmanager", + "base64 0.22.1", + "litellm-auth-aws", + "litellm-core-utils", + "litellm-secrets-types", + "rstest", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", + "veil", + "wiremock", +] + +[[package]] +name = "litellm-secrets-google" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", + "google-cloud-auth", + "google-cloud-gax", + "google-cloud-kms-v1", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "percent-encoding", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + +[[package]] +name = "litellm-secrets-types" +version = "0.1.0" +dependencies = [ + "litellm-auth-types", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "veil", +] + [[package]] name = "litellm-token-counter" version = "0.1.0" @@ -2277,7 +2809,7 @@ dependencies = [ "base64 0.22.1", "rand 0.8.7", "rstest", - "rustc-hash", + "rustc-hash 2.1.3", "serde", "serde_json", "thiserror 2.0.19", @@ -2453,6 +2985,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.5.1" @@ -2487,6 +3029,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2499,7 +3051,7 @@ version = "6.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" dependencies = [ - "bitflags", + "bitflags 2.13.1", "libc", "once_cell", "onig_sys", @@ -2527,6 +3079,42 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "opentelemetry" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" +dependencies = [ + "futures-core", + "futures-sink", + "js-sys", + "pin-project-lite", + "thiserror 2.0.19", + "tracing", +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c913ac17a6c451661ee255f4625d143e51647ae78ebd969b75e41c4442f4fe47" + +[[package]] +name = "opentelemetry_sdk" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b59f80e1ac4d5ff7a2db8fb6c80badb7f0f3f858211fba08dd9aaec750894f9" +dependencies = [ + "futures-channel", + "futures-executor", + "futures-util", + "opentelemetry", + "percent-encoding", + "portable-atomic", + "rand 0.9.5", + "thiserror 2.0.19", +] + [[package]] name = "outref" version = "0.5.2" @@ -2590,6 +3178,44 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phf" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +dependencies = [ + "phf_shared", + "rand 0.8.7", +] + +[[package]] +name = "phf_shared" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -2662,6 +3288,15 @@ version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -2712,7 +3347,7 @@ checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" dependencies = [ "bit-set", "bit-vec", - "bitflags", + "bitflags 2.13.1", "num-traits", "rand 0.9.5", "rand_chacha 0.9.0", @@ -2723,6 +3358,38 @@ dependencies = [ "unarray", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -2821,7 +3488,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.3", "rustls 0.23.42", "socket2 0.6.5", "thiserror 2.0.19", @@ -2843,7 +3510,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash", + "rustc-hash 2.1.3", "rustls 0.23.42", "rustls-pki-types", "slab", @@ -3022,7 +3689,7 @@ dependencies = [ "arcstr", "combine", "itoa", - "num-bigint", + "num-bigint 0.5.1", "percent-encoding", "ryu", "sha1_smol", @@ -3049,7 +3716,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.1", ] [[package]] @@ -3180,6 +3847,9 @@ dependencies = [ "rustls 0.23.42", "rustls-pki-types", "rustls-platform-verifier", + "serde", + "serde_json", + "serde_urlencoded", "sync_wrapper", "tokio", "tokio-rustls 0.26.4", @@ -3248,6 +3918,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.3" @@ -3269,7 +3945,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys", @@ -3295,6 +3971,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -3374,6 +4051,63 @@ dependencies = [ "untrusted", ] +[[package]] +name = "rustpython-ast" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5" +dependencies = [ + "is-macro", + "num-bigint 0.4.8", + "rustpython-parser-core", + "static_assertions", +] + +[[package]] +name = "rustpython-parser" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb" +dependencies = [ + "anyhow", + "is-macro", + "itertools 0.11.0", + "lalrpop-util", + "log", + "num-bigint 0.4.8", + "num-traits", + "phf", + "phf_codegen", + "rustc-hash 1.1.0", + "rustpython-ast", + "rustpython-parser-core", + "tiny-keccak", + "unic-emoji-char", + "unic-ucd-ident", + "unicode_names2", +] + +[[package]] +name = "rustpython-parser-core" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad" +dependencies = [ + "is-macro", + "memchr", + "rustpython-parser-vendored", +] + +[[package]] +name = "rustpython-parser-vendored" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6" +dependencies = [ + "memchr", + "once_cell", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -3462,7 +4196,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "core-foundation", "core-foundation-sys", "libc", @@ -3622,6 +4356,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -3638,6 +4381,15 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -3660,6 +4412,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + [[package]] name = "slab" version = "0.4.12" @@ -3869,6 +4627,15 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tiktoken-rs" version = "0.12.0" @@ -3881,7 +4648,7 @@ dependencies = [ "fancy-regex 0.17.0", "lazy_static", "regex", - "rustc-hash", + "rustc-hash 2.1.3", ] [[package]] @@ -3914,6 +4681,15 @@ dependencies = [ "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -4029,6 +4805,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-tungstenite" version = "0.24.0" @@ -4088,6 +4876,44 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "base64 0.22.1", + "bytes", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "rustls-native-certs", + "sync_wrapper", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -4096,11 +4922,15 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.14.0", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -4110,7 +4940,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "async-compression", - "bitflags", + "bitflags 2.13.1", "bytes", "futures-core", "futures-util", @@ -4167,6 +4997,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", ] [[package]] @@ -4179,6 +5010,31 @@ dependencies = [ "tracing", ] +[[package]] +name = "tracing-opentelemetry" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" +dependencies = [ + "js-sys", + "opentelemetry", + "tracing", + "tracing-core", + "tracing-subscriber", + "web-time", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "sharded-slab", + "thread_local", + "tracing-core", +] + [[package]] name = "try-lock" version = "0.2.5" @@ -4268,6 +5124,58 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" +[[package]] +name = "unic-char-property" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221" +dependencies = [ + "unic-char-range", +] + +[[package]] +name = "unic-char-range" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc" + +[[package]] +name = "unic-common" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc" + +[[package]] +name = "unic-emoji-char" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d" +dependencies = [ + "unic-char-property", + "unic-char-range", + "unic-ucd-version", +] + +[[package]] +name = "unic-ucd-ident" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987" +dependencies = [ + "unic-char-property", + "unic-char-range", + "unic-ucd-version", +] + +[[package]] +name = "unic-ucd-version" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4" +dependencies = [ + "unic-common", +] + [[package]] name = "unicase" version = "2.9.0" @@ -4295,12 +5203,40 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode_categories" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "unicode_names2" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd" +dependencies = [ + "phf", + "unicode_names2_generator", +] + +[[package]] +name = "unicode_names2_generator" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e" +dependencies = [ + "getopts", + "log", + "phf_codegen", + "rand 0.8.7", +] + [[package]] name = "untrusted" version = "0.9.0" @@ -4348,6 +5284,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "veil" version = "0.3.0" @@ -4724,6 +5666,29 @@ dependencies = [ "memchr", ] +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64 0.22.1", + "deadpool", + "futures", + "http 1.4.2", + "http-body-util", + "hyper 1.10.1", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "wit-bindgen" version = "0.57.1" @@ -4817,6 +5782,20 @@ name = "zeroize" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] [[package]] name = "zerotrie" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b781d564805..2f6f5feb4ad 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -18,6 +18,10 @@ litellm-auth-types = { path = "crates/auth-types" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } +litellm-secrets = { path = "crates/secrets" } +litellm-secrets-types = { path = "crates/secrets-types" } +litellm-secrets-aws = { path = "crates/secrets-aws" } +litellm-secrets-google = { path = "crates/secrets-google" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } @@ -32,6 +36,8 @@ litellm-host-python = { path = "crates/host-python" } bytes = "1" http = "1" +google-cloud-auth = { version = "1.16.0", default-features = false } +jsonwebtoken = { version = "11.1.0", default-features = false } hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] } proptest = "1.7.0" pyo3 = "0.29.2" diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index be215cc9016..9e7c6bfab43 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -3,6 +3,8 @@ pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; pub const AWS_REGION: &str = "AWS_REGION"; +pub const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION"; +pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "AWS_BEDROCK_RUNTIME_ENDPOINT"; pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 045320efc80..0c6258a193c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +google-sdk = ["dep:google-cloud-auth", "dep:http"] + [dependencies] litellm-auth-types.workspace = true @@ -14,3 +17,5 @@ sha2.workspace = true tokio.workspace = true gcp_auth = "0.12.7" +google-cloud-auth = { workspace = true, optional = true } +http = { workspace = true, optional = true } diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 1ff9487e3f4..8aeddae9efc 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -8,6 +8,11 @@ use moka::future::Cache; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; +#[cfg(feature = "google-sdk")] +mod sdk; +#[cfg(feature = "google-sdk")] +pub use sdk::GoogleCredentials; + const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; @@ -26,19 +31,31 @@ pub struct VertexConfig { } impl VertexConfig { + pub fn new( + credentials: Option>, + project_id: Option, + location: Option, + ) -> Self { + Self { + credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()), + project_id: project_id.filter(|value| !value.trim().is_empty()), + location: location.filter(|value| !value.trim().is_empty()), + } + } + pub fn from_sourced_optional_params( params: &Map, sources: &BTreeMap, ) -> Result { - Ok(Self { - credentials: optional_credentials( + Ok(Self::new( + optional_credentials( params, sources, &["vertex_credentials", "vertex_ai_credentials"], )?, - project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?, - location: optional_string(params, &["vertex_location", "vertex_ai_location"])?, - }) + optional_string(params, &["vertex_project", "vertex_ai_project"])?, + optional_string(params, &["vertex_location", "vertex_ai_location"])?, + )) } pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self { @@ -469,6 +486,39 @@ mod tests { assert_eq!(config.location(), Some("alias-location")); } + #[test] + fn typed_config_preserves_source_and_empty_value_fallback() { + let configured = VertexConfig::new( + Some(Sourced::new( + SecretValue::new("inline-json"), + InputSource::Request, + )), + Some("project".into()), + Some("location".into()), + ); + assert!(matches!( + credential_source(&configured, &|_| Some("environment-json".into())), + CredentialSource::Inline(value) if value.expose() == "inline-json" + )); + let empty = VertexConfig::new( + Some(Sourced::new(SecretValue::new(" "), InputSource::Request)), + Some(" ".into()), + Some(" ".into()), + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(), + Some("env-location") + ); + } + #[test] fn project_and_location_prefer_input_then_environment() { let configured = diff --git a/litellm-rust/crates/auth-gcp/src/sdk.rs b/litellm-rust/crates/auth-gcp/src/sdk.rs new file mode 100644 index 00000000000..566df291b7d --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/sdk.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag}; +use google_cloud_auth::errors::CredentialsError; +use http::{Extensions, HeaderMap, HeaderName, HeaderValue}; +use litellm_auth_types::Error; + +use crate::{VertexAuth, VertexConfig}; + +type EnvironmentLookup = dyn Fn(&str) -> Option + Send + Sync; + +pub struct GoogleCredentials { + auth: VertexAuth, + config: VertexConfig, + environment: Arc, +} + +impl GoogleCredentials { + pub fn new(config: VertexConfig, environment: Arc) -> Self { + Self { + auth: VertexAuth::default(), + config, + environment, + } + } + + pub async fn request_headers(&self) -> Result { + let response = self + .auth + .validate_environment(Vec::new(), None, &self.config, &|name| { + (self.environment)(name) + }) + .await?; + response + .headers + .into_iter() + .map(|(key, value)| { + let name = + HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidHeader)?; + let value = HeaderValue::from_str(&value).map_err(|_| Error::InvalidHeader)?; + Ok((name, value)) + }) + .collect() + } +} + +impl CredentialsProvider for GoogleCredentials { + async fn headers( + &self, + _: Extensions, + ) -> Result, CredentialsError> { + self.request_headers() + .await + .map(|data| CacheableResource::New { + entity_tag: EntityTag::new(), + data, + }) + .map_err(|_| CredentialsError::from_msg(false, "Google authentication failed")) + } + + async fn universe_domain(&self) -> Option { + None + } +} + +impl std::fmt::Debug for GoogleCredentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GoogleCredentials").finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn sdk_and_http_credentials_share_token_resolution_and_redaction() { + let credentials = GoogleCredentials::new( + VertexConfig::new(None, Some("project".into()), None), + Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private-token".into())), + ); + let direct = credentials.request_headers().await.unwrap(); + let CacheableResource::New { data, .. } = + credentials.headers(Extensions::new()).await.unwrap() + else { + panic!("first request did not return headers"); + }; + assert_eq!(direct, data); + assert_eq!(data[http::header::AUTHORIZATION], "Bearer private-token"); + assert!(!format!("{credentials:?}").contains("private-token")); + } + + #[tokio::test] + async fn invalid_token_headers_return_a_redacted_sdk_error() { + let credentials = GoogleCredentials::new( + VertexConfig::new(None, Some("project".into()), None), + Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private\nvalue".into())), + ); + assert_eq!( + credentials.request_headers().await.unwrap_err(), + Error::InvalidHeader + ); + let error = credentials.headers(Extensions::new()).await.unwrap_err(); + assert!(!format!("{error:?}").contains("private")); + } +} diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml new file mode 100644 index 00000000000..b1dc5b33cda --- /dev/null +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-secrets-aws" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-secrets-types.workspace = true +litellm-core-utils.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tracing = "0.1" +veil.workspace = true +aws-sdk-kms = "1.120.0" +aws-sdk-secretsmanager = "1.117.0" +aws-credential-types = "1.3.0" + +[dev-dependencies] +base64.workspace = true +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs new file mode 100644 index 00000000000..954cfa2f8fd --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -0,0 +1,79 @@ +use std::sync::Arc; + +use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; +use litellm_auth_aws::{ + AwsAuthConfig, + constants::{AWS_DEFAULT_REGION, AWS_REGION, AWS_REGION_NAME}, + resolve_credentials, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; + +use crate::Error; + +#[derive(Clone)] +pub(crate) struct Credentials { + config: AwsAuthConfig, + environment: Arc, +} + +impl Credentials { + pub(crate) fn new( + settings: &KeyManagementSettings, + environment: Arc, + ) -> Self { + Self { + config: AwsAuthConfig { + region_name: region(settings, environment.as_ref()).ok(), + role_name: settings.aws_role_name.clone(), + session_name: settings.aws_session_name.clone(), + external_id: settings + .aws_external_id + .as_ref() + .map(|v| v.expose().to_owned()), + profile_name: settings.aws_profile_name.clone(), + web_identity_token: settings + .aws_web_identity_token + .as_ref() + .map(|v| v.expose().to_owned()), + sts_endpoint: settings.aws_sts_endpoint.clone(), + ..Default::default() + }, + environment, + } + } +} + +impl ProvideCredentials for Credentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::new(async { + resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) + .await + .map_err(|_| { + CredentialsError::provider_error("secret manager authentication failed") + }) + }) + } +} + +pub(crate) fn region( + settings: &KeyManagementSettings, + environment: &dyn Lookup, +) -> Result { + settings + .aws_region_name + .clone() + .or_else(|| environment.get(AWS_REGION_NAME)) + .or_else(|| environment.get(AWS_REGION)) + .or_else(|| environment.get(AWS_DEFAULT_REGION)) + .ok_or(Error::MissingRegion) +} + +impl std::fmt::Debug for Credentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Credentials").finish_non_exhaustive() + } +} diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs new file mode 100644 index 00000000000..3f2c83a6a6d --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -0,0 +1,29 @@ +use aws_sdk_secretsmanager::error::SdkError; + +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("AWS authentication failed")] + Auth(#[from] #[redact] litellm_auth_aws::Error), + #[error("AWS region is not configured")] + MissingRegion, + #[error("KMS response has no plaintext")] + MissingPlaintext, + #[error("AWS request timed out")] + Timeout, + #[error("AWS KMS decrypt failed")] + Decrypt(#[from] #[redact] Box>), + #[error("AWS Secrets Manager request preparation failed")] + Read(#[from] #[redact] Box>), + #[error("AWS Secrets Manager create failed")] + Create(#[from] #[redact] Box>), + #[error("AWS Secrets Manager update failed")] + Put(#[from] #[redact] Box>), + #[error("AWS Secrets Manager delete failed")] + Delete(#[from] #[redact] Box>), + #[error("AWS Secrets Manager replication failed")] + Replicate(#[from] #[redact] Box>), + #[error("primary secret is not a JSON object")] + PrimarySecret, + #[error(transparent)] + Operation(#[from] litellm_secrets_types::Error), +} diff --git a/litellm-rust/crates/secrets-aws/src/kms.rs b/litellm-rust/crates/secrets-aws/src/kms.rs new file mode 100644 index 00000000000..a66b1c4d2fe --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/kms.rs @@ -0,0 +1,63 @@ +use litellm_auth_aws::constants::AWS_REGION_NAME; +use std::sync::Arc; + +use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Region}, + primitives::Blob, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; + +use crate::{Error, auth}; + +#[derive(Clone)] +pub struct AwsKms { + client: Client, +} + +impl AwsKms { + pub fn new(client: Client) -> Self { + Self { client } + } + + pub async fn decrypt(&self, ciphertext: Vec) -> Result, Error> { + let response = self + .client + .decrypt() + .ciphertext_blob(Blob::new(ciphertext)) + .send() + .await + .map_err(|error| Error::Decrypt(Box::new(error)))?; + Ok(response + .plaintext + .ok_or(Error::MissingPlaintext)? + .into_inner()) + } +} + +pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> { + environment + .get(AWS_REGION_NAME) + .map(|_| ()) + .ok_or(Error::MissingRegion) +} + +pub fn load_aws_kms( + use_aws_kms: Option, + settings: &KeyManagementSettings, + environment: Arc, +) -> Result, Error> { + if use_aws_kms != Some(true) { + return Ok(None); + } + if settings.aws_region_name.is_none() { + validate_environment(environment.as_ref())?; + } + let config = aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(auth::region(settings, environment.as_ref())?)) + .credentials_provider(auth::Credentials::new(settings, environment)) + .build(); + Ok(Some(AwsKms::new(Client::from_conf(config)))) +} diff --git a/litellm-rust/crates/secrets-aws/src/lib.rs b/litellm-rust/crates/secrets-aws/src/lib.rs new file mode 100644 index 00000000000..d17eb38c7eb --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/lib.rs @@ -0,0 +1,10 @@ +#![forbid(unsafe_code)] + +mod auth; +mod error; +pub mod kms; +pub mod secret_manager; + +pub use error::Error; +pub use kms::{AwsKms, load_aws_kms}; +pub use secret_manager::{AwsSecretWriteSettings, AwsSecretsManagerV2, RotationResponse}; diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs new file mode 100644 index 00000000000..97645ecac6d --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -0,0 +1,297 @@ +use litellm_auth_aws::constants::AWS_BEDROCK_RUNTIME_ENDPOINT; +use std::{collections::BTreeMap, sync::Arc}; + +use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Region}, + operation::{ + create_secret::CreateSecretOutput, delete_secret::DeleteSecretOutput, + put_secret_value::PutSecretValueOutput, + replicate_secret_to_regions::ReplicateSecretToRegionsOutput, + }, + types::{ReplicaRegionType, Tag}, +}; +use litellm_auth_aws::constants::{ + AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{ + BaseSecretManager, KeyManagementSettings, Secret, SecretValue, async_rotate_secret, +}; +use serde_json::Value; + +use crate::{Error, auth}; + +#[derive(Clone)] +pub struct AwsSecretsManagerV2 { + client: Client, + write_settings: AwsSecretWriteSettings, +} + +#[derive(Clone, Debug, Default)] +pub struct AwsSecretWriteSettings { + pub kms_key_id: Option, + pub tags: Option>, + pub replica_regions: Option>, +} + +impl From<&KeyManagementSettings> for AwsSecretWriteSettings { + fn from(settings: &KeyManagementSettings) -> Self { + Self { + kms_key_id: settings.kms_key_id.clone(), + tags: settings.tags.clone(), + replica_regions: settings.replica_regions.clone(), + } + } +} + +#[derive(Debug)] +pub enum RotationResponse { + Created(CreateSecretOutput), + Updated(PutSecretValueOutput), +} + +impl AwsSecretsManagerV2 { + pub fn new(client: Client, write_settings: AwsSecretWriteSettings) -> Self { + Self { + client, + write_settings, + } + } + + pub fn load_aws_secret_manager( + use_aws_secret_manager: Option, + settings: KeyManagementSettings, + environment: Arc, + ) -> Result, Error> { + if use_aws_secret_manager != Some(true) { + return Ok(None); + } + let builder = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(auth::region(&settings, environment.as_ref())?)) + .credentials_provider(auth::Credentials::new(&settings, environment.clone())); + let config = match environment.get(AWS_BEDROCK_RUNTIME_ENDPOINT) { + Some(url) => builder + .endpoint_url(url.replace("bedrock-runtime", "secretsmanager")) + .build(), + None => builder.build(), + }; + Ok(Some(Self::new( + Client::from_conf(config), + (&settings).into(), + ))) + } + + pub async fn read_secret_for_resolver( + &self, + name: &str, + primary_name: Option<&str>, + environment: &(dyn Lookup + Sync), + ) -> Result, Error> { + if bootstrap_key(name) { + return Ok(environment + .get(name) + .map(SecretValue::new) + .map(Secret::String)); + } + match primary_name.filter(|name| !name.is_empty()) { + None => self + .async_read_secret(name) + .await + .map(|value| value.map(Secret::String)), + Some(primary) => { + let value = if bootstrap_key(primary) { + environment.get(primary).map(SecretValue::new) + } else { + self.async_read_secret(primary).await? + }; + let object: Value = serde_json::from_str( + value + .as_ref() + .map(SecretValue::expose) + .filter(|v| !v.is_empty()) + .unwrap_or("{}"), + ) + .map_err(|_| Error::PrimarySecret)?; + let object = object.as_object().ok_or(Error::PrimarySecret)?; + Ok(object.get(name).cloned().and_then(Secret::from_json)) + } + } + } + + pub async fn async_read_secret(&self, name: &str) -> Result, Error> { + match self.client.get_secret_value().secret_id(name).send().await { + Ok(response) => Ok(response.secret_string.map(SecretValue::new)), + Err(error) + if matches!( + &error, + aws_sdk_secretsmanager::error::SdkError::TimeoutError(_) + ) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) => + { + Err(Error::Timeout) + } + Err(error) if request_preparation_failed(&error) => Err(Error::Read(Box::new(error))), + Err(_) => { + tracing::error!("AWS secret read failed"); + Ok(None) + } + } + } + + pub async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + let response = self + .client + .create_secret() + .name(name) + .secret_string(value.expose()) + .set_description(description.filter(|v| !v.is_empty()).map(str::to_owned)) + .set_kms_key_id( + self.write_settings + .kms_key_id + .clone() + .filter(|v| !v.is_empty()), + ) + .set_tags(self.write_settings.tags.as_ref().map(|tags| { + tags.iter() + .map(|(key, value)| Tag::builder().key(key).value(value).build()) + .collect() + })) + .send() + .await + .map_err(|error| Error::Create(Box::new(error)))?; + if let Some(regions) = &self.write_settings.replica_regions + && !regions.is_empty() + && self.async_replicate_secret(name, regions).await.is_err() + { + tracing::warn!("secret created but replication failed"); + } + Ok(response) + } + + pub async fn async_replicate_secret( + &self, + name: &str, + regions: &[String], + ) -> Result, Error> { + if regions.is_empty() { + return Ok(None); + } + self.client + .replicate_secret_to_regions() + .secret_id(name) + .set_add_replica_regions(Some( + regions + .iter() + .map(|region| ReplicaRegionType::builder().region(region).build()) + .collect(), + )) + .send() + .await + .map(Some) + .map_err(|error| Error::Replicate(Box::new(error))) + } + + pub async fn async_put_secret_value( + &self, + name: &str, + value: &SecretValue, + ) -> Result { + self.client + .put_secret_value() + .secret_id(name) + .secret_string(value.expose()) + .send() + .await + .map_err(|error| Error::Put(Box::new(error))) + } + + pub async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.client + .delete_secret() + .secret_id(name) + .recovery_window_in_days(recovery_window_in_days) + .send() + .await + .map_err(|error| Error::Delete(Box::new(error))) + } + + pub async fn async_rotate_secret( + &self, + current_name: &str, + new_name: &str, + value: &SecretValue, + ) -> Result { + if current_name == new_name { + return self + .async_put_secret_value(current_name, value) + .await + .map(RotationResponse::Updated); + } + async_rotate_secret(self, current_name, new_name, value) + .await + .map(RotationResponse::Created) + } +} + +impl BaseSecretManager for AwsSecretsManagerV2 { + type Error = Error; + type WriteResponse = CreateSecretOutput; + type DeleteResponse = DeleteSecretOutput; + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + self.async_read_secret(name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + self.async_write_secret(name, value, description).await + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.async_delete_secret(name, recovery_window_in_days) + .await + } +} + +fn bootstrap_key(name: &str) -> bool { + matches!( + name, + AWS_ACCESS_KEY_ID + | AWS_SECRET_ACCESS_KEY + | AWS_REGION_NAME + | AWS_REGION + | AWS_BEDROCK_RUNTIME_ENDPOINT + ) +} + +fn request_preparation_failed( + error: &aws_sdk_secretsmanager::error::SdkError< + aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError, + >, +) -> bool { + matches!( + error, + aws_sdk_secretsmanager::error::SdkError::ConstructionFailure(_) + ) || std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |error| { + error.source() + }) + .any(|source| source.is::()) +} diff --git a/litellm-rust/crates/secrets-aws/tests/kms.rs b/litellm-rust/crates/secrets-aws/tests/kms.rs new file mode 100644 index 00000000000..39a50297551 --- /dev/null +++ b/litellm-rust/crates/secrets-aws/tests/kms.rs @@ -0,0 +1,59 @@ +use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, +}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_aws::AwsKms; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header}, +}; + +#[tokio::test] +async fn kms_decrypt_calls_the_sdk_without_applying_lookup_policy() { + let server = MockServer::start().await; + let plaintext = " private-value\n"; + Mock::given(header("x-amz-target", "TrentService.Decrypt")) + .and(body_json( + serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"Plaintext": STANDARD.encode(plaintext)})), + ) + .expect(1) + .mount(&server) + .await; + let client = Client::from_conf( + aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + let manager = AwsKms::new(client); + assert_eq!( + manager.decrypt(b"encrypted".to_vec()).await.unwrap(), + plaintext.as_bytes() + ); +} + +#[test] +fn disabled_kms_loader_does_not_require_environment_configuration() { + use litellm_secrets_aws::load_aws_kms; + use litellm_secrets_types::KeyManagementSettings; + use std::sync::Arc; + for enabled in [None, Some(false)] { + assert!( + load_aws_kms( + enabled, + &KeyManagementSettings::default(), + Arc::new(|_: &str| None) + ) + .unwrap() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs new file mode 100644 index 00000000000..56698b391ec --- /dev/null +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -0,0 +1,292 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, +}; +use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse}; +use litellm_secrets_types::{KeyManagementSettings, SecretValue}; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header}, +}; + +fn manager(server: &MockServer, settings: KeyManagementSettings) -> AwsSecretsManagerV2 { + let client = Client::from_conf( + aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + AwsSecretsManagerV2::new(client, (&settings).into()) +} + +#[rstest::rstest] +#[case::string_value("KEY", Some("value"))] +#[case::missing_value("missing", None)] +#[case::non_string_value("BOOL", None)] +#[tokio::test] +async fn primary_lookup_preserves_read_semantics( + #[case] name: &str, + #[case] expected: Option<&str>, +) { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId":"primary"}))) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"SecretString":json!({"KEY":"value", "BOOL":true}).to_string()}), + ), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, KeyManagementSettings::default()); + assert_eq!( + manager + .read_secret_for_resolver(name, Some("primary"), &|_: &str| None) + .await + .unwrap() + .and_then(|v| v.as_str().map(str::to_owned)) + .as_deref(), + expected + ); +} + +#[rstest::rstest] +#[case::access_key("AWS_ACCESS_KEY_ID")] +#[case::secret_access_key("AWS_SECRET_ACCESS_KEY")] +#[case::region_name("AWS_REGION_NAME")] +#[case::region("AWS_REGION")] +#[case::bedrock_endpoint("AWS_BEDROCK_RUNTIME_ENDPOINT")] +#[tokio::test] +async fn bootstrap_keys_bypass_primary_lookup(#[case] name: &str) { + let server = MockServer::start().await; + let manager = manager(&server, KeyManagementSettings::default()); + assert_eq!( + manager + .read_secret_for_resolver(name, Some("primary"), &|_: &str| Some("bootstrap".into())) + .await + .unwrap() + .unwrap() + .as_str() + .unwrap(), + "bootstrap" + ); +} + +#[tokio::test] +async fn failed_read_returns_none_but_invalid_primary_json_is_an_error() { + let server = MockServer::start().await; + Mock::given(body_partial_json(json!({"SecretId":"missing"}))) + .respond_with( + ResponseTemplate::new(400).set_body_json(json!({"__type":"ResourceNotFoundException"})), + ) + .mount(&server) + .await; + Mock::given(body_partial_json(json!({"SecretId":"invalid"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString":"not-json"}))) + .mount(&server) + .await; + let manager = manager(&server, KeyManagementSettings::default()); + assert!( + manager + .async_read_secret("missing") + .await + .unwrap() + .is_none() + ); + assert!(matches!( + manager + .read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None) + .await, + Err(Error::PrimarySecret) + )); +} + +#[tokio::test] +async fn same_name_rotation_uses_put_and_returns_its_response() { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.PutSecretValue")) + .and(body_partial_json( + json!({"SecretId":"key", "SecretString":"replacement"}), + )) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"Name":"key", "VersionId":"version"})), + ) + .expect(1) + .mount(&server) + .await; + let response = manager(&server, KeyManagementSettings::default()) + .async_rotate_secret("key", "key", &SecretValue::new("replacement")) + .await + .unwrap(); + match response { + RotationResponse::Updated(output) => assert_eq!(output.version_id(), Some("version")), + _ => panic!("rotation created a second secret"), + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn renamed_rotation_reads_creates_verifies_then_deletes() { + let server = MockServer::start().await; + let step = AtomicUsize::new(0); + Mock::given(wiremock::matchers::method("POST")) + .respond_with(move |request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().unwrap(); + let action = request + .headers + .get("x-amz-target") + .unwrap() + .to_str() + .unwrap(); + match step.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert_eq!(action, "secretsmanager.GetSecretValue"); + assert_eq!(body["SecretId"], "old"); + ResponseTemplate::new(200).set_body_json(json!({"SecretString":"old-value"})) + } + 1 => { + assert_eq!(action, "secretsmanager.CreateSecret"); + assert_eq!(body["Name"], "new"); + assert_eq!(body["Description"], "Rotated from old"); + assert_eq!(body["SecretString"], "replacement"); + ResponseTemplate::new(200).set_body_json(json!({"Name":"new"})) + } + 2 => { + assert_eq!(action, "secretsmanager.GetSecretValue"); + assert_eq!(body["SecretId"], "new"); + ResponseTemplate::new(200).set_body_json(json!({"SecretString":"replacement"})) + } + 3 => { + assert_eq!(action, "secretsmanager.DeleteSecret"); + assert_eq!(body["SecretId"], "old"); + assert_eq!(body["RecoveryWindowInDays"], 7); + ResponseTemplate::new(200).set_body_json(json!({"Name":"old"})) + } + _ => panic!("unexpected request"), + } + }) + .expect(4) + .mount(&server) + .await; + assert!(matches!( + manager(&server, KeyManagementSettings::default()) + .async_rotate_secret("old", "new", &SecretValue::new("replacement")) + .await + .unwrap(), + RotationResponse::Created(_) + )); +} + +#[tokio::test] +async fn creation_passes_tags_and_kms_and_survives_replication_failure() { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.CreateSecret")) + .and(body_partial_json(json!({"Name":"key", "SecretString":"value", "KmsKeyId":"kms-key", "Tags":[{"Key":"stage", "Value":"test"}]}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"Name":"key"}))).expect(1).mount(&server).await; + Mock::given(header( + "x-amz-target", + "secretsmanager.ReplicateSecretToRegions", + )) + .and(body_partial_json( + json!({"SecretId":"key", "AddReplicaRegions":[{"Region":"replica-region"}]}), + )) + .respond_with( + ResponseTemplate::new(400).set_body_json(json!({"__type":"InvalidRequestException"})), + ) + .expect(1) + .mount(&server) + .await; + let settings = KeyManagementSettings { + kms_key_id: Some("kms-key".into()), + tags: Some(std::collections::BTreeMap::from([( + "stage".into(), + "test".into(), + )])), + replica_regions: Some(vec!["replica-region".into()]), + ..Default::default() + }; + let manager = manager(&server, settings); + assert_eq!( + manager + .async_write_secret("key", &SecretValue::new("value"), None) + .await + .unwrap() + .name(), + Some("key") + ); + assert!( + manager + .async_replicate_secret("key", &[]) + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn credential_failures_are_not_swallowed_as_missing_secrets() { + use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; + #[derive(Debug)] + struct FailedCredentials; + impl ProvideCredentials for FailedCredentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::ready(Err(CredentialsError::provider_error( + "private-auth-detail", + ))) + } + } + let server = MockServer::start().await; + let config = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(FailedCredentials) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(); + let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); + let error = manager.async_read_secret("key").await.unwrap_err(); + assert!(!format!("{error:?}").contains("private-auth-detail")); + assert!(matches!(error, Error::Read(_))); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { + use std::time::Duration; + let server = MockServer::start().await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(1)) + .set_body_json(json!({"SecretString":"late"})), + ) + .mount(&server) + .await; + let config = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .timeout_config( + aws_sdk_secretsmanager::config::timeout::TimeoutConfig::builder() + .operation_timeout(Duration::from_millis(30)) + .build(), + ) + .build(); + let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); + assert!(matches!( + manager.async_read_secret("key").await, + Err(Error::Timeout) + )); +} diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml new file mode 100644 index 00000000000..daecf20ff9e --- /dev/null +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "litellm-secrets-google" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-gcp = { workspace = true, features = ["google-sdk"] } +litellm-secrets-types.workspace = true +litellm-auth-types.workspace = true +litellm-core-utils.workspace = true +base64.workspace = true +serde_json.workspace = true +thiserror.workspace = true +moka.workspace = true +veil.workspace = true +google-cloud-kms-v1 = "1.14.0" +google-cloud-gax = { version = "1.14.0", default-features = false } +percent-encoding = "2.3" +serde.workspace = true +reqwest.workspace = true + +[dev-dependencies] +google-cloud-auth.workspace = true +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-google/src/auth.rs b/litellm-rust/crates/secrets-google/src/auth.rs new file mode 100644 index 00000000000..45fc99d8d5f --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/auth.rs @@ -0,0 +1,21 @@ +use std::sync::Arc; + +use litellm_auth_gcp::{GoogleCredentials, VertexConfig}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::SecretValue; + +pub(crate) fn credentials( + project: Option, + credentials: Option, + environment: Arc, +) -> GoogleCredentials { + GoogleCredentials::new( + VertexConfig::new( + credentials.map(|value| Sourced::new(value, InputSource::Environment)), + project, + None, + ), + Arc::new(move |name| environment.get(name)), + ) +} diff --git a/litellm-rust/crates/secrets-google/src/error.rs b/litellm-rust/crates/secrets-google/src/error.rs new file mode 100644 index 00000000000..a94cc8c2de6 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/error.rs @@ -0,0 +1,43 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("Google KMS client configuration failed")] + Client( + #[from] + #[redact] + google_cloud_gax::client_builder::Error, + ), + #[error("Google authentication failed")] + Auth( + #[from] + #[redact] + litellm_auth_types::Error, + ), + #[error("Google KMS request failed")] + Kms( + #[from] + #[redact] + google_cloud_gax::error::Error, + ), + #[error("Google Secret Manager HTTP request failed")] + Http( + #[from] + #[redact] + reqwest::Error, + ), + #[error("Google Secret Manager returned HTTP {0}")] + Status(u16), + #[error("Google Secret Manager returned no payload")] + MissingPayload, + #[error("required environment variable is missing: {0}")] + MissingEnvironment(&'static str), + #[error("invalid refresh interval")] + RefreshInterval, + #[error("payload is not valid base64")] + Base64(#[from] base64::DecodeError), + #[error("decrypted value is not UTF-8")] + Utf8, + #[error("invalid Google Secret Manager endpoint")] + Endpoint, + #[error("Google Secret Manager requires an enterprise license")] + EnterpriseRequired, +} diff --git a/litellm-rust/crates/secrets-google/src/kms.rs b/litellm-rust/crates/secrets-google/src/kms.rs new file mode 100644 index 00000000000..3a247edaa35 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/kms.rs @@ -0,0 +1,67 @@ +use std::sync::Arc; + +use google_cloud_kms_v1::client::KeyManagementService; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::SecretValue; + +use crate::{Error, auth}; + +const GOOGLE_APPLICATION_CREDENTIALS: &str = "GOOGLE_APPLICATION_CREDENTIALS"; +const GOOGLE_KMS_RESOURCE_NAME: &str = "GOOGLE_KMS_RESOURCE_NAME"; + +#[derive(Clone)] +pub struct GoogleKms { + client: KeyManagementService, + resource_name: String, +} + +impl GoogleKms { + pub fn new(client: KeyManagementService, resource_name: String) -> Self { + Self { + client, + resource_name, + } + } + + pub async fn decrypt(&self, ciphertext: Vec) -> Result, Error> { + let response = self + .client + .decrypt() + .set_name(&self.resource_name) + .set_ciphertext(ciphertext) + .send() + .await?; + Ok(response.plaintext.to_vec()) + } +} + +pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> { + for key in [GOOGLE_APPLICATION_CREDENTIALS, GOOGLE_KMS_RESOURCE_NAME] { + if environment.get(key).is_none() { + return Err(Error::MissingEnvironment(key)); + } + } + Ok(()) +} + +pub async fn load_google_kms( + use_google_kms: Option, + environment: Arc, +) -> Result, Error> { + if use_google_kms != Some(true) { + return Ok(None); + } + validate_environment(environment.as_ref())?; + let credentials = environment + .get(GOOGLE_APPLICATION_CREDENTIALS) + .ok_or(Error::MissingEnvironment(GOOGLE_APPLICATION_CREDENTIALS))?; + let resource_name = environment + .get(GOOGLE_KMS_RESOURCE_NAME) + .ok_or(Error::MissingEnvironment(GOOGLE_KMS_RESOURCE_NAME))?; + let credentials = auth::credentials(None, Some(SecretValue::new(credentials)), environment); + let client = KeyManagementService::builder() + .with_credentials(credentials) + .build() + .await?; + Ok(Some(GoogleKms::new(client, resource_name))) +} diff --git a/litellm-rust/crates/secrets-google/src/lib.rs b/litellm-rust/crates/secrets-google/src/lib.rs new file mode 100644 index 00000000000..a664f11a018 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/lib.rs @@ -0,0 +1,10 @@ +#![forbid(unsafe_code)] + +mod auth; +mod error; +pub mod kms; +pub mod secret_manager; + +pub use error::Error; +pub use kms::{GoogleKms, load_google_kms}; +pub use secret_manager::GoogleSecretManager; diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs new file mode 100644 index 00000000000..3eb8475a367 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -0,0 +1,169 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{Secret, SecretValue}; +use moka::future::Cache; +use serde::Deserialize; + +use litellm_auth_gcp::GoogleCredentials; + +use crate::{Error, auth}; + +const GOOGLE_SECRET_MANAGER_PROJECT_ID: &str = "GOOGLE_SECRET_MANAGER_PROJECT_ID"; +const GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL: &str = "GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL"; +const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL"; +const GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER: &str = + "GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER"; +const GCS_PATH_SERVICE_ACCOUNT: &str = "GCS_PATH_SERVICE_ACCOUNT"; +const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400); +const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(600); +const CACHE_CAPACITY: u64 = 200; + +#[derive(Clone)] +pub struct GoogleSecretManager { + client: reqwest::Client, + credentials: Arc, + endpoint: reqwest::Url, + project: String, + cache: Cache>, + always_read: bool, +} + +#[derive(Deserialize)] +struct Response { + payload: Option, +} + +#[derive(Deserialize)] +struct Payload { + data: Option, +} + +impl GoogleSecretManager { + pub fn with_client( + client: reqwest::Client, + endpoint: reqwest::Url, + project: String, + environment: Arc, + refresh_interval: Option, + always_read: bool, + ) -> Result { + let credentials = auth::credentials( + Some(project.clone()), + environment + .get(GCS_PATH_SERVICE_ACCOUNT) + .map(SecretValue::new), + environment, + ); + let ttl = refresh_interval + .filter(|ttl| !ttl.is_zero()) + .unwrap_or(DEFAULT_CACHE_TTL); + let cache = Cache::builder() + .max_capacity(CACHE_CAPACITY) + .time_to_live(ttl) + .build(); + Ok(Self { + client, + credentials: Arc::new(credentials), + endpoint, + project, + cache, + always_read, + }) + } + + pub fn new( + environment: Arc, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let project = environment + .get(GOOGLE_SECRET_MANAGER_PROJECT_ID) + .ok_or(Error::MissingEnvironment(GOOGLE_SECRET_MANAGER_PROJECT_ID))?; + let ttl = environment + .get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL) + .filter(|v| !v.is_empty()) + .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .transpose()? + .unwrap_or( + environment + .get(SECRET_MANAGER_REFRESH_INTERVAL) + .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .transpose()? + .unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64), + ); + let always_read = environment + .get(GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER) + .is_some_and(|v| v.eq_ignore_ascii_case("true")); + Self::with_client( + reqwest::Client::new(), + reqwest::Url::parse("https://secretmanager.googleapis.com").expect("static URL"), + project, + environment, + Some(if ttl < 0 { + Duration::from_nanos(1) + } else { + Duration::from_secs(ttl as u64) + }), + always_read, + ) + } + + pub async fn get_secret_from_google_secret_manager( + &self, + name: &str, + ) -> Result, Error> { + if !self.always_read + && let Some(cached) = self.cache.get(name).await + { + return Ok(cached.and_then(cached_secret)); + } + let url = self + .endpoint + .join(&format!( + "/v1/projects/{}/secrets/{}/versions/latest:access", + percent_encoding::utf8_percent_encode( + &self.project, + percent_encoding::NON_ALPHANUMERIC + ), + percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) + )) + .map_err(|_| Error::Endpoint)?; + let response = self + .client + .get(url) + .headers(self.credentials.request_headers().await?) + .send() + .await?; + if response.status() != reqwest::StatusCode::OK { + self.cache.insert(name.to_owned(), None).await; + return Err(Error::Status(response.status().as_u16())); + } + let response: Response = response.json().await?; + let Some(data) = response.payload.and_then(|payload| payload.data) else { + self.cache.insert(name.to_owned(), None).await; + return Err(Error::MissingPayload); + }; + let filtered: String = data + .chars() + .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) + .collect(); + let bytes = STANDARD.decode(filtered)?; + let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?; + let value = SecretValue::new(plaintext); + self.cache + .insert(name.to_owned(), Some(value.clone())) + .await; + Ok(Some(Secret::String(value))) + } +} + +fn cached_secret(value: SecretValue) -> Option { + match serde_json::from_str(value.expose()) { + Ok(json) => Secret::from_json(json), + Err(_) => Some(Secret::String(value)), + } +} diff --git a/litellm-rust/crates/secrets-google/tests/kms.rs b/litellm-rust/crates/secrets-google/tests/kms.rs new file mode 100644 index 00000000000..667ecd268c8 --- /dev/null +++ b/litellm-rust/crates/secrets-google/tests/kms.rs @@ -0,0 +1,49 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use google_cloud_kms_v1::client::KeyManagementService; +use litellm_secrets_google::GoogleKms; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, path}, +}; + +#[tokio::test] +async fn google_kms_decrypts_using_the_configured_resource() { + let server = MockServer::start().await; + let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key"; + Mock::given(path(format!("/v1/{resource}:decrypt"))) + .and(body_json( + serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = KeyManagementService::builder() + .with_endpoint(server.uri()) + .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build()) + .with_retry_policy(google_cloud_gax::retry_policy::NeverRetry) + .build() + .await + .unwrap(); + let manager = GoogleKms::new(client, resource.into()); + assert_eq!( + manager.decrypt(b"encrypted".to_vec()).await.unwrap(), + b" value\n" + ); +} + +#[tokio::test] +async fn disabled_google_kms_loader_does_not_require_environment_configuration() { + use std::sync::Arc; + for enabled in [None, Some(false)] { + assert!( + litellm_secrets_google::load_google_kms(enabled, Arc::new(|_: &str| None)) + .await + .unwrap() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs new file mode 100644 index 00000000000..66cc0068ebb --- /dev/null +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -0,0 +1,176 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_google::{Error, GoogleSecretManager}; +use litellm_secrets_types::Secret; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, path}, +}; + +fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager { + GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())), + Some(ttl), + always_read, + ) + .unwrap() +} + +#[rstest::rstest] +#[case::nonempty("private-value")] +#[case::empty("")] +#[tokio::test] +async fn successful_reads_use_auth_latest_version_and_cache_including_empty_values( + #[case] value: &str, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .and(header("authorization", "Bearer token")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(value)}})), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str() + .unwrap(), + value + ); + } +} + +#[rstest::rstest] +#[case::not_found(ResponseTemplate::new(404))] +#[case::missing_payload( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"payload":{}})) +)] +#[tokio::test] +async fn negative_cache_returns_none_after_initial_error(#[case] response: ResponseTemplate) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with(response) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + assert!(matches!( + manager.get_secret_from_google_secret_manager("key").await, + Err(Error::Status(404) | Error::MissingPayload) + )); + assert!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .is_none() + ); +} + +#[rstest::rstest] +#[case::always_read(true, Duration::from_secs(60))] +#[case::expired_cache(false, Duration::from_millis(1))] +#[tokio::test] +async fn always_read_and_expired_cache_fetch_again( + #[case] always_read: bool, + #[case] ttl: Duration, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("value")}})), + ) + .expect(2) + .mount(&server) + .await; + let manager = manager(&server, always_read, ttl); + for _ in 0..2 { + tokio::time::sleep(Duration::from_millis(5)).await; + assert!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .is_some() + ); + } +} + +#[test] +fn google_manager_requires_host_license_and_project_configuration() { + assert!(matches!( + GoogleSecretManager::new(Arc::new(|_: &str| None), false), + Err(Error::EnterpriseRequired) + )); + assert!(matches!( + GoogleSecretManager::new(Arc::new(|_: &str| None), true), + Err(Error::MissingEnvironment( + "GOOGLE_SECRET_MANAGER_PROJECT_ID" + )) + )); +} + +#[rstest::rstest] +#[case::boolean("true", Some(Secret::Bool(true)))] +#[case::null("null", None)] +#[case::string( + "\"text\"", + Some(Secret::String(litellm_secrets_types::SecretValue::new("text"))) +)] +#[case::object( + "{\"key\":1}", + Secret::from_json(serde_json::json!({"key":1})) +)] +#[tokio::test] +async fn cached_values_preserve_python_json_conversion( + #[case] raw: &str, + #[case] expected: Option, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(raw)}})), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some(raw) + ); + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap(), + expected + ); +} diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml new file mode 100644 index 00000000000..acd29746722 --- /dev/null +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-secrets-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-types.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs new file mode 100644 index 00000000000..d71bce64221 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs @@ -0,0 +1,58 @@ +use crate::{Error, SecretValue}; + +pub fn validate_secret_name(name: &str) -> Result<(), Error> { + if name.split('/').any(|segment| segment == "..") + || name + .chars() + .any(|c| c.is_control() || matches!(c, '\u{2028}' | '\u{2029}')) + { + return Err(Error::UnsafeSecretName); + } + Ok(()) +} + +#[expect( + async_fn_in_trait, + reason = "closed backend dispatch does not require Send bounds on generic rotation" +)] +pub trait BaseSecretManager { + type Error: From; + type WriteResponse; + type DeleteResponse; + + async fn async_read_secret(&self, name: &str) -> Result, Self::Error>; + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result; + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result; +} + +pub async fn async_rotate_secret( + manager: &M, + current_name: &str, + new_name: &str, + value: &SecretValue, +) -> Result { + if manager.async_read_secret(current_name).await?.is_none() { + return Err(Error::CurrentSecretMissing.into()); + } + let response = manager + .async_write_secret( + new_name, + value, + Some(&format!("Rotated from {current_name}")), + ) + .await?; + if manager.async_read_secret(new_name).await?.is_none() { + return Err(Error::NewSecretMissing.into()); + } + manager.async_delete_secret(current_name, 7).await?; + Ok(response) +} diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs new file mode 100644 index 00000000000..36d319311a3 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -0,0 +1,92 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::SecretValue; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum KeyManagementSystem { + GoogleKms, + AzureKeyVault, + AwsSecretManager, + GoogleSecretManager, + HashicorpVault, + Cyberark, + Local, + AwsKms, + Custom, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AccessMode { + #[default] + ReadOnly, + WriteOnly, + ReadAndWrite, +} + +impl AccessMode { + pub fn readable(self) -> bool { + matches!(self, Self::ReadOnly | Self::ReadAndWrite) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(default)] +pub struct KeyManagementSettings { + pub hosted_keys: Option>, + pub store_virtual_keys: Option, + pub prefix_for_stored_virtual_keys: String, + pub access_mode: AccessMode, + pub primary_secret_name: Option, + pub description: Option, + pub tags: Option>, + pub kms_key_id: Option, + pub custom_secret_manager: Option, + pub aws_region_name: Option, + pub aws_role_name: Option, + pub aws_session_name: Option, + #[serde(serialize_with = "serialize_secret")] + pub aws_external_id: Option, + pub aws_profile_name: Option, + #[serde(serialize_with = "serialize_secret")] + pub aws_web_identity_token: Option, + pub aws_sts_endpoint: Option, + pub replica_regions: Option>, +} + +impl Default for KeyManagementSettings { + fn default() -> Self { + Self { + hosted_keys: None, + store_virtual_keys: Some(false), + prefix_for_stored_virtual_keys: "litellm/".into(), + access_mode: AccessMode::ReadOnly, + primary_secret_name: None, + description: None, + tags: None, + kms_key_id: None, + custom_secret_manager: None, + aws_region_name: None, + aws_role_name: None, + aws_session_name: None, + aws_external_id: None, + aws_profile_name: None, + aws_web_identity_token: None, + aws_sts_endpoint: None, + replica_regions: None, + } + } +} + +fn serialize_secret( + value: &Option, + serializer: S, +) -> Result { + value + .as_ref() + .map(SecretValue::expose) + .serialize(serializer) +} diff --git a/litellm-rust/crates/secrets-types/src/error.rs b/litellm-rust/crates/secrets-types/src/error.rs new file mode 100644 index 00000000000..cae9c7f4c69 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/error.rs @@ -0,0 +1,9 @@ +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[error("secret name contains an unsafe path segment or control character")] + UnsafeSecretName, + #[error("current secret was not found")] + CurrentSecretMissing, + #[error("new secret could not be verified")] + NewSecretMissing, +} diff --git a/litellm-rust/crates/secrets-types/src/lib.rs b/litellm-rust/crates/secrets-types/src/lib.rs new file mode 100644 index 00000000000..0823ed13c06 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/lib.rs @@ -0,0 +1,12 @@ +#![forbid(unsafe_code)] + +mod base_secret_manager; +mod config; +mod error; +mod value; + +pub use base_secret_manager::{BaseSecretManager, async_rotate_secret, validate_secret_name}; +pub use config::{AccessMode, KeyManagementSettings, KeyManagementSystem}; +pub use error::Error; +pub use litellm_auth_types::SecretValue; +pub use value::Secret; diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs new file mode 100644 index 00000000000..524045ac007 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -0,0 +1,32 @@ +use crate::SecretValue; + +#[derive(Clone, PartialEq, Eq, veil::Redact)] +pub enum Secret { + String(SecretValue), + Bool(#[redact] bool), + Json(#[redact] serde_json::Value), +} + +impl From for Secret { + fn from(value: SecretValue) -> Self { + Self::String(value) + } +} + +impl Secret { + pub fn from_json(value: serde_json::Value) -> Option { + match value { + serde_json::Value::Null => None, + serde_json::Value::String(value) => Some(Self::String(SecretValue::new(value))), + serde_json::Value::Bool(value) => Some(Self::Bool(value)), + value => Some(Self::Json(value)), + } + } + + pub fn as_str(&self) -> Option<&str> { + match self { + Self::String(value) => Some(value.expose()), + Self::Bool(_) | Self::Json(_) => None, + } + } +} diff --git a/litellm-rust/crates/secrets-types/tests/config.rs b/litellm-rust/crates/secrets-types/tests/config.rs new file mode 100644 index 00000000000..4a5f17bc68a --- /dev/null +++ b/litellm-rust/crates/secrets-types/tests/config.rs @@ -0,0 +1,60 @@ +use litellm_secrets_types::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +use serde_json::json; + +#[test] +fn config_preserves_defaults_nulls_and_serialized_names() { + let empty: KeyManagementSettings = serde_json::from_value(json!({})).unwrap(); + assert_eq!(empty, KeyManagementSettings::default()); + assert_eq!(empty.access_mode, AccessMode::ReadOnly); + assert_eq!(empty.store_virtual_keys, Some(false)); + assert_eq!(empty.prefix_for_stored_virtual_keys, "litellm/"); + let configured: KeyManagementSettings = serde_json::from_value(json!({ + "hosted_keys": [], "store_virtual_keys": null, "access_mode": "write_only", + "aws_web_identity_token": "private-token", "aws_external_id": "private-id", + "tags": {"stage": "test"}, "replica_regions": ["test-region"] + })) + .unwrap(); + assert!(!configured.access_mode.readable()); + assert_eq!(configured.store_virtual_keys, None); + assert_eq!(configured.hosted_keys.as_deref(), Some([].as_slice())); + assert!(!format!("{configured:?}").contains("private-")); + let serialized = serde_json::to_value(&configured).unwrap(); + assert_eq!(serialized["access_mode"], "write_only"); + assert_eq!(serialized["aws_web_identity_token"], "private-token"); + assert_eq!( + serde_json::from_value::(serialized).unwrap(), + configured + ); +} + +#[rstest::rstest] +#[case::aws_kms("aws_kms", KeyManagementSystem::AwsKms)] +#[case::aws_secret_manager("aws_secret_manager", KeyManagementSystem::AwsSecretManager)] +#[case::google_kms("google_kms", KeyManagementSystem::GoogleKms)] +#[case::google_secret_manager("google_secret_manager", KeyManagementSystem::GoogleSecretManager)] +#[case::azure_key_vault("azure_key_vault", KeyManagementSystem::AzureKeyVault)] +#[case::hashicorp_vault("hashicorp_vault", KeyManagementSystem::HashicorpVault)] +#[case::cyberark("cyberark", KeyManagementSystem::Cyberark)] +#[case::custom("custom", KeyManagementSystem::Custom)] +#[case::local("local", KeyManagementSystem::Local)] +fn key_management_system_serialization_round_trips( + #[case] name: &str, + #[case] system: KeyManagementSystem, +) { + assert_eq!( + serde_json::from_value::(json!(name)).unwrap(), + system + ); + assert_eq!(serde_json::to_value(system).unwrap(), name); +} + +#[test] +fn secret_debug_never_exposes_values() { + assert!( + !format!("{:?}", Secret::String(SecretValue::new("sensitive-value"))) + .contains("sensitive-value") + ); + assert!(!format!("{:?}", Secret::Bool(true)).contains("true")); +} diff --git a/litellm-rust/crates/secrets-types/tests/rotation.rs b/litellm-rust/crates/secrets-types/tests/rotation.rs new file mode 100644 index 00000000000..48a5304ece8 --- /dev/null +++ b/litellm-rust/crates/secrets-types/tests/rotation.rs @@ -0,0 +1,105 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use litellm_secrets_types::{ + BaseSecretManager, Error, SecretValue, async_rotate_secret, validate_secret_name, +}; + +struct Manager { + step: AtomicUsize, + absent_at: Option, +} + +impl BaseSecretManager for Manager { + type Error = Error; + type WriteResponse = &'static str; + type DeleteResponse = (); + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + let step = self.step.fetch_add(1, Ordering::SeqCst); + assert_eq!(name, if step == 0 { "old" } else { "new" }); + Ok((self.absent_at != Some(step)).then(|| SecretValue::new("value"))) + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 1); + assert_eq!(name, "new"); + assert_eq!(value.expose(), "replacement"); + assert_eq!(description, Some("Rotated from old")); + Ok("provider-response") + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result<(), Error> { + assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 3); + assert_eq!(name, "old"); + assert_eq!(recovery_window_in_days, 7); + Ok(()) + } +} + +#[tokio::test] +async fn rotation_verifies_before_deleting_and_returns_provider_response() { + let manager = Manager { + step: AtomicUsize::new(0), + absent_at: None, + }; + assert_eq!( + async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement")) + .await + .unwrap(), + "provider-response" + ); + assert_eq!(manager.step.load(Ordering::SeqCst), 4); +} + +#[rstest::rstest] +#[case::current_secret_missing(0, Error::CurrentSecretMissing, 1)] +#[case::new_secret_missing(2, Error::NewSecretMissing, 3)] +#[tokio::test] +async fn missing_old_or_new_value_stops_rotation_before_deletion( + #[case] absent_at: usize, + #[case] expected: Error, + #[case] calls: usize, +) { + let manager = Manager { + step: AtomicUsize::new(0), + absent_at: Some(absent_at), + }; + assert_eq!( + async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement")) + .await + .unwrap_err(), + expected + ); + assert_eq!(manager.step.load(Ordering::SeqCst), calls); +} + +#[rstest::rstest] +#[case::parent("..")] +#[case::parent_prefix("../x")] +#[case::parent_segment("x/../y")] +#[case::parent_suffix("x/..")] +#[case::line_feed("line\n")] +#[case::next_line("\u{85}")] +#[case::line_separator("\u{2028}")] +#[case::paragraph_separator("\u{2029}")] +fn names_reject_path_traversal_and_control_characters(#[case] name: &str) { + assert_eq!(validate_secret_name(name), Err(Error::UnsafeSecretName)); +} + +#[rstest::rstest] +#[case::embedded_double_dot("release-1.0..2")] +#[case::path_separator("folder/key")] +#[case::empty("")] +#[case::three_dots("...")] +fn names_allow_safe_values(#[case] name: &str) { + assert_eq!(validate_secret_name(name), Ok(())); +} diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml new file mode 100644 index 00000000000..3414470d234 --- /dev/null +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "litellm-secrets" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[features] +default = [] +aws = ["dep:litellm-secrets-aws"] +google = ["dep:litellm-secrets-google"] + +[dependencies] +litellm-secrets-types.workspace = true +litellm-secrets-aws = { workspace = true, optional = true } +litellm-secrets-google = { workspace = true, optional = true } +litellm-core-utils.workspace = true +base64.workspace = true +serde.workspace = true +strum.workspace = true +jsonwebtoken.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tracing = "0.1" +reqwest.workspace = true +moka.workspace = true +tokio = { workspace = true, features = ["fs"] } + +rustpython-parser = { version = "0.4.0", default-features = false, features = ["num-bigint"] } + +[dev-dependencies] +rstest.workspace = true +wiremock = "0.6.5" +tempfile = "3" +aws-sdk-kms = "1.120.0" +google-cloud-kms-v1 = "1.14.0" +google-cloud-auth.workspace = true diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs new file mode 100644 index 00000000000..d240adb6a67 --- /dev/null +++ b/litellm-rust/crates/secrets/src/error.rs @@ -0,0 +1,39 @@ +use crate::KeyManagementSystem; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("encrypted environment value is missing")] + MissingCiphertext, + #[error("ciphertext is not valid base64 for the configured manager")] + InvalidCiphertext, + #[error("decrypted value is not UTF-8")] + Utf8, + #[error("secret manager backend is not compiled: {0:?}")] + UnsupportedBackend(KeyManagementSystem), + #[error("configured secret manager does not match its backend")] + BackendMismatch, + #[error("unsupported OIDC provider or missing build feature")] + UnsupportedOidc, + #[error("OIDC reference requires a provider and audience")] + InvalidOidc, + #[error("OIDC environment variable is missing")] + MissingEnvironment, + #[error("OIDC request failed")] + OidcHttp, + #[error("OIDC provider returned HTTP {0}")] + OidcStatus(u16), + #[error("OIDC response is invalid")] + OidcResponse, + #[error("OIDC file path must be absolute and within the credential allowlist")] + UnsafeOidcPath, + #[error("OIDC file could not be read")] + OidcFile, + #[error("secret manager returned no secret")] + MissingSecret, + #[cfg(feature = "aws")] + #[error(transparent)] + Aws(#[from] litellm_secrets_aws::Error), + #[cfg(feature = "google")] + #[error(transparent)] + Google(#[from] litellm_secrets_google::Error), +} diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs new file mode 100644 index 00000000000..6d19965605d --- /dev/null +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -0,0 +1,118 @@ +use litellm_core_utils::settings::Lookup; + +use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue}; + +#[derive(Clone)] +pub enum SecretManager { + Local, + #[cfg(feature = "aws")] + AwsKms(crate::aws::AwsKms), + #[cfg(feature = "aws")] + AwsSecretsManagerV2(crate::aws::AwsSecretsManagerV2), + #[cfg(feature = "google")] + GoogleKms(crate::google::GoogleKms), + #[cfg(feature = "google")] + GoogleSecretManager(crate::google::GoogleSecretManager), +} + +impl SecretManager { + pub fn system(&self) -> KeyManagementSystem { + match self { + Self::Local => KeyManagementSystem::Local, + #[cfg(feature = "aws")] + Self::AwsKms(_) => KeyManagementSystem::AwsKms, + #[cfg(feature = "aws")] + Self::AwsSecretsManagerV2(_) => KeyManagementSystem::AwsSecretManager, + #[cfg(feature = "google")] + Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, + #[cfg(feature = "google")] + Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + } + } +} + +pub async fn get_secret_from_manager( + client: &SecretManager, + secret_name: &str, + _settings: &KeyManagementSettings, + environment: &(dyn Lookup + Send + Sync), +) -> Result, Error> { + match client { + SecretManager::Local => Ok(environment + .get(secret_name) + .map(SecretValue::new) + .map(Secret::String)), + #[cfg(feature = "aws")] + SecretManager::AwsKms(client) => { + let ciphertext = environment + .get(secret_name) + .ok_or(Error::MissingCiphertext)?; + let plaintext = client + .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Permissive)?) + .await?; + let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + Ok(Some(Secret::String(SecretValue::new(value.trim())))) + } + #[cfg(feature = "google")] + SecretManager::GoogleKms(client) => { + let ciphertext = environment + .get(secret_name) + .ok_or(Error::MissingCiphertext)?; + let plaintext = client + .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Canonical)?) + .await?; + let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + Ok(Some(Secret::String(SecretValue::new(value)))) + } + #[cfg(feature = "aws")] + SecretManager::AwsSecretsManagerV2(client) => client + .read_secret_for_resolver( + secret_name, + _settings.primary_secret_name.as_deref(), + environment, + ) + .await + .map_err(Error::from), + #[cfg(feature = "google")] + SecretManager::GoogleSecretManager(client) => client + .get_secret_from_google_secret_manager(secret_name) + .await? + .map(Some) + .ok_or(Error::MissingSecret), + } +} + +#[cfg(any(feature = "aws", feature = "google"))] +#[derive(Clone, Copy)] +enum Base64Mode { + #[cfg(feature = "google")] + Canonical, + #[cfg(feature = "aws")] + Permissive, +} + +#[cfg(any(feature = "aws", feature = "google"))] +fn decode_ciphertext(value: &str, mode: Base64Mode) -> Result, Error> { + use base64::{Engine, engine::general_purpose::STANDARD}; + let canonical = match mode { + #[cfg(feature = "google")] + Base64Mode::Canonical => true, + #[cfg(feature = "aws")] + Base64Mode::Permissive => false, + }; + let encoded = if canonical { + value.to_owned() + } else { + value + .chars() + .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) + .collect() + }; + let ciphertext = STANDARD + .decode(&encoded) + .map_err(|_| Error::InvalidCiphertext)?; + if canonical && STANDARD.encode(&ciphertext) != encoded { + return Err(Error::InvalidCiphertext); + } + Ok(ciphertext) +} diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs new file mode 100644 index 00000000000..c434cfaaf03 --- /dev/null +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -0,0 +1,21 @@ +#![forbid(unsafe_code)] + +mod error; +mod handler; +mod oidc; +mod resolver; +mod state; + +pub use error::Error; +pub use handler::{SecretManager, get_secret_from_manager}; +pub use litellm_secrets_types::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +pub use oidc::{OidcProvider, OidcReference, OidcResolver}; +pub use resolver::SecretResolver; +pub use state::{SecretManagerState, secret_manager_would_be_consulted}; + +#[cfg(feature = "aws")] +pub use litellm_secrets_aws as aws; +#[cfg(feature = "google")] +pub use litellm_secrets_google as google; diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs new file mode 100644 index 00000000000..48e6a9bc0b0 --- /dev/null +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -0,0 +1,264 @@ +use std::{ + path::Path, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use jsonwebtoken::dangerous::insecure_decode_claims; +use litellm_core_utils::settings::Lookup; +use moka::future::Cache; +use serde::Deserialize; + +use crate::{Error, SecretValue}; + +const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540); +const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295); +const TOKEN_EXPIRY_MARGIN_SECONDS: f64 = 60.0; +const CIRCLE_OIDC_TOKEN: &str = "CIRCLE_OIDC_TOKEN"; +const CIRCLE_OIDC_TOKEN_V2: &str = "CIRCLE_OIDC_TOKEN_V2"; +const AZURE_FEDERATED_TOKEN_FILE: &str = "AZURE_FEDERATED_TOKEN_FILE"; +const ACTIONS_ID_TOKEN_REQUEST_URL: &str = "ACTIONS_ID_TOKEN_REQUEST_URL"; +const ACTIONS_ID_TOKEN_REQUEST_TOKEN: &str = "ACTIONS_ID_TOKEN_REQUEST_TOKEN"; +const OIDC_ALLOWED_CREDENTIAL_DIRS: &str = "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS"; +const DEFAULT_CREDENTIAL_DIRS: &str = "/var/run/secrets,/run/secrets"; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::EnumString, strum::AsRefStr)] +#[strum(serialize_all = "snake_case")] +pub enum OidcProvider { + Google, + #[strum(serialize = "circleci")] + CircleCi, + #[strum(serialize = "circleci_v2")] + CircleCiV2, + Github, + Azure, + File, + Env, + EnvPath, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct OidcReference<'a> { + pub provider: OidcProvider, + pub audience: &'a str, +} + +impl<'a> TryFrom<&'a str> for OidcReference<'a> { + type Error = Error; + + fn try_from(reference: &'a str) -> Result { + let (provider, audience) = reference + .strip_prefix("oidc/") + .and_then(|body| body.split_once('/')) + .ok_or(Error::InvalidOidc)?; + Ok(Self { + provider: provider.parse().map_err(|_| Error::UnsupportedOidc)?, + audience, + }) + } +} + +#[derive(Deserialize)] +struct OidcTokenClaims { + exp: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum NumericDate { + Number(f64), + String(String), +} + +impl NumericDate { + fn seconds(self) -> Option { + match self { + Self::Number(value) => Some(value), + Self::String(value) => value.parse().ok(), + } + .filter(|value| value.is_finite()) + } +} + +pub struct OidcResolver { + client: reqwest::Client, + google_identity_endpoint: reqwest::Url, + cache: Cache, + clock: fn() -> SystemTime, +} + +impl Default for OidcResolver { + fn default() -> Self { + Self::new( + reqwest::Client::builder().timeout(Duration::from_secs(600)).connect_timeout(Duration::from_secs(5)).build().expect("HTTP client configuration"), + reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"), + ) + } +} + +impl OidcResolver { + pub fn new(client: reqwest::Client, google_identity_endpoint: reqwest::Url) -> Self { + Self { + client, + google_identity_endpoint, + cache: Cache::builder() + .max_capacity(200) + .time_to_live(GOOGLE_TOKEN_MAX_TTL) + .build(), + clock: SystemTime::now, + } + } + + pub fn with_clock(self, clock: fn() -> SystemTime) -> Self { + Self { clock, ..self } + } + + pub async fn resolve( + &self, + reference: &str, + environment: &(dyn Lookup + Send + Sync), + ) -> Result, Error> { + let OidcReference { provider, audience } = reference.try_into()?; + match provider { + OidcProvider::CircleCi => required_env(environment, CIRCLE_OIDC_TOKEN) + .map(SecretValue::new) + .map(Some), + OidcProvider::CircleCiV2 => required_env(environment, CIRCLE_OIDC_TOKEN_V2) + .map(SecretValue::new) + .map(Some), + OidcProvider::Env => required_env(environment, audience) + .map(SecretValue::new) + .map(Some), + OidcProvider::EnvPath => read_file(&required_env(environment, audience)?) + .await + .map(Some), + OidcProvider::File => read_allowed_file(audience, environment).await.map(Some), + OidcProvider::Azure => { + if let Some(path) = environment.get(AZURE_FEDERATED_TOKEN_FILE) { + return read_file(&path).await.map(Some); + } + Err(Error::UnsupportedOidc) + } + OidcProvider::Github => { + let url = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_URL)?; + let authorization = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_TOKEN)?; + if let Some(value) = self.cached(reference).await { + return Ok(Some(value)); + } + let response = self + .client + .get(url) + .query(&[("audience", audience)]) + .bearer_auth(authorization) + .header("Accept", "application/json; api-version=2.0") + .send() + .await + .map_err(|_| Error::OidcHttp)?; + if response.status() != reqwest::StatusCode::OK { + return Err(Error::OidcStatus(response.status().as_u16())); + } + #[derive(Deserialize)] + struct Token { + value: Option, + } + let token: Token = response.json().await.map_err(|_| Error::OidcResponse)?; + if let Some(value) = &token.value { + self.cache + .insert( + reference.to_owned(), + (value.clone(), (self.clock)() + GITHUB_TOKEN_TTL), + ) + .await; + } + Ok(token.value) + } + OidcProvider::Google => { + if !cfg!(feature = "google") { + return Err(Error::UnsupportedOidc); + } + if let Some(value) = self.cached(reference).await { + return Ok(Some(value)); + } + let response = self + .client + .get(self.google_identity_endpoint.clone()) + .query(&[("audience", audience)]) + .header("Metadata-Flavor", "Google") + .send() + .await + .map_err(|_| Error::OidcHttp)?; + if response.status() != reqwest::StatusCode::OK { + return Err(Error::OidcStatus(response.status().as_u16())); + } + let token = response.text().await.map_err(|_| Error::OidcResponse)?; + let now = (self.clock)(); + let ttl = oidc_token_cache_ttl(&token, now, GOOGLE_TOKEN_MAX_TTL); + let value = SecretValue::new(token); + if let Some(ttl) = ttl.filter(|ttl| !ttl.is_zero()) { + self.cache + .insert(reference.to_owned(), (value.clone(), now + ttl)) + .await; + } + Ok(Some(value)) + } + } + } + + async fn cached(&self, reference: &str) -> Option { + self.cache + .get(reference) + .await + .and_then(|(value, expires)| ((self.clock)() < expires).then_some(value)) + } +} + +fn required_env(environment: &dyn Lookup, name: &str) -> Result { + environment.get(name).ok_or(Error::MissingEnvironment) +} + +async fn read_file(path: &str) -> Result { + tokio::fs::read_to_string(path) + .await + .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) + .map_err(|_| Error::OidcFile) +} + +async fn read_allowed_file( + path: &str, + environment: &(dyn Lookup + Sync), +) -> Result { + if !Path::new(path).is_absolute() { + return Err(Error::UnsafeOidcPath); + } + let resolved = tokio::fs::canonicalize(path) + .await + .map_err(|_| Error::OidcFile)?; + let allowed = environment + .get(OIDC_ALLOWED_CREDENTIAL_DIRS) + .filter(|v| !v.is_empty()) + .unwrap_or_else(|| DEFAULT_CREDENTIAL_DIRS.into()); + for directory in allowed.split(',').map(str::trim).filter(|d| !d.is_empty()) { + if let Ok(directory) = tokio::fs::canonicalize(directory).await + && resolved.starts_with(directory) + { + return tokio::fs::read_to_string(&resolved) + .await + .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) + .map_err(|_| Error::OidcFile); + } + } + Err(Error::UnsafeOidcPath) +} + +fn oidc_token_cache_ttl(token: &str, now: SystemTime, max_ttl: Duration) -> Option { + let fallback = Some(max_ttl); + let Ok(claims) = insecure_decode_claims::(token) else { + return fallback; + }; + let Some(exp) = claims.exp.and_then(NumericDate::seconds) else { + return fallback; + }; + let seconds = exp.trunc() + - now.duration_since(UNIX_EPOCH).ok()?.as_secs() as f64 + - TOKEN_EXPIRY_MARGIN_SECONDS; + (seconds > 0.0).then(|| Duration::from_secs_f64(seconds.min(max_ttl.as_secs_f64()))) +} diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs new file mode 100644 index 00000000000..93daddae0c4 --- /dev/null +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -0,0 +1,140 @@ +use std::sync::Arc; + +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; + +use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; + +use crate::state::{LookupTarget, normalize_secret_name}; + +pub struct SecretResolver { + state: Arc, + environment: Arc, + oidc: OidcResolver, +} + +impl Default for SecretResolver { + fn default() -> Self { + Self::new( + Arc::new(SecretManagerState::default()), + Arc::new(ProcessEnvironment), + OidcResolver::default(), + ) + } +} + +impl SecretResolver { + pub fn new( + state: Arc, + environment: Arc, + oidc: OidcResolver, + ) -> Self { + Self { + state, + environment, + oidc, + } + } + + pub async fn get_secret( + &self, + name: &str, + _default_value: Option, + ) -> Result, Error> { + let name = normalize_secret_name(name); + if name.starts_with("oidc/") { + return self + .oidc + .resolve(name, self.environment.as_ref()) + .await + .map(|value| value.map(Secret::String)); + } + if !self.state.readable() { + return Ok(self + .environment + .get(name) + .map(|value| match str_to_bool(&value) { + Some(value) => Secret::Bool(value), + None => Secret::String(SecretValue::new(value)), + })); + } + let result = match self.state.lookup_target(name) { + LookupTarget::Environment => Ok(self.environment_secret(name)), + LookupTarget::Manager { backend, settings } => { + crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref()) + .await + } + }; + let value = match result { + Ok(value) => value, + Err(_) => { + tracing::error!("secret manager lookup failed; falling back to environment"); + self.environment_secret(name) + } + }; + Ok(value.and_then(managed_secret)) + } + + fn environment_secret(&self, name: &str) -> Option { + self.environment + .get(name) + .map(SecretValue::new) + .map(Secret::String) + } + + pub async fn get_secret_str( + &self, + name: &str, + default_value: Option, + ) -> Result, Error> { + Ok(match self.get_secret(name, default_value).await? { + Some(Secret::String(value)) => Some(value), + Some(Secret::Bool(_) | Secret::Json(_)) | None => None, + }) + } + + pub async fn get_secret_bool( + &self, + name: &str, + default_value: Option, + ) -> Result, Error> { + Ok( + match self + .get_secret(name, default_value.map(Secret::Bool)) + .await? + { + Some(Secret::Bool(value)) => Some(value), + Some(Secret::String(value)) => str_to_bool(value.expose()), + Some(Secret::Json(_)) | None => None, + }, + ) + } +} + +fn str_to_bool(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "true" => Some(true), + "false" => Some(false), + _ => None, + } +} + +fn literal_bool(value: &str) -> Option { + use rustpython_parser::{Parse, ast}; + match ast::Expr::parse(value.trim_start_matches([' ', '\t']), "").ok()? { + ast::Expr::Constant(node) => match node.value { + ast::Constant::Bool(value) => Some(value), + _ => None, + }, + _ => None, + } +} + +fn managed_secret(value: Secret) -> Option { + match value { + Secret::String(value) => Some(match literal_bool(value.expose()) { + Some(boolean) => Secret::Bool(boolean), + None => Secret::String(value), + }), + Secret::Bool(_) | Secret::Json(_) => None, + } +} diff --git a/litellm-rust/crates/secrets/src/state.rs b/litellm-rust/crates/secrets/src/state.rs new file mode 100644 index 00000000000..ca942aa65db --- /dev/null +++ b/litellm-rust/crates/secrets/src/state.rs @@ -0,0 +1,105 @@ +use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; + +pub(crate) enum LookupTarget<'a> { + Environment, + Manager { + backend: &'a SecretManager, + settings: &'a KeyManagementSettings, + }, +} + +pub(crate) fn normalize_secret_name(name: &str) -> &str { + name.strip_prefix("os.environ/").unwrap_or(name) +} + +#[derive(Clone, Default)] +pub struct SecretManagerState { + system: Option, + settings: Option, + backend: Option, +} + +impl SecretManagerState { + pub fn new( + system: Option, + settings: Option, + backend: Option, + ) -> Result { + if let Some(system) = system { + let available = match system { + KeyManagementSystem::Local => true, + KeyManagementSystem::AwsKms | KeyManagementSystem::AwsSecretManager => { + cfg!(feature = "aws") + } + KeyManagementSystem::GoogleKms | KeyManagementSystem::GoogleSecretManager => { + cfg!(feature = "google") + } + KeyManagementSystem::AzureKeyVault + | KeyManagementSystem::HashicorpVault + | KeyManagementSystem::Cyberark + | KeyManagementSystem::Custom => false, + }; + if !available { + return Err(Error::UnsupportedBackend(system)); + } + if let Some(backend) = &backend + && system != backend.system() + { + return Err(Error::BackendMismatch); + } + } + Ok(Self { + system, + settings, + backend, + }) + } + + pub fn system(&self) -> Option { + self.system + } + pub fn settings(&self) -> Option<&KeyManagementSettings> { + self.settings.as_ref() + } + pub fn backend(&self) -> Option<&SecretManager> { + self.backend.as_ref() + } + + pub(crate) fn readable(&self) -> bool { + self.backend.is_some() + && self + .settings + .as_ref() + .is_some_and(|settings| settings.access_mode.readable()) + } + + pub(crate) fn lookup_target(&self, name: &str) -> LookupTarget<'_> { + match (&self.backend, &self.settings) { + (Some(backend), Some(settings)) + if settings.access_mode.readable() + && hosts_secret(settings, name) + && self + .system + .is_some_and(|system| system != KeyManagementSystem::Local) => + { + LookupTarget::Manager { backend, settings } + } + _ => LookupTarget::Environment, + } + } +} + +pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool { + state.readable() + && state + .settings + .as_ref() + .is_some_and(|settings| hosts_secret(settings, normalize_secret_name(name))) +} + +fn hosts_secret(settings: &KeyManagementSettings, name: &str) -> bool { + settings + .hosted_keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == name)) +} diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs new file mode 100644 index 00000000000..a2cbbd843e1 --- /dev/null +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -0,0 +1,107 @@ +#[cfg(feature = "aws")] +#[tokio::test] +async fn aws_handler_reads_ciphertext_decodes_trims_and_redacts() { + use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Credentials, Region}, + }; + use base64::{Engine, engine::general_purpose::STANDARD}; + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, aws::AwsKms, get_secret_from_manager, + }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_json}; + + let server = MockServer::start().await; + Mock::given(body_json( + serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"Plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = Client::from_conf( + aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .build(), + ); + let manager = SecretManager::AwsKms(AwsKms::new(client)); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|name: &str| { + assert_eq!(name, "KEY"); + Some(format!(" {}\n", STANDARD.encode("encrypted"))) + }) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some("value")); + assert!(!format!("{value:?}").contains("value")); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, + Err(Error::MissingCiphertext) + )); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some("abc".into())).await, + Err(Error::InvalidCiphertext) + )); +} + +#[cfg(feature = "google")] +#[tokio::test] +async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whitespace() { + use base64::{Engine, engine::general_purpose::STANDARD}; + use google_cloud_kms_v1::client::KeyManagementService; + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, get_secret_from_manager, google::GoogleKms, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, path}, + }; + + let server = MockServer::start().await; + let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key"; + Mock::given(path(format!("/v1/{resource}:decrypt"))) + .and(body_json( + serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = KeyManagementService::builder() + .with_endpoint(server.uri()) + .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build()) + .build() + .await + .unwrap(); + let manager = SecretManager::GoogleKms(GoogleKms::new(client, resource.into())); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| { + Some(STANDARD.encode("encrypted")) + }) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some(" value\n")); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some(format!( + " {}", + STANDARD.encode("encrypted") + ))) + .await, + Err(Error::InvalidCiphertext) + )); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, + Err(Error::MissingCiphertext) + )); +} diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs new file mode 100644 index 00000000000..b17e7de7f9d --- /dev/null +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -0,0 +1,295 @@ +use std::{collections::BTreeMap, sync::Arc}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets::{Error, OidcResolver, Secret, SecretManagerState, SecretResolver}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, method, path, query_param}, +}; + +fn environment(pairs: &[(&str, &str)]) -> Arc { + let values: BTreeMap = pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + Arc::new(move |name: &str| values.get(name).cloned()) +} + +#[rstest::rstest] +#[case::environment("oidc/env/TOKEN", "true")] +#[case::circleci("oidc/circleci/audience", "circle")] +#[case::circleci_v2("oidc/circleci_v2/audience", "circle-v2")] +#[tokio::test] +async fn environment_sources_resolve_expected_value( + #[case] reference: &str, + #[case] expected: &str, +) { + let env = environment(&[ + ("TOKEN", "true"), + ("CIRCLE_OIDC_TOKEN", "circle"), + ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"), + ]); + assert_eq!( + OidcResolver::default() + .resolve(reference, env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + expected + ); +} + +#[tokio::test] +async fn environment_sources_bypass_boolean_conversion_and_defaults() { + let env = environment(&[("TOKEN", "true")]); + let oidc = OidcResolver::default(); + let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc); + assert_eq!( + resolver + .get_secret_str("os.environ/oidc/env/TOKEN", None) + .await + .unwrap() + .unwrap() + .expose(), + "true" + ); + assert_eq!( + resolver + .get_secret_bool("oidc/env/TOKEN", None) + .await + .unwrap(), + Some(true) + ); + assert!(matches!( + resolver + .get_secret("oidc/env/MISSING", Some(Secret::Bool(true))) + .await, + Err(Error::MissingEnvironment) + )); + assert!(matches!( + resolver.get_secret("oidc/invalid", None).await, + Err(Error::InvalidOidc) + )); +} + +#[tokio::test] +async fn github_requests_are_authenticated_cached_and_revalidate_environment() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/token")) + .and(query_param("audience", "https://service/oidc/path")) + .and(header("authorization", "Bearer request-token")) + .and(header("accept", "application/json; api-version=2.0")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value":"identity-token"})), + ) + .expect(1) + .mount(&server) + .await; + let env = environment(&[ + ( + "ACTIONS_ID_TOKEN_REQUEST_URL", + &format!("{}/token", server.uri()), + ), + ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"), + ]); + let oidc = OidcResolver::default(); + for _ in 0..2 { + assert_eq!( + oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "identity-token" + ); + } + assert!(matches!( + oidc.resolve( + "oidc/github/https://service/oidc/path", + environment(&[]).as_ref() + ) + .await, + Err(Error::MissingEnvironment) + )); +} + +#[tokio::test] +async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explicit() { + let allowed = tempfile::tempdir().unwrap(); + let outside = tempfile::tempdir().unwrap(); + let token = allowed.path().join("token"); + let private = outside.path().join("private"); + std::fs::write(&token, "token\r\n").unwrap(); + std::fs::write(&private, "outside").unwrap(); + let env = environment(&[ + ( + "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", + allowed.path().to_str().unwrap(), + ), + ("PATH_TOKEN", private.to_str().unwrap()), + ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()), + ]); + let oidc = OidcResolver::default(); + assert_eq!( + oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "token\n" + ); + assert!(matches!( + oidc.resolve("oidc/file/relative", env.as_ref()).await, + Err(Error::UnsafeOidcPath) + )); + assert!(matches!( + oidc.resolve(&format!("oidc/file/{}", private.display()), env.as_ref()) + .await, + Err(Error::UnsafeOidcPath) + )); + assert_eq!( + oidc.resolve("oidc/env_path/PATH_TOKEN", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "outside" + ); + assert_eq!( + oidc.resolve("oidc/azure/scope", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "token\n" + ); + #[cfg(unix)] + { + let link = allowed.path().join("link"); + std::os::unix::fs::symlink(&private, &link).unwrap(); + assert!(matches!( + oidc.resolve(&format!("oidc/file/{}", link.display()), env.as_ref()) + .await, + Err(Error::UnsafeOidcPath) + )); + } +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::at_refresh_boundary(serde_json::json!(1060), 2)] +#[case::beyond_refresh_boundary(serde_json::json!(1061), 1)] +#[case::already_expired(serde_json::json!(999), 2)] +#[case::string_expiry(serde_json::json!("999"), 2)] +#[case::fractional_expiry(serde_json::json!(1060.9), 2)] +#[case::negative_expiry(serde_json::json!(-1), 2)] +#[case::null_expiry(serde_json::Value::Null, 1)] +#[case::unreadable_expiry(serde_json::json!("invalid"), 1)] +#[case::nonfinite_expiry(serde_json::json!("NaN"), 1)] +#[tokio::test] +async fn google_expiry_caps_cache_and_preserves_audience( + #[case] expiry: serde_json::Value, + #[case] calls: u64, +) { + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + fn now() -> SystemTime { + UNIX_EPOCH + Duration::from_secs(1000) + } + let server = MockServer::start().await; + let token = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(serde_json::json!({"alg":"RS256","typ":"JWT"}).to_string()), + URL_SAFE_NO_PAD.encode(serde_json::json!({"exp":expiry}).to_string()) + ); + Mock::given(method("GET")) + .and(header("metadata-flavor", "Google")) + .and(query_param("audience", "https://service/oidc/path")) + .respond_with(ResponseTemplate::new(200).set_body_string(&token)) + .expect(calls) + .mount(&server) + .await; + let oidc = + OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()).with_clock(now); + for _ in 0..2 { + assert_eq!( + oidc.resolve( + "oidc/google/https://service/oidc/path", + environment(&[]).as_ref() + ) + .await + .unwrap() + .unwrap() + .expose(), + token + ); + } +} + +#[cfg(not(feature = "google"))] +#[tokio::test] +async fn google_oidc_requires_its_build_feature() { + assert!(matches!( + OidcResolver::default() + .resolve("oidc/google/audience", environment(&[]).as_ref()) + .await, + Err(Error::UnsupportedOidc) + )); +} + +#[tokio::test] +async fn azure_oidc_without_a_token_file_requires_an_unimplemented_backend() { + assert!(matches!( + OidcResolver::default() + .resolve("oidc/azure/scope", environment(&[]).as_ref()) + .await, + Err(Error::UnsupportedOidc) + )); +} + +#[rstest::rstest] +#[case::missing_prefix("env/TOKEN", false)] +#[case::missing_audience_separator("oidc/env", false)] +#[case::unknown_provider("oidc/unknown/TOKEN", true)] +#[tokio::test] +async fn invalid_references_fail_before_environment_lookup( + #[case] reference: &str, + #[case] unsupported: bool, +) { + let error = OidcResolver::default() + .resolve(reference, &|_: &str| { + panic!("invalid reference reached environment lookup") + }) + .await + .unwrap_err(); + assert!(matches!(error, Error::UnsupportedOidc) == unsupported); + assert!(matches!(error, Error::InvalidOidc) != unsupported); +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::opaque("opaque-token")] +#[case::missing_expiry("header.e30.signature")] +#[tokio::test] +async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) { + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_string(token)) + .expect(1) + .mount(&server) + .await; + let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()); + for _ in 0..2 { + assert_eq!( + resolver + .resolve("oidc/google/audience", environment(&[]).as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + token, + ); + } +} diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs new file mode 100644 index 00000000000..e2b056f01cc --- /dev/null +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -0,0 +1,343 @@ +use std::sync::Arc; + +use litellm_secrets::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, OidcResolver, Secret, SecretManager, + SecretManagerState, SecretResolver, SecretValue, secret_manager_would_be_consulted, +}; + +fn resolver(value: Option<&str>, readable: bool) -> SecretResolver { + let state = if readable { + SecretManagerState::new( + Some(KeyManagementSystem::Local), + Some(KeyManagementSettings::default()), + Some(SecretManager::Local), + ) + .unwrap() + } else { + SecretManagerState::default() + }; + let value = value.map(str::to_owned); + SecretResolver::new( + Arc::new(state), + Arc::new(move |_: &str| value.clone()), + OidcResolver::default(), + ) +} + +#[rstest::rstest] +#[case::lowercase_true("true", Some(true), None)] +#[case::whitespace_lowercase_false(" FALSE ", Some(false), None)] +#[case::python_true("True", Some(true), Some(true))] +#[case::python_false("False", Some(false), Some(false))] +#[case::parenthesized_python_true("(True)", None, Some(true))] +#[case::commented_python_false("False # comment", None, Some(false))] +#[case::integer("1", None, None)] +#[case::yes("yes", None, None)] +#[case::plain_string("secret", None, None)] +#[tokio::test] +async fn boolean_conversion_preserves_local_and_manager_differences( + #[case] input: &str, + #[case] local: Option, + #[case] manager: Option, + #[values(false, true)] readable: bool, +) { + let boolean = if readable { manager } else { local }; + let resolver = resolver(Some(input), readable); + let expected = boolean + .map(Secret::Bool) + .unwrap_or_else(|| Secret::String(SecretValue::new(input))); + assert_eq!( + resolver.get_secret("key", None).await.unwrap(), + Some(expected) + ); + assert_eq!( + resolver + .get_secret_str("key", None) + .await + .unwrap() + .map(|v| v.expose().to_owned()), + boolean.is_none().then(|| input.to_owned()) + ); +} + +#[tokio::test] +async fn manager_boolean_conversion_trims_whitespace() { + assert_eq!( + resolver(Some(" true "), true) + .get_secret_bool("key", None) + .await + .unwrap(), + Some(true) + ); +} + +#[tokio::test] +async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() { + let missing = resolver(None, false); + assert_eq!( + missing + .get_secret("missing", Some(Secret::Bool(true))) + .await + .unwrap(), + None + ); + assert_eq!( + missing + .get_secret_bool("missing", Some(true)) + .await + .unwrap(), + None + ); + let resolver = SecretResolver::new( + Arc::new(SecretManagerState::default()), + Arc::new(|name: &str| (name == "KEY").then(|| "value".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str("os.environ/KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[rstest::rstest] +#[case::all_keys(None)] +#[case::no_keys(Some(Vec::new()))] +#[case::allowlisted_key(Some(vec!["KEY".into()]))] +fn manager_gating_requires_client_readable_settings_and_allowlisted_name( + #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)] + access_mode: AccessMode, + #[values(false, true)] client: bool, + #[case] keys: Option>, +) { + let expected = client + && access_mode.readable() + && keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == "KEY")); + let state = SecretManagerState::new( + Some(KeyManagementSystem::Local), + Some(KeyManagementSettings { + access_mode, + hosted_keys: keys, + ..Default::default() + }), + client.then_some(SecretManager::Local), + ) + .unwrap(); + assert_eq!( + secret_manager_would_be_consulted(&state, "os.environ/KEY"), + expected + ); +} + +#[test] +fn manager_gating_requires_settings() { + let no_settings = SecretManagerState::new(None, None, Some(SecretManager::Local)).unwrap(); + assert!(!secret_manager_would_be_consulted(&no_settings, "KEY")); +} + +#[cfg(feature = "aws")] +#[rstest::rstest] +#[case::missing_value(None, None)] +#[case::lookup_error(Some("primary".to_owned()), Some("environment-value"))] +#[tokio::test] +async fn aws_missing_values_do_not_fallback_but_lookup_errors_do( + #[case] primary: Option, + #[case] expected: Option<&str>, +) { + use litellm_secrets::aws::AwsSecretsManagerV2; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_partial_json}; + let server = MockServer::start().await; + Mock::given(body_partial_json(serde_json::json!({"SecretId":"KEY"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(u64::from(primary.is_none())) + .mount(&server) + .await; + Mock::given(body_partial_json(serde_json::json!({"SecretId":"primary"}))) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"SecretString":"invalid-json"})), + ) + .expect(u64::from(primary.is_some())) + .mount(&server) + .await; + let endpoint = server.uri(); + let environment: Arc = + Arc::new(move |name: &str| match name { + "AWS_REGION_NAME" => Some("us-east-1".into()), + "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()), + "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()), + "KEY" => Some("environment-value".into()), + _ => None, + }); + let settings = KeyManagementSettings { + primary_secret_name: primary, + ..Default::default() + }; + let manager = AwsSecretsManagerV2::load_aws_secret_manager( + Some(true), + settings.clone(), + environment.clone(), + ) + .unwrap() + .unwrap(); + let state = SecretManagerState::new( + Some(KeyManagementSystem::AwsSecretManager), + Some(settings), + Some(SecretManager::AwsSecretsManagerV2(manager)), + ) + .unwrap(); + let resolver = SecretResolver::new( + Arc::new(state), + environment.clone(), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str("os.environ/KEY", None) + .await + .unwrap() + .map(|v| v.expose().to_owned()) + .as_deref(), + expected + ); +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::hosted_filter(Some(Vec::new()), Some(KeyManagementSystem::GoogleSecretManager))] +#[case::negative_cache(None, Some(KeyManagementSystem::GoogleSecretManager))] +#[case::missing_system(None, None)] +#[case::hosted_nested_prefix(Some(vec!["os.environ/KEY".into()]), Some(KeyManagementSystem::GoogleSecretManager))] +#[tokio::test] +async fn google_negative_cache_still_falls_back_and_hosted_filter_avoids_io( + #[case] hosted_keys: Option>, + #[case] system: Option, +) { + use litellm_secrets::google::GoogleSecretManager; + use std::time::Duration; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::path}; + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/os%2Eenviron%2FKEY/versions/latest:access", + )) + .respond_with(ResponseTemplate::new(404)) + .expect(u64::from( + hosted_keys.as_ref().is_none_or(|keys| !keys.is_empty()) && system.is_some(), + )) + .mount(&server) + .await; + let environment: Arc = + Arc::new(|name: &str| match name { + "VERTEX_AI_API_KEY" => Some("token".into()), + "os.environ/KEY" => Some("environment-value".into()), + _ => None, + }); + let manager = GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + environment.clone(), + Some(Duration::from_secs(60)), + false, + ) + .unwrap(); + let settings = KeyManagementSettings { + hosted_keys, + ..Default::default() + }; + let state = SecretManagerState::new( + system, + Some(settings), + Some(SecretManager::GoogleSecretManager(manager)), + ) + .unwrap(); + let resolver = SecretResolver::new( + Arc::new(state), + environment.clone(), + OidcResolver::default(), + ); + for _ in 0..2 { + assert_eq!( + resolver + .get_secret_str("os.environ/os.environ/KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "environment-value" + ); + } +} + +#[tokio::test] +async fn resolver_future_can_run_on_a_tokio_worker() { + let resolver = resolver(Some("worker-value"), false); + let result = tokio::spawn(async move { resolver.get_secret_str("KEY", None).await }) + .await + .unwrap() + .unwrap(); + assert_eq!(result.unwrap().expose(), "worker-value"); +} + +#[rstest::rstest] +#[case::nested_true("((True)) # comment", Some(true))] +#[case::commented_false("(False # comment\n)", Some(false))] +#[case::boolean_expression("True and False", None)] +#[case::string_literal("'True'", None)] +#[case::tuple("(True,)", None)] +#[case::unary_expression("not False", None)] +#[case::multiple_expressions("True\nFalse", None)] +#[case::incomplete_expression("(True", None)] +#[tokio::test] +async fn manager_boolean_literals_follow_python_syntax( + #[case] input: &str, + #[case] expected: Option, +) { + let value = resolver(Some(input), true) + .get_secret("key", None) + .await + .unwrap(); + assert_eq!( + value, + Some( + expected + .map(Secret::Bool) + .unwrap_or_else(|| Secret::String(SecretValue::new(input))) + ) + ); +} + +#[tokio::test] +async fn environment_prefix_is_removed_only_once_and_gating_uses_the_same_name() { + let name = "os.environ/folder/os.environ/KEY"; + let state = SecretManagerState::new( + Some(KeyManagementSystem::Local), + Some(KeyManagementSettings { + hosted_keys: Some(vec!["folder/os.environ/KEY".into()]), + ..Default::default() + }), + Some(SecretManager::Local), + ) + .unwrap(); + assert!(secret_manager_would_be_consulted(&state, name)); + let resolver = SecretResolver::new( + Arc::new(state), + Arc::new(|name: &str| (name == "folder/os.environ/KEY").then(|| "value".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str(name, None) + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} From 15f368060ff0d420a86ec21bb818f0573c21e99b Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 23:30:54 +0000 Subject: [PATCH 110/146] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-pro-latest: off_peak_pricing, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-latest: max_tokens, max_output_tokens openrouter/deepseek/deepseek-v4-pro-0813: off_peak_pricing, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 19 ++++++++++--------- model_prices_and_context_window.json | 19 ++++++++++--------- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 39c418de524..172ec45147d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2404e-07, + "input_cost_per_token": 1.32e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7468e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 4.4e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7468e-08, - "input_cost_per_token": 5.2404e-07, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 39c418de524..172ec45147d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2404e-07, + "input_cost_per_token": 1.32e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7468e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 4.4e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7468e-08, - "input_cost_per_token": 5.2404e-07, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 5ee75af90852d76da0bf6532c05368e8b984b609 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 23:34:40 +0000 Subject: [PATCH 111/146] fix(openrouter): keep glm-latest output limit at alias target value Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 172ec45147d..7089a6ffaa9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71594,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 172ec45147d..7089a6ffaa9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71594,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 7bcb01a40c7393b373ff8d4c517edde6ad68c865 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:00:57 +0000 Subject: [PATCH 112/146] chore(prices): sync OpenRouter prices: 4 models openrouter/~deepseek/deepseek-flash-latest: max_tokens, max_output_tokens openrouter/~z-ai/glm-flash-latest: max_tokens, max_output_tokens openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 20 +++++++++---------- model_prices_and_context_window.json | 20 +++++++++---------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7089a6ffaa9..e45e2db275a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 4.22298e-07, + "input_cost_per_token": 9.483e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 8.44596e-07, + "output_cost_per_token": 1.8966e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 3.51915e-08, + "cache_read_input_token_cost": 7.9025e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.556e-08, - "output_cost_per_token": 7.112e-08, - "cache_read_input_token_cost": 7.112e-09, + "input_cost_per_token": 8.9866e-08, + "output_cost_per_token": 1.79732e-07, + "cache_read_input_token_cost": 1.79732e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71300,8 +71300,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "output_cost_per_token": 1.2e-06, @@ -71574,8 +71574,8 @@ "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7089a6ffaa9..e45e2db275a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 4.22298e-07, + "input_cost_per_token": 9.483e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 8.44596e-07, + "output_cost_per_token": 1.8966e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 3.51915e-08, + "cache_read_input_token_cost": 7.9025e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.556e-08, - "output_cost_per_token": 7.112e-08, - "cache_read_input_token_cost": 7.112e-09, + "input_cost_per_token": 8.9866e-08, + "output_cost_per_token": 1.79732e-07, + "cache_read_input_token_cost": 1.79732e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71300,8 +71300,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "output_cost_per_token": 1.2e-06, @@ -71574,8 +71574,8 @@ "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", From a3ffc395b9399cb7b8ce63e9e66887dd376b5ad6 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:31:06 +0000 Subject: [PATCH 113/146] chore(prices): sync OpenRouter prices: 3 models openrouter/ibm-granite/granite-4.2-8b: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/meta-llama/llama-3.1-70b-instruct: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token openrouter/meta-llama/llama-4-maverick: input_cost_per_token, output_cost_per_token --- ...model_prices_and_context_window_backup.json | 18 +++++++++--------- model_prices_and_context_window.json | 18 +++++++++--------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e45e2db275a..a51cbdec882 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -68617,8 +68617,8 @@ "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 1.875e-07, - "output_cost_per_token": 6.525e-07, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 16384, @@ -69010,12 +69010,12 @@ "supports_reasoning": false }, "openrouter/meta-llama/llama-3.1-70b-instruct": { - "input_cost_per_token": 4e-07, - "output_cost_per_token": 4e-07, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_output_tokens": 8192, + "max_tokens": 8192, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -72727,14 +72727,14 @@ "supports_web_search": false }, "openrouter/ibm-granite/granite-4.2-8b": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 6e-08, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 117964, "max_tokens": 117964, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e45e2db275a..a51cbdec882 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -68617,8 +68617,8 @@ "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 1.875e-07, - "output_cost_per_token": 6.525e-07, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 16384, @@ -69010,12 +69010,12 @@ "supports_reasoning": false }, "openrouter/meta-llama/llama-3.1-70b-instruct": { - "input_cost_per_token": 4e-07, - "output_cost_per_token": 4e-07, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_output_tokens": 8192, + "max_tokens": 8192, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -72727,14 +72727,14 @@ "supports_web_search": false }, "openrouter/ibm-granite/granite-4.2-8b": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 6e-08, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 117964, "max_tokens": 117964, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From acd303043b10d53cebb66002048afc06cabbd6f5 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 01:31:07 +0000 Subject: [PATCH 114/146] chore(prices): sync OpenRouter prices: 1 model openrouter/qwen/qwen3.8-27b: output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a51cbdec882..bdae7ede25c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66592,8 +66592,8 @@ }, "openrouter/qwen/qwen3.8-27b": { "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.55e-06, - "cache_read_input_token_cost": 8.5e-08, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a51cbdec882..bdae7ede25c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66592,8 +66592,8 @@ }, "openrouter/qwen/qwen3.8-27b": { "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.55e-06, - "cache_read_input_token_cost": 8.5e-08, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, From e3f69fc4d8e5348fdae1086d33c720b54df9a1be Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 19:08:25 -0700 Subject: [PATCH 115/146] refactor(rust): define consistent secret lookup contracts --- litellm-rust/Cargo.lock | 254 +------- litellm-rust/crates/secrets-aws/src/error.rs | 4 +- .../crates/secrets-aws/src/secret_manager.rs | 44 +- .../secrets-aws/tests/secret_manager.rs | 20 + .../secrets-google/src/secret_manager.rs | 26 +- .../secrets-google/tests/secret_manager.rs | 110 ++-- .../crates/secrets-types/src/value.rs | 9 +- litellm-rust/crates/secrets/Cargo.toml | 3 - litellm-rust/crates/secrets/README.md | 11 + litellm-rust/crates/secrets/src/error.rs | 10 +- litellm-rust/crates/secrets/src/handler.rs | 5 +- litellm-rust/crates/secrets/src/lib.rs | 2 +- litellm-rust/crates/secrets/src/oidc.rs | 7 +- litellm-rust/crates/secrets/src/resolver.rs | 135 ++--- litellm-rust/crates/secrets/src/state.rs | 90 +-- .../crates/secrets/tests/resolution.rs | 567 +++++++++--------- 16 files changed, 522 insertions(+), 775 deletions(-) create mode 100644 litellm-rust/crates/secrets/README.md diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 81fcabaf122..8c35a0be0b4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1515,15 +1515,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "getopts" -version = "0.2.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" -dependencies = [ - "unicode-width", -] - [[package]] name = "getrandom" version = "0.2.17" @@ -2227,27 +2218,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "is-macro" -version = "0.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8267aa6001e25494f3015f9663bbd88a18240c74483afa5f0934a1b3e4c388e9" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.119", -] - -[[package]] -name = "itertools" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.13.0" @@ -2394,12 +2364,6 @@ dependencies = [ "zeroize", ] -[[package]] -name = "lalrpop-util" -version = "0.20.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553" - [[package]] name = "lazy_static" version = "1.5.0" @@ -2716,14 +2680,12 @@ dependencies = [ "moka", "reqwest 0.12.28", "rstest", - "rustpython-parser", "serde", "serde_json", "strum", "tempfile", "thiserror 2.0.19", "tokio", - "tracing", "wiremock", ] @@ -2809,7 +2771,7 @@ dependencies = [ "base64 0.22.1", "rand 0.8.7", "rstest", - "rustc-hash 2.1.3", + "rustc-hash", "serde", "serde_json", "thiserror 2.0.19", @@ -2985,16 +2947,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "num-bigint" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" -dependencies = [ - "num-integer", - "num-traits", -] - [[package]] name = "num-bigint" version = "0.5.1" @@ -3178,44 +3130,6 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" -[[package]] -name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_codegen" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" -dependencies = [ - "phf_generator", - "phf_shared", -] - -[[package]] -name = "phf_generator" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" -dependencies = [ - "phf_shared", - "rand 0.8.7", -] - -[[package]] -name = "phf_shared" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" -dependencies = [ - "siphasher", -] - [[package]] name = "pin-project" version = "1.1.13" @@ -3488,7 +3402,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.3", + "rustc-hash", "rustls 0.23.42", "socket2 0.6.5", "thiserror 2.0.19", @@ -3510,7 +3424,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash 2.1.3", + "rustc-hash", "rustls 0.23.42", "rustls-pki-types", "slab", @@ -3689,7 +3603,7 @@ dependencies = [ "arcstr", "combine", "itoa", - "num-bigint 0.5.1", + "num-bigint", "percent-encoding", "ryu", "sha1_smol", @@ -3918,12 +3832,6 @@ dependencies = [ "syn 2.0.119", ] -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - [[package]] name = "rustc-hash" version = "2.1.3" @@ -4051,63 +3959,6 @@ dependencies = [ "untrusted", ] -[[package]] -name = "rustpython-ast" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5" -dependencies = [ - "is-macro", - "num-bigint 0.4.8", - "rustpython-parser-core", - "static_assertions", -] - -[[package]] -name = "rustpython-parser" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb" -dependencies = [ - "anyhow", - "is-macro", - "itertools 0.11.0", - "lalrpop-util", - "log", - "num-bigint 0.4.8", - "num-traits", - "phf", - "phf_codegen", - "rustc-hash 1.1.0", - "rustpython-ast", - "rustpython-parser-core", - "tiny-keccak", - "unic-emoji-char", - "unic-ucd-ident", - "unicode_names2", -] - -[[package]] -name = "rustpython-parser-core" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad" -dependencies = [ - "is-macro", - "memchr", - "rustpython-parser-vendored", -] - -[[package]] -name = "rustpython-parser-vendored" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6" -dependencies = [ - "memchr", - "once_cell", -] - [[package]] name = "rustversion" version = "1.0.23" @@ -4412,12 +4263,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" -[[package]] -name = "siphasher" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" - [[package]] name = "slab" version = "0.4.12" @@ -4648,7 +4493,7 @@ dependencies = [ "fancy-regex 0.17.0", "lazy_static", "regex", - "rustc-hash 2.1.3", + "rustc-hash", ] [[package]] @@ -4681,15 +4526,6 @@ dependencies = [ "time-core", ] -[[package]] -name = "tiny-keccak" -version = "2.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" -dependencies = [ - "crunchy", -] - [[package]] name = "tinystr" version = "0.8.3" @@ -5124,58 +4960,6 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" -[[package]] -name = "unic-char-property" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221" -dependencies = [ - "unic-char-range", -] - -[[package]] -name = "unic-char-range" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc" - -[[package]] -name = "unic-common" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc" - -[[package]] -name = "unic-emoji-char" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d" -dependencies = [ - "unic-char-property", - "unic-char-range", - "unic-ucd-version", -] - -[[package]] -name = "unic-ucd-ident" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987" -dependencies = [ - "unic-char-property", - "unic-char-range", - "unic-ucd-version", -] - -[[package]] -name = "unic-ucd-version" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4" -dependencies = [ - "unic-common", -] - [[package]] name = "unicase" version = "2.9.0" @@ -5203,40 +4987,12 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" -[[package]] -name = "unicode-width" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" - [[package]] name = "unicode_categories" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" -[[package]] -name = "unicode_names2" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd" -dependencies = [ - "phf", - "unicode_names2_generator", -] - -[[package]] -name = "unicode_names2_generator" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e" -dependencies = [ - "getopts", - "log", - "phf_codegen", - "rand 0.8.7", -] - [[package]] name = "untrusted" version = "0.9.0" diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs index 3f2c83a6a6d..23595397a13 100644 --- a/litellm-rust/crates/secrets-aws/src/error.rs +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -12,7 +12,7 @@ pub enum Error { Timeout, #[error("AWS KMS decrypt failed")] Decrypt(#[from] #[redact] Box>), - #[error("AWS Secrets Manager request preparation failed")] + #[error("AWS Secrets Manager read failed")] Read(#[from] #[redact] Box>), #[error("AWS Secrets Manager create failed")] Create(#[from] #[redact] Box>), @@ -22,6 +22,8 @@ pub enum Error { Delete(#[from] #[redact] Box>), #[error("AWS Secrets Manager replication failed")] Replicate(#[from] #[redact] Box>), + #[error("AWS Secrets Manager response has no string payload")] + MissingString, #[error("primary secret is not a JSON object")] PrimarySecret, #[error(transparent)] diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 97645ecac6d..493cb1d2e8f 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -106,23 +106,24 @@ impl AwsSecretsManagerV2 { } else { self.async_read_secret(primary).await? }; - let object: Value = serde_json::from_str( - value - .as_ref() - .map(SecretValue::expose) - .filter(|v| !v.is_empty()) - .unwrap_or("{}"), - ) - .map_err(|_| Error::PrimarySecret)?; + let Some(value) = value else { + return Ok(None); + }; + let object: Value = + serde_json::from_str(value.expose()).map_err(|_| Error::PrimarySecret)?; let object = object.as_object().ok_or(Error::PrimarySecret)?; - Ok(object.get(name).cloned().and_then(Secret::from_json)) + Ok(object.get(name).cloned().map(Secret::from_json)) } } } pub async fn async_read_secret(&self, name: &str) -> Result, Error> { match self.client.get_secret_value().secret_id(name).send().await { - Ok(response) => Ok(response.secret_string.map(SecretValue::new)), + Ok(response) => response + .secret_string + .map(SecretValue::new) + .map(Some) + .ok_or(Error::MissingString), Err(error) if matches!( &error, @@ -131,11 +132,14 @@ impl AwsSecretsManagerV2 { { Err(Error::Timeout) } - Err(error) if request_preparation_failed(&error) => Err(Error::Read(Box::new(error))), - Err(_) => { - tracing::error!("AWS secret read failed"); + Err(error) + if error + .as_service_error() + .is_some_and(|error| error.is_resource_not_found_exception()) => + { Ok(None) } + Err(error) => Err(Error::Read(Box::new(error))), } } @@ -281,17 +285,3 @@ fn bootstrap_key(name: &str) -> bool { | AWS_BEDROCK_RUNTIME_ENDPOINT ) } - -fn request_preparation_failed( - error: &aws_sdk_secretsmanager::error::SdkError< - aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError, - >, -) -> bool { - matches!( - error, - aws_sdk_secretsmanager::error::SdkError::ConstructionFailure(_) - ) || std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |error| { - error.source() - }) - .any(|source| source.is::()) -} diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs index 56698b391ec..a410767cb5a 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -290,3 +290,23 @@ async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { Err(Error::Timeout) )); } + +#[rstest::rstest] +#[case::denied(400, "AccessDeniedException")] +#[case::throttled(400, "ThrottlingException")] +#[case::unavailable(503, "ServiceUnavailableException")] +#[tokio::test] +async fn service_failures_remain_errors(#[case] status: u16, #[case] code: &str) { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with(ResponseTemplate::new(status).set_body_json(json!({"__type":code}))) + .expect(1) + .mount(&server) + .await; + assert!(matches!( + manager(&server, KeyManagementSettings::default()) + .async_read_secret("key") + .await, + Err(Error::Read(_)) + )); +} diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 3eb8475a367..3c34d9cbcc4 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -26,7 +26,7 @@ pub struct GoogleSecretManager { credentials: Arc, endpoint: reqwest::Url, project: String, - cache: Cache>, + cache: Cache, always_read: bool, } @@ -119,7 +119,7 @@ impl GoogleSecretManager { if !self.always_read && let Some(cached) = self.cache.get(name).await { - return Ok(cached.and_then(cached_secret)); + return Ok(Some(Secret::String(cached))); } let url = self .endpoint @@ -138,32 +138,20 @@ impl GoogleSecretManager { .headers(self.credentials.request_headers().await?) .send() .await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } if response.status() != reqwest::StatusCode::OK { - self.cache.insert(name.to_owned(), None).await; return Err(Error::Status(response.status().as_u16())); } let response: Response = response.json().await?; let Some(data) = response.payload.and_then(|payload| payload.data) else { - self.cache.insert(name.to_owned(), None).await; return Err(Error::MissingPayload); }; - let filtered: String = data - .chars() - .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) - .collect(); - let bytes = STANDARD.decode(filtered)?; + let bytes = STANDARD.decode(data)?; let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?; let value = SecretValue::new(plaintext); - self.cache - .insert(name.to_owned(), Some(value.clone())) - .await; + self.cache.insert(name.to_owned(), value.clone()).await; Ok(Some(Secret::String(value))) } } - -fn cached_secret(value: SecretValue) -> Option { - match serde_json::from_str(value.expose()) { - Ok(json) => Secret::from_json(json), - Err(_) => Some(Secret::String(value)), - } -} diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs index 66cc0068ebb..b3b1d29e62c 100644 --- a/litellm-rust/crates/secrets-google/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_secrets_google::{Error, GoogleSecretManager}; -use litellm_secrets_types::Secret; + use wiremock::{ Mock, MockServer, ResponseTemplate, matchers::{header, path}, @@ -55,32 +55,58 @@ async fn successful_reads_use_auth_latest_version_and_cache_including_empty_valu } #[rstest::rstest] -#[case::not_found(ResponseTemplate::new(404))] -#[case::missing_payload( - ResponseTemplate::new(200).set_body_json(serde_json::json!({"payload":{}})) -)] +#[case::not_found(404, serde_json::json!({}))] +#[case::unauthorized(401, serde_json::json!({}))] +#[case::forbidden(403, serde_json::json!({}))] +#[case::throttled(429, serde_json::json!({}))] +#[case::unavailable(503, serde_json::json!({}))] +#[case::missing_payload(200, serde_json::json!({"payload":{}}))] +#[case::invalid_base64(200, serde_json::json!({"payload":{"data":"%%%"}}))] #[tokio::test] -async fn negative_cache_returns_none_after_initial_error(#[case] response: ResponseTemplate) { +async fn failed_or_missing_reads_are_not_cached( + #[case] status: u16, + #[case] body: serde_json::Value, +) { let server = MockServer::start().await; + let manager = manager(&server, false, Duration::from_secs(60)); + let failing = Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .expect(1) + .mount_as_scoped(&server) + .await; + let result = manager.get_secret_from_google_secret_manager("key").await; + match status { + 404 => assert_eq!(result.unwrap(), None), + 200 => assert!(matches!( + result, + Err(Error::MissingPayload | Error::Base64(_)) + )), + status => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)), + } + drop(failing); Mock::given(path( "/v1/projects/project/secrets/key/versions/latest:access", )) - .respond_with(response) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("recovered")}})), + ) .expect(1) .mount(&server) .await; - let manager = manager(&server, false, Duration::from_secs(60)); - assert!(matches!( - manager.get_secret_from_google_secret_manager("key").await, - Err(Error::Status(404) | Error::MissingPayload) - )); - assert!( - manager - .get_secret_from_google_secret_manager("key") - .await - .unwrap() - .is_none() - ); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some("recovered") + ); + } } #[rstest::rstest] @@ -130,21 +156,12 @@ fn google_manager_requires_host_license_and_project_configuration() { } #[rstest::rstest] -#[case::boolean("true", Some(Secret::Bool(true)))] -#[case::null("null", None)] -#[case::string( - "\"text\"", - Some(Secret::String(litellm_secrets_types::SecretValue::new("text"))) -)] -#[case::object( - "{\"key\":1}", - Secret::from_json(serde_json::json!({"key":1})) -)] +#[case("true")] +#[case("null")] +#[case("\"text\"")] +#[case("{\"key\":1}")] #[tokio::test] -async fn cached_values_preserve_python_json_conversion( - #[case] raw: &str, - #[case] expected: Option, -) { +async fn cache_preserves_raw_values(#[case] raw: &str) { let server = MockServer::start().await; Mock::given(path( "/v1/projects/project/secrets/key/versions/latest:access", @@ -157,20 +174,15 @@ async fn cached_values_preserve_python_json_conversion( .mount(&server) .await; let manager = manager(&server, false, Duration::from_secs(60)); - assert_eq!( - manager - .get_secret_from_google_secret_manager("key") - .await - .unwrap() - .unwrap() - .as_str(), - Some(raw) - ); - assert_eq!( - manager - .get_secret_from_google_secret_manager("key") - .await - .unwrap(), - expected - ); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some(raw) + ); + } } diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs index 524045ac007..087537fb3eb 100644 --- a/litellm-rust/crates/secrets-types/src/value.rs +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -14,12 +14,11 @@ impl From for Secret { } impl Secret { - pub fn from_json(value: serde_json::Value) -> Option { + pub fn from_json(value: serde_json::Value) -> Self { match value { - serde_json::Value::Null => None, - serde_json::Value::String(value) => Some(Self::String(SecretValue::new(value))), - serde_json::Value::Bool(value) => Some(Self::Bool(value)), - value => Some(Self::Json(value)), + serde_json::Value::String(value) => Self::String(SecretValue::new(value)), + serde_json::Value::Bool(value) => Self::Bool(value), + value => Self::Json(value), } } diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 3414470d234..a7e7ec80636 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -21,13 +21,10 @@ strum.workspace = true jsonwebtoken.workspace = true serde_json.workspace = true thiserror.workspace = true -tracing = "0.1" reqwest.workspace = true moka.workspace = true tokio = { workspace = true, features = ["fs"] } -rustpython-parser = { version = "0.4.0", default-features = false, features = ["num-bigint"] } - [dev-dependencies] rstest.workspace = true wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md new file mode 100644 index 00000000000..183a39e15bb --- /dev/null +++ b/litellm-rust/crates/secrets/README.md @@ -0,0 +1,11 @@ +# Secret resolution + +Construct `SecretManagerState::new(backend, settings)` for a configured manager or use `SecretManagerState::default()` for environment lookups. The configured backend determines its provider identity. Write-only settings and names excluded by `hosted_keys` use the environment directly. `secret_manager_would_be_consulted` follows the same routing decision as resolution + +`get_secret` returns `Ok(Some(value))` for a found value, `Ok(None)` when no source contains the value, and `Err(error)` when lookup fails. For managed names, resolution checks the manager, then the environment, then the caller's default. An empty string, `false`, or an explicitly stored JSON null is a found value + +Backend failures propagate by default. To allow fallback during a backend failure, construct the resolver with `.with_failure_policy(FailurePolicy::EnvironmentFallback)`. It then tries the environment and default, in that order. If neither exists, the original error is returned. This policy applies to manager lookups. Explicit OIDC references retain their own authentication errors and never fall back to environment secrets under the reference name + +`get_secret` preserves value types. `get_secret_str` accepts a string default and rejects boolean or JSON values with `Error::TypeMismatch`. `get_secret_bool` accepts a boolean default and converts strings containing `true` or `false`, ignoring surrounding whitespace and ASCII case. Other strings and JSON values produce `Error::TypeMismatch`. Conversion failures never activate fallback or replace a found value with the default + +Provider payloads remain strings unless explicitly selecting a field from an AWS primary JSON secret. Google caches only successfully decoded string payloads, so reads have identical values and types before and after caching. Confirmed absence and failed reads are not cached. AWS resource-not-found responses and Google HTTP 404 responses indicate absence. Other provider errors remain errors, and successful responses without the required payload are malformed responses rather than missing secrets diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index d240adb6a67..0c6e681b8aa 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -1,5 +1,3 @@ -use crate::KeyManagementSystem; - #[derive(Debug, thiserror::Error)] pub enum Error { #[error("encrypted environment value is missing")] @@ -8,10 +6,6 @@ pub enum Error { InvalidCiphertext, #[error("decrypted value is not UTF-8")] Utf8, - #[error("secret manager backend is not compiled: {0:?}")] - UnsupportedBackend(KeyManagementSystem), - #[error("configured secret manager does not match its backend")] - BackendMismatch, #[error("unsupported OIDC provider or missing build feature")] UnsupportedOidc, #[error("OIDC reference requires a provider and audience")] @@ -28,8 +22,8 @@ pub enum Error { UnsafeOidcPath, #[error("OIDC file could not be read")] OidcFile, - #[error("secret manager returned no secret")] - MissingSecret, + #[error("secret cannot be converted to {expected}")] + TypeMismatch { expected: &'static str }, #[cfg(feature = "aws")] #[error(transparent)] Aws(#[from] litellm_secrets_aws::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 6d19965605d..943ffdf6158 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -76,9 +76,8 @@ pub async fn get_secret_from_manager( #[cfg(feature = "google")] SecretManager::GoogleSecretManager(client) => client .get_secret_from_google_secret_manager(secret_name) - .await? - .map(Some) - .ok_or(Error::MissingSecret), + .await + .map_err(Error::from), } } diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index c434cfaaf03..ff2e95f7b2f 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -12,7 +12,7 @@ pub use litellm_secrets_types::{ AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, }; pub use oidc::{OidcProvider, OidcReference, OidcResolver}; -pub use resolver::SecretResolver; +pub use resolver::{FailurePolicy, SecretResolver}; pub use state::{SecretManagerState, secret_manager_would_be_consulted}; #[cfg(feature = "aws")] diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index 48e6a9bc0b0..fd477859bf6 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -88,8 +88,13 @@ pub struct OidcResolver { impl Default for OidcResolver { fn default() -> Self { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(600)) + .connect_timeout(Duration::from_secs(5)) + .build() + .expect("HTTP client configuration"); Self::new( - reqwest::Client::builder().timeout(Duration::from_secs(600)).connect_timeout(Duration::from_secs(5)).build().expect("HTTP client configuration"), + client, reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"), ) } diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 93daddae0c4..89439893852 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -2,14 +2,21 @@ use std::sync::Arc; use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use crate::state::{LookupTarget, normalize_secret_name}; use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; -use crate::state::{LookupTarget, normalize_secret_name}; +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum FailurePolicy { + #[default] + Propagate, + EnvironmentFallback, +} pub struct SecretResolver { state: Arc, environment: Arc, oidc: OidcResolver, + failure_policy: FailurePolicy, } impl Default for SecretResolver { @@ -32,13 +39,21 @@ impl SecretResolver { state, environment, oidc, + failure_policy: FailurePolicy::default(), + } + } + + pub fn with_failure_policy(self, failure_policy: FailurePolicy) -> Self { + Self { + failure_policy, + ..self } } pub async fn get_secret( &self, name: &str, - _default_value: Option, + default_value: Option, ) -> Result, Error> { let name = normalize_secret_name(name); if name.starts_with("oidc/") { @@ -46,32 +61,26 @@ impl SecretResolver { .oidc .resolve(name, self.environment.as_ref()) .await - .map(|value| value.map(Secret::String)); + .map(|value| value.map(Secret::String).or(default_value)); } - if !self.state.readable() { - return Ok(self - .environment - .get(name) - .map(|value| match str_to_bool(&value) { - Some(value) => Secret::Bool(value), - None => Secret::String(SecretValue::new(value)), - })); + let LookupTarget::Manager { backend, settings } = self.state.lookup_target(name) else { + return Ok(self.environment_secret(name).or(default_value)); + }; + match crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref()) + .await + { + Ok(value) => Ok(value + .or_else(|| self.environment_secret(name)) + .or(default_value)), + Err(error) => match self.failure_policy { + FailurePolicy::Propagate => Err(error), + FailurePolicy::EnvironmentFallback => self + .environment_secret(name) + .or(default_value) + .map(Some) + .ok_or(error), + }, } - let result = match self.state.lookup_target(name) { - LookupTarget::Environment => Ok(self.environment_secret(name)), - LookupTarget::Manager { backend, settings } => { - crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref()) - .await - } - }; - let value = match result { - Ok(value) => value, - Err(_) => { - tracing::error!("secret manager lookup failed; falling back to environment"); - self.environment_secret(name) - } - }; - Ok(value.and_then(managed_secret)) } fn environment_secret(&self, name: &str) -> Option { @@ -84,12 +93,18 @@ impl SecretResolver { pub async fn get_secret_str( &self, name: &str, - default_value: Option, + default_value: Option, ) -> Result, Error> { - Ok(match self.get_secret(name, default_value).await? { - Some(Secret::String(value)) => Some(value), - Some(Secret::Bool(_) | Secret::Json(_)) | None => None, - }) + match self + .get_secret(name, default_value.map(Secret::String)) + .await? + { + Some(Secret::String(value)) => Ok(Some(value)), + None => Ok(None), + Some(Secret::Bool(_) | Secret::Json(_)) => { + Err(Error::TypeMismatch { expected: "string" }) + } + } } pub async fn get_secret_bool( @@ -97,44 +112,24 @@ impl SecretResolver { name: &str, default_value: Option, ) -> Result, Error> { - Ok( - match self - .get_secret(name, default_value.map(Secret::Bool)) - .await? - { - Some(Secret::Bool(value)) => Some(value), - Some(Secret::String(value)) => str_to_bool(value.expose()), - Some(Secret::Json(_)) | None => None, - }, - ) - } -} - -fn str_to_bool(value: &str) -> Option { - match value.trim().to_ascii_lowercase().as_str() { - "true" => Some(true), - "false" => Some(false), - _ => None, - } -} - -fn literal_bool(value: &str) -> Option { - use rustpython_parser::{Parse, ast}; - match ast::Expr::parse(value.trim_start_matches([' ', '\t']), "").ok()? { - ast::Expr::Constant(node) => match node.value { - ast::Constant::Bool(value) => Some(value), - _ => None, - }, - _ => None, - } -} - -fn managed_secret(value: Secret) -> Option { - match value { - Secret::String(value) => Some(match literal_bool(value.expose()) { - Some(boolean) => Secret::Bool(boolean), - None => Secret::String(value), - }), - Secret::Bool(_) | Secret::Json(_) => None, + match self + .get_secret(name, default_value.map(Secret::Bool)) + .await? + { + Some(Secret::Bool(value)) => Ok(Some(value)), + Some(Secret::String(value)) => { + match value.expose().trim().to_ascii_lowercase().as_str() { + "true" => Ok(Some(true)), + "false" => Ok(Some(false)), + _ => Err(Error::TypeMismatch { + expected: "boolean", + }), + } + } + Some(Secret::Json(_)) => Err(Error::TypeMismatch { + expected: "boolean", + }), + None => Ok(None), + } } } diff --git a/litellm-rust/crates/secrets/src/state.rs b/litellm-rust/crates/secrets/src/state.rs index ca942aa65db..7854d763ef2 100644 --- a/litellm-rust/crates/secrets/src/state.rs +++ b/litellm-rust/crates/secrets/src/state.rs @@ -1,4 +1,4 @@ -use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; +use crate::{KeyManagementSettings, KeyManagementSystem, SecretManager}; pub(crate) enum LookupTarget<'a> { Environment, @@ -14,73 +14,37 @@ pub(crate) fn normalize_secret_name(name: &str) -> &str { #[derive(Clone, Default)] pub struct SecretManagerState { - system: Option, - settings: Option, - backend: Option, + manager: Option<(SecretManager, KeyManagementSettings)>, } impl SecretManagerState { - pub fn new( - system: Option, - settings: Option, - backend: Option, - ) -> Result { - if let Some(system) = system { - let available = match system { - KeyManagementSystem::Local => true, - KeyManagementSystem::AwsKms | KeyManagementSystem::AwsSecretManager => { - cfg!(feature = "aws") - } - KeyManagementSystem::GoogleKms | KeyManagementSystem::GoogleSecretManager => { - cfg!(feature = "google") - } - KeyManagementSystem::AzureKeyVault - | KeyManagementSystem::HashicorpVault - | KeyManagementSystem::Cyberark - | KeyManagementSystem::Custom => false, - }; - if !available { - return Err(Error::UnsupportedBackend(system)); - } - if let Some(backend) = &backend - && system != backend.system() - { - return Err(Error::BackendMismatch); - } + pub fn new(backend: SecretManager, settings: KeyManagementSettings) -> Self { + Self { + manager: Some((backend, settings)), } - Ok(Self { - system, - settings, - backend, - }) } pub fn system(&self) -> Option { - self.system - } - pub fn settings(&self) -> Option<&KeyManagementSettings> { - self.settings.as_ref() - } - pub fn backend(&self) -> Option<&SecretManager> { - self.backend.as_ref() + self.backend().map(SecretManager::system) } - pub(crate) fn readable(&self) -> bool { - self.backend.is_some() - && self - .settings - .as_ref() - .is_some_and(|settings| settings.access_mode.readable()) + pub fn settings(&self) -> Option<&KeyManagementSettings> { + self.manager.as_ref().map(|(_, settings)| settings) + } + + pub fn backend(&self) -> Option<&SecretManager> { + self.manager.as_ref().map(|(backend, _)| backend) } pub(crate) fn lookup_target(&self, name: &str) -> LookupTarget<'_> { - match (&self.backend, &self.settings) { - (Some(backend), Some(settings)) - if settings.access_mode.readable() - && hosts_secret(settings, name) - && self - .system - .is_some_and(|system| system != KeyManagementSystem::Local) => + match &self.manager { + Some((backend, settings)) + if backend.system() != KeyManagementSystem::Local + && settings.access_mode.readable() + && settings + .hosted_keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == name)) => { LookupTarget::Manager { backend, settings } } @@ -90,16 +54,6 @@ impl SecretManagerState { } pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool { - state.readable() - && state - .settings - .as_ref() - .is_some_and(|settings| hosts_secret(settings, normalize_secret_name(name))) -} - -fn hosts_secret(settings: &KeyManagementSettings, name: &str) -> bool { - settings - .hosted_keys - .as_ref() - .is_none_or(|keys| keys.iter().any(|key| key == name)) + let name = normalize_secret_name(name); + !name.starts_with("oidc/") && matches!(state.lookup_target(name), LookupTarget::Manager { .. }) } diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index e2b056f01cc..3a826092d72 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -1,18 +1,13 @@ use std::sync::Arc; use litellm_secrets::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, OidcResolver, Secret, SecretManager, - SecretManagerState, SecretResolver, SecretValue, secret_manager_would_be_consulted, + Error, KeyManagementSettings, OidcResolver, Secret, SecretManager, SecretManagerState, + SecretResolver, SecretValue, secret_manager_would_be_consulted, }; -fn resolver(value: Option<&str>, readable: bool) -> SecretResolver { - let state = if readable { - SecretManagerState::new( - Some(KeyManagementSystem::Local), - Some(KeyManagementSettings::default()), - Some(SecretManager::Local), - ) - .unwrap() +fn resolver(value: Option<&str>, configured: bool) -> SecretResolver { + let state = if configured { + SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default()) } else { SecretManagerState::default() }; @@ -25,77 +20,107 @@ fn resolver(value: Option<&str>, readable: bool) -> SecretResolver { } #[rstest::rstest] -#[case::lowercase_true("true", Some(true), None)] -#[case::whitespace_lowercase_false(" FALSE ", Some(false), None)] -#[case::python_true("True", Some(true), Some(true))] -#[case::python_false("False", Some(false), Some(false))] -#[case::parenthesized_python_true("(True)", None, Some(true))] -#[case::commented_python_false("False # comment", None, Some(false))] -#[case::integer("1", None, None)] -#[case::yes("yes", None, None)] -#[case::plain_string("secret", None, None)] +#[case("true", Some(true))] +#[case(" FALSE ", Some(false))] +#[case("(True)", None)] +#[case("False # comment", None)] +#[case("1", None)] +#[case("secret", None)] #[tokio::test] -async fn boolean_conversion_preserves_local_and_manager_differences( +async fn conversion_is_explicit_and_independent_of_manager_configuration( #[case] input: &str, - #[case] local: Option, - #[case] manager: Option, - #[values(false, true)] readable: bool, + #[case] boolean: Option, + #[values(false, true)] configured: bool, ) { - let boolean = if readable { manager } else { local }; - let resolver = resolver(Some(input), readable); - let expected = boolean - .map(Secret::Bool) - .unwrap_or_else(|| Secret::String(SecretValue::new(input))); + let resolver = resolver(Some(input), configured); assert_eq!( resolver.get_secret("key", None).await.unwrap(), - Some(expected) + Some(Secret::String(SecretValue::new(input))) ); assert_eq!( resolver .get_secret_str("key", None) .await .unwrap() - .map(|v| v.expose().to_owned()), - boolean.is_none().then(|| input.to_owned()) + .unwrap() + .expose(), + input + ); + match boolean { + Some(value) => assert_eq!( + resolver.get_secret_bool("key", None).await.unwrap(), + Some(value) + ), + None => assert!(matches!( + resolver.get_secret_bool("key", Some(true)).await, + Err(Error::TypeMismatch { + expected: "boolean" + }) + )), + } +} + +#[rstest::rstest] +#[tokio::test] +async fn defaults_apply_only_to_absence(#[values(false, true)] configured: bool) { + let missing = resolver(None, configured); + assert_eq!(missing.get_secret("key", None).await.unwrap(), None); + assert_eq!( + missing.get_secret_bool("key", Some(false)).await.unwrap(), + Some(false) + ); + assert_eq!( + missing + .get_secret_str("key", Some(SecretValue::new("default"))) + .await + .unwrap() + .unwrap() + .expose(), + "default" + ); + for value in [ + Secret::Bool(false), + Secret::from_json(serde_json::json!({"key":1})), + Secret::from_json(serde_json::Value::Null), + ] { + assert_eq!( + missing + .get_secret("key", Some(value.clone())) + .await + .unwrap(), + Some(value) + ); + } + assert_eq!( + resolver(Some(""), configured) + .get_secret_str("key", Some(SecretValue::new("default"))) + .await + .unwrap() + .unwrap() + .expose(), + "" ); } #[tokio::test] -async fn manager_boolean_conversion_trims_whitespace() { +async fn prefix_is_removed_once_and_local_manager_is_not_consulted() { + let state = SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default()); assert_eq!( - resolver(Some(" true "), true) - .get_secret_bool("key", None) - .await - .unwrap(), - Some(true) - ); -} - -#[tokio::test] -async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() { - let missing = resolver(None, false); - assert_eq!( - missing - .get_secret("missing", Some(Secret::Bool(true))) - .await - .unwrap(), - None - ); - assert_eq!( - missing - .get_secret_bool("missing", Some(true)) - .await - .unwrap(), - None + state.system(), + Some(litellm_secrets::KeyManagementSystem::Local) ); + assert!(!secret_manager_would_be_consulted( + &state, + "os.environ/os.environ/KEY" + )); let resolver = SecretResolver::new( - Arc::new(SecretManagerState::default()), - Arc::new(|name: &str| (name == "KEY").then(|| "value".into())), + Arc::new(state), + Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())), OidcResolver::default(), ); assert_eq!( resolver - .get_secret_str("os.environ/KEY", None) + .get_secret_str("os.environ/os.environ/KEY", None) .await .unwrap() .unwrap() @@ -104,177 +129,6 @@ async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() { ); } -#[rstest::rstest] -#[case::all_keys(None)] -#[case::no_keys(Some(Vec::new()))] -#[case::allowlisted_key(Some(vec!["KEY".into()]))] -fn manager_gating_requires_client_readable_settings_and_allowlisted_name( - #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)] - access_mode: AccessMode, - #[values(false, true)] client: bool, - #[case] keys: Option>, -) { - let expected = client - && access_mode.readable() - && keys - .as_ref() - .is_none_or(|keys| keys.iter().any(|key| key == "KEY")); - let state = SecretManagerState::new( - Some(KeyManagementSystem::Local), - Some(KeyManagementSettings { - access_mode, - hosted_keys: keys, - ..Default::default() - }), - client.then_some(SecretManager::Local), - ) - .unwrap(); - assert_eq!( - secret_manager_would_be_consulted(&state, "os.environ/KEY"), - expected - ); -} - -#[test] -fn manager_gating_requires_settings() { - let no_settings = SecretManagerState::new(None, None, Some(SecretManager::Local)).unwrap(); - assert!(!secret_manager_would_be_consulted(&no_settings, "KEY")); -} - -#[cfg(feature = "aws")] -#[rstest::rstest] -#[case::missing_value(None, None)] -#[case::lookup_error(Some("primary".to_owned()), Some("environment-value"))] -#[tokio::test] -async fn aws_missing_values_do_not_fallback_but_lookup_errors_do( - #[case] primary: Option, - #[case] expected: Option<&str>, -) { - use litellm_secrets::aws::AwsSecretsManagerV2; - use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_partial_json}; - let server = MockServer::start().await; - Mock::given(body_partial_json(serde_json::json!({"SecretId":"KEY"}))) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) - .expect(u64::from(primary.is_none())) - .mount(&server) - .await; - Mock::given(body_partial_json(serde_json::json!({"SecretId":"primary"}))) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"SecretString":"invalid-json"})), - ) - .expect(u64::from(primary.is_some())) - .mount(&server) - .await; - let endpoint = server.uri(); - let environment: Arc = - Arc::new(move |name: &str| match name { - "AWS_REGION_NAME" => Some("us-east-1".into()), - "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()), - "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()), - "KEY" => Some("environment-value".into()), - _ => None, - }); - let settings = KeyManagementSettings { - primary_secret_name: primary, - ..Default::default() - }; - let manager = AwsSecretsManagerV2::load_aws_secret_manager( - Some(true), - settings.clone(), - environment.clone(), - ) - .unwrap() - .unwrap(); - let state = SecretManagerState::new( - Some(KeyManagementSystem::AwsSecretManager), - Some(settings), - Some(SecretManager::AwsSecretsManagerV2(manager)), - ) - .unwrap(); - let resolver = SecretResolver::new( - Arc::new(state), - environment.clone(), - OidcResolver::default(), - ); - assert_eq!( - resolver - .get_secret_str("os.environ/KEY", None) - .await - .unwrap() - .map(|v| v.expose().to_owned()) - .as_deref(), - expected - ); -} - -#[cfg(feature = "google")] -#[rstest::rstest] -#[case::hosted_filter(Some(Vec::new()), Some(KeyManagementSystem::GoogleSecretManager))] -#[case::negative_cache(None, Some(KeyManagementSystem::GoogleSecretManager))] -#[case::missing_system(None, None)] -#[case::hosted_nested_prefix(Some(vec!["os.environ/KEY".into()]), Some(KeyManagementSystem::GoogleSecretManager))] -#[tokio::test] -async fn google_negative_cache_still_falls_back_and_hosted_filter_avoids_io( - #[case] hosted_keys: Option>, - #[case] system: Option, -) { - use litellm_secrets::google::GoogleSecretManager; - use std::time::Duration; - use wiremock::{Mock, MockServer, ResponseTemplate, matchers::path}; - let server = MockServer::start().await; - Mock::given(path( - "/v1/projects/project/secrets/os%2Eenviron%2FKEY/versions/latest:access", - )) - .respond_with(ResponseTemplate::new(404)) - .expect(u64::from( - hosted_keys.as_ref().is_none_or(|keys| !keys.is_empty()) && system.is_some(), - )) - .mount(&server) - .await; - let environment: Arc = - Arc::new(|name: &str| match name { - "VERTEX_AI_API_KEY" => Some("token".into()), - "os.environ/KEY" => Some("environment-value".into()), - _ => None, - }); - let manager = GoogleSecretManager::with_client( - reqwest::Client::new(), - server.uri().parse().unwrap(), - "project".into(), - environment.clone(), - Some(Duration::from_secs(60)), - false, - ) - .unwrap(); - let settings = KeyManagementSettings { - hosted_keys, - ..Default::default() - }; - let state = SecretManagerState::new( - system, - Some(settings), - Some(SecretManager::GoogleSecretManager(manager)), - ) - .unwrap(); - let resolver = SecretResolver::new( - Arc::new(state), - environment.clone(), - OidcResolver::default(), - ); - for _ in 0..2 { - assert_eq!( - resolver - .get_secret_str("os.environ/os.environ/KEY", None) - .await - .unwrap() - .unwrap() - .expose(), - "environment-value" - ); - } -} - #[tokio::test] async fn resolver_future_can_run_on_a_tokio_worker() { let resolver = resolver(Some("worker-value"), false); @@ -285,59 +139,230 @@ async fn resolver_future_can_run_on_a_tokio_worker() { assert_eq!(result.unwrap().expose(), "worker-value"); } -#[rstest::rstest] -#[case::nested_true("((True)) # comment", Some(true))] -#[case::commented_false("(False # comment\n)", Some(false))] -#[case::boolean_expression("True and False", None)] -#[case::string_literal("'True'", None)] -#[case::tuple("(True,)", None)] -#[case::unary_expression("not False", None)] -#[case::multiple_expressions("True\nFalse", None)] -#[case::incomplete_expression("(True", None)] -#[tokio::test] -async fn manager_boolean_literals_follow_python_syntax( - #[case] input: &str, - #[case] expected: Option, -) { - let value = resolver(Some(input), true) - .get_secret("key", None) - .await - .unwrap(); - assert_eq!( - value, - Some( - expected - .map(Secret::Bool) - .unwrap_or_else(|| Secret::String(SecretValue::new(input))) +#[cfg(feature = "aws")] +mod aws { + use super::*; + use litellm_secrets::{AccessMode, FailurePolicy, aws::AwsSecretsManagerV2}; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + + fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState { + let endpoint = server.uri(); + let environment = Arc::new(move |name: &str| match name { + "AWS_REGION_NAME" => Some("us-east-1".into()), + "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()), + "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()), + _ => None, + }); + let manager = + AwsSecretsManagerV2::load_aws_secret_manager(Some(true), settings.clone(), environment) + .unwrap() + .unwrap(); + SecretManagerState::new(SecretManager::AwsSecretsManagerV2(manager), settings) + } + + #[rstest::rstest] + #[case::missing(400, serde_json::json!({"__type":"ResourceNotFoundException"}), false)] + #[case::denied(400, serde_json::json!({"__type":"AccessDeniedException"}), true)] + #[case::malformed(200, serde_json::json!({}), true)] + #[tokio::test] + async fn failure_policy_preserves_errors_and_fallback_precedence( + #[case] status: u16, + #[case] body: serde_json::Value, + #[case] fails: bool, + #[values(FailurePolicy::Propagate, FailurePolicy::EnvironmentFallback)] + policy: FailurePolicy, + #[values(None, Some("environment"))] environment: Option<&'static str>, + #[values(None, Some("default"))] default: Option<&str>, + ) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .expect(1) + .mount(&server) + .await; + let resolver = SecretResolver::new( + Arc::new(state(&server, KeyManagementSettings::default())), + Arc::new(move |_: &str| environment.map(str::to_owned)), + OidcResolver::default(), ) - ); + .with_failure_policy(policy); + let result = resolver + .get_secret_str("KEY", default.map(SecretValue::new)) + .await; + let fallback = environment.or(default); + if fails && (policy == FailurePolicy::Propagate || fallback.is_none()) { + assert!(matches!(result, Err(Error::Aws(_)))); + } else { + assert_eq!(result.unwrap().as_ref().map(SecretValue::expose), fallback); + } + } + + #[rstest::rstest] + #[case::boolean(serde_json::json!(false))] + #[case::object(serde_json::json!({"key":1}))] + #[case::null(serde_json::Value::Null)] + #[case::string(serde_json::json!("true"))] + #[tokio::test] + async fn typed_values_survive_resolution_and_accessors_reject_wrong_types( + #[case] value: serde_json::Value, + ) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"SecretString":serde_json::json!({"KEY":value}).to_string()}), + )) + .expect(3) + .mount(&server) + .await; + let settings = KeyManagementSettings { + primary_secret_name: Some("primary".into()), + ..Default::default() + }; + let resolver = SecretResolver::new( + Arc::new(state(&server, settings)), + Arc::new(|_: &str| Some("fallback".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret("KEY", Some(Secret::Bool(true))) + .await + .unwrap(), + Some(Secret::from_json(value.clone())) + ); + match &value { + serde_json::Value::String(text) => assert_eq!( + resolver + .get_secret_str("KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + text + ), + _ => assert!(matches!( + resolver.get_secret_str("KEY", None).await, + Err(Error::TypeMismatch { expected: "string" }) + )), + } + match value { + serde_json::Value::Bool(boolean) => assert_eq!( + resolver.get_secret_bool("KEY", None).await.unwrap(), + Some(boolean) + ), + serde_json::Value::String(_) => assert_eq!( + resolver.get_secret_bool("KEY", None).await.unwrap(), + Some(true) + ), + _ => assert!(matches!( + resolver.get_secret_bool("KEY", None).await, + Err(Error::TypeMismatch { + expected: "boolean" + }) + )), + } + } + + #[rstest::rstest] + #[tokio::test] + async fn gating_prediction_matches_actual_lookup( + #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)] + access_mode: AccessMode, + #[values(None, Some(vec![]), Some(vec!["KEY".into()]))] hosted_keys: Option>, + #[values("os.environ/KEY", "os.environ/oidc/env/KEY")] name: &str, + ) { + let server = MockServer::start().await; + let expected = name == "os.environ/KEY" + && access_mode.readable() + && hosted_keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == "KEY")); + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"SecretString":"remote"})), + ) + .expect(u64::from(expected)) + .mount(&server) + .await; + let state = state( + &server, + KeyManagementSettings { + access_mode, + hosted_keys, + ..Default::default() + }, + ); + assert!(state.backend().is_some()); + assert_eq!(state.settings().unwrap().access_mode, access_mode); + assert_eq!(secret_manager_would_be_consulted(&state, name), expected); + let resolver = SecretResolver::new( + Arc::new(state), + Arc::new(|_: &str| Some("environment".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str(name, None) + .await + .unwrap() + .unwrap() + .expose(), + if expected { "remote" } else { "environment" } + ); + } } +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::missing(404)] +#[case::failure(503)] #[tokio::test] -async fn environment_prefix_is_removed_only_once_and_gating_uses_the_same_name() { - let name = "os.environ/folder/os.environ/KEY"; - let state = SecretManagerState::new( - Some(KeyManagementSystem::Local), - Some(KeyManagementSettings { - hosted_keys: Some(vec!["folder/os.environ/KEY".into()]), - ..Default::default() - }), - Some(SecretManager::Local), +async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) { + use litellm_secrets::{FailurePolicy, google::GoogleSecretManager}; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(status)) + .expect(2) + .mount(&server) + .await; + let environment: Arc = + Arc::new(|name: &str| match name { + "VERTEX_AI_API_KEY" => Some("token".into()), + "KEY" => Some("environment".into()), + _ => None, + }); + let manager = GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + environment.clone(), + None, + false, ) .unwrap(); - assert!(secret_manager_would_be_consulted(&state, name)); - let resolver = SecretResolver::new( - Arc::new(state), - Arc::new(|name: &str| (name == "folder/os.environ/KEY").then(|| "value".into())), - OidcResolver::default(), + let state = SecretManagerState::new( + SecretManager::GoogleSecretManager(manager), + KeyManagementSettings::default(), ); + let resolver = SecretResolver::new(Arc::new(state), environment, OidcResolver::default()); + let result = resolver.get_secret_str("KEY", None).await; + if status == 404 { + assert_eq!(result.unwrap().unwrap().expose(), "environment"); + } else { + assert!( + matches!(result, Err(Error::Google(litellm_secrets::google::Error::Status(actual))) if actual == status) + ); + } assert_eq!( resolver - .get_secret_str(name, None) + .with_failure_policy(FailurePolicy::EnvironmentFallback) + .get_secret_str("KEY", None) .await .unwrap() .unwrap() .expose(), - "value" + "environment" ); } From 2bacd6caa0154c46461ec3f790e4b438052156af Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:30:52 +0000 Subject: [PATCH 116/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~moonshotai/kimi-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/moonshotai/kimi-k3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bdae7ede25c..1095ec9e467 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.7e-06, - "output_cost_per_token": 8.5e-06, - "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.675e-06, + "output_cost_per_token": 9.38e-06, + "cache_read_input_token_cost": 1.943e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.7e-07, - "input_cost_per_token": 1.7e-06, + "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.675e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8.5e-06, + "output_cost_per_token": 9.38e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bdae7ede25c..1095ec9e467 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.7e-06, - "output_cost_per_token": 8.5e-06, - "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.675e-06, + "output_cost_per_token": 9.38e-06, + "cache_read_input_token_cost": 1.943e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.7e-07, - "input_cost_per_token": 1.7e-06, + "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.675e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8.5e-06, + "output_cost_per_token": 9.38e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 4258bd366c7c6342d1e29641d060823c7cea7a40 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:00:39 +0000 Subject: [PATCH 117/146] chore(prices): sync OpenRouter prices: 2 models openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1095ec9e467..ba001cd3923 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.483e-07, + "input_cost_per_token": 9.5526e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8966e-06, + "output_cost_per_token": 1.91052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9025e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.9866e-08, - "output_cost_per_token": 1.79732e-07, - "cache_read_input_token_cost": 1.79732e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1095ec9e467..ba001cd3923 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.483e-07, + "input_cost_per_token": 9.5526e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8966e-06, + "output_cost_per_token": 1.91052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9025e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.9866e-08, - "output_cost_per_token": 1.79732e-07, - "cache_read_input_token_cost": 1.79732e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From d5921713a8105804917134614a28457bc7fc552c Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:30:42 +0000 Subject: [PATCH 118/146] chore(prices): sync OpenRouter prices: 2 models openrouter/~moonshotai/kimi-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/moonshotai/kimi-k3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ba001cd3923..1976437f900 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.675e-06, - "output_cost_per_token": 9.38e-06, - "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 8.5e-06, + "cache_read_input_token_cost": 1.7e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.943e-07, - "input_cost_per_token": 1.675e-06, + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.7e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 9.38e-06, + "output_cost_per_token": 8.5e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ba001cd3923..1976437f900 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.675e-06, - "output_cost_per_token": 9.38e-06, - "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 8.5e-06, + "cache_read_input_token_cost": 1.7e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.943e-07, - "input_cost_per_token": 1.675e-06, + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.7e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 9.38e-06, + "output_cost_per_token": 8.5e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 83223885e6d4f9dd536f73214b446e1f14b64244 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 04:09:59 +0000 Subject: [PATCH 119/146] fix(responses): forward safety_identifier through the chat completion bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 1 + tests/e2e/llm_translation/endpoints_client.py | 3 + .../e2e/llm_translation/test_responses_e2e.py | 89 ++++++++++++++++++- tests/e2e/models.py | 2 + tests/e2e/provider_edge.py | 10 ++- .../test_litellm_completion_responses.py | 10 +++ 6 files changed, 109 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index cf3075ee28d..3ca2cc28c9a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -454,6 +454,7 @@ class LiteLLMCompletionResponsesConfig: "stream": stream, "metadata": kwargs.get("metadata"), "service_tier": kwargs.get("service_tier"), + "safety_identifier": responses_api_request.get("safety_identifier"), "web_search_options": web_search_options, "response_format": response_format, "reasoning_effort": reasoning.effort, diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 4d2c73e7078..eb7bae2220c 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -75,6 +75,7 @@ class ResponsesRequest(BaseModel): stream: bool = False tools: list[ResponsesFunctionTool] | None = None guardrails: list[str] | None = None + safety_identifier: str | None = None cache: dict[str, bool] | None = {"no-cache": True} @@ -316,6 +317,7 @@ class EndpointsClient: *, stream: bool = False, guardrails: list[str] | None = None, + safety_identifier: str | None = None, ) -> StreamingResponse: return self._send( "/v1/responses", @@ -326,6 +328,7 @@ class EndpointsClient: instructions="You are a helpful assistant", stream=stream, guardrails=guardrails, + safety_identifier=safety_identifier, ), stream=stream, ) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 3fcf2d1ac05..083bc66b3f2 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -8,13 +8,18 @@ litellm-regression-tests/tests/test_inference_endpoints.py. from __future__ import annotations import json -from typing import cast +import threading +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Final, cast import pytest -from e2e_config import unique_marker +from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, unique_marker from e2e_http import ( assert_client_error, require_successful_call, + unwrap, ) from endpoints_client import ( EndpointsClient, @@ -26,7 +31,9 @@ from endpoints_client import ( ResponsesStreamEventType, ) from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import ChatBody, ChatMessage, LiteLLMParamsBody +from provider_edge import LiveEdge, start_provider_edge +from provider_edge_bedrock import bedrock_signer from pydantic import BaseModel, ValidationError pytestmark = pytest.mark.e2e @@ -39,6 +46,33 @@ class _OptionalResponsesBody(BaseModel): BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +BEDROCK_EDGE_REGION: Final = "us-east-1" +BEDROCK_EDGE_MOUNT: Final = f"bedrock/{BEDROCK_EDGE_REGION}" + + +class ConverseRequestBody(BaseModel): + additionalModelRequestFields: dict[str, str] | None = None + + +@dataclass(slots=True) +class ConverseRequestCapture: + """The Converse bodies the proxy actually sent upstream, as seen by a live + edge sitting between the proxy and Bedrock.""" + + _bodies: list[ConverseRequestBody] = field(default_factory=list) + _lock: threading.Lock = field(default_factory=threading.Lock) + + def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None: + if body is None or "/converse" not in url: + return + with self._lock: + self._bodies.append(ConverseRequestBody.model_validate_json(body)) + + @property + def bodies(self) -> tuple[ConverseRequestBody, ...]: + with self._lock: + return tuple(self._bodies) + WEATHER_TOOL = ResponsesFunctionTool( name="get_weather", @@ -295,6 +329,55 @@ class TestResponses: arguments = WeatherArguments.model_validate(raw_arguments) assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + @pytest.mark.parametrize("endpoint", ["/v1/responses", "/v1/chat/completions"]) + def test_bedrock_forwards_allowed_safety_identifier_as_additional_model_request_field( + self, endpoints_client: EndpointsClient, resources: ResourceManager, endpoint: str + ) -> None: + capture: Final = ConverseRequestCapture() + edge: Final = start_provider_edge( + LiveEdge(observe_request=capture.observe, sign=bedrock_signer(BEDROCK_EDGE_REGION)), + mounts=MappingProxyType({BEDROCK_EDGE_MOUNT: f"https://bedrock-runtime.{BEDROCK_EDGE_REGION}.amazonaws.com"}), + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + ) + resources.defer(edge.shutdown) + model: Final = f"e2e-responses-{unique_marker()}" + model_id: Final = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model=BEDROCK_CONVERSE_BACKEND, + api_base=edge.edge.api_base(BEDROCK_EDGE_MOUNT), + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name=BEDROCK_EDGE_REGION, + allowed_openai_params=["safety_identifier"], + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key: Final = resources.key() + safety_identifier: Final = f"end-user-{unique_marker()}" + + if endpoint == "/v1/responses": + responses_result: Final = endpoints_client.responses( + key, model, "reply with one word", safety_identifier=safety_identifier + ) + require_successful_call(responses_result) + else: + unwrap( + endpoints_client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="reply with one word")], + safety_identifier=safety_identifier, + ), + ) + ) + + assert [body.additionalModelRequestFields for body in capture.bodies] == [ + {"safety_identifier": safety_identifier} + ], f"{endpoint} did not forward safety_identifier to Bedrock Converse: {capture.bodies}" + @pytest.mark.skip(reason="stage red: product gap, /v1/responses 500s (aresponses TypeError) on missing input instead of 400") @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") def test_missing_input_returns_error( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 4b202e3c663..47ef672ebec 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -298,6 +298,7 @@ class ChatBody(BaseModel): max_completion_tokens: int | None = None temperature: float | None = None user: str | None = None + safety_identifier: str | None = None metadata: ChatMetadata | None = None reasoning_effort: str | None = None thinking: ThinkingParam | None = None @@ -976,6 +977,7 @@ class LiteLLMParamsBody(BaseModel): api_base: str | None = None api_version: str | None = None realtime_protocol: str | None = None + allowed_openai_params: list[str] | None = None aws_access_key_id: str | None = None aws_secret_access_key: str | None = None aws_region_name: str | None = None diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 136b00208f7..fc10dde2a77 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -99,6 +99,7 @@ from provider_cache import ( SIGNATURE_HEADERS, CacheEdge, MountPolicy, + RequestSigner, is_bedrock, scoped_edge_base, split_test_segment, @@ -539,6 +540,7 @@ class ReplayEdge: @dataclass(frozen=True, slots=True) class LiveEdge: observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None + sign: RequestSigner | None = None type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge @@ -788,14 +790,16 @@ def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None, observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None, + sign: RequestSigner | None = None, ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS } if observe_request is not None: observe_request(url, forwarded, body) + outbound: Final = forwarded if sign is None else sign(method, url, forwarded, body) head: Final = ( - forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) + forward_stream(method, url, headers=outbound, body=body, timeout=timeout) if cache is None else cache.forward(mount, method, url, forwarded, body, timeout, test_key=test_key) ) match head: @@ -871,10 +875,10 @@ def handle_edge_request( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, mount, test_key, ) - case LiveEdge(observe_request=observe_request): + case LiveEdge(observe_request=observe_request, sign=sign): return _handle_live( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, - observe_request=observe_request, + observe_request=observe_request, sign=sign, ) case RecordEdge(): return _handle_record( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 6077f281a81..7b9de4644b4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1248,6 +1248,16 @@ class TestFunctionCallTransformation: assert "tool_choice" not in result assert "tools" not in result + def test_safety_identifier_forwarded_to_chat_completion_request(self) -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="bedrock/global.openai.gpt-5.6-luna", + input="hi", + responses_api_request={"safety_identifier": "user-7f3a"}, + custom_llm_provider="bedrock", + ) + + assert result["safety_identifier"] == "user-7f3a" + def test_parallel_tool_calls_dropped_when_no_chat_tools_remain(self) -> None: transform: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request codex_tool_search: Final = { From ba93c7402a029ac7efb41148e31b3a6b1d186f8b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 04:53:10 +0000 Subject: [PATCH 120/146] test(e2e): tolerate provider retries in safety_identifier capture assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/llm_translation/test_responses_e2e.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 083bc66b3f2..a5dc48a015f 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -374,9 +374,11 @@ class TestResponses: ) ) - assert [body.additionalModelRequestFields for body in capture.bodies] == [ - {"safety_identifier": safety_identifier} - ], f"{endpoint} did not forward safety_identifier to Bedrock Converse: {capture.bodies}" + forwarded: Final = tuple(body.additionalModelRequestFields for body in capture.bodies) + assert forwarded, f"{endpoint} produced no Bedrock Converse request" + assert forwarded == ({"safety_identifier": safety_identifier},) * len(forwarded), ( + f"{endpoint} did not forward safety_identifier to Bedrock Converse on every attempt: {capture.bodies}" + ) @pytest.mark.skip(reason="stage red: product gap, /v1/responses 500s (aresponses TypeError) on missing input instead of 400") @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") From 7e0fa40fe3b260efe8447077fa4726e8e0080729 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 06:17:03 +0000 Subject: [PATCH 121/146] test(e2e): gate the Bedrock edge capture behind a provider_edge_host opt-in The Buildkite ephemeral stack runs the gateway in another pod, so it cannot reach the pytest host's provider edge. The GitHub changed-e2e lane runs gateways on the runner and sets E2E_PROVIDER_EDGE_HOST_REACHABLE Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-e2e-changed.yml | 1 + tests/e2e/AGENTS.md | 2 +- tests/e2e/conftest.py | 7 +++++++ tests/e2e/e2e_config.py | 1 + tests/e2e/llm_translation/test_responses_e2e.py | 1 + tests/e2e/pytest.ini | 1 + 6 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index c9f08deb36e..f659afad7e3 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -175,6 +175,7 @@ jobs: env: TESTS: ${{ needs.detect.outputs.tests }} E2E_FIXTURE_MODE: live + E2E_PROVIDER_EDGE_HOST_REACHABLE: '1' run: | umask 077 read -r -a test_files <<< "${TESTS}" diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 78ba2ec4ed1..c25e958242f 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -122,7 +122,7 @@ E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 Point the proxy at bogus provider credentials for the replay run and it still has to pass: that is the whole proof that nothing left the process. Bundles are never committed. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and hard-fails after seven days. CI records and replays this lane on a schedule in `.github/workflows/e2e_record_replay.yml`, publishing the bundle as a private `e2e-fixtures-bundle` artifact instead of committing it, selecting the tests with the `@pytest.mark.replayable` marker, and proving the bogus-credentials replay hermetic by counting provider egress with `.github/scripts/e2e_egress_sentinel.py` -Current limits: Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode +Current limits: Bedrock cannot be mounted in record or replay (SigV4 signs the Host header, so a rewritten api_base fails signature verification); a test that needs to observe the Converse body registers its own `LiveEdge` with `provider_edge_bedrock.bedrock_signer` re-signing the forwarded request, and carries the `provider_edge_host` opt-in marker because the gateway must reach the pytest host, which the Buildkite ephemeral stack cannot (the GitHub changed-e2e lane, whose gateways run on the runner, sets `E2E_PROVIDER_EDGE_HOST_REACHABLE`). Deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode ## Typing diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 52a634693c5..8776d00d502 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -31,6 +31,7 @@ from e2e_config import ( MANAGED_FILES_OPT_IN_ENV, MCP_OAUTH_LIVE_OPT_IN_ENV, PROMPT_CACHING_OPT_IN_ENV, + PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, REDIS_CHAOS_OPT_IN_ENV, WEEKLY_ANOMALY_OPT_IN_ENV, @@ -59,6 +60,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, + "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, } ) @@ -143,6 +145,11 @@ def pytest_configure(config: pytest.Config) -> None: "mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless " "E2E_MCP_OAUTH_LIVE is set", ) + config.addinivalue_line( + "markers", + "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " + "gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index a79c158f9c4..14de4619664 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -146,6 +146,7 @@ PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" +PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index a5dc48a015f..1c2336bd166 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -329,6 +329,7 @@ class TestResponses: arguments = WeatherArguments.model_validate(raw_arguments) assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + @pytest.mark.provider_edge_host @pytest.mark.parametrize("endpoint", ["/v1/responses", "/v1/chat/completions"]) def test_bedrock_forwards_allowed_safety_identifier_as_additional_model_request_field( self, endpoints_client: EndpointsClient, resources: ResourceManager, endpoint: str diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 97acb9ec52b..c6acd449884 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -13,3 +13,4 @@ markers = cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set + provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set From 333fadad6cebe500ff68f702fcb9d9028ad2c6f2 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 21 Sep 2026 06:25:43 +0000 Subject: [PATCH 122/146] fix(helm): render a fixed replicaCount on componentized deployments when HPA is disabled Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm/templates/backend/deployment.yaml | 3 + .../litellm/templates/gateway/deployment.yaml | 3 + helm/litellm/templates/ui/deployment.yaml | 3 + helm/litellm/tests/replica_count_tests.yaml | 84 +++++++++++++++++++ helm/litellm/values.yaml | 8 ++ 5 files changed, 101 insertions(+) create mode 100644 helm/litellm/tests/replica_count_tests.yaml diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 0db2f0b3d43..6370f581d71 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -7,6 +7,9 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: backend spec: + {{- if not .Values.backend.hpa.enabled }} + replicas: {{ .Values.backend.replicaCount }} + {{- end }} {{- with .Values.backend.strategy }} strategy: {{- toYaml . | nindent 4 }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index c06cc9583a0..0c537b1bdeb 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -7,6 +7,9 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: gateway spec: + {{- if not .Values.gateway.hpa.enabled }} + replicas: {{ .Values.gateway.replicaCount }} + {{- end }} {{- with .Values.gateway.strategy }} strategy: {{- toYaml . | nindent 4 }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index b992b347bad..b794418b7e9 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -7,6 +7,9 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: ui spec: + {{- if not .Values.ui.hpa.enabled }} + replicas: {{ .Values.ui.replicaCount }} + {{- end }} {{- with .Values.ui.strategy }} strategy: {{- toYaml . | nindent 4 }} diff --git a/helm/litellm/tests/replica_count_tests.yaml b/helm/litellm/tests/replica_count_tests.yaml new file mode 100644 index 00000000000..e6a28689ff5 --- /dev/null +++ b/helm/litellm/tests/replica_count_tests.yaml @@ -0,0 +1,84 @@ +suite: test fixed replica count when HPA is disabled +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - ui/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: gateway renders replicaCount into spec.replicas when its HPA is disabled + template: gateway/deployment.yaml + set: + gateway.hpa.enabled: false + gateway.replicaCount: 3 + asserts: + - isKind: + of: Deployment + - equal: + path: spec.replicas + value: 3 + + - it: backend renders replicaCount into spec.replicas when its HPA is disabled + template: backend/deployment.yaml + set: + backend.hpa.enabled: false + backend.replicaCount: 2 + asserts: + - equal: + path: spec.replicas + value: 2 + + - it: ui renders replicaCount into spec.replicas when its HPA is disabled + template: ui/deployment.yaml + set: + ui.hpa.enabled: false + ui.replicaCount: 2 + asserts: + - equal: + path: spec.replicas + value: 2 + + - it: replicaCount 0 scales the gateway to zero instead of being treated as unset + template: gateway/deployment.yaml + set: + gateway.hpa.enabled: false + gateway.replicaCount: 0 + asserts: + - equal: + path: spec.replicas + value: 0 + + - it: every component omits spec.replicas when its HPA is enabled, so the autoscaler owns the count + set: + gateway.hpa.enabled: true + gateway.replicaCount: 3 + backend.hpa.enabled: true + backend.replicaCount: 3 + ui.hpa.enabled: true + ui.replicaCount: 3 + asserts: + - notExists: + path: spec.replicas + template: gateway/deployment.yaml + - notExists: + path: spec.replicas + template: backend/deployment.yaml + - notExists: + path: spec.replicas + template: ui/deployment.yaml + + - it: a component with HPA disabled renders replicas while a sibling with HPA enabled does not + set: + gateway.hpa.enabled: false + gateway.replicaCount: 4 + backend.hpa.enabled: true + backend.replicaCount: 4 + asserts: + - equal: + path: spec.replicas + value: 4 + template: gateway/deployment.yaml + - notExists: + path: spec.replicas + template: backend/deployment.yaml diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 4ca54131d6a..d7836b03a34 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -397,6 +397,10 @@ gateway: # failureThreshold: 30 # periodSeconds: 10 startupProbe: {} + # Fixed pod count, rendered into the Deployment's spec.replicas only when + # hpa.enabled is false. With the HPA on, the autoscaler owns the count and + # this value is ignored. + replicaCount: 1 hpa: enabled: true minReplicas: 1 @@ -524,6 +528,8 @@ backend: strategy: {} # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} + # Same semantics as gateway.replicaCount. + replicaCount: 1 hpa: enabled: true minReplicas: 1 @@ -590,6 +596,8 @@ ui: strategy: {} # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} + # Same semantics as gateway.replicaCount. + replicaCount: 1 hpa: enabled: false minReplicas: 1 From d266d7324b49099ba782a90669cf8d6d1d2d1e54 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 21 Sep 2026 06:35:05 +0000 Subject: [PATCH 123/146] fix(helm): leave spec.replicas unset unless replicaCount is explicitly configured Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- helm/litellm/templates/backend/deployment.yaml | 2 +- helm/litellm/templates/gateway/deployment.yaml | 2 +- helm/litellm/templates/ui/deployment.yaml | 2 +- helm/litellm/tests/replica_count_tests.yaml | 16 ++++++++++++++++ helm/litellm/values.yaml | 13 +++++++------ 5 files changed, 26 insertions(+), 9 deletions(-) diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 6370f581d71..3eb64e5528c 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -7,7 +7,7 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: backend spec: - {{- if not .Values.backend.hpa.enabled }} + {{- if and (not .Values.backend.hpa.enabled) (not (kindIs "invalid" .Values.backend.replicaCount)) }} replicas: {{ .Values.backend.replicaCount }} {{- end }} {{- with .Values.backend.strategy }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 0c537b1bdeb..49b452b3053 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -7,7 +7,7 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: gateway spec: - {{- if not .Values.gateway.hpa.enabled }} + {{- if and (not .Values.gateway.hpa.enabled) (not (kindIs "invalid" .Values.gateway.replicaCount)) }} replicas: {{ .Values.gateway.replicaCount }} {{- end }} {{- with .Values.gateway.strategy }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index b794418b7e9..efee2d5fc34 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -7,7 +7,7 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: ui spec: - {{- if not .Values.ui.hpa.enabled }} + {{- if and (not .Values.ui.hpa.enabled) (not (kindIs "invalid" .Values.ui.replicaCount)) }} replicas: {{ .Values.ui.replicaCount }} {{- end }} {{- with .Values.ui.strategy }} diff --git a/helm/litellm/tests/replica_count_tests.yaml b/helm/litellm/tests/replica_count_tests.yaml index e6a28689ff5..791e47ff798 100644 --- a/helm/litellm/tests/replica_count_tests.yaml +++ b/helm/litellm/tests/replica_count_tests.yaml @@ -49,6 +49,22 @@ tests: path: spec.replicas value: 0 + - it: a component with HPA disabled but no replicaCount set keeps omitting spec.replicas, so upgrades do not reset a hand-scaled Deployment + set: + gateway.hpa.enabled: false + backend.hpa.enabled: false + ui.hpa.enabled: false + asserts: + - notExists: + path: spec.replicas + template: gateway/deployment.yaml + - notExists: + path: spec.replicas + template: backend/deployment.yaml + - notExists: + path: spec.replicas + template: ui/deployment.yaml + - it: every component omits spec.replicas when its HPA is enabled, so the autoscaler owns the count set: gateway.hpa.enabled: true diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index d7836b03a34..2c0c7151a32 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -397,10 +397,11 @@ gateway: # failureThreshold: 30 # periodSeconds: 10 startupProbe: {} - # Fixed pod count, rendered into the Deployment's spec.replicas only when - # hpa.enabled is false. With the HPA on, the autoscaler owns the count and - # this value is ignored. - replicaCount: 1 + # Optional fixed pod count, rendered into the Deployment's spec.replicas only + # when hpa.enabled is false. Unset by default so an existing Deployment keeps + # its current count; with the HPA on, the autoscaler owns the count, e.g.: + # replicaCount: 3 + replicaCount: hpa: enabled: true minReplicas: 1 @@ -529,7 +530,7 @@ backend: # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} # Same semantics as gateway.replicaCount. - replicaCount: 1 + replicaCount: hpa: enabled: true minReplicas: 1 @@ -597,7 +598,7 @@ ui: # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} # Same semantics as gateway.replicaCount. - replicaCount: 1 + replicaCount: hpa: enabled: false minReplicas: 1 From 0317903a440117c0ffd4faba3639049147a2645e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:16:02 +0000 Subject: [PATCH 124/146] ci(e2e-changed): surface failed test ids from the pytest log Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-e2e-changed.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index f659afad7e3..6da16a33ea3 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -176,6 +176,7 @@ jobs: TESTS: ${{ needs.detect.outputs.tests }} E2E_FIXTURE_MODE: live E2E_PROVIDER_EDGE_HOST_REACHABLE: '1' + COLUMNS: '400' run: | umask 077 read -r -a test_files <<< "${TESTS}" @@ -190,6 +191,7 @@ jobs: uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}" verified=$? set -e + grep -E '^(FAILED|ERROR) ' "${log}" || true grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1 echo "::endgroup::" if [ "${status}" = "5" ]; then From b38504b5a6f646c7d4eaedb1847b32b4b18647b9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:37:28 +0000 Subject: [PATCH 125/146] test(e2e): assert the forwarded Converse body without requiring the model to accept safety_identifier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../e2e/llm_translation/test_responses_e2e.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 1c2336bd166..525231de917 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -19,7 +19,6 @@ from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, un from e2e_http import ( assert_client_error, require_successful_call, - unwrap, ) from endpoints_client import ( EndpointsClient, @@ -359,20 +358,15 @@ class TestResponses: safety_identifier: Final = f"end-user-{unique_marker()}" if endpoint == "/v1/responses": - responses_result: Final = endpoints_client.responses( - key, model, "reply with one word", safety_identifier=safety_identifier - ) - require_successful_call(responses_result) + endpoints_client.responses(key, model, "reply with one word", safety_identifier=safety_identifier) else: - unwrap( - endpoints_client.proxy.chat( - key, - ChatBody( - model=model, - messages=[ChatMessage(role="user", content="reply with one word")], - safety_identifier=safety_identifier, - ), - ) + endpoints_client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="reply with one word")], + safety_identifier=safety_identifier, + ), ) forwarded: Final = tuple(body.additionalModelRequestFields for body in capture.bodies) From 8795be0a65754221bee0b318ad355843d5e3ca34 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 10:51:46 +0000 Subject: [PATCH 126/146] refactor(types): replace Any with proven types in 32 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/_logging.py | 8 +++++--- litellm/a2a_protocol/streaming_iterator.py | 12 ++++++------ litellm/integrations/focus/focus_logger.py | 4 +++- .../generic_prompt_manager.py | 2 +- .../mavvrik_focus/mavvrik_focus_logger.py | 17 ++++++++++++++--- litellm/integrations/otel/presets/agentops.py | 7 +++++-- litellm/integrations/otel/runtime.py | 13 ++++++++----- .../bounded_prometheus_series_tracker.py | 16 +++++++++++----- .../vector_store_pre_call_hook.py | 2 +- litellm/integrations/weights_biases.py | 9 +++++---- .../usage_object_transformation.py | 6 +++--- litellm/litellm_core_utils/safe_json_dumps.py | 2 +- litellm/llms/azure/fine_tuning/handler.py | 10 +++++----- .../llms/codestral/completion/transformation.py | 2 +- litellm/llms/databricks/common_utils.py | 4 ++-- litellm/llms/gemini/realtime/transformation.py | 4 ++-- .../llms/jina_ai/embedding/transformation.py | 2 +- .../llms/openrouter/embedding/transformation.py | 4 +++- litellm/llms/reducto/common.py | 4 +++- .../embedding/transformation.py | 3 ++- .../vertex_ai/agent_engine/transformation.py | 2 +- litellm/proxy/a2a/discovery.py | 7 ++++--- .../proxy/agent_endpoints/databricks_oauth.py | 11 ++++++----- litellm/proxy/client/cli/main.py | 13 +++++++++++-- .../generic_guardrail_api/__init__.py | 2 +- .../guardrails/guardrail_hooks/onyx/onyx.py | 4 ++-- litellm/proxy/hooks/responses_id_security.py | 6 +++--- .../callback_logs_endpoints.py | 5 +++-- .../management_v1/spend_logs.py | 8 ++++---- .../object_permission_utils.py | 4 ++-- .../gemini_passthrough_logging_handler.py | 2 +- .../proxy/public_endpoints/public_endpoints.py | 10 +++++++--- 32 files changed, 127 insertions(+), 78 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 5ba0c080364..644a79d8cbd 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -5,6 +5,7 @@ import logging import os import re import sys +from collections.abc import Sequence from datetime import datetime from logging import Formatter from typing import Any, Final, TextIO @@ -186,7 +187,8 @@ class SecretRedactionFilter(logging.Filter): record.stack_info = _redact_string(record.stack_info) # rebind-ok: a Filter scrubs records in place # Redact extra fields passed via logger.debug("msg", extra={...}) - for key, value in list(record.__dict__.items()): + record_items: Final[Sequence[tuple[str, object]]] = list(record.__dict__.items()) + for key, value in record_items: if key in _STANDARD_RECORD_ATTRS: continue if isinstance(value, str): @@ -507,7 +509,7 @@ handler.addFilter(_secret_filter) handler.addFilter(_correlation_filter) -def _try_parse_json_message(message: str) -> dict[str, Any] | None: +def _try_parse_json_message(message: str) -> dict[str, object] | None: """ Try to parse a log message as JSON. Returns parsed dict if valid, else None. Handles messages that are entirely valid JSON (e.g. json.dumps output). @@ -585,7 +587,7 @@ class JsonFormatter(Formatter): def format(self, record): message_str: Final = record.getMessage() - json_record: Final[dict[str, Any]] = { + json_record: Final[dict[str, object]] = { "message": message_str, "level": record.levelname, "timestamp": self.formatTime(record), diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index d936caeb75e..8232d7cf2d8 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -5,7 +5,7 @@ A2A Streaming Iterator with token tracking and logging support. import asyncio from collections.abc import AsyncIterator from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import litellm from litellm._logging import verbose_logger @@ -15,7 +15,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj if TYPE_CHECKING: - from a2a.types import SendStreamingMessageRequest, SendStreamingMessageResponse + from a2a.compat.v0_3.types import SendStreamingMessageRequest, SendStreamingMessageResponse class A2AStreamingIterator: @@ -39,9 +39,9 @@ class A2AStreamingIterator: self.start_time = datetime.now() # Collect chunks for token counting - self.chunks: list[Any] = [] + self.chunks: list[SendStreamingMessageResponse] = [] self.collected_text_parts: list[str] = [] - self.final_chunk: Any | None = None + self.final_chunk: SendStreamingMessageResponse | None = None def __aiter__(self): return self @@ -69,7 +69,7 @@ class A2AStreamingIterator: await self._handle_stream_complete() raise - def _collect_text_from_chunk(self, chunk: Any) -> None: + def _collect_text_from_chunk(self, chunk: "SendStreamingMessageResponse") -> None: """Extract text from a streaming chunk and add to collected parts.""" try: chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} @@ -79,7 +79,7 @@ class A2AStreamingIterator: except Exception: verbose_logger.debug("Failed to extract text from A2A streaming chunk") - def _is_completed_chunk(self, chunk: Any) -> bool: + def _is_completed_chunk(self, chunk: "SendStreamingMessageResponse") -> bool: """Check if chunk indicates stream completion.""" try: chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index c9b47835948..dce51b8190b 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -15,6 +15,8 @@ from .destinations import FocusTimeWindow if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from .export_engine import FocusExportEngine else: AsyncIOScheduler = Any @@ -111,7 +113,7 @@ class FocusLogger(CustomLogger): """Entry point for scheduler jobs to run export cycle with locking.""" from litellm.proxy.proxy_server import proxy_logging_obj - pod_lock_manager = None + pod_lock_manager: PodLockManager | None = None if proxy_logging_obj is not None: writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index 77d315d0cee..de01b2bb02c 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -58,7 +58,7 @@ class GenericPromptManager(CustomPromptManagement): api_key: str | None = None, timeout: int = 30, prompt_id: str | None = None, - additional_provider_specific_query_params: dict[str, Any] | None = None, + additional_provider_specific_query_params: Mapping[str, object] | None = None, **kwargs, ): """ diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 3c189b4d53e..7e3c4cc3ce8 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -21,7 +21,7 @@ from __future__ import annotations import os from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import litellm from litellm._logging import verbose_proxy_logger @@ -35,6 +35,17 @@ else: AsyncIOScheduler = Any +class _PodLockManager(Protocol): + """The subset of PodLockManager this logger drives to serialize the export across pods.""" + + @property + def redis_cache(self) -> object: ... + + async def acquire_lock(self, cronjob_id: str) -> bool | None: ... + + async def release_lock(self, cronjob_id: str) -> None: ... + + def _parse_metrics_marker( marker: object | None, ) -> datetime | None: @@ -226,9 +237,9 @@ class MavvrikFocusLogger(FocusLogger): """Scheduler entry point — uses Mavvrik-specific pod-lock key.""" from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415 - pod_lock_manager = None + pod_lock_manager: _PodLockManager | None = None if proxy_logging_obj is not None: - writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) + writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: pod_lock_manager = getattr(writer, "pod_lock_manager", None) diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 965213f2ee4..58123656caa 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,9 +9,12 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ +from collections.abc import Sequence from typing import Any, Final import httpx +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -71,7 +74,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> Any: +def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -106,7 +109,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> Any: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Any) -> Any: + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: self._ensure_authenticated() return super().export(spans) diff --git a/litellm/integrations/otel/runtime.py b/litellm/integrations/otel/runtime.py index c6eaecd108b..13903597e1a 100644 --- a/litellm/integrations/otel/runtime.py +++ b/litellm/integrations/otel/runtime.py @@ -8,13 +8,16 @@ identity unconditionally. """ from collections.abc import Callable, Iterator -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager from functools import cache -from typing import Any, Final +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from opentelemetry.trace import Span @cache -def _otel_runtime() -> "tuple[Callable[[str], Any], Callable[..., None]] | None": +def _otel_runtime() -> "tuple[Callable[[str], AbstractContextManager[Span | None]], Callable[..., None]] | None": """Resolve the SDK-backed hooks once and cache the outcome, absence included. CPython never caches a failed import, so without this memoization every call @@ -29,7 +32,7 @@ def _otel_runtime() -> "tuple[Callable[[str], Any], Callable[..., None]] | None" @contextmanager -def phase_span(name: str) -> "Iterator[Any]": +def phase_span(name: str) -> "Iterator[Span | None]": """Run a request phase inside a live active span so its DB/service calls nest. Yields ``None`` (a plain no-op) when the OTel SDK is unavailable or V2 is not @@ -43,7 +46,7 @@ def phase_span(name: str) -> "Iterator[Any]": yield span -def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None: +def seed_request_identity(user_api_key_dict: object, model: object = None) -> None: """Seed request-identity Baggage at the auth boundary (no-op without V2).""" runtime: Final = _otel_runtime() if runtime is None: diff --git a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py index c1ccf09d5d6..ba7d54fafea 100644 --- a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py +++ b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py @@ -3,7 +3,13 @@ from __future__ import annotations import time from collections import OrderedDict from threading import RLock -from typing import Any, Final +from typing import Final, Protocol + + +class _RemovableMetric(Protocol): + """The one prometheus-client metric method this tracker calls.""" + + def remove(self, *labelvalues: object) -> None: ... class BoundedPrometheusSeriesTracker: @@ -21,7 +27,7 @@ class BoundedPrometheusSeriesTracker: def track_series( self, - metric: Any, + metric: _RemovableMetric, metric_name: str, label_values: tuple[str | None, ...], max_series: int | None, @@ -60,7 +66,7 @@ class BoundedPrometheusSeriesTracker: break del series[tracked_label_values] - def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool: + def remove_series(self, metric: _RemovableMetric, label_values: tuple[str | None, ...]) -> bool: """Drop one child series, True when it is gone (removed or never existed).""" return self._remove_metric_child(metric, label_values) @@ -82,7 +88,7 @@ class BoundedPrometheusSeriesTracker: def _remove_metric_series( self, - metric: Any, + metric: _RemovableMetric, series: OrderedDict[tuple[str | None, ...], float], label_values: tuple[str | None, ...], ) -> None: @@ -90,7 +96,7 @@ class BoundedPrometheusSeriesTracker: series.pop(label_values, None) @staticmethod - def _remove_metric_child(metric: Any, label_values: tuple[str | None, ...]) -> bool: + def _remove_metric_child(metric: _RemovableMetric, label_values: tuple[str | None, ...]) -> bool: """ Remove the Prometheus child for ``label_values`` and report whether the tracker should commit the matching state change. diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index b2243060c6c..74fb8a8d6a3 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -406,7 +406,7 @@ class VectorStorePreCallHook(CustomLogger): request_data: dict, response_chunk: Any, call_type: CallTypes | None, - ) -> Any | None: + ) -> object | None: """ Add search results to the final streaming chunk. diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index d1a8ec098cf..9bc070a1f9a 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -4,6 +4,7 @@ imported_openAIResponse = True try: import io import logging + from collections.abc import Mapping from typing import Any, Literal, Protocol, TypeVar from wandb.sdk.data_types import trace_tree @@ -43,7 +44,7 @@ try: @staticmethod def results_to_trace_tree( - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, results: list[trace_tree.Result], time_elapsed: float, @@ -73,7 +74,7 @@ try: def _resolve_edit( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree: @@ -91,7 +92,7 @@ try: def _resolve_completion( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree: @@ -134,7 +135,7 @@ try: def _request_response_result_to_trace( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, request_str: str, choices: list[str], diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index f11f6d46fb2..a02c40b7611 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -50,7 +50,7 @@ _INTERACTIONS_MODALITY_FIELDS: Final[Mapping[str, str]] = MappingProxyType( ) -def _modality_field(entry: Mapping[str, Any]) -> str | None: +def _modality_field(entry: Mapping[str, object]) -> str | None: return _INTERACTIONS_MODALITY_FIELDS.get(str(entry.get("modality", "")).lower()) @@ -58,7 +58,7 @@ def _token_count(value: object) -> int: return value if isinstance(value, int) else 0 -def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, int]: +def _modality_token_sums(entries: Sequence[Mapping[str, object]]) -> Mapping[str, int]: fields: Final = frozenset(field for entry in entries if (field := _modality_field(entry)) is not None) return MappingProxyType( { @@ -68,7 +68,7 @@ def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, i ) -def _google_search_query_count(usage_object: Mapping[str, Any]) -> int: +def _google_search_query_count(usage_object: Mapping[str, object]) -> int: entries: Final = usage_object.get("grounding_tool_count") if not isinstance(entries, Sequence): return 0 diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 4f9ac82d57d..63242a580e7 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -85,7 +85,7 @@ def safe_json_structure( def safe_dumps( - data: Any, + data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, value_transform: Callable[[str | None, str], str] | None = None, ) -> str: diff --git a/litellm/llms/azure/fine_tuning/handler.py b/litellm/llms/azure/fine_tuning/handler.py index ac1e430e063..36c4fae04c7 100644 --- a/litellm/llms/azure/fine_tuning/handler.py +++ b/litellm/llms/azure/fine_tuning/handler.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine -from typing import Any, Final, cast +from typing import Final, cast import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI @@ -19,7 +19,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): """ @staticmethod - def _ensure_training_type(create_fine_tuning_job_data: dict[str, Any]) -> None: + def _ensure_training_type(create_fine_tuning_job_data: dict[str, object]) -> None: """ Azure requires trainingType in extra_body. Default to 1 (supervised) if omitted. """ @@ -66,7 +66,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: self._ensure_training_type(create_fine_tuning_job_data) openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( @@ -109,7 +109,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -149,7 +149,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, diff --git a/litellm/llms/codestral/completion/transformation.py b/litellm/llms/codestral/completion/transformation.py index e3c3fd1231c..baa134bb398 100644 --- a/litellm/llms/codestral/completion/transformation.py +++ b/litellm/llms/codestral/completion/transformation.py @@ -29,7 +29,7 @@ class CodestralTextCompletionConfig(OpenAITextCompletionConfig): random_seed: int | None = None, stop: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[dict[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index a4ec2c5378b..f2f9df422e8 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -12,7 +12,7 @@ Authentication priority: import os import re -from typing import Any, Final, Literal +from typing import Final, Literal from urllib.parse import urlsplit, urlunsplit from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -48,7 +48,7 @@ class DatabricksBase: ] @classmethod - def redact_sensitive_data(cls, data: Any) -> Any: + def redact_sensitive_data(cls, data: object) -> object: """ Redact sensitive information (tokens, secrets) from data before logging. diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 79985569c5f..e64cbf88d95 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -453,7 +453,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return normalized @staticmethod - def _finalize_gemini_live_setup(model: str, setup: dict[str, Any]) -> dict[str, Any]: + def _finalize_gemini_live_setup(model: str, setup: dict[str, object]) -> dict[str, object]: generation_config: Final = setup.get("generationConfig") if isinstance(generation_config, dict): modalities: Final = generation_config.get("responseModalities") @@ -1172,7 +1172,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): def map_openai_event( self, key: str, - value: Any, + value: object, current_delta_type: ALL_DELTA_TYPES | None, ) -> OpenAIRealtimeEventTypes | ResponsesAPIStreamEvents: if isinstance(value, dict): diff --git a/litellm/llms/jina_ai/embedding/transformation.py b/litellm/llms/jina_ai/embedding/transformation.py index 26f512979e5..260d9e6e494 100644 --- a/litellm/llms/jina_ai/embedding/transformation.py +++ b/litellm/llms/jina_ai/embedding/transformation.py @@ -31,7 +31,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): def __init__( self, ) -> None: - locals_: Final = locals().copy() + locals_: Final[dict[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 29d0c8c1c56..14b0e462ea7 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -170,7 +170,9 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> OpenRouterException: """ Get the error class for OpenRouter errors. """ diff --git a/litellm/llms/reducto/common.py b/litellm/llms/reducto/common.py index 9b9efd24b72..b194590fdb9 100644 --- a/litellm/llms/reducto/common.py +++ b/litellm/llms/reducto/common.py @@ -3,6 +3,8 @@ import binascii from collections import defaultdict from typing import TYPE_CHECKING, Any, Final, NoReturn +import httpx + from litellm.constants import request_timeout REDUCTO_API_BASE: Final = "https://platform.reducto.ai" @@ -62,7 +64,7 @@ def extract_file_id_or_bytes( return None, raw_bytes, mime -def _extract_file_id_from_upload_response(response: Any) -> str: +def _extract_file_id_from_upload_response(response: httpx.Response) -> str: try: payload: Final = response.json() except ValueError as exc: diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index 3f228c0881d..fc9c6bcc19f 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllEmbeddingInputValues @@ -160,7 +161,7 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: + def get_error_class(self, error_message: str, status_code: int, headers: Any) -> BaseLLMException: """ Get the error class for Vercel AI Gateway errors. """ diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index e430d9e2280..c5ca9f38144 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -205,7 +205,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): session_id: Final = self._get_session_id(optional_params) # Build the input - input_data: Final[dict[str, Any]] = { + input_data: Final[dict[str, str]] = { "message": prompt, "user_id": user_id, } diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index e08c938f195..ff58c9c85ec 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -14,6 +14,7 @@ fetcher dispatches by ``discovery_mode``: pure-A2A fallback strategy returns 404 for these deployments. """ +from collections.abc import Mapping from enum import Enum from typing import Any, Final from urllib.parse import urlencode @@ -55,7 +56,7 @@ def _normalize_base_url(base_url: str) -> str: def _build_langgraph_platform_paths( - params: dict[str, Any] | None, + params: Mapping[str, object] | None, ) -> tuple[str, ...]: """Build the paths to try for LangGraph Platform discovery. @@ -71,7 +72,7 @@ def _build_langgraph_platform_paths( return tuple(f"{path}?{query}" for path in AGENT_CARD_WELL_KNOWN_PATHS) -def _paths_for_mode(mode: DiscoveryMode, params: dict[str, Any] | None) -> tuple[str, ...]: +def _paths_for_mode(mode: DiscoveryMode, params: Mapping[str, object] | None) -> tuple[str, ...]: if mode == DiscoveryMode.WELL_KNOWN_FALLBACK: return AGENT_CARD_WELL_KNOWN_PATHS if mode == DiscoveryMode.LANGGRAPH_PLATFORM: @@ -83,7 +84,7 @@ async def fetch_well_known_card( base_url: str, *, discovery_mode: DiscoveryMode = DiscoveryMode.WELL_KNOWN_FALLBACK, - params: dict[str, Any] | None = None, + params: Mapping[str, object] | None = None, timeout: float = DEFAULT_DISCOVERY_TIMEOUT_SECONDS, headers: dict[str, str] | None = None, ) -> dict[str, Any]: diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 4c3b1bc084d..38a76ea6890 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -25,8 +25,9 @@ Config example:: import asyncio import base64 import hashlib +from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Final +from typing import Final import httpx @@ -43,7 +44,7 @@ _TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60 _DEFAULT_TTL_SECONDS: Final = 3600 -def _resolve_secret(value: Any) -> str | None: +def _resolve_secret(value: object) -> str | None: """Resolve a config value, expanding ``os.environ/`` references.""" if not isinstance(value, str): return None @@ -75,7 +76,7 @@ class DatabricksAppOAuthConfig: def parse_databricks_oauth_config( - litellm_params: dict[str, Any] | None, + litellm_params: Mapping[str, object] | None, ) -> DatabricksAppOAuthConfig | None: """Build a Databricks App OAuth config from an agent's ``litellm_params``. @@ -191,7 +192,7 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): except httpx.HTTPError as exc: raise ValueError(f"Databricks App OAuth token request failed: {exc}") from exc - body: Final = response.json() + body: Final[object] = response.json() if not isinstance(body, dict): raise ValueError( f"Databricks App OAuth token response returned non-object JSON (got {type(body).__name__})" @@ -215,7 +216,7 @@ databricks_app_oauth_token_cache: Final = DatabricksAppOAuthTokenCache() async def resolve_databricks_app_auth_header( - litellm_params: dict[str, Any] | None, + litellm_params: Mapping[str, object] | None, ) -> dict[str, str] | None: """Return ``{"Authorization": "Bearer "}`` for a Databricks App agent. diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 63e38c93221..6d63acc7479 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -9,7 +9,15 @@ from litellm._version import version as litellm_version from litellm.proxy.client.health import HealthManagementClient from .commands.agents import agent_commands -from .commands.auth import auth_group, context_secret_vault, get_stored_api_key, login, logout, whoami +from .commands.auth import ( + CliContextObj, + auth_group, + context_secret_vault, + get_stored_api_key, + login, + logout, + whoami, +) from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.config import config_commands, get_config_value, hidden_command_names @@ -126,7 +134,8 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s @click.pass_context def version(ctx: click.Context): """Show the LiteLLM Proxy CLI and server version.""" - print_version(ctx.obj.get("base_url"), ctx.obj.get("api_key")) + ctx_obj: Final[CliContextObj] = ctx.obj + print_version(ctx_obj.get("base_url"), ctx_obj.get("api_key")) # Add authentication commands as top-level commands diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index d1576b68813..e3511d46544 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def _get_config_value(litellm_params: Any, optional_params: Any, attribute_name: str) -> Any | None: +def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> Any | None: if optional_params is not None: value: Final = ( optional_params.get(attribute_name) diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 7529c4ce3f3..1a6feb47215 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -6,7 +6,7 @@ # +-------------------------------------------------------------+ import os import uuid -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional import httpx from fastapi import HTTPException @@ -63,7 +63,7 @@ class OnyxGuardrail(CustomGuardrail): async def _validate_with_guard_server( self, - payload: Any, + payload: object, input_type: Literal["request", "response"], conversation_id: str, ) -> dict: diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index d9050489095..bdf7e2ab53d 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -40,7 +40,7 @@ _UNMANAGED_RESPONSE_ID_DETAIL: Final = ( _PROXY_ADMIN_ROLES: Final = frozenset({LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value}) -def _proxy_general_settings() -> Mapping[str, Any]: +def _proxy_general_settings() -> Mapping[str, object]: from litellm.proxy.proxy_server import general_settings return general_settings @@ -107,7 +107,7 @@ def _is_responses_api_create_route(request_route: str | None) -> bool: class ResponsesIDSecurity(CustomLogger): def __init__( self, - general_settings_reader: Callable[[], Mapping[str, Any]] = _proxy_general_settings, + general_settings_reader: Callable[[], Mapping[str, object]] = _proxy_general_settings, signing_key_reader: Callable[[], str | None] = _proxy_signing_key, ) -> None: self._general_settings_reader: Final = general_settings_reader @@ -307,7 +307,7 @@ class ResponsesIDSecurity(CustomLogger): data: dict, user_api_key_dict: "UserAPIKeyAuth", response: LLMResponseTypes, - ) -> Any: + ) -> LLMResponseTypes: """ Queue response IDs for batch processing instead of writing directly to DB. diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index cecadc03d71..4a1079871b0 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -15,6 +15,7 @@ self-describing `StandardLoggingPayload`, so completions/responses can use it to """ import uuid +from collections.abc import Mapping from datetime import datetime, timezone from typing import Any, Final @@ -48,7 +49,7 @@ class CallbackLogsReplayer: """ @staticmethod - def _epoch_to_datetime(value: Any) -> datetime: + def _epoch_to_datetime(value: object) -> datetime: """`StandardLoggingPayload` stores startTime/endTime as float epoch seconds.""" if isinstance(value, (int, float)): return datetime.fromtimestamp(float(value), tz=timezone.utc) @@ -114,7 +115,7 @@ class CallbackLogsReplayer: return logging_obj @staticmethod - def _response_obj_from_payload(payload: dict[str, Any]) -> dict[str, Any]: + def _response_obj_from_payload(payload: Mapping[str, object]) -> dict[str, object]: """Minimal response object so usage-derived spend-log fields resolve.""" return { "id": payload.get("id"), diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index f6907a7f87a..1cbc454ca5e 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,7 +1,7 @@ """`/management/v1/spend_logs` facets.""" from datetime import datetime, timezone -from typing import Annotated, Any, Final, Literal +from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, Query, Request @@ -39,7 +39,7 @@ async def _spend_log_scope_clause( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, next_param_index: int, -) -> tuple[str | None, tuple[Any, ...]]: +) -> tuple[str | None, tuple[str | list[str], ...]]: """SQL predicate restricting the facet to spend logs this caller may read. Returns ``(None, ())`` for a proxy admin. Mirrors the scoping ``/spend/logs/ui`` @@ -101,8 +101,8 @@ async def _list_spend_log_facet( ) column_sql: Final = "end_user" if column == "end_user" else '"user"' - window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time)) - search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else () + window_params: Final[tuple[datetime, datetime]] = (_as_utc(start_time), _as_utc(end_time)) + search_params: Final[tuple[str, ...]] = (f"%{escape_like(q)}%",) if q else () search_clause: Final = (f"{column_sql} ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () scope_clause, scope_params = await _spend_log_scope_clause( diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index daab38d3662..437e6763502 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -230,7 +230,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index a95ee87fd31..d97ddb9a909 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -147,7 +147,7 @@ class GeminiPassthroughLoggingHandler: - Creates standard logging object - Logs in litellm callbacks """ - kwargs: dict[str, Any] = {} + kwargs: dict[str, object] = {} model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) complete_streaming_response: Final = GeminiPassthroughLoggingHandler._build_complete_streaming_response( all_chunks=all_chunks, diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index e395f56194f..26a5c44fce1 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -199,9 +199,13 @@ def _build_endpoints(raw: _ProvidersFile) -> list[_EndpointEntry]: return result +_PROVIDERS_FILE_ADAPTER: Final = TypeAdapter(_ProvidersFile) +_PROVIDER_CREATE_FIELDS_ADAPTER: Final = TypeAdapter(list[ProviderCreateInfo]) + + def _load_endpoints() -> list[_EndpointEntry]: - raw: Final[_ProvidersFile] = json.loads( - files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8") + raw: Final = _PROVIDERS_FILE_ADAPTER.validate_python( + json.loads(files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8")) ) return _build_endpoints(raw) @@ -398,7 +402,7 @@ async def get_provider_fields() -> list[ProviderCreateInfo]: ) with open(provider_create_fields_path, "r") as f: - provider_create_fields: Final = json.load(f) + provider_create_fields: Final = _PROVIDER_CREATE_FIELDS_ADAPTER.validate_python(json.load(f)) return provider_create_fields From 85e72a2e2fac84e9be99cbac4ede9d4991b834af Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:00:50 +0000 Subject: [PATCH 127/146] chore(prices): sync OpenRouter prices: 2 models openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1976437f900..d3053dee25d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.5526e-07, + "input_cost_per_token": 9.53172e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.91052e-06, + "output_cost_per_token": 1.906344e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9605e-08, + "cache_read_input_token_cost": 7.9431e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 5.852e-08, + "output_cost_per_token": 1.1704e-07, + "cache_read_input_token_cost": 1.1704e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1976437f900..d3053dee25d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.5526e-07, + "input_cost_per_token": 9.53172e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.91052e-06, + "output_cost_per_token": 1.906344e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9605e-08, + "cache_read_input_token_cost": 7.9431e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 5.852e-08, + "output_cost_per_token": 1.1704e-07, + "cache_read_input_token_cost": 1.1704e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From d9a97d74db83d791276e8b8099412a1d87e930bf Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:21:50 +0000 Subject: [PATCH 128/146] fix(model_prices): add groq qwen3.6-27b deprecation date and bedrock qwen3-next regional pricing Groq lists qwen/qwen3.6-27b for shutdown on 2026-09-14. Adds the six regional Bedrock qwen.qwen3-next-80b-a3b entries priced per AWS's published regional rates (absorbs #42191) with a regression test that the regional entry is used instead of the US rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 79 +++++++++++++++++++ model_prices_and_context_window.json | 79 +++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 27 +++++++ 3 files changed, 185 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 176df53f706..856396a1d75 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -35294,6 +35294,7 @@ "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://console.groq.com/docs/model/qwen/qwen3.6-27b", + "deprecation_date": "2026-09-14", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": false, @@ -44299,6 +44300,84 @@ "supports_system_messages": true, "supports_native_structured_output": true }, + "bedrock/ap-northeast-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.45e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/ap-south-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.41e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/ap-southeast-2/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.545e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.236e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/eu-west-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.41e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/eu-west-2/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 2.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.86e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/sa-east-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.45e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, "qwen.qwen3-vl-235b-a22b": { "input_cost_per_token": 5.3e-07, "litellm_provider": "bedrock_converse", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 176df53f706..856396a1d75 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -35294,6 +35294,7 @@ "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://console.groq.com/docs/model/qwen/qwen3.6-27b", + "deprecation_date": "2026-09-14", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": false, @@ -44299,6 +44300,84 @@ "supports_system_messages": true, "supports_native_structured_output": true }, + "bedrock/ap-northeast-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.45e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/ap-south-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.41e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/ap-southeast-2/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.545e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.236e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/eu-west-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.41e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/eu-west-2/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 2.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.86e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, + "bedrock/sa-east-1/qwen.qwen3-next-80b-a3b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.45e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true + }, "qwen.qwen3-vl-235b-a22b": { "input_cost_per_token": 5.3e-07, "litellm_provider": "bedrock_converse", diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index aef17f3d5d0..da3d022d669 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4230,3 +4230,30 @@ def test_completion_cost_prices_responses_websocket_turns_per_service_tier(): assert ws_cost == pytest.approx(_http_cost(100, 40, "default") + _http_cost(60, 10, "priority")) assert ws_cost != pytest.approx(_http_cost(160, 50, "default")) assert ws_cost != pytest.approx(_http_cost(160, 50, "priority")) + + +QWEN3_NEXT_REGIONS: Final = ("ap-northeast-1", "ap-south-1", "ap-southeast-2", "eu-west-1", "eu-west-2", "sa-east-1") + + +@pytest.mark.parametrize("region", QWEN3_NEXT_REGIONS) +def test_cost_per_token_bedrock_qwen3_next_uses_regional_entry_not_us_rate( + monkeypatch: pytest.MonkeyPatch, region: str +) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + regional: Final = litellm.model_cost[f"bedrock/{region}/qwen.qwen3-next-80b-a3b"] + us: Final = litellm.model_cost["qwen.qwen3-next-80b-a3b"] + assert regional["input_cost_per_token"] != us["input_cost_per_token"] + assert regional["output_cost_per_token"] != us["output_cost_per_token"] + + prompt_tokens, completion_tokens = 1000, 500 + prompt_usd, completion_usd = cost_per_token( + model=f"bedrock/{region}/qwen.qwen3-next-80b-a3b", + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + custom_llm_provider="bedrock", + ) + + assert prompt_usd == pytest.approx(prompt_tokens * regional["input_cost_per_token"]) + assert completion_usd == pytest.approx(completion_tokens * regional["output_cost_per_token"]) From 2888b4f5f4d74e246a11d9fda9f01cea9a6309b2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:30:35 +0000 Subject: [PATCH 129/146] fix(bedrock): whitelist regional qwen3-next keys for converse routing check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/local_testing/whitelisted_bedrock_models.txt | 6 ++++++ whitelisted_bedrock_models.txt | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/tests/local_testing/whitelisted_bedrock_models.txt b/tests/local_testing/whitelisted_bedrock_models.txt index 762d655b886..7615a540b23 100644 --- a/tests/local_testing/whitelisted_bedrock_models.txt +++ b/tests/local_testing/whitelisted_bedrock_models.txt @@ -133,3 +133,9 @@ meta.llama3-2-11b-instruct-v1:0 us.meta.llama3-2-11b-instruct-v1:0 meta.llama3-2-90b-instruct-v1:0 us.meta.llama3-2-90b-instruct-v1:0 +bedrock/ap-northeast-1/qwen.qwen3-next-80b-a3b +bedrock/ap-south-1/qwen.qwen3-next-80b-a3b +bedrock/ap-southeast-2/qwen.qwen3-next-80b-a3b +bedrock/eu-west-1/qwen.qwen3-next-80b-a3b +bedrock/eu-west-2/qwen.qwen3-next-80b-a3b +bedrock/sa-east-1/qwen.qwen3-next-80b-a3b diff --git a/whitelisted_bedrock_models.txt b/whitelisted_bedrock_models.txt index 6124cb41044..254842e2714 100644 --- a/whitelisted_bedrock_models.txt +++ b/whitelisted_bedrock_models.txt @@ -44,6 +44,7 @@ bedrock/ap-northeast-1/minimax.minimax-m2.5 bedrock/ap-northeast-1/moonshotai.kimi-k2-thinking bedrock/ap-northeast-1/moonshotai.kimi-k2.5 bedrock/ap-northeast-1/qwen.qwen3-coder-next +bedrock/ap-northeast-1/qwen.qwen3-next-80b-a3b bedrock/moonshotai.kimi-k2-thinking bedrock/moonshotai.kimi-k2.5 bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0 @@ -54,6 +55,7 @@ bedrock/ap-south-1/minimax.minimax-m2.5 bedrock/ap-south-1/moonshotai.kimi-k2-thinking bedrock/ap-south-1/moonshotai.kimi-k2.5 bedrock/ap-south-1/qwen.qwen3-coder-next +bedrock/ap-south-1/qwen.qwen3-next-80b-a3b bedrock/ap-southeast-2/minimax.minimax-m2.5 bedrock/ap-southeast-3/deepseek.v3.2 bedrock/ap-southeast-3/minimax.minimax-m2.1 @@ -83,11 +85,13 @@ bedrock/eu-west-1/meta.llama3-8b-instruct-v1:0 bedrock/eu-west-1/minimax.minimax-m2.1 bedrock/eu-west-1/minimax.minimax-m2.5 bedrock/eu-west-1/qwen.qwen3-coder-next +bedrock/eu-west-1/qwen.qwen3-next-80b-a3b bedrock/eu-west-2/meta.llama3-70b-instruct-v1:0 bedrock/eu-west-2/meta.llama3-8b-instruct-v1:0 bedrock/eu-west-2/minimax.minimax-m2.1 bedrock/eu-west-2/minimax.minimax-m2.5 bedrock/eu-west-2/qwen.qwen3-coder-next +bedrock/eu-west-2/qwen.qwen3-next-80b-a3b bedrock/eu-west-3/mistral.mistral-7b-instruct-v0:2 bedrock/eu-west-3/mistral.mistral-large-2402-v1:0 bedrock/eu-west-3/mistral.mixtral-8x7b-instruct-v0:1 @@ -103,6 +107,7 @@ bedrock/sa-east-1/minimax.minimax-m2.5 bedrock/sa-east-1/moonshotai.kimi-k2-thinking bedrock/sa-east-1/moonshotai.kimi-k2.5 bedrock/sa-east-1/qwen.qwen3-coder-next +bedrock/sa-east-1/qwen.qwen3-next-80b-a3b bedrock/us-east-1/1-month-commitment/anthropic.claude-instant-v1 bedrock/us-east-1/1-month-commitment/anthropic.claude-v1 bedrock/us-east-1/1-month-commitment/anthropic.claude-v2:1 @@ -240,3 +245,4 @@ bedrock/us-gov-east-1/anthropic.claude-sonnet-5 bedrock/us-gov-east-1/anthropic.claude-opus-4-8 bedrock/us-gov-east-1/anthropic.claude-opus-5 bedrock/us-gov-east-1/anthropic.claude-fable-5-1 +bedrock/ap-southeast-2/qwen.qwen3-next-80b-a3b From 6db2bce43ce28f6cb8d32540347b576c379c7e84 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:30:42 +0000 Subject: [PATCH 130/146] chore(prices): sync OpenRouter prices: 1 model openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d3053dee25d..d3f05f77793 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.53172e-07, + "input_cost_per_token": 9.51432e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.906344e-06, + "output_cost_per_token": 1.902864e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9431e-08, + "cache_read_input_token_cost": 7.9286e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d3053dee25d..d3f05f77793 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.53172e-07, + "input_cost_per_token": 9.51432e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.906344e-06, + "output_cost_per_token": 1.902864e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9431e-08, + "cache_read_input_token_cost": 7.9286e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From 6786eb0131c5da5aabd8372726be27c70f0a1712 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 21 Sep 2026 13:39:32 +0000 Subject: [PATCH 131/146] fix(a2a): send message/stream for Bedrock AgentCore streaming requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../providers/bedrock_agentcore/handler.py | 2 +- .../test_bedrock_agentcore_a2a.py | 32 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index 306a8871b12..da5eb522187 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -98,7 +98,7 @@ class BedrockAgentCoreA2AHandler: request_id=request_id, params=params, litellm_params=litellm_params, - method="message/send", + method="message/stream", stream=True, agent_extra_headers=agent_extra_headers, ) diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index a8fe464ec32..256d73e3612 100644 --- a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -605,6 +605,38 @@ class TestNonStreaming: assert result["error"]["message"] == "Bad request" +class TestStreaming: + """Streaming requests must ask AgentCore for a stream, not a single send.""" + + @pytest.mark.asyncio + async def test_streaming_request_uses_message_stream_method_and_yields_sse_events(self, httpx_transport): + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + sse_body = ( + 'data: {"jsonrpc": "2.0", "id": "req-001", "result": {"kind": "task", "id": "t1"}}\n\n' + 'data: {"jsonrpc": "2.0", "id": "req-001", "result": {"kind": "status-update", "final": true}}\n\n' + ) + with respx.mock(assert_all_called=True) as router: + route = router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse_body) + ) + events = [ + event + async for event in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=SAMPLE_PARAMS, + litellm_params=SAMPLE_LITELLM_PARAMS, + ) + ] + + sent_body = json.loads(route.calls.last.request.content) + assert sent_body["method"] == "message/stream", sent_body + assert sent_body["params"]["message"]["messageId"] == "msg-001" + assert [event["result"]["kind"] for event in events] == ["task", "status-update"] + + class TestConfigManager: """Test that config manager routes 'bedrock' correctly.""" From 18ca95c3a7973008f308ba044f3ff4d34402f238 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:00:51 +0000 Subject: [PATCH 132/146] chore(prices): sync OpenRouter prices: 1 model openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d3f05f77793..39ef6510954 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.51432e-07, + "input_cost_per_token": 9.48126e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.902864e-06, + "output_cost_per_token": 1.896252e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9286e-08, + "cache_read_input_token_cost": 7.90105e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d3f05f77793..39ef6510954 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.51432e-07, + "input_cost_per_token": 9.48126e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.902864e-06, + "output_cost_per_token": 1.896252e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9286e-08, + "cache_read_input_token_cost": 7.90105e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From 346ad002c8bc0178e04e4adee654e2a2adc4fdc7 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:30:57 +0000 Subject: [PATCH 133/146] chore(prices): sync OpenRouter prices: 4 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/meta-llama/llama-3.1-70b-instruct: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token --- ...odel_prices_and_context_window_backup.json | 38 +++++++++---------- model_prices_and_context_window.json | 38 +++++++++---------- 2 files changed, 38 insertions(+), 38 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 39ef6510954..2d22376e2ea 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.7948e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.73844e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 4.4e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.8438e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.852e-08, - "output_cost_per_token": 1.1704e-07, - "cache_read_input_token_cost": 1.1704e-08, + "input_cost_per_token": 5.698e-08, + "output_cost_per_token": 1.1396e-07, + "cache_read_input_token_cost": 1.1396e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -69010,12 +69010,12 @@ "supports_reasoning": false }, "openrouter/meta-llama/llama-3.1-70b-instruct": { - "input_cost_per_token": 7.2e-07, - "output_cost_per_token": 7.2e-07, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 4e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -71317,15 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 4.4e-08, - "input_cost_per_token": 1.32e-06, + "cache_read_input_token_cost": 1.8438e-08, + "input_cost_per_token": 5.7948e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, - "output_cost_per_token": 3.96e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, + "output_cost_per_token": 1.73844e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 39ef6510954..2d22376e2ea 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.7948e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.73844e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 4.4e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.8438e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.852e-08, - "output_cost_per_token": 1.1704e-07, - "cache_read_input_token_cost": 1.1704e-08, + "input_cost_per_token": 5.698e-08, + "output_cost_per_token": 1.1396e-07, + "cache_read_input_token_cost": 1.1396e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -69010,12 +69010,12 @@ "supports_reasoning": false }, "openrouter/meta-llama/llama-3.1-70b-instruct": { - "input_cost_per_token": 7.2e-07, - "output_cost_per_token": 7.2e-07, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 4e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -71317,15 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 4.4e-08, - "input_cost_per_token": 1.32e-06, + "cache_read_input_token_cost": 1.8438e-08, + "input_cost_per_token": 5.7948e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, - "output_cost_per_token": 3.96e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, + "output_cost_per_token": 1.73844e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 7b2d3b36b6c7fb205e2b24b0e49220a337905cc8 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 14:40:46 +0000 Subject: [PATCH 134/146] fix(prices): align deepseek-v4-pro-0813 cache hit cost with cache read cost 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 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2d22376e2ea..7a5c28c9df4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41580,7 +41580,7 @@ }, "openrouter/deepseek/deepseek-v4-pro-0813": { "input_cost_per_token": 5.7948e-07, - "input_cost_per_token_cache_hit": 4.4e-08, + "input_cost_per_token_cache_hit": 1.8438e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2d22376e2ea..7a5c28c9df4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41580,7 +41580,7 @@ }, "openrouter/deepseek/deepseek-v4-pro-0813": { "input_cost_per_token": 5.7948e-07, - "input_cost_per_token_cache_hit": 4.4e-08, + "input_cost_per_token_cache_hit": 1.8438e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, From 1ac4d7ae042129f29aaf1a0d05b20b823888f24e Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 21 Sep 2026 09:55:09 -0500 Subject: [PATCH 135/146] fix(anthropic): type safeguards and safeguard_results as the arrays Anthropic sends Driving a real Claude Code 2.1.278 through the proxy, and a direct call to api.anthropic.com, both show these two fields are JSON arrays on the wire rather than objects. The request carries safeguards as [{"type": "dangerous_tool_use", "classifier_context": {...}}] under beta dangerous-tool-use-2026-09-03, and the 200 comes back with safeguard_results as [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {...}}}]. No runtime change: the request filter matches on TypedDict keys and never inspects the value. The test fixtures move to the captured shapes so the regression tests pin what the client and the provider actually exchange. --- litellm/types/llms/anthropic.py | 6 +++--- .../anthropic_messages/anthropic_response.py | 2 +- .../test_handler_output_config_passthrough.py | 2 +- ...experimental_pass_through_messages_handler.py | 16 ++++++++++------ 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index f57591d0262..c59c88698f7 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -411,7 +411,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): output_config: AnthropicOutputConfig | None # Configuration for Claude's output behavior cache_control: dict[str, Any] | None # Automatic prompt caching reasoning_effort: str | None - safeguards: ReadOnly[dict[str, object] | None] + safeguards: ReadOnly[list[dict[str, object]] | None] class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False): @@ -531,7 +531,7 @@ class AnthropicStopDetails(TypedDict, total=False): class MessageDelta(TypedDict, total=False): stop_reason: str | None stop_details: ReadOnly[AnthropicStopDetails] - safeguard_results: ReadOnly[dict[str, object]] + safeguard_results: ReadOnly[list[dict[str, object]]] class ServerToolUsage(TypedDict, total=False): @@ -602,7 +602,7 @@ class MessageChunk(TypedDict, total=False): stop_reason: str | None stop_sequence: str | None usage: UsageDelta - safeguard_results: ReadOnly[dict[str, object]] + safeguard_results: ReadOnly[list[dict[str, object]]] class MessageStartBlock(TypedDict): diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 41060e96d85..1d4c3cdc864 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -97,4 +97,4 @@ class AnthropicMessagesResponse(TypedDict, total=False): type: Literal["message"] | None usage: AnthropicUsage | None context_management: NotRequired[ContextManagementResponse] - safeguard_results: NotRequired[ReadOnly[dict[str, object]]] + safeguard_results: NotRequired[ReadOnly[list[dict[str, object]]]] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py index d6de6372e0b..6246f502344 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py @@ -113,7 +113,7 @@ class TestOutputConfigStrippedFromCompletionKwargs: def test_safeguards_is_stripped_for_non_anthropic_target(self): extra_kwargs = { "custom_llm_provider": "azure", - "safeguards": {"auto_mode": {"enabled": True, "version": "2026-09-01"}}, + "safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}], } result = _call_prepare(extra_kwargs=extra_kwargs) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 0acb9d634a3..e8bfcb86bf6 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1442,10 +1442,12 @@ async def test_anthropic_messages_leaves_non_provider_failures_unmapped(): @pytest.mark.asyncio async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthropic(): + """Shapes are what Claude Code 2.1.278 sends and api.anthropic.com returns, captured 2026-09-21.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler - safeguards = {"auto_mode": {"enabled": True, "version": "2026-09-01"}} - client_betas = "safeguards-2026-09-01,interleaved-thinking-2025-05-14" + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + client_betas = "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14" + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}] captured: dict[str, object] = {} def upstream_records_the_request(request: httpx.Request) -> httpx.Response: @@ -1462,7 +1464,7 @@ async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthro "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 1, "output_tokens": 1}, - "safeguard_results": {"verdict": "allow"}, + "safeguard_results": safeguard_results, }, request=request, ) @@ -1483,15 +1485,17 @@ async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthro assert captured["body"]["safeguards"] == safeguards assert set(captured["anthropic-beta"].split(",")) == set(client_betas.split(",")) - assert response["safeguard_results"] == {"verdict": "allow"} + assert response["safeguard_results"] == safeguard_results @pytest.mark.asyncio async def test_anthropic_messages_streaming_forwards_safeguards_and_keeps_safeguard_results(): + """Shapes are what Claude Code 2.1.278 sends and api.anthropic.com returns, captured 2026-09-21.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler - safeguards = {"auto_mode": {"enabled": True, "version": "2026-09-01"}} - safeguard_results = {"verdict": "allow", "checks": ["shell_command"]} + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] captured: dict[str, object] = {} message_start = { "type": "message_start", From 33e64e53f9d1590b0ea48b45459a6232f098ba65 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:00:53 +0000 Subject: [PATCH 136/146] chore(prices): sync OpenRouter prices: 4 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 36 +++++++++---------- model_prices_and_context_window.json | 36 +++++++++---------- 2 files changed, 36 insertions(+), 36 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7a5c28c9df4..ef581d447eb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.48126e-07, + "input_cost_per_token": 9.46386e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.896252e-06, + "output_cost_per_token": 1.892772e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.90105e-08, + "cache_read_input_token_cost": 7.88655e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7948e-07, + "input_cost_per_token": 5.7816e-07, "input_cost_per_token_cache_hit": 1.8438e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.73844e-06, + "output_cost_per_token": 1.73448e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.8438e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, + "cache_read_input_token_cost": 1.9272e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.698e-08, - "output_cost_per_token": 1.1396e-07, - "cache_read_input_token_cost": 1.1396e-08, + "input_cost_per_token": 5.544e-08, + "output_cost_per_token": 1.1088e-07, + "cache_read_input_token_cost": 1.1088e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71317,15 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.8438e-08, - "input_cost_per_token": 5.7948e-07, + "cache_read_input_token_cost": 1.9272e-08, + "input_cost_per_token": 5.7816e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, - "output_cost_per_token": 1.73844e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, + "output_cost_per_token": 1.73448e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7a5c28c9df4..ef581d447eb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.48126e-07, + "input_cost_per_token": 9.46386e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.896252e-06, + "output_cost_per_token": 1.892772e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.90105e-08, + "cache_read_input_token_cost": 7.88655e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7948e-07, + "input_cost_per_token": 5.7816e-07, "input_cost_per_token_cache_hit": 1.8438e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.73844e-06, + "output_cost_per_token": 1.73448e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.8438e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, + "cache_read_input_token_cost": 1.9272e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.698e-08, - "output_cost_per_token": 1.1396e-07, - "cache_read_input_token_cost": 1.1396e-08, + "input_cost_per_token": 5.544e-08, + "output_cost_per_token": 1.1088e-07, + "cache_read_input_token_cost": 1.1088e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71317,15 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.8438e-08, - "input_cost_per_token": 5.7948e-07, + "cache_read_input_token_cost": 1.9272e-08, + "input_cost_per_token": 5.7816e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7948e-7,"output_cost_per_token":0.00000173844,"cache_read_input_token_cost":1.8438e-8}, - "output_cost_per_token": 1.73844e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, + "output_cost_per_token": 1.73448e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 97fc8220dfd083725f540dda4230702a8bc483b8 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 15:14:19 +0000 Subject: [PATCH 137/146] fix(prices): align deepseek-v4-pro-0813 off-peak and cache-hit rates with the base cache read rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ef581d447eb..1f569ad1fbc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41580,7 +41580,7 @@ }, "openrouter/deepseek/deepseek-v4-pro-0813": { "input_cost_per_token": 5.7816e-07, - "input_cost_per_token_cache_hit": 1.8438e-08, + "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -41594,7 +41594,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "cache_read_input_token_cost": 1.9272e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71324,7 +71324,7 @@ "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, "output_cost_per_token": 1.73448e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ef581d447eb..1f569ad1fbc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41580,7 +41580,7 @@ }, "openrouter/deepseek/deepseek-v4-pro-0813": { "input_cost_per_token": 5.7816e-07, - "input_cost_per_token_cache_hit": 1.8438e-08, + "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -41594,7 +41594,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "cache_read_input_token_cost": 1.9272e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71324,7 +71324,7 @@ "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.8438e-8}, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, "output_cost_per_token": 1.73448e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, From a3dcec463b9e64d424496644b1b9cdc57dbd049c Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:31:08 +0000 Subject: [PATCH 138/146] chore(prices): sync Azure prices: 1 model, 1 deprecated azure_ai/MAI-Image-2.5-Pro: deprecation_date --- 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 3bc8b3c176c..91e869cbb36 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11175,6 +11175,7 @@ "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2.5-Pro": { + "deprecation_date": "2026-10-01", "input_cost_per_image_token": 8e-06, "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3bc8b3c176c..91e869cbb36 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11175,6 +11175,7 @@ "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2.5-Pro": { + "deprecation_date": "2026-10-01", "input_cost_per_image_token": 8e-06, "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", From ade39978a9202d28ad53088207a6143f16c4c0ad Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:31:16 +0000 Subject: [PATCH 139/146] chore(prices): sync AWS Bedrock prices: 25 models [enrichment failed: AWS Bedrock, 66 held] anthropic.claude-fable-5: anthropic.claude-fable-5-1: anthropic.claude-opus-4-7: anthropic.claude-opus-4-8: anthropic.claude-opus-5: anthropic.claude-sonnet-4-6: anthropic.claude-sonnet-5: global.anthropic.claude-fable-5: global.anthropic.claude-fable-5-1: global.anthropic.claude-opus-4-7: global.anthropic.claude-opus-4-8: global.anthropic.claude-opus-5: global.anthropic.claude-sonnet-4-6: global.anthropic.claude-sonnet-5: us-gov.anthropic.claude-fable-5-1: us-gov.anthropic.claude-opus-4-8: us-gov.anthropic.claude-opus-5: us-gov.anthropic.claude-sonnet-5: us.anthropic.claude-fable-5: us.anthropic.claude-fable-5-1: us.anthropic.claude-opus-4-7: us.anthropic.claude-opus-4-8: us.anthropic.claude-opus-5: us.anthropic.claude-sonnet-4-6: us.anthropic.claude-sonnet-5: --- ...odel_prices_and_context_window_backup.json | 50 +++++++++---------- model_prices_and_context_window.json | 50 +++++++++---------- 2 files changed, 50 insertions(+), 50 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3bc8b3c176c..98e570d7874 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1327,7 +1327,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-mythos-preview": { "input_cost_per_token": 0, @@ -1381,7 +1381,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, @@ -1419,7 +1419,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, @@ -1531,7 +1531,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-fable-5-1": { "cache_creation_input_token_cost": 1.25e-05, @@ -1570,7 +1570,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, @@ -1608,7 +1608,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-fable-5-1": { "cache_creation_input_token_cost": 1.25e-05, @@ -1647,7 +1647,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -1685,7 +1685,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-fable-5-1": { "cache_creation_input_token_cost": 1.375e-05, @@ -1724,7 +1724,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -1837,7 +1837,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, @@ -1875,7 +1875,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, @@ -1913,7 +1913,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, @@ -2063,7 +2063,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2102,7 +2102,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2141,7 +2141,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2329,7 +2329,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2368,7 +2368,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2407,7 +2407,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2556,7 +2556,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -2591,7 +2591,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -2626,7 +2626,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -47289,7 +47289,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -47323,7 +47323,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -47356,7 +47356,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -47407,7 +47407,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us-gov.nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3bc8b3c176c..98e570d7874 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1327,7 +1327,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-mythos-preview": { "input_cost_per_token": 0, @@ -1381,7 +1381,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, @@ -1419,7 +1419,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, @@ -1531,7 +1531,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-fable-5-1": { "cache_creation_input_token_cost": 1.25e-05, @@ -1570,7 +1570,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, @@ -1608,7 +1608,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-fable-5-1": { "cache_creation_input_token_cost": 1.25e-05, @@ -1647,7 +1647,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -1685,7 +1685,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-fable-5-1": { "cache_creation_input_token_cost": 1.375e-05, @@ -1724,7 +1724,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -1837,7 +1837,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, @@ -1875,7 +1875,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, @@ -1913,7 +1913,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, @@ -2063,7 +2063,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2102,7 +2102,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2141,7 +2141,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2329,7 +2329,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2368,7 +2368,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2407,7 +2407,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2556,7 +2556,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "global.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -2591,7 +2591,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -2626,7 +2626,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -47289,7 +47289,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -47323,7 +47323,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 1024, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -47356,7 +47356,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -47407,7 +47407,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" + "source": "https://aws.amazon.com/bedrock/pricing/" }, "us-gov.nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, From 739227fefc8eb4824c59988f131c92599b6fdd30 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:31:20 +0000 Subject: [PATCH 140/146] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 30 +++++++++---------- model_prices_and_context_window.json | 30 +++++++++---------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3bc8b3c176c..4d92825b34f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41590,21 +41590,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.46386e-07, + "input_cost_per_token": 9.42906e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.892772e-06, + "output_cost_per_token": 1.885812e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.88655e-08, + "cache_read_input_token_cost": 7.85755e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41632,22 +41632,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7816e-07, + "input_cost_per_token": 5.7684e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.73448e-06, + "output_cost_per_token": 1.73052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.9272e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, + "cache_read_input_token_cost": 1.8354e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71484,15 +71484,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.9272e-08, - "input_cost_per_token": 5.7816e-07, + "cache_read_input_token_cost": 1.8354e-08, + "input_cost_per_token": 5.7684e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, - "output_cost_per_token": 1.73448e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, + "output_cost_per_token": 1.73052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3bc8b3c176c..4d92825b34f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41590,21 +41590,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.46386e-07, + "input_cost_per_token": 9.42906e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.892772e-06, + "output_cost_per_token": 1.885812e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.88655e-08, + "cache_read_input_token_cost": 7.85755e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41632,22 +41632,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7816e-07, + "input_cost_per_token": 5.7684e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.73448e-06, + "output_cost_per_token": 1.73052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.9272e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, + "cache_read_input_token_cost": 1.8354e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71484,15 +71484,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.9272e-08, - "input_cost_per_token": 5.7816e-07, + "cache_read_input_token_cost": 1.8354e-08, + "input_cost_per_token": 5.7684e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7816e-7,"output_cost_per_token":0.00000173448,"cache_read_input_token_cost":1.9272e-8}, - "output_cost_per_token": 1.73448e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, + "output_cost_per_token": 1.73052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From a233ba910d536d39cf8ba376aa97ea2822380a3d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 21 Sep 2026 08:37:23 -0700 Subject: [PATCH 141/146] feat(auto-router): configure heuristic v2 success threshold --- .../complexity_router/README.md | 16 +++- .../complexity_router/complexity_router.py | 5 +- .../complexity_router/config.py | 12 +++ .../complexity_router/tier_predictor.py | 7 +- .../router_strategy/test_complexity_router.py | 89 +++++++++++++++++-- .../add_model/AutoRouterRoutingTest.test.tsx | 19 +++- .../add_model/AutoRouterRoutingTest.tsx | 8 +- .../add_model/ClassificationMethodConfig.tsx | 85 +++++++++++++++++- .../add_model/ComplexityRouterConfig.test.tsx | 45 ++++++++++ .../add_model/ComplexityRouterConfig.tsx | 5 +- .../add_model/add_auto_router_tab.test.tsx | 75 ++++++++++++++++ .../add_model/add_auto_router_tab.tsx | 3 + .../build_complexity_router_config.test.ts | 34 +++++++ .../build_complexity_router_config.ts | 13 +++ ...d_updated_complexity_router_config.test.ts | 29 ++++++ ...dit_auto_router_modal.integration.test.tsx | 68 ++++++++++++++ .../edit_auto_router_modal.tsx | 9 ++ .../src/lib/autorouter_presets.test.ts | 11 +++ .../src/lib/autorouter_presets.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 ++ 20 files changed, 518 insertions(+), 21 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 6505746bca1..f023d5001d9 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -191,6 +191,7 @@ model_list: model: auto_router/complexity_router complexity_router_config: classifier_type: heuristic_v2 + heuristic_v2_success_threshold: 0.9 tiers: SIMPLE: luna MEDIUM: terra @@ -201,9 +202,18 @@ model_list: No classifier model call or per-model training data is required. The classifier uses global tier quality, request-type quality, and similar-request cohorts from the bundled UltraFeedback artifact. It estimates success at every tier, enforces -monotonic probabilities, and returns the first tier meeting the trained 0.75 -threshold. The existing complexity-router tier pool then selects and dispatches -a model from that tier +monotonic probabilities, and returns the first tier meeting the success threshold, +or REASONING if no tier meets it. The existing complexity-router tier pool then +selects and dispatches a model from that tier + +Set `heuristic_v2_success_threshold` to a value from 0 to 1 to override the +artifact's threshold. For example, `0.9` requires a predicted success probability +of at least 90%. Higher thresholds favor more capable tiers. Omit the setting or +set it to `null` to use the artifact's `routing_threshold`, which is `0.75` for +the bundled artifact. The override leaves the predicted probabilities unchanged + +In the dashboard, select Heuristic v2 under Advanced: Classification Method and +set Success threshold. Clear the field to restore the artifact's default Spend logs record `routing_decision.cause: heuristic_v2`, the detected request type, and all four predicted probabilities. Existing `classifier_type: heuristic` diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a3d6ccbd437..83fcfdfc329 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1429,7 +1429,10 @@ class ComplexityRouter(CustomLogger): _ClassifierCircuitBreaker(circuit_breaker_cooldown) if circuit_breaker_cooldown is not None else None ) self._tier_success_predictor: TierSuccessPredictor | None = ( - TierSuccessPredictor(resolve_tier_artifact(self.config.heuristic_v2_artifact)) + TierSuccessPredictor( + resolve_tier_artifact(self.config.heuristic_v2_artifact), + routing_threshold=self.config.heuristic_v2_success_threshold, + ) if self.config.classifier_type == "heuristic_v2" else None ) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index aa39dff8c53..0b2caa93665 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1036,6 +1036,18 @@ class ComplexityRouterConfig(BaseModel): "UltraFeedback artifact is selected by default; an inline trained artifact may replace it" ), ) + heuristic_v2_success_threshold: float | None = Field( + default=None, + strict=True, + ge=0.0, + le=1.0, + description=( + "Minimum predicted success probability for classifier_type 'heuristic_v2' to select a tier. " + "The first tier meeting this threshold is selected, or REASONING if none meets it. " + "When omitted or null, uses the artifact's routing_threshold (0.75 for the bundled artifact). " + "Other classifier types ignore this setting" + ), + ) classifier_llm_config: ClassifierLLMConfig | None = Field( default=None, description=( diff --git a/litellm/router_strategy/complexity_router/tier_predictor.py b/litellm/router_strategy/complexity_router/tier_predictor.py index 764f6e6ad56..7775c36e795 100644 --- a/litellm/router_strategy/complexity_router/tier_predictor.py +++ b/litellm/router_strategy/complexity_router/tier_predictor.py @@ -108,8 +108,9 @@ class TierPrediction: class TierSuccessPredictor: - def __init__(self, artifact: TrainedTierArtifact) -> None: + def __init__(self, artifact: TrainedTierArtifact, *, routing_threshold: float | None = None) -> None: self._artifact = artifact + self._routing_threshold: Final = artifact.routing_threshold if routing_threshold is None else routing_threshold self._global: Mapping[int, TierGlobalStatistic] = MappingProxyType( {stat.tier: stat for stat in artifact.global_statistics} ) @@ -122,7 +123,7 @@ class TierSuccessPredictor: @property def routing_threshold(self) -> float: - return self._artifact.routing_threshold + return self._routing_threshold def predict(self, prompt: str, request_type: RequestType) -> TierPrediction: cohort: Final = similarity_cohort(prompt, request_type) @@ -132,7 +133,7 @@ class TierSuccessPredictor: {int(tier): probability for tier, probability in zip(_TIERS, monotonic)} ) required_tier: Final = next( - (tier for tier in _TIERS if probabilities[tier] >= self._artifact.routing_threshold), + (tier for tier in _TIERS if probabilities[tier] >= self.routing_threshold), 4, ) return TierPrediction(probabilities=probabilities, required_tier=required_tier) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index ecd25ff654f..90ab39f601c 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3722,16 +3722,37 @@ class TestLLMClassifier: @pytest.mark.asyncio @pytest.mark.parametrize("redact", (False, True)) + @pytest.mark.parametrize( + "override,threshold,tier,model", + ( + ({}, 0.8, "COMPLEX", "complex-model"), + ({"heuristic_v2_success_threshold": None}, 0.8, "COMPLEX", "complex-model"), + ({"heuristic_v2_success_threshold": 0.0}, 0.0, "SIMPLE", "simple-model"), + ({"heuristic_v2_success_threshold": 21 / 102}, 21 / 102, "MEDIUM", "medium-model"), + ({"heuristic_v2_success_threshold": 0.95}, 0.95, "REASONING", "reasoning-model"), + ({"heuristic_v2_success_threshold": 1.0}, 1.0, "REASONING", "reasoning-model"), + ), + ids=("omitted", "null", "zero", "inclusive", "higher", "no-tier-passes"), + ) async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier( - self, mock_router_instance: MagicMock, redact: bool, monkeypatch: pytest.MonkeyPatch + self, + mock_router_instance: MagicMock, + redact: bool, + monkeypatch: pytest.MonkeyPatch, + override: Mapping[str, float | None], + threshold: float, + tier: str, + model: str, ) -> None: monkeypatch.setattr(litellm, "turn_off_message_logging", redact) - router = ComplexityRouter( + artifact: Final = _heuristic_v2_artifact() + router: Final = ComplexityRouter( model_name="tier-router", litellm_router_instance=mock_router_instance, complexity_router_config={ "classifier_type": "heuristic_v2", - "heuristic_v2_artifact": _heuristic_v2_artifact(), + "heuristic_v2_artifact": artifact, + **override, "tiers": { "SIMPLE": "simple-model", "MEDIUM": "medium-model", @@ -3741,15 +3762,15 @@ class TestLLMClassifier: }, ) - response = await router.async_pre_routing_hook( + response: Final = await router.async_pre_routing_hook( model="tier-router", request_kwargs={}, messages=[{"role": "user", "content": "Handle this new request"}], ) assert response is not None - assert response.model == "complex-model" - assert response.routing_decision["tier"] == "COMPLEX" + assert response.model == model + assert response.routing_decision["tier"] == tier assert response.routing_decision["cause"] == "heuristic_v2" assert response.routing_decision["signals"] == [ "request-type:general", @@ -3769,10 +3790,62 @@ class TestLLMClassifier: "COMPLEX": 91 / 102, "REASONING": 100 / 102, }, - "threshold": 0.8, - "predicted_tier": "COMPLEX", + "threshold": threshold, + "predicted_tier": tier, "request_type": "general", } + assert artifact.routing_threshold == 0.8 + + @pytest.mark.parametrize("threshold", (-0.01, 1.01, math.nan, math.inf, -math.inf, True, "0.95")) + def test_heuristic_v2_success_threshold_rejects_invalid_values(self, threshold: float | bool | str) -> None: + with pytest.raises(ValidationError, match="heuristic_v2_success_threshold"): + ComplexityRouterConfig.model_validate( + {"classifier_type": "heuristic_v2", "heuristic_v2_success_threshold": threshold} + ) + + @pytest.mark.asyncio + async def test_heuristic_v2_threshold_reload_and_rejected_update_keep_router_isolated(self) -> None: + artifact: Final = _heuristic_v2_artifact() + + def deployment(threshold: float, name: str = "editable") -> Deployment: + return Deployment( + model_name=name, + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_config={ + "classifier_type": "heuristic_v2", + "heuristic_v2_artifact": artifact.model_dump(), + "heuristic_v2_success_threshold": threshold, + "session_affinity": False, + "tiers": {"SIMPLE": "simple-model", "REASONING": "reasoning-model"}, + }, + ), + model_info={"id": name}, + ) + + router: Final = Router( + model_list=[ + deployment(0.95).model_dump(exclude_none=True), + deployment(0.95, "unchanged").model_dump(exclude_none=True), + ], + ignore_invalid_deployments=True, + ) + + async def routed_threshold(name: str) -> tuple[str, float]: + response: Final = await router.async_pre_routing_hook( + model=name, + request_kwargs={}, + messages=[{"role": "user", "content": "Handle this new request"}], + ) + assert response is not None and response.routing_decision is not None + return response.model, response.routing_decision["heuristic_v2_forecast"]["threshold"] + + assert await routed_threshold("editable") == ("reasoning-model", 0.95) + assert router.upsert_deployment(deployment(0.0)) is not None + assert await routed_threshold("editable") == ("simple-model", 0.0) + assert await routed_threshold("unchanged") == ("reasoning-model", 0.95) + assert router.upsert_deployment(deployment(1.01)) is None + assert await routed_threshold("editable") == ("simple-model", 0.0) def test_heuristic_v2_needs_no_classifier_model(self): config = ComplexityRouterConfig(classifier_type="heuristic_v2") diff --git a/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.test.tsx b/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.test.tsx index 74e7193c8b0..86d161d1c8e 100644 --- a/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.test.tsx @@ -15,7 +15,8 @@ vi.mock("../networking", () => ({ const CONFIG = { tiers: { SIMPLE: ["cheap"], MEDIUM: ["mid"], COMPLEX: ["strong"], REASONING: ["o3"] }, - classifier_type: "heuristic", + classifier_type: "heuristic_v2", + heuristic_v2_success_threshold: 0, } as unknown as ComplexityRouterConfigPayload; const Harness = () => ( @@ -62,6 +63,22 @@ describe("AutoRouterRoutingTest", () => { expect(screen.getByTestId("auto-router-routing-test-send")).toBeDisabled(); }); + it("blocks previewing an invalid success threshold instead of sending NaN as null", () => { + renderWithProviders( + , + ); + fireEvent.change(screen.getByTestId("auto-router-routing-test-prompt"), { target: { value: "hello" } }); + expect(screen.getByTestId("auto-router-routing-test-send")).toBeDisabled(); + expect(screen.getByText("Success threshold must be a number between 0 and 1")).toBeVisible(); + expect(testAutoRouterRouting).not.toHaveBeenCalled(); + }); + it("routes the typed prompt through the config being edited and shows where it landed", async () => { const user = userEvent.setup(); vi.mocked(testAutoRouterRouting).mockResolvedValue(successResponse); diff --git a/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.tsx b/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.tsx index 00b2e75dfd1..2b6c06e9a96 100644 --- a/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AutoRouterRoutingTest.tsx @@ -5,7 +5,7 @@ import { Button } from "@/components/ui/button"; import { Textarea } from "@/components/ui/textarea"; import RoutingDecisionCard from "@/components/view_logs/LogDetailsDrawer/RoutingDecisionCard"; import { AutoRouterRoutingTestResult, testAutoRouterRouting } from "../networking"; -import { ComplexityRouterConfigPayload } from "./build_complexity_router_config"; +import { ComplexityRouterConfigPayload, getHeuristicV2SuccessThresholdError } from "./build_complexity_router_config"; import { buildAutoRouterRoutingTestRequest } from "./build_auto_router_routing_test_request"; interface AutoRouterRoutingTestProps { @@ -31,8 +31,10 @@ const AutoRouterRoutingTest: React.FC = ({ }) => { const [prompt, setPrompt] = React.useState(""); const [state, setState] = React.useState({ status: "idle" }); + const configError = getHeuristicV2SuccessThresholdError(config.heuristic_v2_success_threshold); const send = async () => { + if (configError) return; setState({ status: "running" }); const params = { prompt, config, defaultModel, routerName, teamId }; const request = buildAutoRouterRoutingTestRequest(params); @@ -62,13 +64,15 @@ const AutoRouterRoutingTest: React.FC = ({
+ {configError &&

{configError}

} + {state.status === "failed" && (
> = ({ + value, + onChange, +}) => { + const threshold = value.heuristic_v2_success_threshold; + if (effectiveClassifierType(value) === "heuristic_v2" || threshold === undefined) return null; + const error = getHeuristicV2SuccessThresholdError(threshold); + return ( +
+

+ Heuristic v2 success threshold (inactive):{" "} + + {Number.isFinite(threshold) ? threshold : "Invalid value"} + +

+

Only used when Heuristic v2 is selected

+ {error && ( +

+ {error} +

+ )} + +
+ ); +}; + const ClassifierTypeRadios: React.FC<{ value: ComplexityRouterConfigValue; classifierType: ClassifierType; @@ -259,6 +296,12 @@ const ClassificationMethodConfig: React.FC = ({ const classifierModel = value.classifier_llm_config?.model ?? ""; const classifierReasoningEffort = value.classifier_llm_config?.reasoning_effort; const explicitlySupportedClassifierEfforts = effortOptionsByModel[classifierModel]; + const successThresholdError = getHeuristicV2SuccessThresholdError(value.heuristic_v2_success_threshold); + const successThresholdDraft = + draft?.id === HEURISTIC_V2_SUCCESS_THRESHOLD_ID && + Object.is(value.heuristic_v2_success_threshold, draft.raw.trim() === "" ? undefined : Number(draft.raw)) + ? draft.raw + : null; const handleClassifierTypeChange = (classifierType: ClassifierType) => { onChange(transitionClassifierType(value, classifierType)); @@ -275,6 +318,14 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, hybrid_boundary_margin: Math.min(1, Math.max(0, parsed)) }); }; + const handleSuccessThresholdChange = (raw: string) => { + setDraft({ id: HEURISTIC_V2_SUCCESS_THRESHOLD_ID, raw }); + onChange({ + ...value, + heuristic_v2_success_threshold: raw.trim() === "" ? undefined : Number(raw), + }); + }; + // One write for everything the prompt dialog owns. The rubric arrives here rather than through the // rubric handler because two onChange calls in one tick would both spread this render's `value`, // so whichever landed second would drop the other's edit. @@ -407,6 +458,36 @@ const ClassificationMethodConfig: React.FC = ({ <> + {classifierType === "heuristic_v2" && ( +
+ + handleSuccessThresholdChange(event.target.value)} + onBlur={() => { + if (!successThresholdError) setDraft(null); + }} + aria-invalid={Boolean(successThresholdError)} + aria-describedby={`${HEURISTIC_V2_SUCCESS_THRESHOLD_ID}-help${successThresholdError ? ` ${HEURISTIC_V2_SUCCESS_THRESHOLD_ID}-error` : ""}`} + /> +

+ Minimum predicted success probability, from 0 to 1. Higher values favor more capable tiers. Leave blank to + use the artifact default +

+ {successThresholdError && ( + + )} +
+ )} + {classifierType === "heuristic_first" && (
Decide locally up to 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 e91ff1d59c1..70658b787f0 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -153,6 +153,51 @@ describe("ComplexityRouterConfig", () => { expect(screen.queryByText(/Score < 0.15/)).not.toBeInTheDocument(); }); + it.each<[string, Partial]>([ + ["heuristic", { classifier_type: "heuristic" }], + ["LLM", { classifier_type: "llm" }], + ["heuristic first", { classifier_type: "heuristic_first" }], + ["hybrid", { classifier_type: "hybrid" }], + ["Capability", { classifier_type: "capability" }], + ["Fuse v2", { classifier_type: "llm_v2" }], + [ + "custom tiers", + { + classifier_type: "heuristic_v2", + custom_tier_set: { + tiers: [{ id: "review", name: "REVIEW", definition: "Review code", models: ["gpt-4"] }], + fallback_tier_id: "review", + }, + }, + ], + ])("shows and clears an invalid inactive threshold under %s", (_label, overrides) => { + const value = { ...defaultValue, ...overrides, heuristic_v2_success_threshold: Number.NaN }; + const onChange = vi.fn(); + renderWithProviders(); + const retained = screen.getByRole("region", { name: "Inactive Heuristic v2 threshold" }); + expect(within(retained).getByRole("status", { name: "Retained Heuristic v2 threshold" })).toHaveTextContent( + "Invalid value", + ); + expect(within(retained).getByRole("alert")).toHaveTextContent("Success threshold must be a number between 0 and 1"); + fireEvent.click(within(retained).getByRole("button", { name: "Clear Heuristic v2 threshold" })); + expect(onChange).toHaveBeenCalledWith({ ...value, heuristic_v2_success_threshold: undefined }); + }); + + it("shows an inactive zero threshold until explicitly cleared and hides the summary for active or absent values", () => { + const onChange = vi.fn(); + const value = { ...defaultValue, heuristic_v2_success_threshold: 0 }; + const { rerender } = renderWithProviders( + , + ); + expect(screen.getByRole("status", { name: "Retained Heuristic v2 threshold" })).toHaveTextContent("0"); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + expect(onChange).not.toHaveBeenCalled(); + rerender(); + expect(screen.queryByRole("region", { name: "Inactive Heuristic v2 threshold" })).not.toBeInTheDocument(); + rerender(); + expect(screen.queryByRole("region", { name: "Inactive Heuristic v2 threshold" })).not.toBeInTheDocument(); + }); + it("should show classifier fields and use the configured values when classifier_type is llm", () => { const llmValue: ComplexityRouterConfigValue = { ...defaultValue, diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index f6b50ce20bc..8216df139aa 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -37,7 +37,7 @@ import { import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; -import ClassificationMethodConfig from "./ClassificationMethodConfig"; +import ClassificationMethodConfig, { InactiveHeuristicV2Threshold } from "./ClassificationMethodConfig"; import ContextWindowEscalationConfig from "./ContextWindowEscalationConfig"; import ResponseFormatControls from "./ResponseFormatControls"; import StallEscalationConfig from "./StallEscalationConfig"; @@ -374,6 +374,7 @@ export interface ComplexityRouterConfigValue { /** An explicit pin. Unset means the default tracks the tiers - see resolveComplexityDefaultModel. */ default_model?: string; classifier_type: ClassifierType; + heuristic_v2_success_threshold?: number; capability_classifier_config?: CapabilitySettings; llm_v2_config?: FuseSettings; classifier_llm_config?: ClassifierLLMConfig; @@ -618,6 +619,8 @@ const ComplexityRouterConfig: React.FC = ({ )}
+ + {forecast ? ( <> { }); }); + it("blocks invalid success thresholds and creates a heuristic v2 router with explicit zero", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + renderWithProviders(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "threshold-router" } }); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(screen.getByRole("radio", { name: /^Heuristic v2/ })); + + const threshold = screen.getByRole("textbox", { name: "Success threshold" }); + expect(threshold).toHaveValue(""); + fireEvent.change(threshold, { target: { value: "invalid" } }); + fireEvent.blur(threshold); + expect(threshold).toHaveValue("invalid"); + expect(threshold).toHaveAttribute("aria-invalid", "true"); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + expect(screen.getByTestId("auto-router-test-routing-btn")).toBeDisabled(); + + await user.click(screen.getByRole("radio", { name: /^Heuristic \(default\)/ })); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + await user.click(screen.getByRole("radio", { name: /^Heuristic v2/ })); + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "1.01" } }); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "0" } }); + await user.click(screen.getByRole("button", { name: "Add Auto Router" })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + classifier_type: "heuristic_v2", + heuristic_v2_success_threshold: 0, + }); + }); + + it("clears an invalid threshold draft when automatic setup replaces the configuration", async () => { + const user = userEvent.setup(); + mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); + renderWithProviders(); + const automaticSetup = await screen.findByRole("button", { name: "Configure automatically" }); + await waitFor(() => expect(automaticSetup).toBeEnabled()); + await user.click(automaticSetup); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "reset-threshold-router" } }); + await user.click(screen.getByText("Advanced: Classification Method")); + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "1.1" } }); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + + await user.click(automaticSetup); + expect(screen.getByRole("textbox", { name: "Success threshold" })).toHaveValue(""); + expect(screen.getByRole("textbox", { name: "Success threshold" })).toHaveAttribute("aria-invalid", "false"); + await user.click(screen.getByRole("button", { name: "Add Auto Router" })); + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).not.toHaveProperty( + "heuristic_v2_success_threshold", + ); + }); + + it("clears an invalid inactive threshold before creating the router", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + renderWithProviders(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "clear-threshold-router" } }); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(screen.getByRole("radio", { name: /^Heuristic v2/ })); + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "invalid" } }); + await user.click(screen.getByRole("radio", { name: /^Heuristic \(default\)/ })); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + await user.click(screen.getByRole("button", { name: "Clear Heuristic v2 threshold" })); + expect(screen.queryByRole("region", { name: "Inactive Heuristic v2 threshold" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Add Auto Router" })); + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).not.toHaveProperty( + "heuristic_v2_success_threshold", + ); + }); + it("carries a context-window escalation opt-out through to the create payload", 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 126d9ba2311..57a6201bc7b 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 @@ -47,6 +47,7 @@ import { buildComplexityRouterConfig, getKeywordTierRulesError, getClassifierModelError, + getHeuristicV2SuccessThresholdError, getClassifierReasoningEffortError, getMissingTiersError, getPlanModeTierError, @@ -146,6 +147,7 @@ export const getSubmitBlockedReason = ( getPlanModeTierError(config.plan_mode_min_tier, activeTierRows(config)) ?? getKeywordTierRulesError(keywordTierRules, activeTierRows(config)) ?? getClassifierModelError(config) ?? + getHeuristicV2SuccessThresholdError(config.heuristic_v2_success_threshold) ?? (heuristicScoringRole(config) === "decides" ? customDimensionsError(config.custom_dimensions) : null) ?? getClassifierReasoningEffortError(config, modelInfo) ?? getReferencedModelsError(referencedModelsParams, availability) @@ -405,6 +407,7 @@ const AddAutoRouterTab: React.FC = ({ classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, + heuristicV2SuccessThreshold: complexityRouterConfig.heuristic_v2_success_threshold, capabilityClassifierConfig: complexityRouterConfig.capability_classifier_config, llmV2Config: complexityRouterConfig.llm_v2_config, 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 2990878d086..63fed7c7175 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 @@ -4,6 +4,7 @@ import { normalizeClassifierLlmConfig, getKeywordTierRulesError, getClassifierModelError, + getHeuristicV2SuccessThresholdError, getClassifierReasoningEffortError, getMissingTiersError, hydrateCustomTierSet, @@ -211,8 +212,28 @@ describe("buildComplexityRouterConfig", () => { expect(config.classifier_llm_config).toBeUndefined(); expect(config.classifier_context_window_size).toBeUndefined(); expect(config.classifier_fallback).toBeUndefined(); + expect(config).not.toHaveProperty("heuristic_v2_success_threshold"); }); + it.each([0, 0.95, 1])("serializes a heuristic v2 success threshold of %s", (heuristicV2SuccessThreshold) => { + const config = buildComplexityRouterConfig({ + ...baseParams, + classifierType: "heuristic_v2", + heuristicV2SuccessThreshold, + }); + expect(config.heuristic_v2_success_threshold).toBe(heuristicV2SuccessThreshold); + }); + + it.each(["heuristic", "llm", "heuristic_first", "hybrid", "capability", "llm_v2"] as const)( + "retains the inactive success threshold under %s", + (classifierType) => { + expect( + buildComplexityRouterConfig({ ...baseParams, classifierType, heuristicV2SuccessThreshold: 0.91 }) + .heuristic_v2_success_threshold, + ).toBe(0.91); + }, + ); + it("includes classifier_context_window_size and classifier_context_budget_chars only when classifier_type is llm", () => { const params: BuildComplexityRouterConfigParams = { ...baseParams, @@ -884,6 +905,19 @@ describe("buildComplexityRouterConfig tier model params", () => { }); }); +describe("getHeuristicV2SuccessThresholdError", () => { + it.each([undefined, 0, 0.95, 1])("accepts the optional probability %s", (threshold) => { + expect(getHeuristicV2SuccessThresholdError(threshold)).toBeNull(); + }); + + it.each([-0.01, 1.01, Number.NaN, Number.POSITIVE_INFINITY, Number.NEGATIVE_INFINITY])( + "rejects invalid success threshold %s", + (threshold) => { + expect(getHeuristicV2SuccessThresholdError(threshold)).toBe("Success threshold must be a number between 0 and 1"); + }, + ); +}); + describe("getClassifierModelError", () => { it("stays quiet for a heuristic router, which needs no classifier model", () => { expect(getClassifierModelError({ classifier_type: "heuristic" })).toBeNull(); 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 8a377c17ad7..05dc327968c 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 @@ -144,6 +144,7 @@ export interface StoredComplexityRouterConfig { hybrid_boundary_margin?: unknown; tier_labels?: unknown; classifier_type?: ClassifierType; + heuristic_v2_success_threshold?: unknown; capability_classifier_config?: unknown; llm_v2_config?: unknown; classifier_llm_config?: ClassifierLLMConfig; @@ -182,6 +183,7 @@ export interface BuildComplexityRouterConfigParams { planModeMinTier: string | undefined; tierLabels: ComplexityTierLabels | undefined; classifierType: ClassifierType; + heuristicV2SuccessThreshold?: number; capabilityClassifierConfig?: CapabilitySettings; llmV2Config?: FuseSettings; classifierLlmConfig: ClassifierLLMConfigWire | undefined; @@ -248,6 +250,7 @@ export interface ComplexityRouterConfigPayload { plan_mode_min_tier?: string; tier_labels?: ComplexityTierLabels; classifier_type: ClassifierType; + heuristic_v2_success_threshold?: number; capability_classifier_config?: CapabilitySettings; llm_v2_config?: FuseSettings; classifier_llm_config?: ClassifierLLMConfig; @@ -356,6 +359,12 @@ export const getKeywordTierRulesError = ( return `Keyword rule(s) ${orphaned.join(", ")} route to a tier this router no longer has`; }; +export const getHeuristicV2SuccessThresholdError = (threshold: number | undefined): string | null => { + if (threshold === undefined) return null; + const validProbability = Number.isFinite(threshold) && threshold >= 0 && threshold <= 1; + return validProbability ? null : "Success threshold must be a number between 0 and 1"; +}; + // An edited tier set forces the LLM classifier, so the model requirement follows the EFFECTIVE type. // Both forms' submit gates and their submit handlers read this one answer so they cannot drift. export const getClassifierModelError = ( @@ -557,6 +566,7 @@ export const buildComplexityRouterConfig = ({ planModeMinTier, tierLabels, classifierType, + heuristicV2SuccessThreshold, capabilityClassifierConfig, llmV2Config, classifierLlmConfig, @@ -640,6 +650,9 @@ export const buildComplexityRouterConfig = ({ ...(planModeMinTier?.trim() && { plan_mode_min_tier: planModeMinTier }), ...(cleanedTierLabels && { tier_labels: cleanedTierLabels }), classifier_type: classifierType, + ...(heuristicV2SuccessThreshold !== undefined && { + heuristic_v2_success_threshold: heuristicV2SuccessThreshold, + }), ...classifierWireFields(effectiveType, classifierInputs), ...(effectiveType === "capability" && capabilityClassifierConfig && { capability_classifier_config: capabilityClassifierConfig }), 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 4ae6efbb12d..e4b4cbafdf6 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 @@ -46,6 +46,34 @@ const hydratedState: KeywordMatchingState = { }; describe("buildUpdatedComplexityRouterConfig keyword matching", () => { + it.each([0, 0.92, 1])("hydrates and saves a success threshold of %s without changing the artifact", (threshold) => { + const stored = { + ...STORED, + classifier_type: "heuristic_v2" as const, + heuristic_v2_success_threshold: threshold, + heuristic_v2_artifact: { routing_threshold: 0.82, custom_metadata: "retained" }, + }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + expect(hydrated.heuristic_v2_success_threshold).toBe(threshold); + const saved = buildUpdatedComplexityRouterConfig(stored, hydrated); + expect(saved.heuristic_v2_success_threshold).toBe(threshold); + expect(saved.heuristic_v2_artifact).toEqual(stored.heuristic_v2_artifact); + + const cleared = buildUpdatedComplexityRouterConfig(stored, { + ...hydrated, + heuristic_v2_success_threshold: undefined, + }); + expect(cleared).not.toHaveProperty("heuristic_v2_success_threshold"); + expect(cleared.heuristic_v2_artifact).toEqual(stored.heuristic_v2_artifact); + }); + + it.each([undefined, null])("keeps an inherited success threshold %s omitted after saving", (threshold) => { + const stored = { ...STORED, heuristic_v2_success_threshold: threshold }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + expect(hydrated.heuristic_v2_success_threshold).toBeUndefined(); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated)).not.toHaveProperty("heuristic_v2_success_threshold"); + }); + it.each(["capability", "llm_v2", "heuristic"] as const)( "handles enabled stored overrides when editing %s with or without keyword form state", (classifier_type) => { @@ -669,6 +697,7 @@ describe("managed keys survive an untouched open-and-save", () => { plan_mode_min_tier: "COMPLEX", tier_labels: { SIMPLE: "Cheap" }, classifier_type: "heuristic_first", + heuristic_v2_success_threshold: 0.89, heuristic_first_max_tier: "SIMPLE", classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000, reasoning_effort: "low" }, classifier_context_window_size: 5, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 0bb3340ac09..34db61483cf 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -132,6 +132,74 @@ describe("EditAutoRouterModal keyword matching", () => { expect(await screen.findByText(/Keyword\/Semantic Matching/i)).toBeInTheDocument(); }); + it.each(["0", ""])("hydrates the saved threshold and saves an edit to '%s'", async (raw) => { + const user = userEvent.setup(); + renderModal({ + modelData: { + ...MODEL_DATA, + litellm_params: { + ...MODEL_DATA.litellm_params, + complexity_router_config: { + ...STORED_CONFIG, + classifier_type: "heuristic_v2", + heuristic_v2_success_threshold: 0.91, + }, + }, + }, + }); + await user.click(await screen.findByText("Advanced: Classification Method")); + const threshold = screen.getByRole("textbox", { name: "Success threshold" }); + expect(threshold).toHaveValue("0.91"); + fireEvent.change(threshold, { target: { value: raw } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + if (raw === "") expect(savedConfig()).not.toHaveProperty("heuristic_v2_success_threshold"); + else expect(savedConfig().heuristic_v2_success_threshold).toBe(0); + }); + + it("blocks an invalid threshold edit and retains a corrected value when switching classifiers", async () => { + const user = userEvent.setup(); + renderModal({ + modelData: { + ...MODEL_DATA, + litellm_params: { + ...MODEL_DATA.litellm_params, + complexity_router_config: { + ...STORED_CONFIG, + classifier_type: "heuristic_v2", + heuristic_v2_success_threshold: 0.91, + }, + }, + }, + }); + await user.click(await screen.findByText("Advanced: Classification Method")); + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "-0.1" } }); + expect(screen.getByRole("button", { name: "Save Changes" })).toBeDisabled(); + expect(modelPatchUpdateCall).not.toHaveBeenCalled(); + + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "0.88" } }); + await user.click(screen.getByRole("radio", { name: /^Heuristic \(default\)/ })); + expect(screen.queryByRole("textbox", { name: "Success threshold" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + expect(savedConfig()).toMatchObject({ classifier_type: "heuristic", heuristic_v2_success_threshold: 0.88 }); + }); + + it("clears an invalid inactive threshold before saving the router", async () => { + const user = userEvent.setup(); + renderModal(); + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(screen.getByRole("radio", { name: /^Heuristic v2/ })); + fireEvent.change(screen.getByRole("textbox", { name: "Success threshold" }), { target: { value: "1.1" } }); + await user.click(screen.getByRole("radio", { name: /^Heuristic \(default\)/ })); + expect(screen.getByRole("button", { name: "Save Changes" })).toBeDisabled(); + await user.click(screen.getByRole("button", { name: "Clear Heuristic v2 threshold" })); + expect(screen.queryByRole("region", { name: "Inactive Heuristic v2 threshold" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + expect(savedConfig()).not.toHaveProperty("heuristic_v2_success_threshold"); + }); + // These keys are rewritten from form state on save, so if the modal renders the controls // without hydrating them, an untouched save silently wipes the stored configuration. This // drives the real component; a test of the payload builder alone cannot see that bug. 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 e25c7f07dd7..5991049f4fc 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 @@ -42,6 +42,7 @@ import { type BuildComplexityRouterConfigParams, buildComplexityRouterConfig, getClassifierModelError, + getHeuristicV2SuccessThresholdError, getClassifierReasoningEffortError, getKeywordTierRulesError, getMissingTiersError, @@ -127,6 +128,10 @@ export const hydrateComplexityRouterConfig = ( plan_mode_min_tier: hydratePlanModeMinTier(parsedConfig.plan_mode_min_tier, custom_tier_set), tier_labels: hydrateTierLabels(parsedConfig.tier_labels), classifier_type: parsedConfig.classifier_type || "heuristic", + heuristic_v2_success_threshold: + typeof parsedConfig.heuristic_v2_success_threshold === "number" + ? parsedConfig.heuristic_v2_success_threshold + : undefined, capability_classifier_config: capabilitySettingsSchema.safeParse(parsedConfig.capability_classifier_config).data, llm_v2_config: fuseSettingsSchema.safeParse(parsedConfig.llm_v2_config).data, classifier_llm_config: parsedConfig.classifier_llm_config, @@ -227,6 +232,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classification_examples", "heuristic_first_max_tier", "hybrid_boundary_margin", + "heuristic_v2_success_threshold", "classification_mode", "session_affinity", "session_affinity_ttl_seconds", @@ -329,6 +335,7 @@ export const buildUpdatedComplexityRouterConfig = ( classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, + heuristicV2SuccessThreshold: value.heuristic_v2_success_threshold, capabilityClassifierConfig: value.capability_classifier_config, llmV2Config: value.llm_v2_config, classifierLlmConfig: value.classifier_llm_config, @@ -427,6 +434,7 @@ const EditAutoRouterModal: React.FC = ({ getPlanModeTierError(complexityRouterConfig.plan_mode_min_tier, activeTierRows(complexityRouterConfig)) ?? getKeywordTierRulesError(keywordTierRules, activeTierRows(complexityRouterConfig)) ?? getClassifierModelError(complexityRouterConfig) ?? + getHeuristicV2SuccessThresholdError(complexityRouterConfig.heuristic_v2_success_threshold) ?? getForecastConfigError(complexityRouterConfig) ?? (heuristicScoringRole(complexityRouterConfig) === "decides" ? customDimensionsError(complexityRouterConfig.custom_dimensions) @@ -559,6 +567,7 @@ const EditAutoRouterModal: React.FC = ({ } const classifierError = getClassifierModelError(complexityRouterConfig) ?? + getHeuristicV2SuccessThresholdError(complexityRouterConfig.heuristic_v2_success_threshold) ?? getForecastConfigError(complexityRouterConfig) ?? (heuristicScoringRole(complexityRouterConfig) === "decides" ? customDimensionsError(complexityRouterConfig.custom_dimensions) diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index fed11454c23..7c49b15e279 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -680,6 +680,17 @@ describe("autorouter_presets", () => { }); describe("buildPresetPrefill", () => { + it.each([undefined, 0, 0.95])("carries a preset's success threshold %s into the form", (threshold) => { + const preset = getPresetByKey("anthropic_family")!; + const config = { + ...preset.complexity_router_config, + classifier_type: "heuristic_v2" as const, + heuristic_v2_success_threshold: threshold, + }; + const prefill = buildPresetPrefill(config, groupsOnly(getRequiredModelsInPreset(preset))); + expect(prefill.complexityRouterConfig.heuristic_v2_success_threshold).toBe(threshold); + }); + it("prefills a real bundled preset's tiers into the config", () => { const preset = getPresetByKey("anthropic_family")!; const prefill = buildPresetPrefill( diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 02096cada41..f085b4760a9 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -284,6 +284,7 @@ export const buildPresetPrefill = ( tier_model_params: resolveParamKeys(hydrateTierModelParams(config.tiers, config.tier_model_configs)), tier_labels: hydrateTierLabels(config.tier_labels), classifier_type: config.classifier_type, + heuristic_v2_success_threshold: config.heuristic_v2_success_threshold, classifier_llm_config: config.classifier_llm_config && { ...config.classifier_llm_config, model: resolve(config.classifier_llm_config.model), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d62a758e3a8..aa71adfad42 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36684,6 +36684,11 @@ export interface components { * @default ultrafeedback */ heuristic_v2_artifact: components["schemas"]["TrainedTierArtifact"] | "ultrafeedback"; + /** + * Heuristic V2 Success Threshold + * @description Minimum predicted success probability for classifier_type 'heuristic_v2' to select a tier. The first tier meeting this threshold is selected, or REASONING if none meets it. When omitted or null, uses the artifact's routing_threshold (0.75 for the bundled artifact). Other classifier types ignore this setting + */ + heuristic_v2_success_threshold?: number | null; /** * Housekeeping Patterns * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. From 58729ac69fb66f7f0f4ad2cc5af92c52f0b21e9a Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 16:10:40 +0000 Subject: [PATCH 142/146] test(a2a): sort imports in merged bedrock agentcore test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index cc314351fc2..1c87fb7564d 100644 --- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -10,12 +10,11 @@ Verifies that: """ import json +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from unittest.mock import AsyncMock, MagicMock, patch - SAMPLE_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/my_agent" SAMPLE_MODEL = f"bedrock/agentcore/{SAMPLE_ARN}" From 3d15f08fdad59c5f1c6a020dc99c5216333a5ae5 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:01:30 +0000 Subject: [PATCH 143/146] chore(prices): sync Together AI prices: 2 models together_ai/Qwen/Qwen3.7-Max: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost together_ai/Qwen/Qwen3.8-Flash: input_cost_per_token, output_cost_per_token --- litellm/model_prices_and_context_window_backup.json | 10 +++++----- model_prices_and_context_window.json | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5b686f0a4b5..7c2b5d450b2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -46584,13 +46584,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 7.5e-06, + "output_cost_per_token": 4.5e-06, "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, @@ -64470,12 +64470,12 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3.8-Flash": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_token": 9e-08, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.7e-07, + "output_cost_per_token": 2.82e-07, "source": "https://api.together.ai/v1/models" }, "together_ai/moonshotai/Kimi-K2.6": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5b686f0a4b5..7c2b5d450b2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -46584,13 +46584,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 7.5e-06, + "output_cost_per_token": 4.5e-06, "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, @@ -64470,12 +64470,12 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3.8-Flash": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_token": 9e-08, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.7e-07, + "output_cost_per_token": 2.82e-07, "source": "https://api.together.ai/v1/models" }, "together_ai/moonshotai/Kimi-K2.6": { From 43b81e448d10da4420f74bc44ef6ad5bacceac12 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:01:42 +0000 Subject: [PATCH 144/146] chore(prices): sync OpenRouter prices: 6 models, 1 new openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/~x-ai/grok-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, input_cost_per_token_above_200k_tokens, output_cost_per_token_above_200k_tokens, cache_read_input_token_cost_above_200k_tokens openrouter/anthropic/claude-sonnet-4: max_input_tokens openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/x-ai/grok-4.7: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_pdf_input, supports_reasoning, supports_web_search, supports_audio_input, supports_tool_choice, supports_prompt_caching, supports_response_schema, supports_function_calling, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, input_cost_per_token_above_200k_tokens, output_cost_per_token_above_200k_tokens, cache_read_input_token_cost_above_200k_tokens --- ...odel_prices_and_context_window_backup.json | 59 +++++++++++++------ model_prices_and_context_window.json | 59 +++++++++++++------ 2 files changed, 82 insertions(+), 36 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5b686f0a4b5..1eedbbde6e1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41204,7 +41204,7 @@ "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "litellm_provider": "openrouter", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -41591,21 +41591,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.42906e-07, + "input_cost_per_token": 9.34554e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.885812e-06, + "output_cost_per_token": 1.869108e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.85755e-08, + "cache_read_input_token_cost": 7.78795e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41633,22 +41633,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7684e-07, + "input_cost_per_token": 5.7156e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.73052e-06, + "output_cost_per_token": 1.71468e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.8354e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, + "cache_read_input_token_cost": 1.8186e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71485,15 +71485,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.8354e-08, - "input_cost_per_token": 5.7684e-07, + "cache_read_input_token_cost": 1.8186e-08, + "input_cost_per_token": 5.7156e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, - "output_cost_per_token": 1.73052e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, + "output_cost_per_token": 1.71468e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71715,17 +71715,17 @@ "supports_web_search": true }, "openrouter/~x-ai/grok-latest": { - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_200k_tokens": 8e-07, + "input_cost_per_token": 1.6e-06, + "input_cost_per_token_above_200k_tokens": 3.2e-06, "litellm_provider": "openrouter", "max_input_tokens": 500000, "max_output_tokens": 450000, "max_tokens": 450000, "mode": "chat", - "output_cost_per_token": 6e-06, - "output_cost_per_token_above_200k_tokens": 1.2e-05, + "output_cost_per_token": 4.8e-06, + "output_cost_per_token_above_200k_tokens": 9.6e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -75338,5 +75338,28 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false + }, + "openrouter/x-ai/grok-4.7": { + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_200k_tokens": 8e-07, + "input_cost_per_token": 1.6e-06, + "input_cost_per_token_above_200k_tokens": 3.2e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 500000, + "max_output_tokens": 450000, + "max_tokens": 450000, + "mode": "chat", + "output_cost_per_token": 4.8e-06, + "output_cost_per_token_above_200k_tokens": 9.6e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5b686f0a4b5..1eedbbde6e1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41204,7 +41204,7 @@ "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "litellm_provider": "openrouter", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -41591,21 +41591,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.42906e-07, + "input_cost_per_token": 9.34554e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.885812e-06, + "output_cost_per_token": 1.869108e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.85755e-08, + "cache_read_input_token_cost": 7.78795e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41633,22 +41633,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7684e-07, + "input_cost_per_token": 5.7156e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.73052e-06, + "output_cost_per_token": 1.71468e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.8354e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, + "cache_read_input_token_cost": 1.8186e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71485,15 +71485,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.8354e-08, - "input_cost_per_token": 5.7684e-07, + "cache_read_input_token_cost": 1.8186e-08, + "input_cost_per_token": 5.7156e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7684e-7,"output_cost_per_token":0.00000173052,"cache_read_input_token_cost":1.8354e-8}, - "output_cost_per_token": 1.73052e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, + "output_cost_per_token": 1.71468e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71715,17 +71715,17 @@ "supports_web_search": true }, "openrouter/~x-ai/grok-latest": { - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_200k_tokens": 8e-07, + "input_cost_per_token": 1.6e-06, + "input_cost_per_token_above_200k_tokens": 3.2e-06, "litellm_provider": "openrouter", "max_input_tokens": 500000, "max_output_tokens": 450000, "max_tokens": 450000, "mode": "chat", - "output_cost_per_token": 6e-06, - "output_cost_per_token_above_200k_tokens": 1.2e-05, + "output_cost_per_token": 4.8e-06, + "output_cost_per_token_above_200k_tokens": 9.6e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -75338,5 +75338,28 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false + }, + "openrouter/x-ai/grok-4.7": { + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_200k_tokens": 8e-07, + "input_cost_per_token": 1.6e-06, + "input_cost_per_token_above_200k_tokens": 3.2e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 500000, + "max_output_tokens": 450000, + "max_tokens": 450000, + "mode": "chat", + "output_cost_per_token": 4.8e-06, + "output_cost_per_token_above_200k_tokens": 9.6e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true } } From bed94a48d964f8216fb65a80c6259e057c7057b2 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:31:40 +0000 Subject: [PATCH 145/146] chore(prices): sync OpenRouter prices: 7 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/~moonshotai/kimi-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/ibm-granite/granite-4.2-8b: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/moonshotai/kimi-k3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/qwen/qwen3.8-27b: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 54 +++++++++---------- model_prices_and_context_window.json | 54 +++++++++---------- 2 files changed, 54 insertions(+), 54 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ff57a10cd9a..ef0fd09382e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41591,21 +41591,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.34554e-07, + "input_cost_per_token": 9.27768e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.869108e-06, + "output_cost_per_token": 1.855536e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.78795e-08, + "cache_read_input_token_cost": 7.7314e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41633,22 +41633,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7156e-07, + "input_cost_per_token": 5.7024e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.71468e-06, + "output_cost_per_token": 1.71072e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.8186e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, + "cache_read_input_token_cost": 1.9008e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7024e-7,"output_cost_per_token":0.00000171072,"cache_read_input_token_cost":1.9008e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -66759,9 +66759,9 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.5e-06, - "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.2e-07, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 8.5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, @@ -66942,9 +66942,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.7e-06, - "output_cost_per_token": 8.5e-06, - "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71485,15 +71485,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.8186e-08, - "input_cost_per_token": 5.7156e-07, + "cache_read_input_token_cost": 1.9008e-08, + "input_cost_per_token": 5.7024e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, - "output_cost_per_token": 1.71468e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7024e-7,"output_cost_per_token":0.00000171072,"cache_read_input_token_cost":1.9008e-8}, + "output_cost_per_token": 1.71072e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71575,14 +71575,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.7e-07, - "input_cost_per_token": 1.7e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8.5e-06, + "output_cost_per_token": 1.5e-05, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -72895,14 +72895,14 @@ "supports_web_search": false }, "openrouter/ibm-granite/granite-4.2-8b": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 1e-07, + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 6e-08, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 117964, "max_tokens": 117964, "mode": "chat", - "output_cost_per_token": 1.5e-07, + "output_cost_per_token": 2.5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ff57a10cd9a..ef0fd09382e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41591,21 +41591,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.34554e-07, + "input_cost_per_token": 9.27768e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.869108e-06, + "output_cost_per_token": 1.855536e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.78795e-08, + "cache_read_input_token_cost": 7.7314e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -41633,22 +41633,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.7156e-07, + "input_cost_per_token": 5.7024e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.71468e-06, + "output_cost_per_token": 1.71072e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.8186e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, + "cache_read_input_token_cost": 1.9008e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7024e-7,"output_cost_per_token":0.00000171072,"cache_read_input_token_cost":1.9008e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -66759,9 +66759,9 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.5e-06, - "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.2e-07, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 8.5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, @@ -66942,9 +66942,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.7e-06, - "output_cost_per_token": 8.5e-06, - "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71485,15 +71485,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.8186e-08, - "input_cost_per_token": 5.7156e-07, + "cache_read_input_token_cost": 1.9008e-08, + "input_cost_per_token": 5.7024e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7156e-7,"output_cost_per_token":0.00000171468,"cache_read_input_token_cost":1.8186e-8}, - "output_cost_per_token": 1.71468e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.7024e-7,"output_cost_per_token":0.00000171072,"cache_read_input_token_cost":1.9008e-8}, + "output_cost_per_token": 1.71072e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71575,14 +71575,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.7e-07, - "input_cost_per_token": 1.7e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8.5e-06, + "output_cost_per_token": 1.5e-05, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -72895,14 +72895,14 @@ "supports_web_search": false }, "openrouter/ibm-granite/granite-4.2-8b": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 1e-07, + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 6e-08, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 117964, "max_tokens": 117964, "mode": "chat", - "output_cost_per_token": 1.5e-07, + "output_cost_per_token": 2.5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 00b298a36ddb329dcc85be9cc91b1ed843c016b4 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:31:49 +0000 Subject: [PATCH 146/146] chore(prices): sync AWS Bedrock prices: 13 models, 1 new [1 with gaps, enrichment failed: AWS Bedrock, 38 held] deepseek.v3.2: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_audio_input, supports_response_schema global.moonshotai.kimi-k3: supports_vision, max_input_tokens, supports_audio_input, supports_response_schema, supports_function_calling, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, cache_creation_input_token_cost, supports_tool_choice, supports_prompt_caching google.gemma-3-12b-it: max_tokens, max_output_tokens, supports_audio_input, supports_response_schema, supports_function_calling google.gemma-3-4b-it: max_tokens, max_output_tokens, supports_audio_input, supports_function_calling mistral.devstral-2-123b: max_tokens, supports_vision, max_output_tokens, supports_audio_input, supports_response_schema mistral.magistral-small-2509: max_tokens, supports_vision, max_output_tokens, supports_audio_input, supports_response_schema mistral.ministral-3-14b-instruct: max_tokens, supports_vision, max_output_tokens, supports_audio_input, supports_response_schema mistral.ministral-3-8b-instruct: max_tokens, supports_vision, max_output_tokens, supports_audio_input, supports_response_schema mistral.mistral-large-3-675b-instruct: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_audio_input, supports_response_schema moonshotai.kimi-k2.5: max_tokens, max_input_tokens, max_output_tokens, supports_audio_input, supports_response_schema nvidia.nemotron-nano-3-30b: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_audio_input, supports_response_schema nvidia.nemotron-nano-9b-v2: max_tokens, supports_vision, max_output_tokens, supports_audio_input, supports_response_schema, supports_function_calling nvidia.nemotron-super-3-120b: max_tokens, supports_vision, max_output_tokens, supports_audio_input, supports_response_schema --- ...odel_prices_and_context_window_backup.json | 134 +++++++++++++----- model_prices_and_context_window.json | 134 +++++++++++++----- 2 files changed, 192 insertions(+), 76 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ff57a10cd9a..a40f7bd46d7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22019,16 +22019,19 @@ "deepseek.v3.2": { "input_cost_per_token": 6.2e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "max_tokens": 163840, + "max_input_tokens": 164000, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "dolphin": { "input_cost_per_token": 5e-07, @@ -30353,10 +30356,14 @@ "input_cost_per_token": 9e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.9e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true }, @@ -30375,10 +30382,13 @@ "input_cost_per_token": 4e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 8e-08, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, "supports_system_messages": true, "supports_vision": true }, @@ -36997,38 +37007,49 @@ "input_cost_per_token": 4e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "mistral.magistral-small-2509": { "input_cost_per_token": 5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 40000, + "max_tokens": 40000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, - "supports_system_messages": true + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true }, "mistral.ministral-3-14b-instruct": { "input_cost_per_token": 2e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": true }, "mistral.ministral-3-3b-instruct": { "input_cost_per_token": 1e-07, @@ -37046,13 +37067,17 @@ "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.5e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": true }, "mistral.mistral-7b-instruct-v0:2": { "input_cost_per_token": 1.5e-07, @@ -37088,14 +37113,18 @@ "mistral.mistral-large-3-675b-instruct": { "input_cost_per_token": 5e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 256000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": true }, "mistral.mistral-small-2402-v1:0": { "input_cost_per_token": 1e-06, @@ -38314,16 +38343,18 @@ "moonshotai.kimi-k2.5": { "input_cost_per_token": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "max_tokens": 262144, + "max_input_tokens": 256000, + "max_output_tokens": 16000, + "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true }, "moonshot/kimi-k2-0711-preview": { "cache_read_input_token_cost": 1.5e-07, @@ -39589,39 +39620,50 @@ "input_cost_per_token": 6e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.3e-07, - "supports_system_messages": true + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false }, "nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 6e-08, "litellm_provider": "bedrock_converse", - "max_input_tokens": 262144, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.4e-07, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/", - "supports_native_structured_output": true + "supports_audio_input": false, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": false }, "nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, - "max_output_tokens": 32768, - "max_tokens": 32768, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 6.5e-07, "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": false }, "o1": { "cache_read_input_token_cost": 7.5e-06, @@ -75361,5 +75403,21 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true + }, + "global.moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ff57a10cd9a..a40f7bd46d7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22019,16 +22019,19 @@ "deepseek.v3.2": { "input_cost_per_token": 6.2e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "max_tokens": 163840, + "max_input_tokens": 164000, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "dolphin": { "input_cost_per_token": 5e-07, @@ -30353,10 +30356,14 @@ "input_cost_per_token": 9e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.9e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true }, @@ -30375,10 +30382,13 @@ "input_cost_per_token": 4e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 8e-08, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, "supports_system_messages": true, "supports_vision": true }, @@ -36997,38 +37007,49 @@ "input_cost_per_token": 4e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "mistral.magistral-small-2509": { "input_cost_per_token": 5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 40000, + "max_tokens": 40000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, - "supports_system_messages": true + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true }, "mistral.ministral-3-14b-instruct": { "input_cost_per_token": 2e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": true }, "mistral.ministral-3-3b-instruct": { "input_cost_per_token": 1e-07, @@ -37046,13 +37067,17 @@ "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.5e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": true }, "mistral.mistral-7b-instruct-v0:2": { "input_cost_per_token": 1.5e-07, @@ -37088,14 +37113,18 @@ "mistral.mistral-large-3-675b-instruct": { "input_cost_per_token": 5e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 256000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": true }, "mistral.mistral-small-2402-v1:0": { "input_cost_per_token": 1e-06, @@ -38314,16 +38343,18 @@ "moonshotai.kimi-k2.5": { "input_cost_per_token": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "max_tokens": 262144, + "max_input_tokens": 256000, + "max_output_tokens": 16000, + "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true }, "moonshot/kimi-k2-0711-preview": { "cache_read_input_token_cost": 1.5e-07, @@ -39589,39 +39620,50 @@ "input_cost_per_token": 6e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.3e-07, - "supports_system_messages": true + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false }, "nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 6e-08, "litellm_provider": "bedrock_converse", - "max_input_tokens": 262144, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.4e-07, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/", - "supports_native_structured_output": true + "supports_audio_input": false, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_vision": false }, "nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, - "max_output_tokens": 32768, - "max_tokens": 32768, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 6.5e-07, "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": false }, "o1": { "cache_read_input_token_cost": 7.5e-06, @@ -75361,5 +75403,21 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true + }, + "global.moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } }