From d94b227906a67d6f283c890f4080259e0706160d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 12:27:24 +0000 Subject: [PATCH 001/109] fix(mcp): stable ordering for MCP servers list in Admin UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_management_endpoints.py | 9 ++- .../test_mcp_management_endpoints.py | 70 +++++++++++++++++++ .../_components/mcp_servers.test.tsx | 29 +++++++- .../mcp-servers/_components/mcp_servers.tsx | 49 ++++++------- 4 files changed, 127 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 4c97bbaf5de..1dac2fd3a77 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1024,6 +1024,9 @@ if MCP_AVAILABLE: return {"servers": registry_servers} ## FastAPI Routes + def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]: + return ((server.server_name or server.alias or server.server_id).lower(), server.server_id) + def _get_user_mcp_management_mode() -> UserMCPManagementMode: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, @@ -1174,10 +1177,12 @@ if MCP_AVAILABLE: detail="You do not have permission to view MCP servers for this team.", ) - redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id) + redacted_mcp_servers = sorted( + await _get_team_scoped_mcp_server_list(sanitized_team_id), key=_mcp_server_display_order + ) else: servers: Final = await _resolve_accessible_mcp_servers(user_api_key_dict) - redacted_mcp_servers = _redact_mcp_credentials_list(servers) + redacted_mcp_servers = sorted(_redact_mcp_credentials_list(servers), key=_mcp_server_display_order) if connected_app_view is True and is_ui_session_credential(user_api_key_dict): reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 5e00e7d75be..96bb0e35f0f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1535,6 +1535,76 @@ class TestTeamScopedMCPServerAccess: result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id") assert len(result) == 1 + @pytest.mark.asyncio + async def test_team_scoped_list_is_sorted_by_display_name(self): + """Set-derived resolution order must not leak to the client.""" + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + unsorted = [ + generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"), + generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"), + generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"), + ] + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list", + AsyncMock(return_value=unsorted), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id") + assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"] + + +class TestFetchAllMCPServersOrdering: + @pytest.mark.asyncio + async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(self): + """The registry resolves ids through a set, so the response must impose its own order.""" + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + first_order = [ + generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"), + generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"), + generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"), + ] + second_order = list(reversed(first_order)) + + for resolved in (first_order, second_order): + mock_manager = MagicMock() + mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved) + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth) + assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"] + @pytest.mark.asyncio async def test_restricted_virtual_key_cannot_use_team_id_filter(self): """Restricted virtual keys must not bypass access limits via team_id.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 8abb8855e3d..390c42f1ea3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -3,7 +3,8 @@ import { render, waitFor, screen, act, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import MCPServers from "./mcp_servers"; +import MCPServers, { compareServers } from "./mcp_servers"; +import type { MCPServer } from "@/components/mcp_tools/types"; import * as networking from "@/components/networking"; // Mock the networking module @@ -29,6 +30,32 @@ const createQueryClient = () => }, }); +describe("compareServers", () => { + const server = (server_id: string, name: string, created_at = ""): MCPServer => + ({ server_id, server_name: name, created_at, updated_at: created_at }) as MCPServer; + + const shuffled = [server("c", "github"), server("a", "slack"), server("b", "Jira")]; + + it("orders servers without timestamps by name so config.yaml servers render in a stable order", () => { + const byCreated = [...shuffled].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id); + const byUpdated = [...shuffled].sort((a, b) => compareServers(a, b, "updated_desc")).map((s) => s.server_id); + const byHealth = [...shuffled].sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id); + + expect(byCreated).toEqual(["c", "b", "a"]); + expect(byUpdated).toEqual(["c", "b", "a"]); + expect(byHealth).toEqual(["c", "b", "a"]); + }); + + it("keeps newest-first when timestamps differ", () => { + const newest = server("new", "zzz", "2026-02-01T00:00:00Z"); + const oldest = server("old", "aaa", "2026-01-01T00:00:00Z"); + expect([oldest, newest].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id)).toEqual([ + "new", + "old", + ]); + }); +}); + describe("MCPServers", () => { const defaultProps = { accessToken: "123", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index e6148d5d997..902e8811996 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -45,7 +45,7 @@ import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; import UserEnvVarsModal from "./UserEnvVarsModal"; import { listMCPUserEnvVarStatus } from "@/components/networking"; -type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health"; +export type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health"; const SORT_OPTIONS: { value: SortKey; label: string }[] = [ { value: "created_desc", label: "Recently created" }, @@ -60,32 +60,33 @@ const HEALTH_RANK: Record = { healthy: 2, }; -const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => { +const compareByName = (a: MCPServer, b: MCPServer): number => { + const nameA = (a.server_name || a.alias || a.server_id).toLowerCase(); + const nameB = (b.server_name || b.alias || b.server_id).toLowerCase(); + return nameA.localeCompare(nameB) || a.server_id.localeCompare(b.server_id); +}; + +const compareByTimestampDesc = (a: string | null | undefined, b: string | null | undefined): number => { + const ta = a ? new Date(a).getTime() : 0; + const tb = b ? new Date(b).getTime() : 0; + return tb - ta; +}; + +export const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => { switch (sort) { - case "name_asc": { - const nameA = (a.server_name || a.alias || a.server_id).toLowerCase(); - const nameB = (b.server_name || b.alias || b.server_id).toLowerCase(); - return nameA.localeCompare(nameB); - } - case "updated_desc": { - const ta = a.updated_at ? new Date(a.updated_at).getTime() : 0; - const tb = b.updated_at ? new Date(b.updated_at).getTime() : 0; - return tb - ta; - } + case "name_asc": + return compareByName(a, b); + case "updated_desc": + return compareByTimestampDesc(a.updated_at, b.updated_at) || compareByName(a, b); case "health": { const ra = HEALTH_RANK[a.status ?? "unknown"] ?? 1; const rb = HEALTH_RANK[b.status ?? "unknown"] ?? 1; if (ra !== rb) return ra - rb; - const ta = a.created_at ? new Date(a.created_at).getTime() : 0; - const tb = b.created_at ? new Date(b.created_at).getTime() : 0; - return tb - ta; + return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b); } case "created_desc": - default: { - const ta = a.created_at ? new Date(a.created_at).getTime() : 0; - const tb = b.created_at ? new Date(b.created_at).getTime() : 0; - return tb - ta; - } + default: + return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b); } }; @@ -297,13 +298,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)), ); } - const sorted = [...filtered].sort((a, b) => { - if (!a.created_at && !b.created_at) return 0; - if (!a.created_at) return 1; - if (!b.created_at) return -1; - return new Date(b.created_at).getTime() - new Date(a.created_at).getTime(); - }); - setFilteredServers(sorted); + setFilteredServers(filtered); }, [serversWithHealth], ); From aec592b734a91c3d068f89f0977561b733bc926f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 12:43:50 +0000 Subject: [PATCH 002/109] test(mcp): drop internal patches from ordering tests and regenerate API types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_mcp_management_endpoints.py | 47 ++++++------------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 2 files changed, 14 insertions(+), 35 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 96bb0e35f0f..6381bb7b9fc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1535,38 +1535,19 @@ class TestTeamScopedMCPServerAccess: result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id") assert len(result) == 1 - @pytest.mark.asyncio - async def test_team_scoped_list_is_sorted_by_display_name(self): - """Set-derived resolution order must not leak to the client.""" - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN, - user_id="admin_user", - ) - unsorted = [ - generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"), - generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"), - generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"), - ] - - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list", - AsyncMock(return_value=unsorted), - ), - ): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - fetch_all_mcp_servers, - ) - - result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id") - assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"] - class TestFetchAllMCPServersOrdering: + def test_display_order_is_case_insensitive_name_then_id(self): + servers = [ + generate_mock_mcp_server_db_record(server_id="s-2", alias="github"), + generate_mock_mcp_server_db_record(server_id="s-1", alias="github"), + generate_mock_mcp_server_db_record(server_id="s-0", alias="Slack"), + generate_mock_mcp_server_db_record(server_id="s-3", alias="confluence"), + ] + + ordered = sorted(servers, key=mgmt_endpoints._mcp_server_display_order) + assert [s.server_id for s in ordered] == ["s-3", "s-1", "s-2", "s-0"] + @pytest.mark.asyncio async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(self): """The registry resolves ids through a set, so the response must impose its own order.""" @@ -1585,15 +1566,15 @@ class TestFetchAllMCPServersOrdering: mock_manager = MagicMock() mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved) with ( - patch( + patch( # test-quality-ok: the route reads a module-global manager with no injection seam "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", mock_manager, ), - patch( + patch( # test-quality-ok: admin view is derived from module-global proxy settings "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", return_value=True, ), - patch( + patch( # test-quality-ok: auth contexts need a live prisma client "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", AsyncMock(return_value=[mock_user_auth]), ), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From db17841b3d6ff41900276de6eac8c66fbc924801 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 22:37:16 -0700 Subject: [PATCH 003/109] test(model_management): cover actor edges and wildcard models --- .../test_model_management_endpoints.py | 340 +++++++++++++++++- .../handle_add_model_submit.test.tsx | 19 + 2 files changed, 356 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index d1fe88df26c..302585e42c4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1224,11 +1224,11 @@ class TestUpdateModel: "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), ), - patch( + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value: value, ), - patch( + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock( return_value=ReconcileOutcome(still_desired=None, live_after=None) @@ -4021,7 +4021,7 @@ class TestPatchModelBlockedAuthGate: "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), ), - patch( + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock( return_value=ReconcileOutcome(still_desired=None, live_after=None) @@ -6631,3 +6631,337 @@ class TestTeamMemberAutoRouterWrites: assert json.loads(written["model_info"])["member_auto_router"] is True assert appended.await_args.kwargs["data"].models == ["new-personal-router"] assert appended.await_args.kwargs["data"].team_id == "member-team" + + +class TestModelManagementActorEdges: + @pytest.mark.asyncio + async def test_add_model_rejects_non_team_internal_user(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER) + prisma: Final = MagicMock() + deployment: Final = Deployment( + model_name="internal-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="internal-model-id"), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model(model_params=deployment, user_api_key_dict=actor) + + assert str(exc_info.value.code) == "403" + assert "permission" in str(exc_info.value).lower() + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_add_model_rejects_proxy_admin_viewer(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + actor: Final = UserAPIKeyAuth( + user_id="view-only-user", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + prisma: Final = MagicMock() + deployment: Final = Deployment( + model_name="view-only-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="view-only-model-id"), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model(model_params=deployment, user_api_key_dict=actor) + + assert str(exc_info.value.code) == "403" + assert "view-only" in str(exc_info.value).lower() + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_add_model_requires_database_storage(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + prisma: Final = MagicMock() + deployment: Final = Deployment( + model_name="database-disabled-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="database-disabled-model-id"), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", False), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model(model_params=deployment, user_api_key_dict=actor) + + assert str(exc_info.value.code) == "500" + assert "STORE_MODEL_IN_DB" in str(exc_info.value) + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_legacy_model_update_persists_changed_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + + model_id: Final = "legacy-update-model-id" + existing_row: Final = MagicMock() + existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30} + existing_row.model_dump.return_value = { + "model_name": "legacy-update-model", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": model_id}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + router: Final = MagicMock() + router.get_model_ids.return_value = [model_id] + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(timeout=42), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=actor, + ) + + written: Final = json.loads( + prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) + assert written["timeout"] == 42 + assert written["model"] == "openai/test-model" + + @pytest.mark.asyncio + async def test_legacy_model_update_explicit_null_preserves_existing_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + + model_id: Final = "legacy-null-model-id" + existing_row: Final = MagicMock() + existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30} + existing_row.model_dump.return_value = { + "model_name": "legacy-null-model", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": model_id}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + router: Final = MagicMock() + router.get_model_ids.return_value = [model_id] + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(timeout=None), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=actor, + ) + + written: Final = json.loads( + prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) + assert written["timeout"] == 30 + + @pytest.mark.asyncio + async def test_patch_model_rejects_config_file_model(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + model_id: Final = "config-model-id" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_proxymodeltable.update = AsyncMock() + router: Final = MagicMock() + router.get_deployment.return_value = Deployment( + model_name="config-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id=model_id), + ) + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(timeout=42), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=actor, + ) + + assert str(exc_info.value.code) == "400" + assert "Cannot edit config-based model" in str(exc_info.value) + prisma.db.litellm_proxymodeltable.update.assert_not_awaited() + + @contextlib.contextmanager + def _client_for(self, actor: UserAPIKeyAuth) -> Iterator[TestClient]: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.proxy_server import app + + app.dependency_overrides[proxy_server.user_api_key_auth] = lambda: actor + try: + yield TestClient(app) + finally: + app.dependency_overrides.pop(proxy_server.user_api_key_auth, None) + + def test_post_model_new_binds_to_actor_guard(self): + actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER) + prisma: Final = MagicMock() + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + self._client_for(actor) as client, + ): + response: Final = client.post( + "/model/new", + json={ + "model_name": "internal-model", + "litellm_params": {"model": "openai/test-model"}, + "model_info": {"id": "internal-model-id"}, + }, + ) + + assert response.status_code == 403 + assert "permission" in response.text.lower() + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + def test_post_legacy_model_update_binds_to_persistence(self): + model_id: Final = "legacy-route-model-id" + existing_row: Final = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="legacy-route-model", + litellm_params={"model": "openai/test-model", "timeout": 30}, + model_info={"id": model_id}, + created_by="admin", + updated_by="admin", + ) + updated_row: Final = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="legacy-route-model", + litellm_params={"model": "openai/test-model", "timeout": 42}, + model_info={"id": model_id}, + created_by="admin", + updated_by="admin", + ) + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + router: Final = MagicMock() + router.get_model_ids.return_value = [model_id] + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch( # test-quality-ok: [TQ008] audit logging is outside the persistence contract + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", + new=AsyncMock(return_value=None), + ), + self._client_for(actor) as client, + ): + response: Final = client.post( + "/model/update", + json={ + "litellm_params": {"timeout": 42}, + "model_info": {"id": model_id}, + }, + ) + + assert response.status_code == 200, response.text + written: Final = json.loads( + prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) + assert written["timeout"] == 42 + + def test_patch_config_model_binds_to_patch_route(self): + model_id: Final = "config-route-model-id" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_proxymodeltable.update = AsyncMock() + router: Final = MagicMock() + router.get_deployment.return_value = Deployment( + model_name="config-route-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id=model_id), + ) + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + self._client_for(actor) as client, + ): + response: Final = client.patch( + f"/model/{model_id}/update", + json={ + "litellm_params": {"timeout": 42}, + "model_info": {"id": model_id}, + }, + ) + + assert response.status_code == 400 + assert "Cannot edit config-based model" in response.text + prisma.db.litellm_proxymodeltable.update.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx index 9d792480c9f..923cf2aa0e3 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx @@ -101,4 +101,23 @@ describe("prepareModelAddRequest", () => { expect(deployment.litellmParamsObj.litellm_credential_name).toBe("from-json"); expect(deployment.litellmParamsObj.timeout).toBe(5); }); + + it.each([ + ["OpenAI", "openai/*"], + ["Azure_AI_Studio", "azure_ai/*"], + ["Petals", "petals/*"], + ])("composes wildcard names for the all-model selection", async (custom_llm_provider, wildcardModel) => { + const formValues = { + model_mappings: [], + model: "all-wildcard", + custom_llm_provider, + }; + + const deployments = await prepareModelAddRequest({ ...formValues }, "token", null); + + expect(deployments).toHaveLength(1); + const [deployment] = deployments!; + expect(deployment.modelName).toBe(wildcardModel); + expect(deployment.litellmParamsObj.model).toBe(wildcardModel); + }); }); From c34adb4ab2bd1b6579cc3eb9a3922cf9703aae58 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 22:44:03 -0700 Subject: [PATCH 004/109] test(ui): cover narrowed dashboard form journeys --- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 55 ++++++++++++++ .../tests/tagManagement/tagManagement.spec.ts | 76 +++++++++++++++++++ ...PaginatedSearchSelect.integration.test.tsx | 45 +++++++++++ .../shared/SearchSelect.integration.test.tsx | 33 ++++++++ .../view_logs/RequestLogsFilters.test.tsx | 11 +++ 5 files changed, 220 insertions(+) create mode 100644 tests/e2e/ui/tests/prompts/addPrompt.spec.ts create mode 100644 tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts new file mode 100644 index 00000000000..cd87e4d3b56 --- /dev/null +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -0,0 +1,55 @@ +import { test, expect } from "@playwright/test"; + +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page as DashboardPage } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { masterKey, uniqueSuffix } from "../../helpers/traffic"; + +test.use({ storageState: ADMIN_STORAGE_PATH }); + +test.describe("Prompt upload form", () => { + test("uploads a prompt file and reads the created prompt back", async ({ + page, + }) => { + const promptId = `e2e-prompt-${uniqueSuffix()}`; + await navigateToPage(page, DashboardPage.Prompts); + await page.getByRole("button", { name: "Upload .prompt File" }).click(); + + try { + await expect( + page.getByRole("dialog", { name: "Add New Prompt" }), + ).toBeVisible(); + await page.getByLabel("Prompt ID").fill(promptId); + await page.locator('input[type="file"]').setInputFiles({ + name: "e2e.prompt", + mimeType: "text/plain", + buffer: Buffer.from( + 'model: fake-openai-gpt-4\ntemplate: "Hello {{name}}"\n', + ), + }); + await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); + await page.getByRole("button", { name: "Create Prompt" }).click(); + + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }) + .toBe(true); + await expect(page.getByText(promptId, { exact: true })).toBeVisible(); + } finally { + await page.request.delete( + `/prompts/${encodeURIComponent(promptId)}?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + } + }); +}); diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts new file mode 100644 index 00000000000..785211463dd --- /dev/null +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -0,0 +1,76 @@ +import { test, expect } from "@playwright/test"; + +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page as DashboardPage } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { masterKey, uniqueSuffix } from "../../helpers/traffic"; + +test.use({ storageState: ADMIN_STORAGE_PATH }); + +test.describe("Tag management", () => { + test("creates, edits, reopens, and reads back a tag", async ({ page }) => { + const tagName = `e2e-tag-${uniqueSuffix()}`; + const description = "synthetic tag description"; + const updatedDescription = `${description} updated`; + + await navigateToPage(page, DashboardPage.TagManagement); + await page.getByRole("button", { name: "+ Create New Tag" }).click(); + + try { + await expect( + page.getByRole("dialog", { name: "Create New Tag" }), + ).toBeVisible(); + await page.getByLabel("Tag Name").fill(tagName); + await page.getByLabel("Description").fill(description); + await page.getByRole("button", { name: "Create Tag" }).click(); + + await expect + .poll(async () => { + const response = await readBack< + Record> + >(page, "/tag/list"); + return Object.values(response).some((tag) => tag.name === tagName); + }) + .toBe(true); + await expect(page.getByText(tagName, { exact: true })).toBeVisible(); + + await page.getByText(tagName, { exact: true }).click(); + await expect(page.getByText("Tag Name:")).toBeVisible(); + await page.getByRole("button", { name: "Edit Tag" }).click(); + await page.getByLabel("Description").fill(updatedDescription); + const updateBody = await captureRequestBody( + page, + { method: "POST", urlIncludes: "/tag/update" }, + () => page.getByRole("button", { name: "Save Changes" }).click(), + ); + expect(updateBody).toMatchObject({ + name: tagName, + description: updatedDescription, + }); + + await expect + .poll(async () => { + const infoResponse = await page.request.post("/tag/info", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { names: [tagName] }, + }); + expect(infoResponse.ok()).toBe(true); + const info = (await infoResponse.json()) as Record< + string, + { description?: string } + >; + return info[tagName]?.description; + }) + .toBe(updatedDescription); + } finally { + await page.request.post("/tag/delete", { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { name: tagName }, + }); + } + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx index 2b948ca8420..6d30f1513ef 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx @@ -1,5 +1,6 @@ import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; +import { useQuery } from "@tanstack/react-query"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -406,4 +407,48 @@ describe("PaginatedSearchSelect", () => { expect(input).toHaveValue("aliasalpha"); await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("aliasalpha")); }); + + it("keeps the latest query results when an earlier response resolves last", async () => { + const pending = new Map void>(); + + function QueryBackedSelect() { + const [query, setQuery] = useState(""); + const result = useQuery({ + queryKey: ["paginated-select-race", query], + queryFn: () => + new Promise((resolve) => { + pending.set(query, resolve); + }), + enabled: query.length > 0, + }); + return ( + <> + + + + + ); + } + + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("button", { name: "Search A" })); + await user.click(screen.getByRole("button", { name: "Search B" })); + await waitFor(() => { + expect(pending.has("A")).toBe(true); + expect(pending.has("B")).toBe(true); + }); + + pending.get("B")?.([{ label: "B result", value: "b" }]); + await user.click(screen.getByRole("combobox")); + expect(await screen.findByText("B result")).toBeInTheDocument(); + + pending.get("A")?.([{ label: "A result", value: "a" }]); + await waitFor(() => expect(screen.queryByText("A result")).not.toBeInTheDocument()); + expect(screen.getByText("B result")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index c981010dff9..ed83acef14a 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -111,4 +111,37 @@ describe("SearchSelect", () => { expect(screen.queryByText("Growth")).not.toBeInTheDocument(); expect(onValueChange).not.toHaveBeenCalled(); }); + + it("supports keyboard select, clear, escape, blur, and reopen", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + function Controlled() { + const [value, setValue] = useState(null); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + + render(); + const input = screen.getByRole("combobox"); + await user.tab(); + await user.keyboard("{Enter}"); + await user.keyboard("{ArrowDown}{Enter}"); + expect(onValueChange).toHaveBeenLastCalledWith("team-1"); + const clear = screen.getByRole("button", { name: "Clear" }); + clear.focus(); + await user.keyboard("{Enter}"); + expect(onValueChange).toHaveBeenLastCalledWith(null); + await user.keyboard("{Escape}"); + await user.tab(); + await user.tab({ shift: true }); + expect(input).toHaveFocus(); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index 5a35c7ae16b..8d1847e0121 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -344,4 +344,15 @@ describe("RequestLogsFilters", () => { expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.CACHE_STATUS, undefined); }); + + it("clears the raw Error Code combobox through the undefined filter contract", async () => { + const user = userEvent.setup(); + const { set } = renderFilters({ [LOG_FILTER_IDS.ERROR_CODE]: "429" }); + const input = await screen.findByPlaceholderText("Select or type an error code"); + + await user.click(input); + await user.click(screen.getByRole("button", { name: "Clear", hidden: true })); + + expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.ERROR_CODE, undefined); + }); }); From 8d1ca16652056512999a61c87e96ae3aaf9c236d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:06:17 -0700 Subject: [PATCH 005/109] test(ui): strengthen dashboard journey assertions --- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 14 ++++++++--- .../tests/tagManagement/tagManagement.spec.ts | 3 ++- ...PaginatedSearchSelect.integration.test.tsx | 25 ++++++------------- .../shared/SearchSelect.integration.test.tsx | 9 +++---- 4 files changed, 24 insertions(+), 27 deletions(-) diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts index cd87e4d3b56..f9868f7b04b 100644 --- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -13,6 +13,7 @@ test.describe("Prompt upload form", () => { page, }) => { const promptId = `e2e-prompt-${uniqueSuffix()}`; + const promptContent = "Hello {{name}}"; await navigateToPage(page, DashboardPage.Prompts); await page.getByRole("button", { name: "Upload .prompt File" }).click(); @@ -25,7 +26,7 @@ test.describe("Prompt upload form", () => { name: "e2e.prompt", mimeType: "text/plain", buffer: Buffer.from( - 'model: fake-openai-gpt-4\ntemplate: "Hello {{name}}"\n', + `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, ), }); await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); @@ -39,17 +40,22 @@ test.describe("Prompt upload form", () => { headers: { Authorization: `Bearer ${masterKey()}` }, }, ); - return response.ok(); + if (!response.ok()) return undefined; + const promptInfo = (await response.json()) as { + raw_prompt_template?: { content?: string }; + }; + return promptInfo.raw_prompt_template?.content; }) - .toBe(true); + .toContain(promptContent); await expect(page.getByText(promptId, { exact: true })).toBeVisible(); } finally { - await page.request.delete( + const deleteResponse = await page.request.delete( `/prompts/${encodeURIComponent(promptId)}?environment=development`, { headers: { Authorization: `Bearer ${masterKey()}` }, }, ); + expect(deleteResponse.ok()).toBe(true); } }); }); diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts index 785211463dd..e1d5138ea90 100644 --- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -64,13 +64,14 @@ test.describe("Tag management", () => { }) .toBe(updatedDescription); } finally { - await page.request.post("/tag/delete", { + const deleteResponse = await page.request.post("/tag/delete", { headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json", }, data: { name: tagName }, }); + expect(deleteResponse.ok()).toBe(true); } }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx index 6d30f1513ef..905610a77e6 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx @@ -421,27 +421,18 @@ describe("PaginatedSearchSelect", () => { }), enabled: query.length > 0, }); - return ( - <> - - - - - ); + return ; } const user = userEvent.setup(); render(); - await user.click(screen.getByRole("button", { name: "Search A" })); - await user.click(screen.getByRole("button", { name: "Search B" })); - await waitFor(() => { - expect(pending.has("A")).toBe(true); - expect(pending.has("B")).toBe(true); - }); + const input = screen.getByRole("combobox"); + await user.click(input); + await user.type(input, "A"); + await waitFor(() => expect(pending.has("A")).toBe(true)); + await user.clear(input); + await user.type(input, "B"); + await waitFor(() => expect(pending.has("B")).toBe(true)); pending.get("B")?.([{ label: "B result", value: "b" }]); await user.click(screen.getByRole("combobox")); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index ed83acef14a..9f320049b09 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -112,7 +112,7 @@ describe("SearchSelect", () => { expect(onValueChange).not.toHaveBeenCalled(); }); - it("supports keyboard select, clear, escape, blur, and reopen", async () => { + it("supports keyboard select, clear, and reselect", async () => { const onValueChange = vi.fn(); const user = userEvent.setup(); function Controlled() { @@ -139,9 +139,8 @@ describe("SearchSelect", () => { clear.focus(); await user.keyboard("{Enter}"); expect(onValueChange).toHaveBeenLastCalledWith(null); - await user.keyboard("{Escape}"); - await user.tab(); - await user.tab({ shift: true }); - expect(input).toHaveFocus(); + input.focus(); + await user.keyboard("{Enter}{ArrowDown}{Enter}"); + expect(onValueChange).toHaveBeenLastCalledWith("team-1"); }); }); From b4c3adc37d1be33550803bce9c88bc190c8f4ec6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:12:33 -0700 Subject: [PATCH 006/109] test(ui): assert dashboard form cleanup --- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 39 ++++++++++++------- .../tests/tagManagement/tagManagement.spec.ts | 30 +++++++------- 2 files changed, 40 insertions(+), 29 deletions(-) diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts index f9868f7b04b..a4a6cf62b7e 100644 --- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -17,21 +17,32 @@ test.describe("Prompt upload form", () => { await navigateToPage(page, DashboardPage.Prompts); await page.getByRole("button", { name: "Upload .prompt File" }).click(); - try { - await expect( - page.getByRole("dialog", { name: "Add New Prompt" }), - ).toBeVisible(); - await page.getByLabel("Prompt ID").fill(promptId); - await page.locator('input[type="file"]').setInputFiles({ - name: "e2e.prompt", - mimeType: "text/plain", - buffer: Buffer.from( - `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, - ), - }); - await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); - await page.getByRole("button", { name: "Create Prompt" }).click(); + await expect( + page.getByRole("dialog", { name: "Add New Prompt" }), + ).toBeVisible(); + await page.getByLabel("Prompt ID").fill(promptId); + await page.locator('input[type="file"]').setInputFiles({ + name: "e2e.prompt", + mimeType: "text/plain", + buffer: Buffer.from( + `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, + ), + }); + await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); + await page.getByRole("button", { name: "Create Prompt" }).click(); + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }) + .toBe(true); + try { await expect .poll(async () => { const response = await page.request.get( diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts index e1d5138ea90..1104031263e 100644 --- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -17,22 +17,22 @@ test.describe("Tag management", () => { await navigateToPage(page, DashboardPage.TagManagement); await page.getByRole("button", { name: "+ Create New Tag" }).click(); - try { - await expect( - page.getByRole("dialog", { name: "Create New Tag" }), - ).toBeVisible(); - await page.getByLabel("Tag Name").fill(tagName); - await page.getByLabel("Description").fill(description); - await page.getByRole("button", { name: "Create Tag" }).click(); + await expect( + page.getByRole("dialog", { name: "Create New Tag" }), + ).toBeVisible(); + await page.getByLabel("Tag Name").fill(tagName); + await page.getByLabel("Description").fill(description); + await page.getByRole("button", { name: "Create Tag" }).click(); - await expect - .poll(async () => { - const response = await readBack< - Record> - >(page, "/tag/list"); - return Object.values(response).some((tag) => tag.name === tagName); - }) - .toBe(true); + await expect + .poll(async () => { + const response = await readBack< + Record> + >(page, "/tag/list"); + return Object.values(response).some((tag) => tag.name === tagName); + }) + .toBe(true); + try { await expect(page.getByText(tagName, { exact: true })).toBeVisible(); await page.getByText(tagName, { exact: true }).click(); From 035271b510d5f4f4053685cf156410775d8e5d46 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:15:31 -0700 Subject: [PATCH 007/109] test(ui): preserve cleanup on failed readback --- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 53 +++++++++---------- .../tests/tagManagement/tagManagement.spec.ts | 33 ++++++------ 2 files changed, 42 insertions(+), 44 deletions(-) diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts index a4a6cf62b7e..2b254c78b10 100644 --- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -17,32 +17,32 @@ test.describe("Prompt upload form", () => { await navigateToPage(page, DashboardPage.Prompts); await page.getByRole("button", { name: "Upload .prompt File" }).click(); - await expect( - page.getByRole("dialog", { name: "Add New Prompt" }), - ).toBeVisible(); - await page.getByLabel("Prompt ID").fill(promptId); - await page.locator('input[type="file"]').setInputFiles({ - name: "e2e.prompt", - mimeType: "text/plain", - buffer: Buffer.from( - `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, - ), - }); - await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); - await page.getByRole("button", { name: "Create Prompt" }).click(); - - await expect - .poll(async () => { - const response = await page.request.get( - `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, - { - headers: { Authorization: `Bearer ${masterKey()}` }, - }, - ); - return response.ok(); - }) - .toBe(true); try { + await expect( + page.getByRole("dialog", { name: "Add New Prompt" }), + ).toBeVisible(); + await page.getByLabel("Prompt ID").fill(promptId); + await page.locator('input[type="file"]').setInputFiles({ + name: "e2e.prompt", + mimeType: "text/plain", + buffer: Buffer.from( + `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, + ), + }); + await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); + await page.getByRole("button", { name: "Create Prompt" }).click(); + + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }) + .toBe(true); await expect .poll(async () => { const response = await page.request.get( @@ -60,13 +60,12 @@ test.describe("Prompt upload form", () => { .toContain(promptContent); await expect(page.getByText(promptId, { exact: true })).toBeVisible(); } finally { - const deleteResponse = await page.request.delete( + await page.request.delete( `/prompts/${encodeURIComponent(promptId)}?environment=development`, { headers: { Authorization: `Bearer ${masterKey()}` }, }, ); - expect(deleteResponse.ok()).toBe(true); } }); }); diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts index 1104031263e..785211463dd 100644 --- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -17,22 +17,22 @@ test.describe("Tag management", () => { await navigateToPage(page, DashboardPage.TagManagement); await page.getByRole("button", { name: "+ Create New Tag" }).click(); - await expect( - page.getByRole("dialog", { name: "Create New Tag" }), - ).toBeVisible(); - await page.getByLabel("Tag Name").fill(tagName); - await page.getByLabel("Description").fill(description); - await page.getByRole("button", { name: "Create Tag" }).click(); - - await expect - .poll(async () => { - const response = await readBack< - Record> - >(page, "/tag/list"); - return Object.values(response).some((tag) => tag.name === tagName); - }) - .toBe(true); try { + await expect( + page.getByRole("dialog", { name: "Create New Tag" }), + ).toBeVisible(); + await page.getByLabel("Tag Name").fill(tagName); + await page.getByLabel("Description").fill(description); + await page.getByRole("button", { name: "Create Tag" }).click(); + + await expect + .poll(async () => { + const response = await readBack< + Record> + >(page, "/tag/list"); + return Object.values(response).some((tag) => tag.name === tagName); + }) + .toBe(true); await expect(page.getByText(tagName, { exact: true })).toBeVisible(); await page.getByText(tagName, { exact: true }).click(); @@ -64,14 +64,13 @@ test.describe("Tag management", () => { }) .toBe(updatedDescription); } finally { - const deleteResponse = await page.request.post("/tag/delete", { + await page.request.post("/tag/delete", { headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json", }, data: { name: tagName }, }); - expect(deleteResponse.ok()).toBe(true); } }); }); From 43f096dde8ef6c0a0c20036f0a595a7608ea2ea1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:18:43 -0700 Subject: [PATCH 008/109] test(ui): preserve form failure evidence --- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 116 +++++++++------- .../tests/tagManagement/tagManagement.spec.ts | 124 ++++++++++-------- 2 files changed, 136 insertions(+), 104 deletions(-) diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts index 2b254c78b10..b601b7e8c06 100644 --- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -3,7 +3,6 @@ import { test, expect } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page as DashboardPage } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; -import { readBack } from "../../helpers/roundTrip"; import { masterKey, uniqueSuffix } from "../../helpers/traffic"; test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -14,58 +13,75 @@ test.describe("Prompt upload form", () => { }) => { const promptId = `e2e-prompt-${uniqueSuffix()}`; const promptContent = "Hello {{name}}"; - await navigateToPage(page, DashboardPage.Prompts); - await page.getByRole("button", { name: "Upload .prompt File" }).click(); + const cleanup = async (): Promise => { + try { + const response = await page.request.delete( + `/prompts/${encodeURIComponent(promptId)}?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + } catch { + return false; + } + }; + const testOutcome = await (async () => { + try { + await navigateToPage(page, DashboardPage.Prompts); + await page.getByRole("button", { name: "Upload .prompt File" }).click(); + await expect( + page.getByRole("dialog", { name: "Add New Prompt" }), + ).toBeVisible(); + await page.getByLabel("Prompt ID").fill(promptId); + await page.locator('input[type="file"]').setInputFiles({ + name: "e2e.prompt", + mimeType: "text/plain", + buffer: Buffer.from( + `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, + ), + }); + await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); + await page.getByRole("button", { name: "Create Prompt" }).click(); + + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }) + .toBe(true); + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + if (!response.ok()) return undefined; + const promptInfo = (await response.json()) as { + raw_prompt_template?: { content?: string }; + }; + return promptInfo.raw_prompt_template?.content; + }) + .toContain(promptContent); + await expect(page.getByText(promptId, { exact: true })).toBeVisible(); + return { passed: true as const }; + } catch (error) { + return { passed: false as const, error }; + } + })(); try { - await expect( - page.getByRole("dialog", { name: "Add New Prompt" }), - ).toBeVisible(); - await page.getByLabel("Prompt ID").fill(promptId); - await page.locator('input[type="file"]').setInputFiles({ - name: "e2e.prompt", - mimeType: "text/plain", - buffer: Buffer.from( - `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, - ), - }); - await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); - await page.getByRole("button", { name: "Create Prompt" }).click(); - - await expect - .poll(async () => { - const response = await page.request.get( - `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, - { - headers: { Authorization: `Bearer ${masterKey()}` }, - }, - ); - return response.ok(); - }) - .toBe(true); - await expect - .poll(async () => { - const response = await page.request.get( - `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, - { - headers: { Authorization: `Bearer ${masterKey()}` }, - }, - ); - if (!response.ok()) return undefined; - const promptInfo = (await response.json()) as { - raw_prompt_template?: { content?: string }; - }; - return promptInfo.raw_prompt_template?.content; - }) - .toContain(promptContent); - await expect(page.getByText(promptId, { exact: true })).toBeVisible(); + if (!testOutcome.passed) throw testOutcome.error; } finally { - await page.request.delete( - `/prompts/${encodeURIComponent(promptId)}?environment=development`, - { - headers: { Authorization: `Bearer ${masterKey()}` }, - }, - ); + const cleanupSucceeded = await cleanup(); + if (testOutcome.passed) expect(cleanupSucceeded).toBe(true); } }); }); diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts index 785211463dd..4223324ff09 100644 --- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -13,64 +13,80 @@ test.describe("Tag management", () => { const tagName = `e2e-tag-${uniqueSuffix()}`; const description = "synthetic tag description"; const updatedDescription = `${description} updated`; + const cleanup = async (): Promise => { + try { + const response = await page.request.post("/tag/delete", { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { name: tagName }, + }); + return response.ok(); + } catch { + return false; + } + }; + const testOutcome = await (async () => { + try { + await navigateToPage(page, DashboardPage.TagManagement); + await page.getByRole("button", { name: "+ Create New Tag" }).click(); + await expect( + page.getByRole("dialog", { name: "Create New Tag" }), + ).toBeVisible(); + await page.getByLabel("Tag Name").fill(tagName); + await page.getByLabel("Description").fill(description); + await page.getByRole("button", { name: "Create Tag" }).click(); - await navigateToPage(page, DashboardPage.TagManagement); - await page.getByRole("button", { name: "+ Create New Tag" }).click(); + await expect + .poll(async () => { + const response = await readBack< + Record> + >(page, "/tag/list"); + return Object.values(response).some((tag) => tag.name === tagName); + }) + .toBe(true); + await expect(page.getByText(tagName, { exact: true })).toBeVisible(); + + await page.getByText(tagName, { exact: true }).click(); + await expect(page.getByText("Tag Name:")).toBeVisible(); + await page.getByRole("button", { name: "Edit Tag" }).click(); + await page.getByLabel("Description").fill(updatedDescription); + const updateBody = await captureRequestBody( + page, + { method: "POST", urlIncludes: "/tag/update" }, + () => page.getByRole("button", { name: "Save Changes" }).click(), + ); + expect(updateBody).toMatchObject({ + name: tagName, + description: updatedDescription, + }); + + await expect + .poll(async () => { + const infoResponse = await page.request.post("/tag/info", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { names: [tagName] }, + }); + expect(infoResponse.ok()).toBe(true); + const info = (await infoResponse.json()) as Record< + string, + { description?: string } + >; + return info[tagName]?.description; + }) + .toBe(updatedDescription); + return { passed: true as const }; + } catch (error) { + return { passed: false as const, error }; + } + })(); try { - await expect( - page.getByRole("dialog", { name: "Create New Tag" }), - ).toBeVisible(); - await page.getByLabel("Tag Name").fill(tagName); - await page.getByLabel("Description").fill(description); - await page.getByRole("button", { name: "Create Tag" }).click(); - - await expect - .poll(async () => { - const response = await readBack< - Record> - >(page, "/tag/list"); - return Object.values(response).some((tag) => tag.name === tagName); - }) - .toBe(true); - await expect(page.getByText(tagName, { exact: true })).toBeVisible(); - - await page.getByText(tagName, { exact: true }).click(); - await expect(page.getByText("Tag Name:")).toBeVisible(); - await page.getByRole("button", { name: "Edit Tag" }).click(); - await page.getByLabel("Description").fill(updatedDescription); - const updateBody = await captureRequestBody( - page, - { method: "POST", urlIncludes: "/tag/update" }, - () => page.getByRole("button", { name: "Save Changes" }).click(), - ); - expect(updateBody).toMatchObject({ - name: tagName, - description: updatedDescription, - }); - - await expect - .poll(async () => { - const infoResponse = await page.request.post("/tag/info", { - headers: { Authorization: `Bearer ${masterKey()}` }, - data: { names: [tagName] }, - }); - expect(infoResponse.ok()).toBe(true); - const info = (await infoResponse.json()) as Record< - string, - { description?: string } - >; - return info[tagName]?.description; - }) - .toBe(updatedDescription); + if (!testOutcome.passed) throw testOutcome.error; } finally { - await page.request.post("/tag/delete", { - headers: { - Authorization: `Bearer ${masterKey()}`, - "Content-Type": "application/json", - }, - data: { name: tagName }, - }); + const cleanupSucceeded = await cleanup(); + if (testOutcome.passed) expect(cleanupSucceeded).toBe(true); } }); }); From 0dc2f0b1c1ff86356d6e9ab9180f708402131df8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:20:54 -0700 Subject: [PATCH 009/109] test(ui): protect dashboard form cleanup --- tests/e2e/ui/helpers/roundTrip.ts | 28 ++++++++++- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 42 ++++++---------- .../tests/tagManagement/tagManagement.spec.ts | 49 ++++++++----------- 3 files changed, 61 insertions(+), 58 deletions(-) diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts index 8d6e264e622..ee484b8d512 100644 --- a/tests/e2e/ui/helpers/roundTrip.ts +++ b/tests/e2e/ui/helpers/roundTrip.ts @@ -12,17 +12,41 @@ export async function captureRequestBody( match: { method: string; urlIncludes: string }, action: () => Promise, ): Promise> { - const pending = page.waitForRequest((req) => req.method() === match.method && req.url().includes(match.urlIncludes)); + const pending = page.waitForRequest( + (req) => + req.method() === match.method && req.url().includes(match.urlIncludes), + ); await action(); const request = await pending; return JSON.parse(request.postData() ?? "{}") as Record; } /** Reads an endpoint as the master key, so a failure is bad data and not an expired UI token. */ -export async function readBack(page: Page, endpoint: string): Promise { +export async function readBack( + page: Page, + endpoint: string, +): Promise { const res = await page.request.get(endpoint, { headers: { Authorization: `Bearer ${masterKey()}` }, }); expect(res.ok(), `GET ${endpoint}`).toBe(true); return (await res.json()) as T; } + +export async function runWithCleanup( + action: () => Promise, + cleanup: () => Promise, +): Promise { + const outcome = await action().then( + () => ({ status: "success" as const }), + (error: unknown) => ({ status: "failure" as const, error }), + ); + try { + if (outcome.status === "failure") throw outcome.error; + } finally { + const cleanupSucceeded = await cleanup().catch(() => false); + if (outcome.status === "success" && !cleanupSucceeded) { + throw new Error("Failed to clean up UI E2E resource"); + } + } +} diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts index b601b7e8c06..891739fda28 100644 --- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -3,6 +3,7 @@ import { test, expect } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page as DashboardPage } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; +import { runWithCleanup } from "../../helpers/roundTrip"; import { masterKey, uniqueSuffix } from "../../helpers/traffic"; test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -13,21 +14,9 @@ test.describe("Prompt upload form", () => { }) => { const promptId = `e2e-prompt-${uniqueSuffix()}`; const promptContent = "Hello {{name}}"; - const cleanup = async (): Promise => { - try { - const response = await page.request.delete( - `/prompts/${encodeURIComponent(promptId)}?environment=development`, - { - headers: { Authorization: `Bearer ${masterKey()}` }, - }, - ); - return response.ok(); - } catch { - return false; - } - }; - const testOutcome = await (async () => { - try { + + await runWithCleanup( + async () => { await navigateToPage(page, DashboardPage.Prompts); await page.getByRole("button", { name: "Upload .prompt File" }).click(); await expect( @@ -71,17 +60,16 @@ test.describe("Prompt upload form", () => { }) .toContain(promptContent); await expect(page.getByText(promptId, { exact: true })).toBeVisible(); - return { passed: true as const }; - } catch (error) { - return { passed: false as const, error }; - } - })(); - - try { - if (!testOutcome.passed) throw testOutcome.error; - } finally { - const cleanupSucceeded = await cleanup(); - if (testOutcome.passed) expect(cleanupSucceeded).toBe(true); - } + }, + async () => { + const response = await page.request.delete( + `/prompts/${encodeURIComponent(promptId)}?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }, + ); }); }); diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts index 4223324ff09..bf46b5ea161 100644 --- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -3,7 +3,11 @@ import { test, expect } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page as DashboardPage } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; -import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { + captureRequestBody, + readBack, + runWithCleanup, +} from "../../helpers/roundTrip"; import { masterKey, uniqueSuffix } from "../../helpers/traffic"; test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -13,22 +17,9 @@ test.describe("Tag management", () => { const tagName = `e2e-tag-${uniqueSuffix()}`; const description = "synthetic tag description"; const updatedDescription = `${description} updated`; - const cleanup = async (): Promise => { - try { - const response = await page.request.post("/tag/delete", { - headers: { - Authorization: `Bearer ${masterKey()}`, - "Content-Type": "application/json", - }, - data: { name: tagName }, - }); - return response.ok(); - } catch { - return false; - } - }; - const testOutcome = await (async () => { - try { + + await runWithCleanup( + async () => { await navigateToPage(page, DashboardPage.TagManagement); await page.getByRole("button", { name: "+ Create New Tag" }).click(); await expect( @@ -76,17 +67,17 @@ test.describe("Tag management", () => { return info[tagName]?.description; }) .toBe(updatedDescription); - return { passed: true as const }; - } catch (error) { - return { passed: false as const, error }; - } - })(); - - try { - if (!testOutcome.passed) throw testOutcome.error; - } finally { - const cleanupSucceeded = await cleanup(); - if (testOutcome.passed) expect(cleanupSucceeded).toBe(true); - } + }, + async () => { + const response = await page.request.post("/tag/delete", { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { name: tagName }, + }); + return response.ok(); + }, + ); }); }); From a8ab1187ca67f59432beed8b655e833f622a4055 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:23:48 -0700 Subject: [PATCH 010/109] test(ui): clean up synchronous browser failures --- tests/e2e/ui/helpers/roundTrip.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts index ee484b8d512..55eb5d6ad9c 100644 --- a/tests/e2e/ui/helpers/roundTrip.ts +++ b/tests/e2e/ui/helpers/roundTrip.ts @@ -37,10 +37,12 @@ export async function runWithCleanup( action: () => Promise, cleanup: () => Promise, ): Promise { - const outcome = await action().then( - () => ({ status: "success" as const }), - (error: unknown) => ({ status: "failure" as const, error }), - ); + const outcome = await Promise.resolve() + .then(action) + .then( + () => ({ status: "success" as const }), + (error: unknown) => ({ status: "failure" as const, error }), + ); try { if (outcome.status === "failure") throw outcome.error; } finally { From 1c08c78ad598f0fa1277305f5a32aa7ac45a2022 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:27:54 -0700 Subject: [PATCH 011/109] test(ui): retain primary cleanup failures --- tests/e2e/ui/helpers/roundTrip.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts index 55eb5d6ad9c..4175ba6ec72 100644 --- a/tests/e2e/ui/helpers/roundTrip.ts +++ b/tests/e2e/ui/helpers/roundTrip.ts @@ -46,7 +46,9 @@ export async function runWithCleanup( try { if (outcome.status === "failure") throw outcome.error; } finally { - const cleanupSucceeded = await cleanup().catch(() => false); + const cleanupSucceeded = await Promise.resolve() + .then(cleanup) + .catch(() => false); if (outcome.status === "success" && !cleanupSucceeded) { throw new Error("Failed to clean up UI E2E resource"); } From eb831d956ccb328411ebb86a0161ac7a23b4aba8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 23:46:53 -0700 Subject: [PATCH 012/109] test(ui): address review feedback --- tests/e2e/ui/helpers/roundTrip.ts | 23 +++++++++--- tests/e2e/ui/tests/prompts/addPrompt.spec.ts | 4 +-- .../tests/tagManagement/tagManagement.spec.ts | 9 ++--- ...PaginatedSearchSelect.integration.test.tsx | 36 ------------------- 4 files changed, 26 insertions(+), 46 deletions(-) diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts index 4175ba6ec72..1fc2d0aec1a 100644 --- a/tests/e2e/ui/helpers/roundTrip.ts +++ b/tests/e2e/ui/helpers/roundTrip.ts @@ -46,11 +46,26 @@ export async function runWithCleanup( try { if (outcome.status === "failure") throw outcome.error; } finally { - const cleanupSucceeded = await Promise.resolve() + const cleanupOutcome = await Promise.resolve() .then(cleanup) - .catch(() => false); - if (outcome.status === "success" && !cleanupSucceeded) { - throw new Error("Failed to clean up UI E2E resource"); + .then( + (succeeded) => + succeeded + ? { status: "success" as const } + : { + status: "failure" as const, + error: new Error("Failed to clean up UI E2E resource"), + }, + (error: unknown) => ({ status: "failure" as const, error }), + ); + if (cleanupOutcome.status === "failure") { + if (outcome.status === "failure") { + throw new AggregateError( + [outcome.error, cleanupOutcome.error], + "Action and cleanup failed", + ); + } + throw cleanupOutcome.error; } } } diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts index 891739fda28..9d85236c4a6 100644 --- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -27,7 +27,7 @@ test.describe("Prompt upload form", () => { name: "e2e.prompt", mimeType: "text/plain", buffer: Buffer.from( - `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`, + `---\nmodel: fake-openai-gpt-4\n---\n${promptContent}\n`, ), }); await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); @@ -58,7 +58,7 @@ test.describe("Prompt upload form", () => { }; return promptInfo.raw_prompt_template?.content; }) - .toContain(promptContent); + .toBe(promptContent); await expect(page.getByText(promptId, { exact: true })).toBeVisible(); }, async () => { diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts index bf46b5ea161..fe659080eab 100644 --- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -31,10 +31,11 @@ test.describe("Tag management", () => { await expect .poll(async () => { - const response = await readBack< - Record> - >(page, "/tag/list"); - return Object.values(response).some((tag) => tag.name === tagName); + const response = await readBack>( + page, + "/tag/list", + ); + return response.some((tag) => tag.name === tagName); }) .toBe(true); await expect(page.getByText(tagName, { exact: true })).toBeVisible(); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx index 905610a77e6..2b948ca8420 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx @@ -1,6 +1,5 @@ import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; -import { useQuery } from "@tanstack/react-query"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -407,39 +406,4 @@ describe("PaginatedSearchSelect", () => { expect(input).toHaveValue("aliasalpha"); await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("aliasalpha")); }); - - it("keeps the latest query results when an earlier response resolves last", async () => { - const pending = new Map void>(); - - function QueryBackedSelect() { - const [query, setQuery] = useState(""); - const result = useQuery({ - queryKey: ["paginated-select-race", query], - queryFn: () => - new Promise((resolve) => { - pending.set(query, resolve); - }), - enabled: query.length > 0, - }); - return ; - } - - const user = userEvent.setup(); - render(); - const input = screen.getByRole("combobox"); - await user.click(input); - await user.type(input, "A"); - await waitFor(() => expect(pending.has("A")).toBe(true)); - await user.clear(input); - await user.type(input, "B"); - await waitFor(() => expect(pending.has("B")).toBe(true)); - - pending.get("B")?.([{ label: "B result", value: "b" }]); - await user.click(screen.getByRole("combobox")); - expect(await screen.findByText("B result")).toBeInTheDocument(); - - pending.get("A")?.([{ label: "A result", value: "a" }]); - await waitFor(() => expect(screen.queryByText("A result")).not.toBeInTheDocument()); - expect(screen.getByText("B result")).toBeInTheDocument(); - }); }); From 5a8d1f5ecaedac9348f93fd87873e7a16b9fba0f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 00:05:43 -0700 Subject: [PATCH 013/109] feat(proxy): report sources in config read endpoints --- litellm/proxy/_types.py | 3 + litellm/proxy/config_resolvers/__init__.py | 11 +- .../proxy/config_resolvers/settings_store.py | 12 +- .../router_settings_endpoints.py | 12 +- litellm/proxy/proxy_server.py | 168 ++++++++++-------- .../proxy_setting_endpoints.py | 65 +++++-- .../test_router_settings_endpoints.py | 30 ++++ .../proxy/proxy_server/test_routes_config.py | 137 ++++++++++++++ .../proxy_server/test_routes_model_metrics.py | 38 ++++ .../test_proxy_setting_endpoints.py | 39 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 29 +++ 11 files changed, 450 insertions(+), 94 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b0d31df92ce..8574f2d8bf4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2410,6 +2410,7 @@ class FieldDetail(BaseModel): field_description: str field_default_value: Any = None stored_in_db: bool | None + source: Literal["config", "db", "default", "unset"] = "unset" class ConfigList(LiteLLMPydanticObjectBase): @@ -2418,6 +2419,7 @@ class ConfigList(LiteLLMPydanticObjectBase): field_description: str field_value: Any stored_in_db: bool | None + source: Literal["config", "db", "default", "unset"] = "unset" field_default_value: Any premium_field: bool = False nested_fields: list[FieldDetail] | None = None # For nested dictionary or Pydantic fields @@ -3693,6 +3695,7 @@ class InvitationClaim(LiteLLMPydanticObjectBase): class ConfigFieldInfo(LiteLLMPydanticObjectBase): field_name: str field_value: Any + source: Literal["config", "db", "default", "unset"] = "unset" class CallbackOnUI(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/config_resolvers/__init__.py b/litellm/proxy/config_resolvers/__init__.py index ebd339b34c3..77da2c413c2 100644 --- a/litellm/proxy/config_resolvers/__init__.py +++ b/litellm/proxy/config_resolvers/__init__.py @@ -5,6 +5,13 @@ from litellm.proxy.config_resolvers._descriptors import ( FieldSource, resolve_fields, ) -from litellm.proxy.config_resolvers.settings_store import SettingsStore +from litellm.proxy.config_resolvers.settings_store import SettingsSource, SettingsStore, source_for -__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "resolve_fields") +__all__ = ( + "FieldDescriptor", + "FieldSource", + "SettingsSource", + "SettingsStore", + "resolve_fields", + "source_for", +) diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 3fe869ee2ce..8f400853fa9 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Iterator, Mapping, MutableMapping from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias from litellm.proxy.config_resolvers._descriptors import FieldSource from litellm.proxy.config_resolvers.settings_rules import ( @@ -19,6 +19,7 @@ from litellm.proxy.config_resolvers.settings_rules import ( _EMPTY_VALUES: Final[Mapping[str, JsonValue]] = MappingProxyType({}) _EMPTY_ROWS: Final[Mapping[DbRow, Mapping[str, JsonValue]]] = MappingProxyType({}) +SettingsSource: TypeAlias = Literal["config", "db", "default", "unset"] class SettingsStore(MutableMapping[str, JsonValue]): @@ -109,3 +110,12 @@ class SettingsStore(MutableMapping[str, JsonValue]): yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT) db_value: Final[SettingValue] = self._database_rows.get(rule.db_row, _EMPTY_VALUES).get(key, ABSENT) return resolve(rule, yaml_value, db_value) + + +def source_for(settings: SettingsStore, key: str, default: object = None) -> SettingsSource: + source: Final = settings.source(key) + if source == "unset": + return "default" if default is not None else "unset" + if source in ("config", "db", "default"): + return source + return "unset" diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index fc000b1638b..5d3b6d40601 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -8,7 +8,7 @@ GET /router/fields - Get router settings field definitions without values (for U """ import inspect -from typing import Any, Final, get_args +from typing import Any, Final, cast, get_args from fastapi import APIRouter, Depends from pydantic import BaseModel, Field @@ -16,6 +16,7 @@ from pydantic import BaseModel, Field from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers import SettingsSource, source_for from litellm.router import Router from litellm.types.management_endpoints import ( ROUTER_SETTINGS_FIELDS, @@ -30,6 +31,7 @@ class RouterSettingsResponse(BaseModel): fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata") current_values: dict[str, Any] = Field(description="Current values of router settings") routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") + source: dict[str, SettingsSource] = Field(description="Source of each current router setting") class RouterFieldsResponse(BaseModel): @@ -109,15 +111,21 @@ async def get_router_settings( # Merge with config values (config takes precedence) current_values.update(router_settings_from_config) - # Update field values with current values for field in router_fields: if field.field_name in current_values: field.field_value = current_values[field.field_name] + field_defaults: Final[dict[str, object]] = { + field.field_name: cast(object, field.field_default) for field in router_fields + } + source: Final[dict[str, SettingsSource]] = { + key: source_for(proxy_config.router_settings, key, field_defaults.get(key)) for key in current_values + } return RouterSettingsResponse( fields=router_fields, current_values=current_values, routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, + source=source, ) except Exception as e: verbose_proxy_logger.error("Error fetching router settings: %s", e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6a720a066b4..1131d7bea86 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -50,6 +50,7 @@ import anyio import websockets import websockets.exceptions from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError +from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, assert_never from litellm._uuid import uuid @@ -431,7 +432,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( model_access_group_spend_counter_key, tag_cache_key, ) -from litellm.proxy.config_resolvers import SettingsStore, resolve_fields +from litellm.proxy.config_resolvers import SettingsStore, resolve_fields, source_for from litellm.proxy.config_resolvers.alerting import ( EMAIL_DESCRIPTORS, MS_TEAMS_DESCRIPTORS, @@ -4816,6 +4817,12 @@ def _as_settings_mapping(value: object) -> Mapping[str, SettingsJsonValue]: return _SETTINGS_MAPPING.validate_python(value) +def _get_field_default(field_info: FieldInfo) -> JsonValue: + if field_info.default is PydanticUndefined: + return None + return cast(JsonValue, field_info.default) + + def _bind_general_settings_store(settings: SettingsStore) -> None: global general_settings general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings @@ -15635,17 +15642,16 @@ async def alerting_settings( where={"param_name": "general_settings"} ) - if db_general_settings is not None and db_general_settings.param_value is not None: - db_general_settings_dict: Final = dict(db_general_settings.param_value) - alerting_args_dict: dict = cast( # cast-ok: ConfigGeneralSettings validates alerting_args as a dict on write - dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {}) - ) - alerting_values: list | None = cast( # cast-ok: ConfigGeneralSettings validates alerting as a list on write - list[JsonValue] | None, db_general_settings_dict.get("alerting") - ) - else: - alerting_args_dict = {} - alerting_values = None + db_general_settings_dict: Final[Mapping[str, JsonValue]] = ( + dict(db_general_settings.param_value) + if db_general_settings is not None and db_general_settings.param_value is not None + else {} + ) + alerting_args_dict: Final = cast(dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {})) + alerting_values: Final = cast(list[JsonValue] | None, db_general_settings_dict.get("alerting")) + + settings: Final = proxy_config.settings + settings.apply_db_row("general_settings", db_general_settings_dict) allowed_args: Final = MappingProxyType( { @@ -15674,9 +15680,9 @@ async def alerting_settings( is_slack_enabled = False - if general_settings.get("alerting") and isinstance(general_settings["alerting"], list): - if "slack" in general_settings["alerting"]: - is_slack_enabled = True + alerting: Final = settings.get("alerting") + if isinstance(alerting, list) and "slack" in alerting: + is_slack_enabled = True _response_obj = ConfigList( field_name="slack_alerting", @@ -15684,6 +15690,7 @@ async def alerting_settings( field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.", field_value=is_slack_enabled, stored_in_db=True if alerting_values is not None else False, + source=source_for(settings, "alerting"), field_default_value=None, premium_field=False, ) @@ -15691,6 +15698,7 @@ async def alerting_settings( for field_name, field_info in SlackAlertingArgs.model_fields.items(): if field_name in allowed_args: + field_default: JsonValue = _get_field_default(field_info) _stored_in_db: bool | None = None if field_name in alerting_args_dict: _stored_in_db = True @@ -15701,9 +15709,10 @@ async def alerting_settings( field_name=field_name, field_type=allowed_args[field_name], field_description=field_info.description or "", - field_value=_slack_alerting_args_dict.get(field_name, None), + field_value=_slack_alerting_args_dict.get(field_name, field_default), stored_in_db=_stored_in_db, - field_default_value=field_info.default, + source=source_for(settings, "alerting_args", field_default), + field_default_value=field_default, premium_field=(True if field_name == "region_outage_alert_ttl" else False), ) return_val.append(_response_obj) @@ -17390,20 +17399,6 @@ async def get_config_general_settings( field_name: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - global prisma_client - - ## VALIDATION ## - """ - - Check if prisma_client is None - - Check if user allowed to call this endpoint (admin-only) - - Check if param in general settings - """ - if prisma_client is None: - raise HTTPException( - status_code=400, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, @@ -17416,37 +17411,47 @@ async def get_config_general_settings( detail={"error": f"Invalid field={field_name} passed in."}, ) - ## get general settings from db - db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( - where={"param_name": "general_settings"} - ) - ### pop the value + field_info: Final = ConfigGeneralSettings.model_fields[field_name] + field_default: JsonValue = _get_field_default(field_info) + settings: Final = proxy_config.settings + db_values: Mapping[str, JsonValue] + if prisma_client is None: + db_values = {} + else: + db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( + where={"param_name": "general_settings"} + ) + db_values = ( + dict(db_general_settings.param_value) + if db_general_settings is not None and db_general_settings.param_value is not None + else {} + ) + settings.apply_db_row("general_settings", db_values) - if db_general_settings is None or db_general_settings.param_value is None: + if field_name not in settings and field_default is None: raise HTTPException( status_code=400, detail={"error": f"Field name={field_name} not in DB"}, ) - else: - general_settings = dict(db_general_settings.param_value) - if field_name in general_settings: - field_value = _redact_general_setting_value( - field_name, - general_settings[field_name], - user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, - ) - if field_name == "plugins" and isinstance(field_value, list): - field_value = [ - ({k: ("***" if k == "plugin_key" else v) for k, v in p.items()} if isinstance(p, dict) else p) - for p in field_value - ] - return ConfigFieldInfo(field_name=field_name, field_value=field_value) - else: - raise HTTPException( - status_code=400, - detail={"error": f"Field name={field_name} not in DB"}, - ) + redacted_field_value: Final = _redact_general_setting_value( + field_name, + settings.get(field_name, field_default), + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + ) + field_value: Final = ( + [ + ({k: ("***" if k == "plugin_key" else v) for k, v in p.items()} if isinstance(p, dict) else p) + for p in redacted_field_value + ] + if field_name == "plugins" and isinstance(redacted_field_value, list) + else redacted_field_value + ) + return ConfigFieldInfo( + field_name=field_name, + field_value=field_value, + source=source_for(settings, field_name, field_default), + ) GeneralSettingsUILiteLLMValue = float | bool | str | None @@ -17600,7 +17605,7 @@ async def get_config_list( """ List the available fields + current values for a given type of setting (currently just 'general_settings'user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),) """ - global prisma_client, general_settings + global prisma_client ## VALIDATION ## """ @@ -17627,10 +17632,16 @@ async def get_config_list( where={"param_name": "general_settings"} ) - if db_general_settings is not None and db_general_settings.param_value is not None: - db_general_settings_dict: Mapping[str, JsonValue] = dict(db_general_settings.param_value) - else: - db_general_settings_dict = {} + db_general_settings_dict: Final[Mapping[str, JsonValue]] = ( + dict(db_general_settings.param_value) + if db_general_settings is not None and db_general_settings.param_value is not None + else {} + ) + settings: Final = proxy_config.settings + settings.apply_db_row("general_settings", db_general_settings_dict) + runtime_settings: Final[Mapping[str, JsonValue]] = ( + cast(Mapping[str, JsonValue], general_settings) if not isinstance(general_settings, SettingsStore) else settings + ) allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES @@ -17638,6 +17649,7 @@ async def get_config_list( for field_name, field_info in ConfigGeneralSettings.model_fields.items(): if field_name in allowed_args: + field_default: JsonValue = _get_field_default(field_info) ## HANDLE TYPED DICT typed_dict_type = allowed_args[field_name] @@ -17657,10 +17669,11 @@ async def get_config_list( field_description="", # Add custom logic if descriptions are available field_default_value=_redact_general_setting_value( sub_field, - general_settings.get(sub_field, None), + runtime_settings.get(sub_field, None), is_full_admin, ), stored_in_db=None, + source=source_for(settings, field_name), ) for sub_field, sub_field_type in pydantic_class.__annotations__.items() ] @@ -17677,7 +17690,7 @@ async def get_config_list( _stored_in_db = None if field_name in db_general_settings_dict: _stored_in_db = True - elif field_name in general_settings: + elif field_name in runtime_settings: _stored_in_db = False _response_obj = ConfigList( @@ -17686,11 +17699,12 @@ async def get_config_list( field_description=field_info.description or "", field_value=_redact_general_setting_value( field_name, - general_settings.get(field_name, None), + runtime_settings.get(field_name, field_default), is_full_admin, ), stored_in_db=_stored_in_db, - field_default_value=field_info.default, + source=source_for(settings, field_name, field_default), + field_default_value=field_default, nested_fields=nested_fields, ) return_val.append(_response_obj) @@ -17701,12 +17715,10 @@ async def get_config_list( _stored_in_db = None if field_name in db_general_settings_dict: _stored_in_db = True - elif field_name in general_settings: + elif field_name in runtime_settings: _stored_in_db = False - _field_value = general_settings.get(field_name, None) - if _field_value is None and field_name in db_general_settings_dict: - _field_value = db_general_settings_dict[field_name] + _field_value: JsonValue = runtime_settings.get(field_name, field_default) _response_obj = ConfigList( field_name=field_name, @@ -17714,7 +17726,8 @@ async def get_config_list( field_description=field_info.description or "", field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin), stored_in_db=_stored_in_db, - field_default_value=field_info.default, + source=source_for(settings, field_name, field_default), + field_default_value=field_default, nested_fields=nested_fields, ) return_val.append(_response_obj) @@ -17722,18 +17735,24 @@ async def get_config_list( db_litellm_settings_row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( where={"param_name": "litellm_settings"} ) - db_litellm_settings: Final[dict] = ( + db_litellm_settings: Final[Mapping[str, JsonValue]] = ( dict(db_litellm_settings_row.param_value) if db_litellm_settings_row is not None and db_litellm_settings_row.param_value is not None else {} ) + litellm_settings_store: Final = proxy_config.litellm_settings + litellm_settings_store.apply_db_row("litellm_settings", db_litellm_settings) for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): - current_value: GeneralSettingsUILiteLLMValue = getattr(litellm, litellm_field_name, None) - default_value = _general_settings_ui_litellm_default(spec) + default_value: GeneralSettingsUILiteLLMValue = _general_settings_ui_litellm_default(spec) + current_value: GeneralSettingsUILiteLLMValue = cast( + GeneralSettingsUILiteLLMValue, + litellm_settings_store.get(litellm_field_name, default_value), + ) + source = source_for(litellm_settings_store, litellm_field_name, default_value) stored_in_db_litellm: bool | None if litellm_field_name in db_litellm_settings: stored_in_db_litellm = True - elif current_value != default_value: + elif source == "config": stored_in_db_litellm = False else: stored_in_db_litellm = None @@ -17744,6 +17763,7 @@ async def get_config_list( field_description=spec["description"], field_value=current_value, stored_in_db=stored_in_db_litellm, + source=source, field_default_value=default_value, field_options=list(spec.get("options", ())) or None, field_tab=spec.get("tab"), diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index fd160636d46..4cd031b8780 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -14,8 +14,8 @@ from typing import ( from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile -from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model -from pydantic.fields import FieldInfo +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model +from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers import SettingsSource, source_for from litellm.proxy.config_resolvers.sso import ( SSO_FIELD_ENV_VARS, SSO_SECRET_FIELDS, @@ -197,6 +198,11 @@ class SettingsResponse(BaseModel): """Schema information including descriptions and property types for UI display""" +class _SettingsWithSchema(BaseModel): + values: dict[str, object] + field_schema: dict[str, object] + + class SSOSettingsResponse(SettingsResponse): """Response model for SSO settings""" @@ -327,6 +333,8 @@ class UISettings(BaseModel): class UISettingsResponse(SettingsResponse): """Response model for UI settings""" + source: dict[str, SettingsSource] + # Allowlist of UI settings that can be stored ALLOWED_UI_SETTINGS_FIELDS: Final = { @@ -658,6 +666,13 @@ def _root_schema(settings_class: type[BaseModel]) -> _RootSchema: ) +def _model_field_default(settings_class: type[BaseModel], field_name: str) -> object: + field_info: Final = settings_class.model_fields.get(field_name) + if field_info is None or field_info.default is PydanticUndefined: + return None + return cast(object, field_info.default) + + async def _get_settings_with_schema( settings_key: str, settings_class: type[BaseModel], @@ -1527,7 +1542,7 @@ async def get_ui_settings(): Get UI-specific configuration flags. All authenticated users can fetch these settings for client-side behavior. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, proxy_config if prisma_client is None: raise HTTPException( @@ -1546,26 +1561,46 @@ async def get_ui_settings(): ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS} apply_runtime_general_settings_flags(ui_settings) + proxy_config.settings.apply_db_row("ui_settings", ui_settings) # Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values from litellm.proxy.proxy_server import user_api_key_cache await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) - # Build config-like object for schema helper - config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} - - settings: Final = await _get_settings_with_schema( - settings_key="ui_settings", - settings_class=_get_effective_ui_settings_class(), - config=config, + effective_ui_settings: Final = { + **{key: proxy_config.settings[key] for key in ALLOWED_UI_SETTINGS_FIELDS if key in proxy_config.settings}, + **ui_settings, + } + config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": effective_ui_settings}} + settings_class: Final = _get_effective_ui_settings_class() + resolved_settings: Final = _SettingsWithSchema.model_validate( + await _get_settings_with_schema( + settings_key="ui_settings", + settings_class=settings_class, + config=config, + ) ) + values: Final = { + **resolved_settings.values, + ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), + } + source: Final[dict[str, SettingsSource]] = { + key: ( + "db" + if key in ui_settings + else source_for( + proxy_config.settings, + key, + _model_field_default(settings_class, key), + ) + ) + for key in values + } return UISettingsResponse( - values={ - **settings["values"], - ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), - }, - field_schema=settings["field_schema"], + values=values, + field_schema=resolved_settings.field_schema, + source=source, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py index 308f4d88f02..148f30f517b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py @@ -75,6 +75,36 @@ class TestRouterSettingsEndpoints: assert isinstance(routing_strategy_field["options"], list) assert len(routing_strategy_field["options"]) > 0 + @pytest.mark.asyncio + async def test_get_router_settings_reports_sources(self, monkeypatch): + from litellm.proxy.config_resolvers import SettingsStore + + store = SettingsStore("router_settings") + store.load_yaml({"routing_strategy": "simple-shuffle"}) + store.apply_db_row("router_settings", {"num_retries": 3}) + monkeypatch.setattr(proxy_server.proxy_config, "router_settings", store) + monkeypatch.setattr(proxy_server, "llm_router", None) + + async def fake_get_config(self, config_file_path=None): + return { + "router_settings": { + "routing_strategy": "simple-shuffle", + "num_retries": 3, + } + } + + monkeypatch.setattr( + proxy_server.ProxyConfig, "get_config", fake_get_config, raising=True + ) + + admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-x" + ) + response = await get_router_settings(user_api_key_dict=admin_user) + + assert response.source["routing_strategy"] == "config" + assert response.source["num_retries"] == "db" + @pytest.mark.asyncio async def test_get_router_settings_includes_routing_groups_from_live_router( self, monkeypatch diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index dd3914e3ad5..d6c63ec9a78 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -15,10 +15,15 @@ from __future__ import annotations import asyncio import json +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.proxy.config_resolvers import SettingsStore +from litellm.proxy.config_resolvers.settings_rules import JsonValue + from .conftest import VOLATILE_KEYS, normalize @@ -37,6 +42,21 @@ def _install_litellm_config(mock_prisma: MagicMock) -> MagicMock: return table +def _install_settings_store( + monkeypatch: pytest.MonkeyPatch, + config_values: Mapping[str, JsonValue], + db_values: Mapping[str, JsonValue], +) -> SettingsStore: + from litellm.proxy import proxy_server + + store: Final = SettingsStore("general_settings") + store.load_yaml(config_values) + store.apply_db_row("general_settings", db_values) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + return store + + # --------------------------------------------------------------------------- # POST /config/update # --------------------------------------------------------------------------- @@ -338,6 +358,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch assert normalize(response.json()) == { "field_name": "max_parallel_requests", "field_value": 7, + "source": "db", } @@ -566,6 +587,122 @@ def test_config_list_happy_admin(client, auth_as, mock_prisma, monkeypatch): } +def test_config_read_routes_report_effective_values_and_sources(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {"max_parallel_requests": 7, "max_file_size_mb": 222} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + _install_settings_store( + monkeypatch, + { + "max_parallel_requests": 5, + "max_file_size_mb": 111, + "pass_through_endpoints": [{"path": "/synthetic"}], + }, + {"max_parallel_requests": 7, "max_file_size_mb": 222}, + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + list_response = client.get("/config/list", params={"config_type": "general_settings"}) + config_only_response = client.get( + "/config/field/info", params={"field_name": "max_file_size_mb"} + ) + db_wins_response = client.get( + "/config/field/info", params={"field_name": "max_parallel_requests"} + ) + + assert list_response.status_code == 200 + by_name: Final = {entry["field_name"]: entry for entry in list_response.json()} + assert by_name["max_file_size_mb"]["field_value"] == 111 + assert by_name["max_file_size_mb"]["source"] == "config" + assert by_name["pass_through_endpoints"]["source"] == "config" + assert by_name["pass_through_endpoints"]["nested_fields"][0]["source"] == "config" + assert by_name["max_parallel_requests"]["field_value"] == 7 + assert by_name["max_parallel_requests"]["source"] == "db" + + assert config_only_response.status_code == 200 + assert config_only_response.json() == { + "field_name": "max_file_size_mb", + "field_value": 111, + "source": "config", + } + assert db_wins_response.status_code == 200 + assert db_wins_response.json() == { + "field_name": "max_parallel_requests", + "field_value": 7, + "source": "db", + } + + +def test_config_read_routes_report_default_source(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + _install_settings_store(monkeypatch, {}, {}) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + list_response = client.get("/config/list", params={"config_type": "general_settings"}) + field_response = client.get( + "/config/field/info", params={"field_name": "proxy_config_reload_interval_seconds"} + ) + + assert list_response.status_code == 200 + by_name: Final = {entry["field_name"]: entry for entry in list_response.json()} + assert by_name["proxy_config_reload_interval_seconds"]["field_value"] == 30 + assert by_name["proxy_config_reload_interval_seconds"]["source"] == "default" + assert field_response.status_code == 200 + assert field_response.json() == { + "field_name": "proxy_config_reload_interval_seconds", + "field_value": 30, + "source": "default", + } + + +def test_config_field_info_uses_store_without_db(client, auth_as, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr(ps, "prisma_client", None) + _install_settings_store(monkeypatch, {"max_file_size_mb": 111}, {}) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/config/field/info", params={"field_name": "max_file_size_mb"}) + + assert response.status_code == 200 + assert response.json() == { + "field_name": "max_file_size_mb", + "field_value": 111, + "source": "config", + } + + +def test_config_field_info_unset_source_remains_an_error(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + _install_settings_store(monkeypatch, {}, {}) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) + + assert response.status_code == 400 + assert "not in" in response.json()["detail"]["error"] + + def test_config_list_exposes_config_reload_interval(client, auth_as, mock_prisma, monkeypatch): """proxy_config_reload_interval_seconds must surface in the admin UI general-settings list as an Integer field defaulting to 30, so operators can tune multi-pod convergence diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py index 246e2cbba54..6a57ad1636d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py @@ -179,6 +179,44 @@ def test_model_settings_method_not_allowed(client, auth_as): # --------------------------------------------------------------------------- +def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): + from litellm.proxy.config_resolvers import SettingsStore + + pc = MagicMock() + row = MagicMock() + row.param_value = {"alerting_args": {"daily_report_frequency": 7}} + pc.db.litellm_config.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", pc) + + logging_obj = MagicMock() + args_model = MagicMock() + args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 7}) + logging_obj.slack_alerting_instance.alerting_args = args_model + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) + + store = SettingsStore("general_settings") + store.load_yaml( + { + "alerting": ["slack"], + "alerting_args": {"daily_report_frequency": 3}, + } + ) + store.apply_db_row( + "general_settings", + {"alerting_args": {"daily_report_frequency": 7}}, + ) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + assert by_name["slack_alerting"]["source"] == "config" + assert by_name["daily_report_frequency"]["source"] == "db" + + def test_alerting_settings_no_db_error(client, auth_as, no_prisma): """Pins ``GET /alerting/settings`` (error: db not connected).""" with auth_as(LitellmUserRoles.PROXY_ADMIN): diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 93fb54f84eb..faf5b336410 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1342,6 +1342,45 @@ class TestProxySettingEndpoints: where={"id": "ui_settings"} ) + def test_get_ui_settings_reports_sources(self, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy.config_resolvers import SettingsStore + + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.ui_settings = { + "disable_model_add_for_internal_users": True, + } + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=mock_db_record + ) + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + + store = SettingsStore("general_settings") + store.load_yaml( + { + "disable_model_add_for_internal_users": False, + "forward_client_headers_to_llm_api": True, + } + ) + store.apply_db_row( + "ui_settings", + {"disable_model_add_for_internal_users": True}, + ) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + data = response.json() + assert data["values"]["disable_model_add_for_internal_users"] is True + assert data["values"]["forward_client_headers_to_llm_api"] is True + assert data["source"]["disable_model_add_for_internal_users"] == "db" + assert data["source"]["forward_client_headers_to_llm_api"] == "config" + def test_get_ui_settings_schema_description_preserved_with_extensions( self, mock_auth, monkeypatch ): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fd882937e79..790ddb93546 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26393,6 +26393,12 @@ export interface components { field_name: string; /** Field Value */ field_value: unknown; + /** + * Source + * @default unset + * @enum {string} + */ + source: "config" | "db" | "default" | "unset"; }; /** ConfigFieldUpdate */ ConfigFieldUpdate: { @@ -26895,6 +26901,12 @@ export interface components { * @default false */ premium_field: boolean; + /** + * Source + * @default unset + * @enum {string} + */ + source: "config" | "db" | "default" | "unset"; /** Stored In Db */ stored_in_db: boolean | null; }; @@ -28286,6 +28298,12 @@ export interface components { field_name: string; /** Field Type */ field_type: string; + /** + * Source + * @default unset + * @enum {string} + */ + source: "config" | "db" | "default" | "unset"; /** Stored In Db */ stored_in_db: boolean | null; }; @@ -36439,6 +36457,13 @@ export interface components { routing_strategy_descriptions: { [key: string]: string; }; + /** + * Source + * @description Source of each current router setting + */ + source: { + [key: string]: "config" | "db" | "default" | "unset"; + }; }; /** * RoutingGroup @@ -38856,6 +38881,10 @@ export interface components { field_schema: { [key: string]: unknown; }; + /** Source */ + source: { + [key: string]: "config" | "db" | "default" | "unset"; + }; /** Values */ values: { [key: string]: unknown; From 8dcf9e8b78509e0392eaef34abb43d2754e2a43f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 01:26:10 -0700 Subject: [PATCH 014/109] fix(proxy): correct settings source provenance --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../proxy/config_resolvers/settings_store.py | 20 ++++++ .../router_settings_endpoints.py | 25 ++++++- litellm/proxy/proxy_server.py | 65 +++++++++++++----- .../proxy_setting_endpoints.py | 30 ++++++--- .../config_resolvers/test_settings_store.py | 33 +++++++++ .../test_router_settings_endpoints.py | 2 + .../proxy/proxy_server/test_routes_config.py | 44 +++++++++--- .../proxy_server/test_routes_model_metrics.py | 67 ++++++++++++++++--- .../test_proxy_setting_endpoints.py | 42 ++++++++++++ 10 files changed, 280 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b244678e201..fa046ef0a72 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 8f400853fa9..73bca3222ea 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -28,6 +28,7 @@ class SettingsStore(MutableMapping[str, JsonValue]): self._yaml_values: Mapping[str, JsonValue] = _EMPTY_VALUES self._database_rows: Mapping[DbRow, Mapping[str, JsonValue]] = _EMPTY_ROWS self._runtime_values: Mapping[str, JsonValue] = _EMPTY_VALUES + self._runtime_sources: Mapping[str, FieldSource] = MappingProxyType({}) self._deleted_runtime_keys: frozenset[str] = frozenset() def load_yaml(self, mapping: Mapping[str, JsonValue]) -> None: @@ -39,11 +40,22 @@ class SettingsStore(MutableMapping[str, JsonValue]): self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))}) self._clear_runtime_keys(frozenset((*previous_row, *db_row))) + def without_db(self) -> SettingsStore: + copy: Final = SettingsStore(self._section) + copy.load_yaml(self._yaml_values) + runtime_values: Final = { + key: value for key, value in self._runtime_values.items() if self._runtime_sources.get(key) != "db" + } + copy.apply_runtime_values(runtime_values) + copy._deleted_runtime_keys = self._deleted_runtime_keys + return copy + def resolved(self) -> Mapping[str, JsonValue]: return MappingProxyType(dict(self)) def apply_runtime_values(self, values: Mapping[str, JsonValue]) -> None: self._runtime_values = MappingProxyType(dict(values)) + self._runtime_sources = MappingProxyType({key: self.source(key) for key in values}) self._deleted_runtime_keys = frozenset() def source(self, key: str) -> FieldSource: @@ -61,6 +73,7 @@ class SettingsStore(MutableMapping[str, JsonValue]): def __setitem__(self, key: str, value: JsonValue) -> None: self._runtime_values = MappingProxyType({**self._runtime_values, key: value}) + self._runtime_sources = MappingProxyType({**self._runtime_sources, key: self.source(key)}) self._deleted_runtime_keys = self._deleted_runtime_keys - frozenset((key,)) def __delitem__(self, key: str) -> None: @@ -69,6 +82,9 @@ class SettingsStore(MutableMapping[str, JsonValue]): self._runtime_values = MappingProxyType( {key_: value for key_, value in self._runtime_values.items() if key_ != key} ) + self._runtime_sources = MappingProxyType( + {key_: source for key_, source in self._runtime_sources.items() if key_ != key} + ) self._deleted_runtime_keys = self._deleted_runtime_keys | frozenset((key,)) def __iter__(self) -> Iterator[str]: @@ -84,6 +100,7 @@ class SettingsStore(MutableMapping[str, JsonValue]): def _clear_runtime(self) -> None: self._runtime_values = _EMPTY_VALUES + self._runtime_sources = MappingProxyType({}) self._deleted_runtime_keys = frozenset() def _clear_runtime_keys(self, keys: frozenset[str]) -> None: @@ -92,6 +109,9 @@ class SettingsStore(MutableMapping[str, JsonValue]): self._runtime_values = MappingProxyType( {key: value for key, value in self._runtime_values.items() if key not in keys} ) + self._runtime_sources = MappingProxyType( + {key: source for key, source in self._runtime_sources.items() if key not in keys} + ) self._deleted_runtime_keys = self._deleted_runtime_keys - keys def _keys(self) -> tuple[str, ...]: diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 5d3b6d40601..1869be88d1b 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, Field from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.config_resolvers import SettingsSource, source_for +from litellm.proxy.config_resolvers import SettingsSource, SettingsStore, source_for from litellm.router import Router from litellm.types.management_endpoints import ( ROUTER_SETTINGS_FIELDS, @@ -41,6 +41,18 @@ class RouterFieldsResponse(BaseModel): routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") +def _router_setting_source( + settings: SettingsStore, + key: str, + current_value: object, + field_default: object, +) -> SettingsSource: + source: Final = source_for(settings, key, field_default) + if source != "unset": + return source + return "default" if current_value is not None else "unset" + + def _get_routing_strategies_from_router_class() -> list[str]: """ Dynamically extract routing strategies from the Router class __init__ method. @@ -116,10 +128,17 @@ async def get_router_settings( field.field_value = current_values[field.field_name] field_defaults: Final[dict[str, object]] = { - field.field_name: cast(object, field.field_default) for field in router_fields + field.field_name: cast(object, field.field_default) # cast-ok: Pydantic field defaults are untyped + for field in router_fields } source: Final[dict[str, SettingsSource]] = { - key: source_for(proxy_config.router_settings, key, field_defaults.get(key)) for key in current_values + key: _router_setting_source( + proxy_config.router_settings, + key, + cast(object, current_values[key]), # cast-ok: current values are stored in a typed response map + field_defaults.get(key), + ) + for key in current_values } return RouterSettingsResponse( fields=router_fields, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1131d7bea86..3898f522feb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -432,7 +432,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( model_access_group_spend_counter_key, tag_cache_key, ) -from litellm.proxy.config_resolvers import SettingsStore, resolve_fields, source_for +from litellm.proxy.config_resolvers import SettingsSource, SettingsStore, resolve_fields, source_for from litellm.proxy.config_resolvers.alerting import ( EMAIL_DESCRIPTORS, MS_TEAMS_DESCRIPTORS, @@ -4820,7 +4820,7 @@ def _as_settings_mapping(value: object) -> Mapping[str, SettingsJsonValue]: def _get_field_default(field_info: FieldInfo) -> JsonValue: if field_info.default is PydanticUndefined: return None - return cast(JsonValue, field_info.default) + return cast(JsonValue, field_info.default) # cast-ok: Pydantic field defaults are JSON values at runtime def _bind_general_settings_store(settings: SettingsStore) -> None: @@ -15605,6 +15605,22 @@ async def model_settings(): #### ALERTING MANAGEMENT ENDPOINTS #### +def _nested_setting_source( + settings: SettingsStore, + db_values: Mapping[str, JsonValue], + parent_key: str, + field_name: str, + field_default: JsonValue, +) -> SettingsSource: + db_value: Final = db_values.get(field_name) + if db_value is not None and db_value != []: + return "db" + parent_value: Final = settings.without_db().get(parent_key) + if isinstance(parent_value, Mapping) and field_name in parent_value: + return "config" + return "default" if field_default is not None else "unset" + + @router.get( "/alerting/settings", description="Return the configurable alerting param, description, and current value", @@ -15647,8 +15663,13 @@ async def alerting_settings( if db_general_settings is not None and db_general_settings.param_value is not None else {} ) - alerting_args_dict: Final = cast(dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {})) - alerting_values: Final = cast(list[JsonValue] | None, db_general_settings_dict.get("alerting")) + alerting_args_value: Final = db_general_settings_dict.get("alerting_args") + alerting_args_dict: Final[Mapping[str, JsonValue]] = ( + alerting_args_value if isinstance(alerting_args_value, dict) else {} + ) + alerting_values: Final = cast( # cast-ok: alerting is stored as a JSON list when present + list[JsonValue] | None, db_general_settings_dict.get("alerting") + ) settings: Final = proxy_config.settings settings.apply_db_row("general_settings", db_general_settings_dict) @@ -15711,7 +15732,13 @@ async def alerting_settings( field_description=field_info.description or "", field_value=_slack_alerting_args_dict.get(field_name, field_default), stored_in_db=_stored_in_db, - source=source_for(settings, "alerting_args", field_default), + source=_nested_setting_source( + settings, + alerting_args_dict, + "alerting_args", + field_name, + field_default, + ), field_default_value=field_default, premium_field=(True if field_name == "region_outage_alert_ttl" else False), ) @@ -17414,21 +17441,19 @@ async def get_config_general_settings( field_info: Final = ConfigGeneralSettings.model_fields[field_name] field_default: JsonValue = _get_field_default(field_info) settings: Final = proxy_config.settings - db_values: Mapping[str, JsonValue] - if prisma_client is None: - db_values = {} - else: + if prisma_client is not None: db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( where={"param_name": "general_settings"} ) - db_values = ( + db_values: Final[Mapping[str, JsonValue]] = ( dict(db_general_settings.param_value) if db_general_settings is not None and db_general_settings.param_value is not None else {} ) settings.apply_db_row("general_settings", db_values) + effective_settings: Final = settings.without_db() if prisma_client is None else settings - if field_name not in settings and field_default is None: + if field_name not in effective_settings and field_default is None: raise HTTPException( status_code=400, detail={"error": f"Field name={field_name} not in DB"}, @@ -17436,7 +17461,7 @@ async def get_config_general_settings( redacted_field_value: Final = _redact_general_setting_value( field_name, - settings.get(field_name, field_default), + effective_settings.get(field_name, field_default), user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, ) field_value: Final = ( @@ -17450,7 +17475,7 @@ async def get_config_general_settings( return ConfigFieldInfo( field_name=field_name, field_value=field_value, - source=source_for(settings, field_name, field_default), + source=source_for(effective_settings, field_name, field_default), ) @@ -17640,7 +17665,11 @@ async def get_config_list( settings: Final = proxy_config.settings settings.apply_db_row("general_settings", db_general_settings_dict) runtime_settings: Final[Mapping[str, JsonValue]] = ( - cast(Mapping[str, JsonValue], general_settings) if not isinstance(general_settings, SettingsStore) else settings + cast( # cast-ok: legacy general_settings remains a mapping at this route boundary + Mapping[str, JsonValue], general_settings + ) + if not isinstance(general_settings, SettingsStore) + else settings ) allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES @@ -17744,9 +17773,11 @@ async def get_config_list( litellm_settings_store.apply_db_row("litellm_settings", db_litellm_settings) for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): default_value: GeneralSettingsUILiteLLMValue = _general_settings_ui_litellm_default(spec) - current_value: GeneralSettingsUILiteLLMValue = cast( - GeneralSettingsUILiteLLMValue, - litellm_settings_store.get(litellm_field_name, default_value), + current_value: GeneralSettingsUILiteLLMValue = ( + cast( # cast-ok: UI field defaults are validated by the field spec + GeneralSettingsUILiteLLMValue, + litellm_settings_store.get(litellm_field_name, default_value), + ) ) source = source_for(litellm_settings_store, litellm_field_name, default_value) stored_in_db_litellm: bool | None diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 4cd031b8780..aba65719f9b 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -24,7 +24,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.config_resolvers import SettingsSource, source_for +from litellm.proxy.config_resolvers import SettingsSource, SettingsStore, source_for from litellm.proxy.config_resolvers.sso import ( SSO_FIELD_ENV_VARS, SSO_SECRET_FIELDS, @@ -34,7 +34,10 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, ) -from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled +from litellm.proxy.spend_tracking.ptu_feature_flag import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.organization_repository import OrganizationRepository @@ -44,6 +47,7 @@ from litellm.repositories.table_repositories import ( UISettingsRepository, ) from litellm.repositories.team_repository import TeamRepository +from litellm.secret_managers.main import get_secret from litellm.types.mcp import MCPToolSearchSettings from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, @@ -670,7 +674,19 @@ def _model_field_default(settings_class: type[BaseModel], field_name: str) -> ob field_info: Final = settings_class.model_fields.get(field_name) if field_info is None or field_info.default is PydanticUndefined: return None - return cast(object, field_info.default) + return cast(object, field_info.default) # cast-ok: Pydantic field defaults are untyped + + +def _ui_setting_source( + key: str, + value: object, + settings: SettingsStore, + settings_class: type[BaseModel], +) -> SettingsSource: + if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: + configured_value: Final = get_secret(PTU_COST_ATTRIBUTION_ENV_VAR, None) + return "config" if configured_value is not None or value is True else "default" + return source_for(settings, key, _model_field_default(settings_class, key)) async def _get_settings_with_schema( @@ -1587,13 +1603,7 @@ async def get_ui_settings(): } source: Final[dict[str, SettingsSource]] = { key: ( - "db" - if key in ui_settings - else source_for( - proxy_config.settings, - key, - _model_field_default(settings_class, key), - ) + "db" if key in ui_settings else _ui_setting_source(key, values[key], proxy_config.settings, settings_class) ) for key in values } diff --git a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py index e41c852eacb..efff9ad24c8 100644 --- a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py +++ b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py @@ -82,6 +82,39 @@ def test_settings_store_keeps_unaffected_runtime_values_on_a_db_row_refresh() -> assert store.source("changed") == "db" +def test_settings_store_without_db_uses_yaml_without_mutating_runtime_values() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"max_parallel_requests": 5}) + store.apply_db_row("general_settings", {"max_parallel_requests": 7}) + store.apply_runtime_values({"max_parallel_requests": 7}) + + without_db: Final = store.without_db() + + assert without_db["max_parallel_requests"] == 5 + assert without_db.source("max_parallel_requests") == "config" + assert store["max_parallel_requests"] == 7 + assert store.source("max_parallel_requests") == "db" + + +def test_settings_store_without_db_preserves_non_db_runtime_values() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"max_parallel_requests": "os.environ/MAX_PARALLEL_REQUESTS"}) + store.apply_runtime_values({"max_parallel_requests": 7}) + + without_db: Final = store.without_db() + + assert without_db["max_parallel_requests"] == 7 + assert without_db.source("max_parallel_requests") == "config" + + +def test_settings_store_without_db_preserves_runtime_deletions() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"deleted": 1}) + del store["deleted"] + + assert "deleted" not in store.without_db() + + def test_settings_store_removes_only_runtime_values_affected_by_a_cleared_db_row() -> None: store: Final = SettingsStore("general_settings") store.load_yaml({"template": "os.environ/SETTING"}) diff --git a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py index 148f30f517b..3af7de62abe 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py @@ -146,6 +146,8 @@ class TestRouterSettingsEndpoints: response = await get_router_settings(user_api_key_dict=admin_user) assert response.current_values.get("routing_groups") == groups + assert response.current_values["timeout"] is not None + assert response.source["timeout"] == "default" rg_field = next(f for f in response.fields if f.field_name == "routing_groups") assert rg_field.field_value == groups diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index d6c63ec9a78..fdaa2476219 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -15,8 +15,11 @@ from __future__ import annotations import asyncio import json -from collections.abc import Mapping +from collections.abc import Callable, Mapping +from contextlib import AbstractContextManager from typing import Final + +from fastapi.testclient import TestClient from unittest.mock import AsyncMock, MagicMock import pytest @@ -362,6 +365,33 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch } +def test_config_field_info_clears_stale_db_source_without_connection( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + store = SettingsStore("general_settings") + store.load_yaml({"max_parallel_requests": 5}) + store.apply_db_row("general_settings", {"max_parallel_requests": 7}) + store.apply_runtime_values({"max_parallel_requests": 7}) + monkeypatch.setattr(ps.proxy_config, "settings", store) + monkeypatch.setattr(ps, "prisma_client", None) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) + + assert response.status_code == 200 + assert normalize(response.json()) == { + "field_name": "max_parallel_requests", + "field_value": 5, + "source": "config", + } + assert store["max_parallel_requests"] == 7 + + def test_config_field_info_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin (INTERNAL_USER) is denied — admin-view gate fires.""" from litellm.proxy import proxy_server as ps @@ -608,12 +638,8 @@ def test_config_read_routes_report_effective_values_and_sources(client, auth_as, with auth_as(LitellmUserRoles.PROXY_ADMIN): list_response = client.get("/config/list", params={"config_type": "general_settings"}) - config_only_response = client.get( - "/config/field/info", params={"field_name": "max_file_size_mb"} - ) - db_wins_response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + config_only_response = client.get("/config/field/info", params={"field_name": "max_file_size_mb"}) + db_wins_response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert list_response.status_code == 200 by_name: Final = {entry["field_name"]: entry for entry in list_response.json()} @@ -651,9 +677,7 @@ def test_config_read_routes_report_default_source(client, auth_as, mock_prisma, with auth_as(LitellmUserRoles.PROXY_ADMIN): list_response = client.get("/config/list", params={"config_type": "general_settings"}) - field_response = client.get( - "/config/field/info", params={"field_name": "proxy_config_reload_interval_seconds"} - ) + field_response = client.get("/config/field/info", params={"field_name": "proxy_config_reload_interval_seconds"}) assert list_response.status_code == 200 by_name: Final = {entry["field_name"]: entry for entry in list_response.json()} diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py index 6a57ad1636d..97ca3b3dbbe 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py @@ -11,13 +11,17 @@ Pins (PR2): from __future__ import annotations +from collections.abc import Callable +from contextlib import AbstractContextManager from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi.testclient import TestClient import litellm from litellm.proxy import proxy_server from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.config_resolvers.settings_rules import JsonValue from .conftest import normalize # type: ignore[import-not-found] @@ -53,9 +57,7 @@ def test_model_streaming_metrics_happy(client, auth_as, prisma_with_query_raw): pin can rely on the exact response shape. """ with auth_as(): - response = client.get( - "/model/streaming_metrics", params={"_selected_model_group": "gpt-4"} - ) + response = client.get("/model/streaming_metrics", params={"_selected_model_group": "gpt-4"}) assert response.status_code == 200 assert normalize(response.json()) == {"data": [], "all_api_bases": []} @@ -94,9 +96,7 @@ def test_model_metrics_no_prisma_error(client, auth_as, no_prisma): # --------------------------------------------------------------------------- -def test_model_metrics_slow_responses_happy( - client, auth_as, prisma_with_query_raw, monkeypatch -): +def test_model_metrics_slow_responses_happy(client, auth_as, prisma_with_query_raw, monkeypatch): """Pins ``GET /model/metrics/slow_responses`` (happy: empty list).""" logging_obj = MagicMock() logging_obj.slack_alerting_instance.alerting_threshold = 30 @@ -184,7 +184,12 @@ def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): pc = MagicMock() row = MagicMock() - row.param_value = {"alerting_args": {"daily_report_frequency": 7}} + row.param_value = { + "alerting_args": { + "daily_report_frequency": 7, + "report_check_interval": None, + } + } pc.db.litellm_config.find_first = AsyncMock(return_value=row) monkeypatch.setattr(proxy_server, "prisma_client", pc) @@ -198,12 +203,20 @@ def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): store.load_yaml( { "alerting": ["slack"], - "alerting_args": {"daily_report_frequency": 3}, + "alerting_args": { + "daily_report_frequency": 3, + "report_check_interval": 300, + }, } ) store.apply_db_row( "general_settings", - {"alerting_args": {"daily_report_frequency": 7}}, + { + "alerting_args": { + "daily_report_frequency": 7, + "report_check_interval": None, + } + }, ) monkeypatch.setattr(proxy_server.proxy_config, "settings", store) monkeypatch.setattr(proxy_server, "general_settings", store) @@ -215,6 +228,42 @@ def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): by_name = {entry["field_name"]: entry for entry in response.json()} assert by_name["slack_alerting"]["source"] == "config" assert by_name["daily_report_frequency"]["source"] == "db" + assert by_name["report_check_interval"]["source"] == "config" + assert by_name["budget_alert_ttl"]["source"] == "default" + + +@pytest.mark.parametrize("db_alerting_args", [None, []]) +def test_alerting_settings_handles_empty_db_args( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + db_alerting_args: JsonValue, +): + from litellm.proxy.config_resolvers import SettingsStore + + pc = MagicMock() + row = MagicMock() + row.param_value = {"alerting_args": db_alerting_args} + pc.db.litellm_config.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", pc) + + logging_obj = MagicMock() + args_model = MagicMock() + args_model.model_dump = MagicMock(return_value={}) + logging_obj.slack_alerting_instance.alerting_args = args_model + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) + + store = SettingsStore("general_settings") + store.load_yaml({"alerting_args": {"report_check_interval": 300}}) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + assert by_name["report_check_interval"]["source"] == "config" def test_alerting_settings_no_db_error(client, auth_as, no_prisma): diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index faf5b336410..d8308d7831f 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3198,6 +3198,7 @@ class TestPtuCostAttributionUISetting: assert response.status_code == 200 assert response.json()["values"]["enable_ptu_cost_attribution"] is False + assert response.json()["source"]["enable_ptu_cost_attribution"] == "default" def test_reported_true_once_the_env_var_is_set(self, mock_auth, monkeypatch): from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR @@ -3209,6 +3210,47 @@ class TestPtuCostAttributionUISetting: assert response.status_code == 200 assert response.json()["values"]["enable_ptu_cost_attribution"] is True + assert response.json()["source"]["enable_ptu_cost_attribution"] == "config" + + def test_reported_config_when_secret_manager_enables_the_flag( + self, mock_auth: None, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.is_ptu_cost_attribution_enabled", + lambda: True, + ) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is True + assert response.json()["source"]["enable_ptu_cost_attribution"] == "config" + + def test_reported_config_when_secret_manager_disables_the_flag( + self, mock_auth: None, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.is_ptu_cost_attribution_enabled", + lambda: False, + ) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_secret", + lambda *_args: False, + ) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is False + assert response.json()["source"]["enable_ptu_cost_attribution"] == "config" def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch): """A row written before the allowlist existed must not be able to turn the feature on.""" From 066cc1883afe3a69a49638df4c8d7cfbc1fd0f2d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 02:39:58 -0700 Subject: [PATCH 015/109] test(router): cover legacy lowest TPM selection --- .../router_strategy/test_lowest_tpm_rpm.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py diff --git a/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py b/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py new file mode 100644 index 00000000000..7b13b196d5b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py @@ -0,0 +1,54 @@ +from datetime import datetime, timedelta +from typing import Final + +from litellm import Router +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict + +MODEL_GROUP: Final = "lowest-tpm-router" +HIGH_USAGE_DEPLOYMENT_ID: Final = "highest-usage" +LOW_USAGE_DEPLOYMENT_ID: Final = "lowest-usage" + + +def _deployment(deployment_id: str) -> DeploymentTypedDict: + params: LiteLLMParamsTypedDict = { + "model": "gpt-4o", + "api_key": "key", + "mock_response": f"from {deployment_id}", + } + return { + "model_name": MODEL_GROUP, + "litellm_params": params, + "model_info": {"id": deployment_id}, + } + + +def test_usage_based_routing_v1_selects_the_lowest_recorded_tpm() -> None: + router: Final = Router( + model_list=[ + _deployment(HIGH_USAGE_DEPLOYMENT_ID), + _deployment(LOW_USAGE_DEPLOYMENT_ID), + ], + routing_strategy="usage-based-routing", + num_retries=0, + ) + usage_by_deployment: Final = { + HIGH_USAGE_DEPLOYMENT_ID: 100, + LOW_USAGE_DEPLOYMENT_ID: 1, + } + now: Final = datetime.now() + cache_keys: Final = tuple( + f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" + for offset in range(60) + ) + + for cache_key in cache_keys: + router.cache.set_cache( + key=cache_key, value=usage_by_deployment, ttl=float("inf") + ) + + deployment: Final = router.get_available_deployment( + model=MODEL_GROUP, + messages=[{"role": "user", "content": "test"}], + ) + + assert deployment["model_info"]["id"] == LOW_USAGE_DEPLOYMENT_ID From da3bd9e31c72e29a2a75ef82107d05e7f91cc4db Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 04:08:34 -0700 Subject: [PATCH 016/109] test(ui): model E2E cleanup failures as values --- tests/e2e/ui/helpers/roundTrip.ts | 111 ++++++++++++++++++++++-------- 1 file changed, 81 insertions(+), 30 deletions(-) diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts index 1fc2d0aec1a..91e48b7087c 100644 --- a/tests/e2e/ui/helpers/roundTrip.ts +++ b/tests/e2e/ui/helpers/roundTrip.ts @@ -33,39 +33,90 @@ export async function readBack( return (await res.json()) as T; } -export async function runWithCleanup( - action: () => Promise, - cleanup: () => Promise, -): Promise { - const outcome = await Promise.resolve() +type OperationOutcome = + | { readonly status: "success" } + | { readonly status: "failure"; readonly error: unknown }; + +type RunFailure = + | { readonly status: "action_failure"; readonly error: unknown } + | { readonly status: "cleanup_failure"; readonly error: unknown } + | { + readonly status: "action_and_cleanup_failure"; + readonly actionError: unknown; + readonly cleanupError: unknown; + }; + +function toRunFailure( + actionOutcome: OperationOutcome, + cleanupOutcome: OperationOutcome, +): RunFailure | null { + if ( + actionOutcome.status === "failure" && + cleanupOutcome.status === "failure" + ) { + return { + status: "action_and_cleanup_failure", + actionError: actionOutcome.error, + cleanupError: cleanupOutcome.error, + }; + } + if (actionOutcome.status === "failure") { + return { status: "action_failure", error: actionOutcome.error }; + } + if (cleanupOutcome.status === "failure") { + return { status: "cleanup_failure", error: cleanupOutcome.error }; + } + return null; +} + +function raiseRunFailure(failure: RunFailure): never { + switch (failure.status) { + case "action_failure": + throw failure.error; + case "cleanup_failure": + throw failure.error; + case "action_and_cleanup_failure": + throw new AggregateError( + [failure.actionError, failure.cleanupError], + "Action and cleanup failed", + ); + } +} + +async function runAction( + action: () => void | Promise, +): Promise { + return Promise.resolve() .then(action) .then( () => ({ status: "success" as const }), (error: unknown) => ({ status: "failure" as const, error }), ); - try { - if (outcome.status === "failure") throw outcome.error; - } finally { - const cleanupOutcome = await Promise.resolve() - .then(cleanup) - .then( - (succeeded) => - succeeded - ? { status: "success" as const } - : { - status: "failure" as const, - error: new Error("Failed to clean up UI E2E resource"), - }, - (error: unknown) => ({ status: "failure" as const, error }), - ); - if (cleanupOutcome.status === "failure") { - if (outcome.status === "failure") { - throw new AggregateError( - [outcome.error, cleanupOutcome.error], - "Action and cleanup failed", - ); - } - throw cleanupOutcome.error; - } - } +} + +async function runCleanup( + cleanup: () => boolean | Promise, +): Promise { + return Promise.resolve() + .then(cleanup) + .then( + (succeeded) => + succeeded + ? { status: "success" as const } + : { + status: "failure" as const, + error: new Error("Failed to clean up UI E2E resource"), + }, + (error: unknown) => ({ status: "failure" as const, error }), + ); +} + +export async function runWithCleanup( + action: () => void | Promise, + cleanup: () => boolean | Promise, +): Promise { + const actionOutcome = await runAction(action); + const cleanupOutcome = await runCleanup(cleanup); + const failure = toRunFailure(actionOutcome, cleanupOutcome); + if (failure !== null) raiseRunFailure(failure); } From c1c566db875f31e28e07f662be39207dca9d8344 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 01:18:18 -0700 Subject: [PATCH 017/109] fix(proxy): record aborted outcome when spend-log cleanup is cancelled at shutdown cleanup_old_spend_logs only caught Exception, so a run cut short by CancelledError recorded no outcome and logged nothing. Under uvicorn the job was never cancelled at all: uvicorn re-raises the captured SIGTERM as soon as the lifespan shutdown returns, before asyncio cancels outstanding tasks, so an in-flight scheduler job simply died with the process. The cleanup now handles CancelledError by logging elapsed time, rows deleted and batch count at error level, recording outcome="aborted", and re-raising. The lifespan shutdown stops the scheduler and awaits the jobs it cancels while the database is still connected, so that handler runs under uvicorn too, and the pod lock is released instead of orphaned. Resolves LIT-6990 --- .../db_transaction_queue/spend_log_cleanup.py | 16 +++ litellm/proxy/proxy_server.py | 18 ++- litellm/proxy/shutdown/scheduled_jobs.py | 70 ++++++++++ .../proxy/shutdown/test_scheduled_jobs.py | 125 ++++++++++++++++++ .../proxy/test_spend_log_cleanup.py | 92 +++++++++++++ 5 files changed, 318 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/shutdown/scheduled_jobs.py create mode 100644 tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index e97e9f6e683..1a14210dbec 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -96,6 +96,8 @@ class SpendLogCleanup: self.general_settings = general_settings or default_settings self._refresh_bounds() + self._run_rows_deleted: int = 0 + self._run_batches: int = 0 from litellm.proxy.proxy_server import proxy_logging_obj pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager @@ -422,6 +424,8 @@ class SpendLogCleanup: total_deleted += deleted_count run_count += 1 + self._run_rows_deleted += deleted_count + self._run_batches += 1 # Add a small sleep to prevent overwhelming the database await asyncio.sleep(0.1) @@ -590,6 +594,9 @@ class SpendLogCleanup: If no pod_lock_manager, runs cleanup without distributed locking. """ lock_acquired = False + run_started_at: Final = time.monotonic() + self._run_rows_deleted = 0 + self._run_batches = 0 try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() @@ -670,6 +677,15 @@ class SpendLogCleanup: self._run_outcome(spend_log_results + session_results + health_check_results) ) + except asyncio.CancelledError: + verbose_proxy_logger.error( + "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here", + time.monotonic() - run_started_at, + self._run_rows_deleted, + self._run_batches, + ) + SpendLogCleanupMetrics.record_run("aborted") + raise except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB # timeout is often empty and gives operators no signal to diagnose. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b5f4236d22..cab0f4d0733 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -680,6 +680,10 @@ from litellm.proxy.route_llm_request import route_request from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager +from litellm.proxy.shutdown.scheduled_jobs import ( + AwaitableAsyncIOExecutor, + cancel_in_flight_scheduler_jobs, +) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, @@ -1456,6 +1460,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await proxy_config.stop_auth_cache_invalidation_subscriber() + # Shutdown event - cancel and await in-flight scheduled jobs while the DB is still connected + if scheduler is not None and scheduler_executor is not None: + try: + await cancel_in_flight_scheduler_jobs(scheduler, scheduler_executor) + except Exception as e: + verbose_proxy_logger.error("Error cancelling in-flight scheduled jobs: %s", e) + await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) if prometheus_multiproc_dir: @@ -2451,6 +2462,7 @@ celery_app_conn: Final = None celery_fn: Final = None # Redis Queue for handling requests scheduler = None +scheduler_executor: AwaitableAsyncIOExecutor | None = None # rebind-ok: bound once the scheduler is built at startup # Global variable for anthropic beta headers reload scheduling last_anthropic_beta_headers_reload = None @@ -9763,7 +9775,7 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" - global heuristic_v1_tuning_baselines, store_model_in_db, scheduler # rebind-ok: startup publishes the one read-only baseline snapshot + global heuristic_v1_tuning_baselines, store_model_in_db, scheduler, scheduler_executor # rebind-ok: startup publishes the one read-only baseline snapshot # MEMORY LEAK FIX: Configure scheduler with optimized settings # Memray analysis showed APScheduler's normalize() and _apply_jitter() causing @@ -9772,9 +9784,9 @@ class ProxyStartupEvent: # 1. Remove/minimize jitter to avoid normalize() memory explosion # 2. Use larger misfire_grace_time to prevent backlog calculations # 3. Set replace_existing=True to avoid duplicate jobs - from apscheduler.executors.asyncio import AsyncIOExecutor from apscheduler.jobstores.memory import MemoryJobStore + scheduler_executor = AwaitableAsyncIOExecutor() # rebind-ok: shutdown awaits the jobs this executor runs scheduler = AsyncIOScheduler( job_defaults={ "coalesce": APSCHEDULER_COALESCE, @@ -9787,7 +9799,7 @@ class ProxyStartupEvent: jobstores={"default": MemoryJobStore()}, # explicitly use memory job store # Use simple executor to minimize overhead executors={ - "default": AsyncIOExecutor(), + "default": scheduler_executor, }, # Disable timezone awareness to reduce computation timezone=None, diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py new file mode 100644 index 00000000000..3c9e791f51c --- /dev/null +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -0,0 +1,70 @@ +""" +Cancel the proxy's in-flight scheduled jobs at shutdown so they can record how they ended. + +APScheduler's ``AsyncIOExecutor.shutdown`` cancels the job tasks it has in flight but cannot +wait for them, because it is not a coroutine, and under uvicorn nothing else ever will: uvicorn +re-raises the SIGTERM it captured as soon as the ASGI lifespan shutdown returns, so the process +dies before ``asyncio.run`` reaches its cancel-all-tasks step. A job mid-run at that point is +killed without ever observing cancellation, which is how a spend-log cleanup interrupted by a +rolling restart left no outcome metric and no log line behind. Cancelling here and awaiting the +cancelled tasks while the database is still connected is what lets a job's own +``CancelledError`` handler run. + +The wait is bounded by ``JOB_CANCEL_TIMEOUT_SECONDS``. A job that has just been cancelled has only +its own cleanup left to do, so the bound is there for a job that swallows cancellation, not one +that honours it, and it keeps shutdown well inside a Kubernetes termination grace period. +""" + +# pyright: reportMissingTypeStubs=false # apscheduler ships no type information + +import asyncio +from collections.abc import Collection +from typing import Final, Protocol + +from apscheduler.executors.asyncio import AsyncIOExecutor + +from litellm._logging import verbose_proxy_logger + +JOB_CANCEL_TIMEOUT_SECONDS: Final = 5.0 + + +class StoppableScheduler(Protocol): + """The slice of ``AsyncIOScheduler`` shutdown uses, which ships no type information""" + + @property + def running(self) -> bool: ... + + def shutdown(self, wait: bool = ...) -> None: ... + + +class AwaitableAsyncIOExecutor(AsyncIOExecutor): # pyright: ignore[reportUntypedBaseClass] # apscheduler ships no type information and is absent from the type-check env + """``AsyncIOExecutor`` whose in-flight job tasks can be awaited after ``shutdown`` cancels them""" + + _pending_futures: Collection["asyncio.Future[object]"] + + def in_flight_jobs(self) -> tuple["asyncio.Future[object]", ...]: + """The job tasks that are running right now, as a snapshot""" + return tuple(future for future in self._pending_futures if not future.done()) + + +async def cancel_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: + """ + Stop the scheduler and wait, bounded by JOB_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. + + Must run before the database is disconnected: a job's cancellation handler is what records + the run's outcome, and it needs the connection the job was using. + """ + if not scheduler.running: + return + in_flight: Final = executor.in_flight_jobs() + scheduler.shutdown(wait=False) + if not in_flight: + return + verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(in_flight)) + _done, pending = await asyncio.wait(in_flight, timeout=JOB_CANCEL_TIMEOUT_SECONDS) + if pending: + verbose_proxy_logger.warning( + "%d scheduled job(s) did not finish within %ss of cancellation; giving up on them", + len(pending), + JOB_CANCEL_TIMEOUT_SECONDS, + ) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py new file mode 100644 index 00000000000..b77d7c4ae50 --- /dev/null +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -0,0 +1,125 @@ +""" +Tests for cancelling in-flight scheduled jobs at proxy shutdown. + +These drive a real AsyncIOScheduler: the point of the helper is the hand-off +between APScheduler's fire-and-forget cancellation and the lifespan shutdown +that has to outlive it, and a mocked scheduler would not exercise that. +""" + +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import datetime + +import pytest +from apscheduler.schedulers.asyncio import AsyncIOScheduler + +import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.proxy.shutdown.scheduled_jobs import ( + AwaitableAsyncIOExecutor, + cancel_in_flight_scheduler_jobs, +) + + +class _Job: + """A scheduled job that blocks until cancelled and records what it observed.""" + + def __init__(self, swallow_cancellation: bool = False) -> None: + self.started = asyncio.Event() + self.events: list[str] = [] + self.swallow_cancellation = swallow_cancellation + + async def run(self) -> None: + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.events.append("cancelled") + if self.swallow_cancellation: + await asyncio.Event().wait() + raise + finally: + self.events.append("finished") + + +@asynccontextmanager +async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOScheduler, AwaitableAsyncIOExecutor]]: + """A started scheduler with every job in flight; stopped on the way out whatever the test did.""" + executor = AwaitableAsyncIOExecutor() + scheduler = AsyncIOScheduler(executors={"default": executor}) + for index, job in enumerate(jobs): + scheduler.add_job(job.run, id=f"job-{index}", next_run_time=datetime.now()) + scheduler.start() + try: + for job in jobs: + await asyncio.wait_for(job.started.wait(), timeout=5) + yield scheduler, executor + finally: + if scheduler.running: + scheduler.shutdown(wait=False) + stragglers = executor.in_flight_jobs() + for straggler in stragglers: + straggler.cancel() + await asyncio.gather(*stragglers, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): + """ + The job's own CancelledError handler is what records how a run ended, so + shutdown must not return until that handler has run. + """ + job = _Job() + async with _running_scheduler(job) as (scheduler, executor): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert job.events == ["cancelled", "finished"] + assert scheduler.running is False + assert executor.in_flight_jobs() == () + + +@pytest.mark.asyncio +async def test_every_in_flight_job_is_cancelled_not_only_the_first(): + first, second = _Job(), _Job() + async with _running_scheduler(first, second) as (scheduler, executor): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert first.events == ["cancelled", "finished"] + assert second.events == ["cancelled", "finished"] + + +@pytest.mark.asyncio +async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(monkeypatch, caplog): + """ + A job that swallows CancelledError must not hold the pod past its + termination grace period, so shutdown gives up on it and says so. + """ + monkeypatch.setattr(scheduled_jobs, "JOB_CANCEL_TIMEOUT_SECONDS", 0.05) + job = _Job(swallow_cancellation=True) + async with _running_scheduler(job) as (scheduler, executor): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert job.events == ["cancelled"] + assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text + + +@pytest.mark.asyncio +async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler(): + async with _running_scheduler() as (scheduler, executor): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + await asyncio.sleep(0) + + assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_a_scheduler_that_never_started_is_left_alone(): + """The proxy runs without a scheduler when it has no database; shutdown must not trip on that.""" + executor = AwaitableAsyncIOExecutor() + scheduler = AsyncIOScheduler(executors={"default": executor}) + + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert scheduler.running is False diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index bf1538183ab..ed35af7ee38 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -7,6 +7,7 @@ import math import time from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -1417,3 +1418,94 @@ def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(st """ results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons) assert SpendLogCleanup._run_outcome(results) == expected + + +_OTHER_OUTCOMES: Final = ("completed", "budget_exhausted", "batch_cap_reached", "skipped_locked", "skipped_disabled") + + +def _runs_recorded(outcome: str) -> float: + """The real ``litellm_spend_log_cleanup_runs_total`` sample for one outcome, 0 when unset""" + from prometheus_client import REGISTRY + + return REGISTRY.get_sample_value("litellm_spend_log_cleanup_runs_total", {"outcome": outcome}) or 0.0 + + +@pytest.mark.asyncio +async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch): + """ + Shutdown cancels a run by throwing CancelledError into whichever batch is in + flight. That is a BaseException, so the Exception handler never saw it and + an interrupted run left no outcome metric and no log line; operators could + not tell that cleanup stopped early, let alone how far it got. + """ + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + aborted_runs_before = _runs_recorded("aborted") + other_runs_before = {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} + + third_batch_reached = asyncio.Event() + + async def _execute_raw(sql, *args): + if third_batch_reached.is_set(): + raise AssertionError("no batch may be issued after the cancelled one") + if _execute_raw.calls < 2: + _execute_raw.calls += 1 + return 150 + third_batch_reached.set() + await asyncio.Event().wait() + + _execute_raw.calls = 0 + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = _execute_raw + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = MagicMock() + cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + cleaner.pod_lock_manager.release_lock = AsyncMock() + + run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client)) + await asyncio.wait_for(third_batch_reached.wait(), timeout=5) + run.cancel() + with pytest.raises(asyncio.CancelledError): + await run + + assert _runs_recorded("aborted") == aborted_runs_before + 1 + assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before + cleaner.pod_lock_manager.release_lock.assert_awaited_once() + mock_logger.exception.assert_not_called() + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert rendered.startswith("Spend log cleanup cancelled after ") + assert "s (rows_deleted=300, batches=2)" in rendered + + +@pytest.mark.asyncio +async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch): + """ + The scheduler holds one cleaner for the life of the process, so the + progress counters must start from zero on every run rather than carrying + an earlier run's totals into the cancellation line. + """ + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0]) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()]) + with pytest.raises(asyncio.CancelledError): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert "(rows_deleted=150, batches=1)" in rendered From 8ce8dd9b3b184a8e8e8884cde29f07e07063f6f7 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 02:25:36 -0700 Subject: [PATCH 018/109] fix(proxy): pause the scheduler at shutdown start and keep cleanup progress per run Review follow-ups on #41213: - Pause the scheduler as the first shutdown step so a job whose fire time falls inside the shutdown window does not start only to be cancelled. Jobs already running keep the whole window and are cancelled and awaited before the database disconnects, as before. - Keep the cleanup run's progress in a task-scoped ContextVar rather than on the cleaner instance, so two runs overlapping on one cleaner (APSCHEDULER_MAX_INSTANCES above 1 without a Redis lock) each report their own rows and batches on cancellation. - Drop the module docstrings the repository comment policy does not allow; the rationale lives in the PR description. --- .../db_transaction_queue/spend_log_cleanup.py | 37 +++++++++--- litellm/proxy/proxy_server.py | 5 ++ litellm/proxy/shutdown/scheduled_jobs.py | 25 +++----- .../proxy/shutdown/test_scheduled_jobs.py | 57 ++++++++++++------- .../proxy/test_spend_log_cleanup.py | 53 +++++++++++++---- 5 files changed, 121 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 1a14210dbec..34213c0d2ce 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1,5 +1,6 @@ import asyncio import time +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Final, Literal, TypeAlias @@ -40,6 +41,28 @@ class TableCleanupResult: stop_reason: StopReason +class _RunProgress: + """How far one cleanup run has got, reported if that run is cancelled""" + + def __init__(self) -> None: + self.rows_deleted: int = 0 + self.batches: int = 0 + + def record_batch(self, rows_deleted: int) -> None: + self.rows_deleted += rows_deleted + self.batches += 1 + + +_run_progress: ContextVar[_RunProgress] = ContextVar("spend_log_cleanup_run_progress") + + +def _record_run_batch(rows_deleted: int) -> None: + """Count a batch towards the run in progress, if a run is what issued it""" + progress: Final = _run_progress.get(None) + if progress is not None: + progress.record_batch(rows_deleted) + + class _RemainingRow(BaseModel): """One row of the capped outstanding-rows probe, validated out of prisma's untyped result.""" @@ -96,8 +119,6 @@ class SpendLogCleanup: self.general_settings = general_settings or default_settings self._refresh_bounds() - self._run_rows_deleted: int = 0 - self._run_batches: int = 0 from litellm.proxy.proxy_server import proxy_logging_obj pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager @@ -424,8 +445,7 @@ class SpendLogCleanup: total_deleted += deleted_count run_count += 1 - self._run_rows_deleted += deleted_count - self._run_batches += 1 + _record_run_batch(deleted_count) # Add a small sleep to prevent overwhelming the database await asyncio.sleep(0.1) @@ -595,8 +615,8 @@ class SpendLogCleanup: """ lock_acquired = False run_started_at: Final = time.monotonic() - self._run_rows_deleted = 0 - self._run_batches = 0 + progress: Final = _RunProgress() + progress_token: Final = _run_progress.set(progress) try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() @@ -681,8 +701,8 @@ class SpendLogCleanup: verbose_proxy_logger.error( "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here", time.monotonic() - run_started_at, - self._run_rows_deleted, - self._run_batches, + progress.rows_deleted, + progress.batches, ) SpendLogCleanupMetrics.record_run("aborted") raise @@ -697,6 +717,7 @@ class SpendLogCleanup: SpendLogCleanupMetrics.record_run("aborted") return # Return after error handling finally: + _run_progress.reset(progress_token) # Only release the lock if it was actually acquired if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cab0f4d0733..1a7b3a6ccfb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -683,6 +683,7 @@ from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownMan from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, cancel_in_flight_scheduler_jobs, + pause_scheduled_jobs, ) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_counter_batch import ( @@ -1419,6 +1420,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: if model_info_scheduler is not scheduler: model_info_scheduler.shutdown(wait=False) + # Shutdown event - stop starting scheduled jobs; the ones already running keep the drain window + if scheduler is not None: + pause_scheduled_jobs(scheduler) + # Shutdown event - drain in-flight requests before tearing down dependencies # so SIGTERM (rolling update, scale-down, liveness kill) doesn't drop them. GracefulShutdownManager.start_shutdown() diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index 3c9e791f51c..46c57e1a608 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -1,20 +1,3 @@ -""" -Cancel the proxy's in-flight scheduled jobs at shutdown so they can record how they ended. - -APScheduler's ``AsyncIOExecutor.shutdown`` cancels the job tasks it has in flight but cannot -wait for them, because it is not a coroutine, and under uvicorn nothing else ever will: uvicorn -re-raises the SIGTERM it captured as soon as the ASGI lifespan shutdown returns, so the process -dies before ``asyncio.run`` reaches its cancel-all-tasks step. A job mid-run at that point is -killed without ever observing cancellation, which is how a spend-log cleanup interrupted by a -rolling restart left no outcome metric and no log line behind. Cancelling here and awaiting the -cancelled tasks while the database is still connected is what lets a job's own -``CancelledError`` handler run. - -The wait is bounded by ``JOB_CANCEL_TIMEOUT_SECONDS``. A job that has just been cancelled has only -its own cleanup left to do, so the bound is there for a job that swallows cancellation, not one -that honours it, and it keeps shutdown well inside a Kubernetes termination grace period. -""" - # pyright: reportMissingTypeStubs=false # apscheduler ships no type information import asyncio @@ -34,6 +17,8 @@ class StoppableScheduler(Protocol): @property def running(self) -> bool: ... + def pause(self) -> None: ... + def shutdown(self, wait: bool = ...) -> None: ... @@ -47,6 +32,12 @@ class AwaitableAsyncIOExecutor(AsyncIOExecutor): # pyright: ignore[reportUntype return tuple(future for future in self._pending_futures if not future.done()) +def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: + """Stop the scheduler from starting jobs that shutdown would only cancel; running jobs continue""" + if scheduler.running: + scheduler.pause() + + async def cancel_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: """ Stop the scheduler and wait, bounded by JOB_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index b77d7c4ae50..3301ce34cd6 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -1,16 +1,8 @@ -""" -Tests for cancelling in-flight scheduled jobs at proxy shutdown. - -These drive a real AsyncIOScheduler: the point of the helper is the hand-off -between APScheduler's fire-and-forget cancellation and the lifespan shutdown -that has to outlive it, and a mocked scheduler would not exercise that. -""" - import asyncio import logging from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from datetime import datetime +from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler @@ -19,11 +11,12 @@ import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, cancel_in_flight_scheduler_jobs, + pause_scheduled_jobs, ) class _Job: - """A scheduled job that blocks until cancelled and records what it observed.""" + """A scheduled job that blocks until cancelled and records what it observed""" def __init__(self, swallow_cancellation: bool = False) -> None: self.started = asyncio.Event() @@ -45,7 +38,7 @@ class _Job: @asynccontextmanager async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOScheduler, AwaitableAsyncIOExecutor]]: - """A started scheduler with every job in flight; stopped on the way out whatever the test did.""" + """A started scheduler with every job in flight, stopped on the way out whatever the test did""" executor = AwaitableAsyncIOExecutor() scheduler = AsyncIOScheduler(executors={"default": executor}) for index, job in enumerate(jobs): @@ -66,10 +59,7 @@ async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOSchedule @pytest.mark.asyncio async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): - """ - The job's own CancelledError handler is what records how a run ended, so - shutdown must not return until that handler has run. - """ + """The job's own CancelledError handler records how a run ended, so shutdown must wait for it""" job = _Job() async with _running_scheduler(job) as (scheduler, executor): await cancel_in_flight_scheduler_jobs(scheduler, executor) @@ -91,10 +81,7 @@ async def test_every_in_flight_job_is_cancelled_not_only_the_first(): @pytest.mark.asyncio async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(monkeypatch, caplog): - """ - A job that swallows CancelledError must not hold the pod past its - termination grace period, so shutdown gives up on it and says so. - """ + """A job that swallows CancelledError must not hold the pod past its termination grace period""" monkeypatch.setattr(scheduled_jobs, "JOB_CANCEL_TIMEOUT_SECONDS", 0.05) job = _Job(swallow_cancellation=True) async with _running_scheduler(job) as (scheduler, executor): @@ -116,10 +103,40 @@ async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler(): @pytest.mark.asyncio async def test_a_scheduler_that_never_started_is_left_alone(): - """The proxy runs without a scheduler when it has no database; shutdown must not trip on that.""" + """The proxy runs without a scheduler when it has no database""" executor = AwaitableAsyncIOExecutor() scheduler = AsyncIOScheduler(executors={"default": executor}) await cancel_in_flight_scheduler_jobs(scheduler, executor) assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_pausing_stops_new_jobs_from_starting_but_leaves_running_ones_alone(): + """A job due during the shutdown drain would only be cancelled, so it must not start at all""" + running = _Job() + async with _running_scheduler(running) as (scheduler, executor): + late = _Job() + scheduler.add_job(late.run, id="late", next_run_time=datetime.now() + timedelta(seconds=0.1)) + + pause_scheduled_jobs(scheduler) + await asyncio.sleep(0.3) + + assert late.started.is_set() is False + assert running.events == [] + assert scheduler.running is True + + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert running.events == ["cancelled", "finished"] + assert late.started.is_set() is False + + +@pytest.mark.asyncio +async def test_pausing_a_scheduler_that_never_started_is_a_no_op(): + scheduler = AsyncIOScheduler(executors={"default": AwaitableAsyncIOExecutor()}) + + pause_scheduled_jobs(scheduler) + + assert scheduler.running is False diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index ed35af7ee38..1691b2d174a 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -1432,12 +1432,7 @@ def _runs_recorded(outcome: str) -> float: @pytest.mark.asyncio async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch): - """ - Shutdown cancels a run by throwing CancelledError into whichever batch is in - flight. That is a BaseException, so the Exception handler never saw it and - an interrupted run left no outcome metric and no log line; operators could - not tell that cleanup stopped early, let alone how far it got. - """ + """A run cut short by shutdown must leave its outcome and how far it got behind""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module mock_logger = MagicMock() @@ -1485,11 +1480,7 @@ async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_r @pytest.mark.asyncio async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch): - """ - The scheduler holds one cleaner for the life of the process, so the - progress counters must start from zero on every run rather than carrying - an earlier run's totals into the cancellation line. - """ + """The scheduler holds one cleaner for the life of the process, so progress must not carry over""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module mock_logger = MagicMock() @@ -1509,3 +1500,43 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc (error_call,) = mock_logger.error.call_args_list rendered = error_call[0][0] % error_call[0][1:] assert "(rows_deleted=150, batches=1)" in rendered + + +@pytest.mark.asyncio +async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch): + """With APSCHEDULER_MAX_INSTANCES above one, two runs share the cleaner but not their progress""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + + first_batch_done = asyncio.Event() + second_run_done = asyncio.Event() + + async def _slow_execute_raw(sql, *args): + first_batch_done.set() + await second_run_done.wait() + return 100 + + slow_client = MagicMock() + _wire_tx(slow_client.db) + slow_client.db.execute_raw = _slow_execute_raw + fast_client = MagicMock() + _wire_tx(fast_client.db) + fast_client.db.execute_raw = AsyncMock(side_effect=[150, 150, 0, 0]) + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + + slow_run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(slow_client)) + await asyncio.wait_for(first_batch_done.wait(), timeout=5) + await cleaner.cleanup_old_spend_logs(fast_client) + second_run_done.set() + await asyncio.sleep(0) + slow_run.cancel() + with pytest.raises(asyncio.CancelledError): + await slow_run + + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert "(rows_deleted=100, batches=1)" in rendered From 39a199d9a2285a0647fd0d70b3f2a7e2d72120d1 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 03:13:35 -0700 Subject: [PATCH 019/109] fix(proxy): let in-flight scheduled jobs finish before cancelling them at shutdown Cancelling every in-flight job the moment shutdown reached the scheduler dropped the rows a write job had already popped: flush_gateway_requests drains its accumulator before committing and does not restore it on CancelledError, and update_spend requeues its batch only after the shutdown drain had already run. Shutdown now waits up to JOB_FINISH_TIMEOUT_SECONDS for in-flight jobs to finish on their own, cancels the ones still running, and does both before the shutdown flushes so a requeued batch is still written. The cleanup run never finishes inside the grace, so it is still cancelled and still records outcome="aborted". Resolves LIT-6990 --- litellm/proxy/proxy_server.py | 16 ++++---- litellm/proxy/shutdown/scheduled_jobs.py | 22 +++++++---- .../proxy/shutdown/test_scheduled_jobs.py | 39 ++++++++++++++----- 3 files changed, 52 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1a7b3a6ccfb..7505714b418 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -682,8 +682,8 @@ from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - cancel_in_flight_scheduler_jobs, pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, ) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_counter_batch import ( @@ -1457,6 +1457,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await _drain_spend_event_producer_on_shutdown() + # Shutdown event - finish or cancel in-flight scheduled jobs before the shutdown flushes and the DB disconnect + if scheduler is not None and scheduler_executor is not None: + try: + await stop_in_flight_scheduler_jobs(scheduler, scheduler_executor) + except Exception as e: + verbose_proxy_logger.error("Error stopping in-flight scheduled jobs: %s", e) + await flush_spend_counters_on_shutdown() await _flush_spend_logs_queue_on_shutdown() @@ -1465,13 +1472,6 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await proxy_config.stop_auth_cache_invalidation_subscriber() - # Shutdown event - cancel and await in-flight scheduled jobs while the DB is still connected - if scheduler is not None and scheduler_executor is not None: - try: - await cancel_in_flight_scheduler_jobs(scheduler, scheduler_executor) - except Exception as e: - verbose_proxy_logger.error("Error cancelling in-flight scheduled jobs: %s", e) - await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) if prometheus_multiproc_dir: diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index 46c57e1a608..e7625a73b47 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -8,6 +8,7 @@ from apscheduler.executors.asyncio import AsyncIOExecutor from litellm._logging import verbose_proxy_logger +JOB_FINISH_TIMEOUT_SECONDS: Final = 5.0 JOB_CANCEL_TIMEOUT_SECONDS: Final = 5.0 @@ -38,21 +39,28 @@ def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: scheduler.pause() -async def cancel_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: +async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: """ - Stop the scheduler and wait, bounded by JOB_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. + Let in-flight jobs finish for up to JOB_FINISH_TIMEOUT_SECONDS, then stop the scheduler and + wait, bounded by JOB_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. - Must run before the database is disconnected: a job's cancellation handler is what records - the run's outcome, and it needs the connection the job was using. + Must run before the database is disconnected: a write job that finishes needs its connection, + and a job's cancellation handler is what records the run's outcome. """ if not scheduler.running: return in_flight: Final = executor.in_flight_jobs() + still_running: set[asyncio.Future[object]] = set() + if in_flight: + verbose_proxy_logger.info( + "Waiting up to %ss for %d in-flight scheduled job(s) to finish", JOB_FINISH_TIMEOUT_SECONDS, len(in_flight) + ) + _done, still_running = await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS) scheduler.shutdown(wait=False) - if not in_flight: + if not still_running: return - verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(in_flight)) - _done, pending = await asyncio.wait(in_flight, timeout=JOB_CANCEL_TIMEOUT_SECONDS) + verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running)) + _done, pending = await asyncio.wait(still_running, timeout=JOB_CANCEL_TIMEOUT_SECONDS) if pending: verbose_proxy_logger.warning( "%d scheduled job(s) did not finish within %ss of cancellation; giving up on them", diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 3301ce34cd6..7defd6cef6c 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -10,23 +10,28 @@ from apscheduler.schedulers.asyncio import AsyncIOScheduler import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - cancel_in_flight_scheduler_jobs, + stop_in_flight_scheduler_jobs, pause_scheduled_jobs, ) class _Job: - """A scheduled job that blocks until cancelled and records what it observed""" + """A scheduled job that blocks until cancelled, or for ``work_seconds``, and records what it observed""" - def __init__(self, swallow_cancellation: bool = False) -> None: + def __init__(self, swallow_cancellation: bool = False, work_seconds: float | None = None) -> None: self.started = asyncio.Event() self.events: list[str] = [] self.swallow_cancellation = swallow_cancellation + self.work_seconds = work_seconds async def run(self) -> None: self.started.set() try: - await asyncio.Event().wait() + if self.work_seconds is None: + await asyncio.Event().wait() + else: + await asyncio.sleep(self.work_seconds) + self.events.append("committed") except asyncio.CancelledError: self.events.append("cancelled") if self.swallow_cancellation: @@ -62,18 +67,32 @@ async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): """The job's own CancelledError handler records how a run ended, so shutdown must wait for it""" job = _Job() async with _running_scheduler(job) as (scheduler, executor): - await cancel_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor) assert job.events == ["cancelled", "finished"] assert scheduler.running is False assert executor.in_flight_jobs() == () +@pytest.mark.asyncio +async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(monkeypatch): + """A spend write cancelled mid-commit drops the rows it popped, so short jobs get to finish first""" + monkeypatch.setattr(scheduled_jobs, "JOB_FINISH_TIMEOUT_SECONDS", 2.0) + write = _Job(work_seconds=0.2) + stuck = _Job() + async with _running_scheduler(write, stuck) as (scheduler, executor): + await stop_in_flight_scheduler_jobs(scheduler, executor) + + assert write.events == ["committed", "finished"] + assert stuck.events == ["cancelled", "finished"] + assert scheduler.running is False + + @pytest.mark.asyncio async def test_every_in_flight_job_is_cancelled_not_only_the_first(): first, second = _Job(), _Job() async with _running_scheduler(first, second) as (scheduler, executor): - await cancel_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor) assert first.events == ["cancelled", "finished"] assert second.events == ["cancelled", "finished"] @@ -86,7 +105,7 @@ async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(mo job = _Job(swallow_cancellation=True) async with _running_scheduler(job) as (scheduler, executor): with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - await cancel_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor) assert job.events == ["cancelled"] assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text @@ -95,7 +114,7 @@ async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(mo @pytest.mark.asyncio async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler(): async with _running_scheduler() as (scheduler, executor): - await cancel_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor) await asyncio.sleep(0) assert scheduler.running is False @@ -107,7 +126,7 @@ async def test_a_scheduler_that_never_started_is_left_alone(): executor = AwaitableAsyncIOExecutor() scheduler = AsyncIOScheduler(executors={"default": executor}) - await cancel_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor) assert scheduler.running is False @@ -127,7 +146,7 @@ async def test_pausing_stops_new_jobs_from_starting_but_leaves_running_ones_alon assert running.events == [] assert scheduler.running is True - await cancel_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor) assert running.events == ["cancelled", "finished"] assert late.started.is_set() is False From e9109ddf4a563c7d72b30db9bbcc1d334111fec8 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 18 Sep 2026 15:11:12 -0500 Subject: [PATCH 020/109] fix(router): make context-window escalation opt-in --- .../complexity_router/README.md | 13 ++++ .../complexity_router/config.py | 5 +- .../router_strategy/test_complexity_router.py | 70 ++++++++++++++----- .../ContextWindowEscalationConfig.tsx | 5 +- .../add_model/add_auto_router_tab.test.tsx | 11 +-- .../build_complexity_router_config.test.ts | 19 +++-- .../build_complexity_router_config.ts | 6 +- ...d_updated_complexity_router_config.test.ts | 23 ++++-- .../src/lib/autorouter_presets.test.ts | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 10 files changed, 111 insertions(+), 51 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 6505746bca1..2c2aea333f9 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -68,6 +68,19 @@ still resolve to a deployment in `model_list`; this configuration does not creat - abc ``` +### Context-window escalation + +Context-window escalation is opt-in. Omit `enable_context_window_escalation` or set it to +`false` to keep the complexity-selected model without context-window replacement or filtering + +Set `enable_context_window_escalation: true` inside `complexity_router_config` to restrict the +selected tier to models whose declared windows fit the prompt, or move to the lowest configured +tier with a fitting model when none in the selected tier fit. Unknown windows do not justify +moving a request. `context_window_escalation_buffer` defaults to `0.95` + +Existing saved configurations with explicit `true` keep escalation enabled. Configurations that +omit the setting now default to disabled; set it to `true` to retain their previous behavior + ### Capability forecasting Set `classifier_type: capability` to use diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index aa39dff8c53..213d3864dc0 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1305,7 +1305,7 @@ class ComplexityRouterConfig(BaseModel): ) enable_context_window_escalation: bool = Field( - default=True, + default=False, description=( "Escalate a request off a tier whose models provably cannot hold its prompt, before " "dispatch. The classifier scores complexity and never prompt size, so a long agentic " @@ -1315,7 +1315,8 @@ class ComplexityRouterConfig(BaseModel): "moves to the lowest configured tier with a model whose declared window fits; when " "only some of the tier's models fit, the pick is restricted to those and the tier " "keeps the request. Models with no resolvable window are never escalated away from " - "and never escalated onto. Set false to dispatch on complexity alone, as before." + "and never escalated onto. Disabled by default: omit or set false to dispatch on " + "complexity alone; set true to enable context-window escalation." ), ) context_window_escalation_buffer: float = Field( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 9b25c869f1c..10d674f1ab8 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -13,6 +13,7 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from copy import deepcopy from functools import partial +from types import MappingProxyType from typing import Dict, Final, List, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -90,6 +91,7 @@ from litellm.types.router import ( TaggedPreRoutingStrategy, ) from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.management_endpoints.auto_router_endpoints import RequestComplexityRouterConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -6558,6 +6560,7 @@ class TestTierModelAffinity: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config={ "tiers": {"SIMPLE": ["small-model", "big-model"]}, + "enable_context_window_escalation": True, "adaptive": adaptive, "deployment_affinity": True, "session_affinity": False, @@ -13723,8 +13726,12 @@ _CJK_TURNS = [ ] -def _tier_config(**overrides) -> Dict: - return {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, **overrides} +def _tier_config(**overrides: object) -> dict[str, object]: + return { + "tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, + "enable_context_window_escalation": True, + **overrides, + } class TestContextWindowEscalation: @@ -13783,7 +13790,7 @@ class TestContextWindowEscalation: router = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG), - complexity_router_config={"tiers": {"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}}, + complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -13820,7 +13827,7 @@ class TestContextWindowEscalation: }, ] ), - complexity_router_config={"tiers": {"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}}, + complexity_router_config=_tier_config(tiers={"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -13871,7 +13878,7 @@ class TestContextWindowEscalation: router = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(*deployments), - complexity_router_config={"tiers": tiers}, + complexity_router_config=_tier_config(tiers=tiers), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -13880,19 +13887,37 @@ class TestContextWindowEscalation: assert result.model == expected_model @pytest.mark.asyncio - async def test_the_disabled_gate_dispatches_on_complexity_alone(self): - """The escape hatch: enable_context_window_escalation false restores today's behavior.""" - router = ComplexityRouter( + @pytest.mark.parametrize("enabled", (None, False, True), ids=("omitted", "disabled", "enabled")) + @pytest.mark.parametrize("serialized", (False, True), ids=("config", "http-json")) + async def test_context_window_escalation_requires_opt_in(self, enabled: bool | None, serialized: bool) -> None: + setting: Final = ( + MappingProxyType({"enable_context_window_escalation": enabled}) + if enabled is not None + else MappingProxyType({}) + ) + raw_config: Final = RequestComplexityRouterConfig.model_validate( + MappingProxyType( + {"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting} + ) + ) + config: Final = ( + RequestComplexityRouterConfig.model_validate_json(raw_config.model_dump_json()) + if serialized + else raw_config + ) + router: Final = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(_SMALL, _BIG), - complexity_router_config=_tier_config(enable_context_window_escalation=False), + complexity_router_config=config.model_dump(exclude_unset=not serialized, exclude_none=True), ) - result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + result: Final = await router.async_pre_routing_hook( + model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS + ) assert result is not None - assert result.model == "small-model" - assert "context_escalated" not in result.routing_decision + assert result.model == ("big-model" if enabled else "small-model") + assert result.routing_decision.get("context_escalated", False) is (enabled is True) @pytest.mark.asyncio async def test_out_of_band_system_and_tools_count_against_the_window(self): @@ -14001,7 +14026,7 @@ class TestContextWindowEscalation: }, ] ), - complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["small-model", "mid-model"]}}, + complexity_router_config=_tier_config(adaptive=True, tiers={"SIMPLE": ["small-model", "mid-model"]}), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -14031,7 +14056,7 @@ class TestContextWindowEscalation: }, ] ), - complexity_router_config={"tiers": {"SIMPLE": "cop-pool", "COMPLEX": "big-model"}}, + complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}), ) real_get_llm_provider = litellm.get_llm_provider copilot_resolutions: List = [] @@ -14064,7 +14089,7 @@ class TestContextWindowEscalation: "model_name": "smart-router", "litellm_params": { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}}, + "complexity_router_config": _tier_config(), }, }, { @@ -14894,7 +14919,12 @@ class TestHealthFallbackDispatch: ) -> None: from litellm.types.router import RouterRateLimitError - router: Final = self._router(config={"tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"}}) + router: Final = self._router( + config={ + "tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"}, + "enable_context_window_escalation": True, + } + ) router.add_deployment( Deployment( model_name="large", @@ -14971,7 +15001,13 @@ class TestHealthFallbackDispatch: @pytest.mark.asyncio @pytest.mark.parametrize("default_fits", [True, False]) async def test_modality_default_must_also_fit_context(self, default_fits: bool) -> None: - router: Final = self._router(config={"modality_routing": True, "tiers": {"SIMPLE": "primary"}}) + router: Final = self._router( + config={ + "modality_routing": True, + "tiers": {"SIMPLE": "primary"}, + "enable_context_window_escalation": True, + } + ) for deployment in router.model_list: deployment["model_info"]["supports_vision"] = deployment["model_name"] == "fallback" deployment["model_info"]["max_input_tokens"] = 10000 if default_fits else 10 diff --git a/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx index c0a65076d20..ad09efd8059 100644 --- a/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx @@ -7,7 +7,7 @@ const ContextWindowEscalationConfig: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; }> = ({ value, onChange }) => { - const enabled = value.enable_context_window_escalation ?? true; + const enabled = value.enable_context_window_escalation ?? false; // A number input renders Number("0.") as "0", so a decimal cannot be typed without a local draft. const [bufferDraft, setBufferDraft] = React.useState(null); const commitBuffer = (raw: string) => { @@ -32,7 +32,8 @@ const ContextWindowEscalationConfig: React.FC<{ When a prompt provably cannot fit the decided tier's context windows, route it to the lowest tier whose - window holds it instead of letting the provider reject it. Off means requests dispatch on complexity alone. + window holds it instead of letting the provider reject it. Disabled by default. Off means requests dispatch on + complexity alone. {enabled && (
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 48903d585ff..578b455d2b8 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -669,7 +669,7 @@ describe("AddAutoRouterTab", () => { }); }); - it("carries a context-window escalation opt-out through to the create payload", async () => { + it("starts context-window escalation disabled and carries an explicit opt-in to the create payload", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); @@ -679,14 +679,15 @@ describe("AddAutoRouterTab", () => { expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Context Window Escalation")); const toggle = await screen.findByRole("switch", { name: "Escalate oversized prompts to a tier that fits" }); - expect(toggle).toBeChecked(); + expect(toggle).not.toBeChecked(); + expect(screen.queryByLabelText("Window fit buffer")).not.toBeInTheDocument(); await user.click(toggle); await user.click(screen.getByRole("button", { name: /add auto router/i })); await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ - enable_context_window_escalation: false, + enable_context_window_escalation: true, }); }); @@ -699,6 +700,7 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-buffer-router"); expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Context Window Escalation")); + await user.click(screen.getByRole("switch", { name: "Escalate oversized prompts to a tier that fits" })); const buffer = await screen.findByLabelText("Window fit buffer"); fireEvent.change(buffer, { target: { value: "1.5" } }); fireEvent.blur(buffer, { target: { value: "1.5" } }); @@ -708,7 +710,7 @@ describe("AddAutoRouterTab", () => { await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); const config = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config; expect(config).toMatchObject({ context_window_escalation_buffer: 1 }); - expect(config).not.toHaveProperty("enable_context_window_escalation"); + expect(config).toHaveProperty("enable_context_window_escalation", true); }); it("clearing the buffer removes it from the payload so the router tracks the backend default", async () => { @@ -720,6 +722,7 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-clear-router"); expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Context Window Escalation")); + await user.click(screen.getByRole("switch", { name: "Escalate oversized prompts to a tier that fits" })); const buffer = await screen.findByLabelText("Window fit buffer"); fireEvent.change(buffer, { target: { value: "0.8" } }); fireEvent.blur(buffer, { target: { value: "0.8" } }); 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 6e6e7a3c6cd..0d70b18cd94 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 @@ -49,7 +49,7 @@ const baseParams: BuildComplexityRouterConfigParams = { describe("buildComplexityRouterConfig", () => { it.each(["capability", "llm_v2", "heuristic"] as const)( - "disables the removed overrides only for forecast creates: %s", + "preserves explicit context-window opt-in beside forecast restrictions: %s", (classifierType) => { const forecast = classifierType !== "heuristic"; const params = { @@ -61,14 +61,10 @@ describe("buildComplexityRouterConfig", () => { }; const config = buildComplexityRouterConfig(params); expect(config.adaptive).toBe(!forecast); - expect(config.enable_context_window_escalation).toBe(!forecast); + expect(config.enable_context_window_escalation).toBe(true); + expect(config.context_window_escalation_buffer).toBe(0.9); expect(config.escalation_keywords).toEqual(forecast ? [] : ["LITELLM ESCALATE"]); - for (const key of [ - "adaptive_weights", - "adaptive_eligible", - "tier_distance_penalty", - "context_window_escalation_buffer", - ]) { + for (const key of ["adaptive_weights", "adaptive_eligible", "tier_distance_penalty"]) { expect(Object.hasOwn(config, key)).toBe(!forecast); } if (forecast) { @@ -107,13 +103,14 @@ describe("buildComplexityRouterConfig", () => { expect(config).toEqual(expected); }); - it("carries an explicit context-window escalation opt-out and buffer, false included", () => { + it.each([undefined, false, true])("preserves the context-window escalation setting: %s", (enabled) => { const config = buildComplexityRouterConfig({ ...baseParams, - enableContextWindowEscalation: false, + enableContextWindowEscalation: enabled, contextWindowEscalationBuffer: 0.9, }); - expect(config.enable_context_window_escalation).toBe(false); + expect(config.enable_context_window_escalation).toBe(enabled); + expect(Object.hasOwn(config, "enable_context_window_escalation")).toBe(enabled !== undefined); expect(config.context_window_escalation_buffer).toBe(0.9); }); 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..d1cea6d48a9 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 @@ -629,6 +629,7 @@ export const buildComplexityRouterConfig = ({ // the form never rewrote. The UI gates the same controls on this, not on the raw value. const effectiveType: ClassifierType = customTierSet ? "llm" : classifierType; const forecast = isForecastClassifier(effectiveType); + const preserveContextWindowBuffer = !forecast || enableContextWindowEscalation === true; const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType); const payload: ComplexityRouterConfigPayload = { @@ -682,11 +683,10 @@ export const buildComplexityRouterConfig = ({ adaptive_eligible: adaptiveEligible, }), ...(returnRawModelName && { return_raw_model_name: true }), - // Omission enables the backend default, so hidden forecast controls need an explicit opt-out. ...((forecast || enableContextWindowEscalation !== undefined) && { - enable_context_window_escalation: forecast ? false : enableContextWindowEscalation, + enable_context_window_escalation: enableContextWindowEscalation ?? false, }), - ...(!forecast && + ...(preserveContextWindowBuffer && contextWindowEscalationBuffer !== undefined && { context_window_escalation_buffer: contextWindowEscalationBuffer, }), 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..31f2b8ef68a 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 @@ -64,14 +64,10 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { const saved = buildUpdatedComplexityRouterConfig(stored, value, undefined, keywordState); const forecast = classifier_type !== "heuristic"; expect(saved.adaptive).toBe(!forecast); - expect(saved.enable_context_window_escalation).toBe(!forecast); + expect(saved.enable_context_window_escalation).toBe(true); + expect(saved.context_window_escalation_buffer).toBe(0.9); expect(saved.escalation_keywords).toEqual(forecast ? [] : stored.escalation_keywords); - for (const key of [ - "adaptive_weights", - "adaptive_eligible", - "tier_distance_penalty", - "context_window_escalation_buffer", - ]) { + for (const key of ["adaptive_weights", "adaptive_eligible", "tier_distance_penalty"]) { expect(Object.hasOwn(saved, key)).toBe(!forecast); } expect(saved.keyword_tier_rules).toEqual(STORED.keyword_tier_rules); @@ -83,6 +79,19 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { }, ); + it.each([undefined, false, true])("preserves stored context-window escalation on save: %s", (enabled) => { + const stored = { + ...STORED, + ...(enabled !== undefined && { enable_context_window_escalation: enabled }), + }; + const value = hydrateComplexityRouterConfig(stored, undefined); + const saved = buildUpdatedComplexityRouterConfig(stored, value, undefined, hydratedState); + const serialized: typeof saved = JSON.parse(JSON.stringify(saved)); + expect(value.enable_context_window_escalation).toBe(enabled); + expect(serialized.enable_context_window_escalation).toBe(enabled); + expect(Object.hasOwn(serialized, "enable_context_window_escalation")).toBe(enabled !== undefined); + }); + it("round-trips an untouched edit without changing any keyword-matching value", () => { // Opening the modal hydrates state from STORED; saving with nothing changed must be a // no-op. These keys are now MANAGED, so a hydration bug silently wipes them. diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index fed11454c23..cda9ba104e5 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -708,7 +708,7 @@ describe("autorouter_presets", () => { expect(prefill.escalationKeywords).toEqual([]); }); - it("carries a preset's context-window escalation opt-out and buffer through the prefill", () => { + it.each([undefined, false, true])("preserves a preset's context-window escalation setting: %s", (enabled) => { const prefill = buildPresetPrefill( { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, @@ -716,12 +716,12 @@ describe("autorouter_presets", () => { classification_mode: "every_request", session_affinity: false, deployment_affinity: true, - enable_context_window_escalation: false, + enable_context_window_escalation: enabled, context_window_escalation_buffer: 0.9, }, groupsOnly(["gpt-5-nano"]), ); - expect(prefill.complexityRouterConfig.enable_context_window_escalation).toBe(false); + expect(prefill.complexityRouterConfig.enable_context_window_escalation).toBe(enabled); expect(prefill.complexityRouterConfig.context_window_escalation_buffer).toBe(0.9); }); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d43adfe1ae4..a18b02646ee 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36502,8 +36502,8 @@ export interface components { embedding_model?: string | null; /** * Enable Context Window Escalation - * @description Escalate a request off a tier whose models provably cannot hold its prompt, before dispatch. The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial lands on a small-window tier and the provider rejects it with a context-window 400 that nothing retries. When every model of the decided tier has a declared window smaller than the estimated prompt, the request moves to the lowest configured tier with a model whose declared window fits; when only some of the tier's models fit, the pick is restricted to those and the tier keeps the request. Models with no resolvable window are never escalated away from and never escalated onto. Set false to dispatch on complexity alone, as before. - * @default true + * @description Escalate a request off a tier whose models provably cannot hold its prompt, before dispatch. The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial lands on a small-window tier and the provider rejects it with a context-window 400 that nothing retries. When every model of the decided tier has a declared window smaller than the estimated prompt, the request moves to the lowest configured tier with a model whose declared window fits; when only some of the tier's models fit, the pick is restricted to those and the tier keeps the request. Models with no resolvable window are never escalated away from and never escalated onto. Disabled by default: omit or set false to dispatch on complexity alone; set true to enable context-window escalation. + * @default false */ enable_context_window_escalation: boolean; /** From 5f722bc19559a82df160b748f26746dac7b55488 Mon Sep 17 00:00:00 2001 From: Tin Date: Sat, 19 Sep 2026 10:26:42 -0700 Subject: [PATCH 021/109] chore(router): remove in-repo escalation docs --- litellm/router_strategy/complexity_router/README.md | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 2c2aea333f9..6505746bca1 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -68,19 +68,6 @@ still resolve to a deployment in `model_list`; this configuration does not creat - abc ``` -### Context-window escalation - -Context-window escalation is opt-in. Omit `enable_context_window_escalation` or set it to -`false` to keep the complexity-selected model without context-window replacement or filtering - -Set `enable_context_window_escalation: true` inside `complexity_router_config` to restrict the -selected tier to models whose declared windows fit the prompt, or move to the lowest configured -tier with a fitting model when none in the selected tier fit. Unknown windows do not justify -moving a request. `context_window_escalation_buffer` defaults to `0.95` - -Existing saved configurations with explicit `true` keep escalation enabled. Configurations that -omit the setting now default to disabled; set it to `true` to retain their previous behavior - ### Capability forecasting Set `classifier_type: capability` to use From 0068df5a8beac2b137fc9412586048c2f42a0af3 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 19 Sep 2026 14:46:58 -0700 Subject: [PATCH 022/109] feat(ui): add internal-user savings and auto-router usage --- .../migration.sql | 42 +++ .../litellm_proxy_extras/schema.prisma | 41 ++ litellm/proxy/db/autorouter_session_rollup.py | 170 ++++++--- litellm/proxy/db/baseline_accounting.py | 40 +- .../db_transaction_queue/spend_log_cleanup.py | 24 +- .../auto_router_endpoints.py | 11 +- litellm/proxy/schema.prisma | 41 ++ schema.prisma | 41 ++ .../spend/test_autorouter_session_rollup.py | 171 ++++++++- .../spend/test_baseline_accounting.py | 67 +++- .../db/test_autorouter_session_rollup.py | 115 +++++- .../test_auto_router_endpoints.py | 42 ++- .../proxy/test_spend_log_cleanup.py | 13 +- .../AutoRouterBenchmarksTab.test.tsx | 4 +- .../_components/AutoRouterBenchmarksTab.tsx | 19 +- .../_components/useAutoRouterBenchmarks.ts | 9 +- .../useDailyActivityRange.test.tsx | 2 + .../_components/useDailyActivityRange.ts | 6 +- .../user_info_view.integration.test.tsx | 357 +++++++++++++++++- .../_components/view_users/user_info_view.tsx | 60 ++- .../components/shared/ScopedSavingsTab.tsx | 133 +++++++ .../components/templates/KeySavingsTab.tsx | 133 +------ ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 +- 23 files changed, 1321 insertions(+), 229 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql create mode 100644 ui/litellm-dashboard/src/components/shared/ScopedSavingsTab.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql new file mode 100644 index 00000000000..2b864131ab2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql @@ -0,0 +1,42 @@ +CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterUserSession" ( + "user_id" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "session_id" TEXT NOT NULL, + "router_name" TEXT NOT NULL, + "router_type" TEXT NOT NULL, + "first_turn_at" TIMESTAMP(3) NOT NULL, + "last_turn_at" TIMESTAMP(3) NOT NULL, + "last_model" TEXT NOT NULL, + "models" JSONB NOT NULL DEFAULT '{}', + "turns" INTEGER NOT NULL DEFAULT 0, + "unordered_turns" INTEGER NOT NULL DEFAULT 0, + "covered_turns" INTEGER NOT NULL DEFAULT 0, + "cache_hits" INTEGER NOT NULL DEFAULT 0, + "same_model_turns" INTEGER NOT NULL DEFAULT 0, + "same_model_hits" INTEGER NOT NULL DEFAULT 0, + "first_visit_turns" INTEGER NOT NULL DEFAULT 0, + "first_visit_hits" INTEGER NOT NULL DEFAULT 0, + "return_turns" INTEGER NOT NULL DEFAULT 0, + "return_hits" INTEGER NOT NULL DEFAULT 0, + "return_expired_misses" INTEGER NOT NULL DEFAULT 0, + "return_within_ttl_misses" INTEGER NOT NULL DEFAULT 0, + "ttl_5m_turns" INTEGER NOT NULL DEFAULT 0, + "ttl_1h_turns" INTEGER NOT NULL DEFAULT 0, + "total_tokens" BIGINT NOT NULL DEFAULT 0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "savings_estimated_turns" INTEGER NOT NULL DEFAULT 0, + "savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}', + "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0, + "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0, + "tier_turns" JSONB NOT NULL DEFAULT '{}', + "baseline_models" JSONB NOT NULL DEFAULT '{}', + + CONSTRAINT "LiteLLM_AutoRouterUserSession_pkey" PRIMARY KEY ("user_id", "api_key", "session_id", "router_name") +); + +CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_last_turn" ON "LiteLLM_AutoRouterUserSession"("last_turn_at"); + +CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_user_last_turn" ON "LiteLLM_AutoRouterUserSession"("user_id", "last_turn_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d2032cec0d0..f4015ed9277 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1620,6 +1620,47 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } +model LiteLLM_AutoRouterUserSession { + user_id String + api_key String + session_id String + router_name String + router_type String + first_turn_at DateTime + last_turn_at DateTime + last_model String + models Json @default("{}") + turns Int @default(0) + unordered_turns Int @default(0) + covered_turns Int @default(0) + cache_hits Int @default(0) + same_model_turns Int @default(0) + same_model_hits Int @default(0) + first_visit_turns Int @default(0) + first_visit_hits Int @default(0) + return_turns Int @default(0) + return_hits Int @default(0) + return_expired_misses Int @default(0) + return_within_ttl_misses Int @default(0) + ttl_5m_turns Int @default(0) + ttl_1h_turns Int @default(0) + total_tokens BigInt @default(0) + spend Float @default(0) + saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) + tier_turns Json @default("{}") + baseline_models Json @default("{}") + + @@id([user_id, api_key, session_id, router_name]) + @@index([last_turn_at], map: "idx_autorouter_user_session_last_turn") + @@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn") +} + // Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in // either direction. forward duplicates the requests the keys did not route through the // router through it, answering whether they should adopt it; reverse duplicates the diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 0d812ee812a..dd08cfd1bef 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -4,7 +4,7 @@ Per-session auto-router benchmarks rollup. At request time the spend writer builds one AutoRouterTurnTransaction per successful auto-routed request (a request whose metadata carries a routing_decision) and queues it on the prisma client. The spend-log flush job drains the queue into -LiteLLM_AutoRouterSession with one conditional upsert per turn: the statement classifies +key and user session rollups with one atomic statement per turn: each upsert classifies the turn (same model, first visit, return to a model the session already used, out of order) against the row's own columns, so nothing is read before the write and concurrent pods compose. The benchmarks endpoint aggregates these rows and never touches @@ -35,10 +35,27 @@ if TYPE_CHECKING: CACHE_TTL_5M_SECONDS: Final = 300 CACHE_TTL_1H_SECONDS: Final = 3600 -AUTOROUTER_BENCHMARKS_SQL: Final = """ +_SESSION_COLUMNS: Final = """ + api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, + last_model, models, turns, unordered_turns, covered_turns, cache_hits, + same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, + return_turns, return_hits, return_expired_misses, return_within_ttl_misses, + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, + baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, + savings_estimated_baseline_models +""" + +AUTOROUTER_BENCHMARKS_SQL: Final = f""" WITH windowed AS ( - SELECT * FROM "LiteLLM_AutoRouterSession" - WHERE last_turn_at >= $1::timestamp + SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterSession" + WHERE $4::text IS NULL + AND last_turn_at >= $1::timestamp + AND first_turn_at < $2::timestamp + AND ($3::text IS NULL OR api_key = $3::text) + UNION ALL + SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterUserSession" + WHERE (($4::text IS NOT NULL AND user_id = $4::text) OR ($4::text IS NULL AND api_key = '')) + AND last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp AND ($3::text IS NULL OR api_key = $3::text) ), @@ -53,7 +70,7 @@ tier_maps AS ( ) SELECT agg.*, - COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns + COALESCE(tier_maps.tier_turns, '{{}}'::jsonb) AS tier_turns FROM ( SELECT router_name, @@ -111,6 +128,7 @@ class AutoRouterTurnTransaction: savings_estimated_turns: int = 0 savings_estimated_actual_spend: float = 0.0 savings_estimated_saved_spend: float = 0.0 + user_id: str = "" class TurnCacheFacts(NamedTuple): @@ -214,10 +232,11 @@ def build_autorouter_turn_transaction( if not isinstance(routing_decision, Mapping) or not routing_decision: return None router_name: Final = routing_decision.get("router_model_name") or payload.get("model_group") - api_key: Final = payload.get("api_key") + api_key: Final = payload.get("api_key") or "" + user_id: Final = payload.get("user") or "" session_id: Final = payload.get("session_id") model: Final = payload.get("model") - if not (isinstance(router_name, str) and router_name and api_key and session_id and model): + if not (isinstance(router_name, str) and router_name and (api_key or user_id) and session_id and model): return None turn_at: Final = _turn_time_utc(str(payload.get("startTime") or "")) if turn_at is None: @@ -236,6 +255,7 @@ def build_autorouter_turn_transaction( estimated_savings: Final = recorded_estimated_autorouter_savings(metadata) return AutoRouterTurnTransaction( api_key=api_key, + user_id=user_id, session_id=bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), @@ -293,18 +313,18 @@ _RETURN_MISS: Final = ( _IDLE_SECONDS: Final = f"EXTRACT(EPOCH FROM {_TURN_AT}::timestamp) - (t.models -> {_MODEL} ->> 'at')::float8" _CACHE_TOUCHED: Final = f"{_TOUCHED}::int = 1" -UPSERT_AUTOROUTER_SESSION_SQL: Final = f""" -INSERT INTO "LiteLLM_AutoRouterSession" AS t ( - api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, - last_model, models, turns, unordered_turns, covered_turns, cache_hits, - same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, - return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, - baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, - savings_estimated_baseline_models + +def _session_upsert_sql(*, user_scoped: bool) -> str: + table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession" + user_column: Final = "user_id, " if user_scoped else "" + user_value: Final = f"{_p('user_id')}::text, " if user_scoped else "" + required_identity: Final = _p("user_id" if user_scoped else "api_key") + return f""" +INSERT INTO "{table_name}" AS t ( + {user_column}{_SESSION_COLUMNS} ) -VALUES ( - {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, +SELECT + {user_value}{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, {_MODEL}, jsonb_build_object({_MODEL}, jsonb_build_object('at', EXTRACT(EPOCH FROM {_TURN_AT}::timestamp), 'ttl', {_CACHE_TTL}::int)), 1, 0, {_COVERED}::int, {_CACHE_HIT}::int, 0, 0, 1, {_CACHE_HIT}::int, @@ -315,8 +335,8 @@ VALUES ( {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA}, {_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8, {_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA} -) -ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET +WHERE {required_identity}::text <> '' +ON CONFLICT ({user_column}api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, total_tokens = t.total_tokens + EXCLUDED.total_tokens, spend = t.spend + EXCLUDED.spend, @@ -365,6 +385,17 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET """ +UPSERT_AUTOROUTER_SESSION_SQL: Final = f""" +WITH key_rollup AS ( + {_session_upsert_sql(user_scoped=False)} + RETURNING 1 +) +{_session_upsert_sql(user_scoped=True)} +""" + +UPSERT_AUTOROUTER_USER_SESSION_SQL: Final = _session_upsert_sql(user_scoped=True) + + def _as_sql_param(value: str | float | bool | datetime | None) -> str | float | None: if isinstance(value, bool): return int(value) @@ -377,18 +408,23 @@ def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS) -async def write_autorouter_turn(db: SupportsExecuteRaw, transaction: AutoRouterTurnTransaction) -> None: - await db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction)) +async def write_autorouter_turn( + db: SupportsExecuteRaw, + transaction: AutoRouterTurnTransaction, + statement: str = UPSERT_AUTOROUTER_SESSION_SQL, +) -> None: + await db.execute_raw(statement, *_upsert_params(transaction)) async def _upsert_turn_with_retry( prisma_client: PrismaClient, transaction: AutoRouterTurnTransaction, n_retry_times: int, + statement: str, ) -> None: for attempt in range(n_retry_times + 1): try: - await write_autorouter_turn(prisma_client.db, transaction) + await write_autorouter_turn(prisma_client.db, transaction, statement) except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise @@ -397,6 +433,58 @@ async def _upsert_turn_with_retry( return +def _session_partition(transaction: AutoRouterTurnTransaction) -> tuple[str, str, str, str]: + identity: Final = ("key", transaction.api_key) if transaction.api_key else ("user", transaction.user_id) + return (*identity, transaction.session_id, transaction.router_name) + + +async def _drain_session_partition( + prisma_client: PrismaClient, + transactions: tuple[AutoRouterTurnTransaction, ...], + n_retry_times: int, + statement: str, +) -> tuple[AutoRouterTurnTransaction, ...]: + for position, transaction in enumerate(transactions): + try: + await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times, statement) + except Exception as flush_err: # noqa: BLE001 # stop dependent turns without retrying an ambiguous write + verbose_proxy_logger.error( + "Spend tracking - auto-router session rollup flush failed for router %s; " + "%s of %s turn writes stopped in this partition: %s", + transaction.router_name, + len(transactions) - position, + len(transactions), + flush_err, + ) + return transactions[position:] + return () + + +async def _flush_session_partition( + prisma_client: PrismaClient, + transactions: tuple[AutoRouterTurnTransaction, ...], + n_retry_times: int, +) -> None: + failed_suffix: Final = await _drain_session_partition( + prisma_client, transactions, n_retry_times, UPSERT_AUTOROUTER_SESSION_SQL + ) + if not failed_suffix or not failed_suffix[0].api_key: + return + failed_user: Final = failed_suffix[0].user_id + other_users: Final = sorted( + ( + transaction + for transaction in failed_suffix[1:] + if transaction.user_id and transaction.user_id != failed_user + ), + key=lambda transaction: transaction.user_id, + ) + for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id): + await _drain_session_partition( + prisma_client, tuple(user_turns), n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL + ) + + async def flush_autorouter_turn_transactions( prisma_client: PrismaClient, transactions: Sequence[AutoRouterTurnTransaction], @@ -407,38 +495,20 @@ async def flush_autorouter_turn_transactions( Statements run sequentially in per-session event order: a turn's classification depends on the turns before it, and Postgres rejects one multi-row INSERT touching the same key twice. Only ConnectError is retried, per statement, because it proves - that statement never reached the database. Any other failure drops the remaining - turns of THAT session only, with an error log, and the flush continues with the - next session: sessions are independent state machines, so one poisoned statement - must not discard unrelated sessions, and a repeated increment is worse than an - undercount. Callers must not add their own retry around this function. + that statement never reached the database. A failed write stops its key and user + histories for this batch. Other users sharing that key can still advance their + independent user histories, with the key projection disabled and the real key + identity preserved. The failed turn is never replayed. Callers must not add their + own retry around this function. """ if not transactions: return ordered: Final = sorted( transactions, - key=lambda transaction: ( - transaction.api_key, - transaction.session_id, - transaction.router_name, - transaction.turn_at, - ), + key=lambda transaction: (*_session_partition(transaction), transaction.turn_at), ) - for session_key, session_group in groupby( + for _, session_group in groupby( ordered, - key=lambda transaction: (transaction.api_key, transaction.session_id, transaction.router_name), + key=_session_partition, ): - session_turns = tuple(session_group) - for position, transaction in enumerate(session_turns): - try: - await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times) - except Exception as flush_err: # noqa: BLE001 # a statement failure drops only its session's remainder by design - verbose_proxy_logger.error( - "Spend tracking - auto-router session rollup flush failed for router %s; " - "%s of %s turn transactions dropped for one session: %s", - session_key[2], - len(session_turns) - position, - len(session_turns), - flush_err, - ) - break + await _flush_session_partition(prisma_client, tuple(session_group), n_retry_times) diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 8622cb9e481..4219102d9aa 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -171,6 +171,7 @@ class _Change(BaseModel): request_id: str publication: BaselinePublication api_key: str + user_id: str = "" session_id: str router_name: str baseline_model: str @@ -256,42 +257,54 @@ SET publication = x.publication::text FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb) WHERE observations.request_id = x.request_id """ -_UPDATE_SESSIONS: Final = """ + + +def _session_correction_sql(*, user_scoped: bool) -> str: + table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession" + identity_columns: Final = ("user_id, " if user_scoped else "") + "api_key, session_id, router_name" + user_filter: Final = "WHERE user_id <> ''" if user_scoped else "" + user_match: Final = "session.user_id = totals.user_id AND " if user_scoped else "" + return f""" WITH changes AS ( SELECT * FROM jsonb_to_recordset($1::jsonb) AS x( - api_key text, session_id text, router_name text, baseline_model text, + user_id text, api_key text, session_id text, router_name text, baseline_model text, covered_delta int, actual_delta float8, savings_delta float8 ) + {user_filter} ), totals AS ( - SELECT api_key, session_id, router_name, SUM(covered_delta)::int AS covered_delta, + SELECT {identity_columns}, SUM(covered_delta)::int AS covered_delta, SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta - FROM changes GROUP BY api_key, session_id, router_name + FROM changes GROUP BY {identity_columns} ), models AS ( - SELECT api_key, session_id, router_name, jsonb_object_agg(baseline_model, delta) AS deltas + SELECT {identity_columns}, jsonb_object_agg(baseline_model, delta) AS deltas FROM ( - SELECT api_key, session_id, router_name, baseline_model, SUM(covered_delta)::int AS delta - FROM changes GROUP BY api_key, session_id, router_name, baseline_model - ) grouped GROUP BY api_key, session_id, router_name + SELECT {identity_columns}, baseline_model, SUM(covered_delta)::int AS delta + FROM changes GROUP BY {identity_columns}, baseline_model + ) grouped GROUP BY {identity_columns} ) -UPDATE "LiteLLM_AutoRouterSession" AS session +UPDATE "{table_name}" AS session SET saved_spend = session.saved_spend + totals.savings_delta, savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta, savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta, savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta, savings_estimated_baseline_models = ( - SELECT COALESCE(jsonb_object_agg(key, value), '{}'::jsonb) FROM ( + SELECT COALESCE(jsonb_object_agg(key, value), '{{}}'::jsonb) FROM ( SELECT key, SUM(value::int)::int AS value FROM ( SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models) UNION ALL SELECT * FROM jsonb_each_text(models.deltas) ) combined GROUP BY key HAVING SUM(value::int) > 0 ) counts ) -FROM totals JOIN models USING (api_key, session_id, router_name) -WHERE session.api_key = totals.api_key AND session.session_id = totals.session_id +FROM totals JOIN models USING ({identity_columns}) +WHERE {user_match}session.api_key = totals.api_key AND session.session_id = totals.session_id AND session.router_name = totals.router_name """ +_UPDATE_SESSIONS: Final = _session_correction_sql(user_scoped=False) +_UPDATE_USER_SESSIONS: Final = _session_correction_sql(user_scoped=True) + + def _primary_transaction(client: PrismaClient) -> _TransactionManager: primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db)) return primary.tx(timeout=_TRANSACTION_TIMEOUT) @@ -308,6 +321,7 @@ def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, n request_id=record.observation.request_id, publication=new, api_key=record.api_key, + user_id=record.turn.user_id if record.turn is not None else "", session_id=record.session_id, router_name=record.router_name, baseline_model=record.baseline_model, @@ -357,6 +371,8 @@ async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None: serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":")) await db.execute_raw(_UPDATE_LOGS, serialized) await db.execute_raw(_UPDATE_SESSIONS, serialized) + if any(change.user_id for change in changes): + await db.execute_raw(_UPDATE_USER_SESSIONS, serialized) for entity, table in DAILY_SPEND_TABLES.items(): if adjustments := tuple( change.daily.adjustment(target, change.savings_delta, change.request_id) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index b28a653c9aa..db6045071a8 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -492,6 +492,18 @@ class SpendLogCleanup: deadline=deadline, ) + async def _delete_old_autorouter_user_session_rows( + self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float + ) -> TableCleanupResult: + return await self._delete_old_rows_batched( + prisma_client, + cutoff_date, + table_name="LiteLLM_AutoRouterUserSession", + key_columns=("user_id", "api_key", "session_id", "router_name"), + time_column="last_turn_at", + deadline=deadline, + ) + async def _delete_old_health_check_rows( self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float ) -> TableCleanupResult: @@ -560,9 +572,17 @@ class SpendLogCleanup: ) except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job verbose_proxy_logger.warning("Auto-router baseline retention remains pending") - sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline) + sessions_result: Final = await self._delete_old_autorouter_session_rows( + prisma_client, session_cutoff, self._group_deadline(deadline, 2) + ) verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted) - return (sessions_result,) + user_sessions_result: Final = await self._delete_old_autorouter_user_session_rows( + prisma_client, session_cutoff, deadline + ) + verbose_proxy_logger.info( + "Deleted %s expired auto-router user session rollup rows", user_sessions_result.rows_deleted + ) + return (sessions_result, user_sessions_result) async def _clean_health_checks( self, prisma_client: PrismaClient, retention_seconds: int, deadline: float diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index a6d5a17d73e..32f3bf9accb 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -746,14 +746,18 @@ async def get_auto_router_benchmarks( ] = None, end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None, api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None, + user_id: Annotated[ + str | None, Query(min_length=1, description="Filter to one canonical internal user recorded on each turn") + ] = None, ) -> AutoRouterBenchmarksResponse: """ Benchmarks for the auto-router dashboard: session shape, savings against the configured baseline, and prompt-caching behaviour bucketed by what the router did. - Reads the LiteLLM_AutoRouterSession rollup, folded once per request at spend-write time, - so this endpoint never scans LiteLLM_SpendLogs. A session is in the window when it - overlaps it: its last turn is on or after start_date and its first turn is on or before + Reads session rollups folded once per request at spend-write time, so this endpoint + never scans LiteLLM_SpendLogs. A user filter selects only turns attributed to that + internal user when written; older key-only history remains outside user views. A session + is in the window when it overlaps it: its last turn is on or after start_date and its first turn is on or before end_date. Overall hit rate is over telemetry-bearing turns; each bucket's hit rate is over that bucket's turns. @@ -783,6 +787,7 @@ async def get_auto_router_benchmarks( start_day.isoformat(), (end_day + timedelta(days=1)).isoformat(), api_key, + user_id, ) rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ()) groups: Final = ( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d2032cec0d0..f4015ed9277 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1620,6 +1620,47 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } +model LiteLLM_AutoRouterUserSession { + user_id String + api_key String + session_id String + router_name String + router_type String + first_turn_at DateTime + last_turn_at DateTime + last_model String + models Json @default("{}") + turns Int @default(0) + unordered_turns Int @default(0) + covered_turns Int @default(0) + cache_hits Int @default(0) + same_model_turns Int @default(0) + same_model_hits Int @default(0) + first_visit_turns Int @default(0) + first_visit_hits Int @default(0) + return_turns Int @default(0) + return_hits Int @default(0) + return_expired_misses Int @default(0) + return_within_ttl_misses Int @default(0) + ttl_5m_turns Int @default(0) + ttl_1h_turns Int @default(0) + total_tokens BigInt @default(0) + spend Float @default(0) + saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) + tier_turns Json @default("{}") + baseline_models Json @default("{}") + + @@id([user_id, api_key, session_id, router_name]) + @@index([last_turn_at], map: "idx_autorouter_user_session_last_turn") + @@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn") +} + // Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in // either direction. forward duplicates the requests the keys did not route through the // router through it, answering whether they should adopt it; reverse duplicates the diff --git a/schema.prisma b/schema.prisma index d2032cec0d0..f4015ed9277 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1620,6 +1620,47 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } +model LiteLLM_AutoRouterUserSession { + user_id String + api_key String + session_id String + router_name String + router_type String + first_turn_at DateTime + last_turn_at DateTime + last_model String + models Json @default("{}") + turns Int @default(0) + unordered_turns Int @default(0) + covered_turns Int @default(0) + cache_hits Int @default(0) + same_model_turns Int @default(0) + same_model_hits Int @default(0) + first_visit_turns Int @default(0) + first_visit_hits Int @default(0) + return_turns Int @default(0) + return_hits Int @default(0) + return_expired_misses Int @default(0) + return_within_ttl_misses Int @default(0) + ttl_5m_turns Int @default(0) + ttl_1h_turns Int @default(0) + total_tokens BigInt @default(0) + spend Float @default(0) + saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) + tier_turns Json @default("{}") + baseline_models Json @default("{}") + + @@id([user_id, api_key, session_id, router_name]) + @@index([last_turn_at], map: "idx_autorouter_user_session_last_turn") + @@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn") +} + // Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in // either direction. forward duplicates the requests the keys did not route through the // router through it, answering whether they should adopt it; reverse duplicates the diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index f3c68b489a5..77549b527d8 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -6,17 +6,24 @@ tests/test_litellm/proxy/db/test_autorouter_session_rollup.py. """ import asyncio +import time import uuid from datetime import datetime, timedelta, timezone -from typing import Final +from types import SimpleNamespace +from typing import Final, TypedDict, cast import pytest from prisma import Prisma +from prisma.errors import RawQueryError +from typing_extensions import ReadOnly from litellm.proxy.db.autorouter_session_rollup import ( AUTOROUTER_BENCHMARKS_SQL, UPSERT_AUTOROUTER_SESSION_SQL, + AutoRouterTurnTransaction, + flush_autorouter_turn_transactions, ) +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -45,6 +52,7 @@ async def _turn( tier: "str | None" = None, baseline: "str | None" = None, estimated: bool = True, + user_id: str = "", ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( @@ -68,6 +76,7 @@ async def _turn( int(estimated), spend if estimated else 0.0, saved if estimated else 0.0, + user_id, ) @@ -217,7 +226,7 @@ async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers assert row["savings_estimated_actual_spend"] == pytest.approx(0.01 * sum(writers)) assert row["savings_estimated_saved_spend"] == pytest.approx(0.02 * sum(writers)) groups: Final = await db.query_raw( - AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key, None ) assert len(groups) == 1 assert groups[0]["classifier_cost"] == row["classifier_cost"] @@ -242,7 +251,7 @@ async def test_unknown_and_legacy_turns_preserve_actual_spend_without_entering_t assert row["saved_spend"] == pytest.approx(-0.03) assert row["savings_estimated_baseline_models"] == {"opus": 1} groups: Final = await db.query_raw( - AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key, None ) assert len(groups) == 1 for actual in (row, groups[0]): @@ -277,6 +286,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) matching = [row for row in rows if row["router_name"] == router] assert len(matching) == 1 @@ -304,6 +314,7 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), first_key, + None, ) matching = [row for row in rows if row["router_name"] == router] assert len(matching) == 1 @@ -317,10 +328,160 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), f"k-{uuid.uuid4()}", + None, ) assert [row for row in unknown_key_rows if row["router_name"] == router] == [] +class _BenchmarkRow(TypedDict): + sessions: ReadOnly[int] + turns: ReadOnly[int] + same_model_turns: ReadOnly[int] + first_visit_turns: ReadOnly[int] + spend: ReadOnly[float] + saved_spend: ReadOnly[float] + tier_turns: ReadOnly[dict[str, int]] + cache_hits: ReadOnly[int] + savings_estimated_turns: ReadOnly[int] + savings_estimated_actual_spend: ReadOnly[float] + savings_estimated_saved_spend: ReadOnly[float] + + +async def _scoped_benchmarks( + db: Prisma, router: str, user_id: str | None = None, key: str | None = None +) -> tuple[_BenchmarkRow, ...]: + rows: Final = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + key, + user_id, + ) + return tuple(cast(_BenchmarkRow, row) for row in rows if row["router_name"] == router) + + +async def test_users_keep_written_identity_across_shared_keys_and_keyless_sessions(db: Prisma) -> None: + router: Final = f"r-{uuid.uuid4()}" + alice: Final = f"u-{uuid.uuid4()}" + bob: Final = f"u-{uuid.uuid4()}" + first_key: Final = f"k-{uuid.uuid4()}" + second_key: Final = f"k-{uuid.uuid4()}" + await _legacy_turn(db, first_key, T0, router=router) + await _turn(db, first_key, "A", T0 + timedelta(seconds=10), router=router, user_id=alice, tier="simple") + await _turn( + db, first_key, "B", T0 + timedelta(seconds=20), router=router, user_id=bob, spend=0.03, saved=0.06, tier="complex" + ) + await _turn(db, second_key, "C", T0, router=router, user_id=alice, spend=0.02, saved=0.04) + await _turn(db, "", "A", T0, router=router, user_id=alice, ttl=300) + await _turn(db, "", "A", T0 + timedelta(seconds=1), router=router, user_id=alice, hit=1) + await _turn(db, "", "B", T0, router=router, user_id=bob, spend=0.04, saved=0.08) + await _turn(db, second_key, "C", T0 - timedelta(days=40), router=router, user_id=alice, session_id="expired") + + alice_rows: Final = await _scoped_benchmarks(db, router, user_id=alice) + bob_rows: Final = await _scoped_benchmarks(db, router, user_id=bob) + global_rows: Final = await _scoped_benchmarks(db, router) + key_rows: Final = await _scoped_benchmarks(db, router, key=first_key) + intersection: Final = await _scoped_benchmarks(db, router, user_id=alice, key=first_key) + assert len(alice_rows) == len(bob_rows) == len(global_rows) == len(key_rows) == len(intersection) == 1 + assert (alice_rows[0]["sessions"], alice_rows[0]["turns"], alice_rows[0]["same_model_turns"]) == (3, 4, 1) + assert (bob_rows[0]["sessions"], bob_rows[0]["turns"], bob_rows[0]["first_visit_turns"]) == (2, 2, 2) + assert alice_rows[0]["spend"] == pytest.approx(0.05) + assert bob_rows[0]["spend"] == pytest.approx(0.07) + assert alice_rows[0]["tier_turns"] == {"simple": 1} + assert bob_rows[0]["tier_turns"] == {"complex": 1} + assert (alice_rows[0]["cache_hits"], bob_rows[0]["cache_hits"]) == (1, 0) + assert (global_rows[0]["sessions"], global_rows[0]["turns"]) == (4, 7) + assert (alice_rows[0]["savings_estimated_turns"], bob_rows[0]["savings_estimated_turns"]) == (4, 2) + assert global_rows[0]["savings_estimated_turns"] == 6 + for scoped in (alice_rows[0], bob_rows[0]): + assert scoped["savings_estimated_actual_spend"] == pytest.approx(scoped["spend"]) + assert scoped["savings_estimated_saved_spend"] == pytest.approx(scoped["saved_spend"]) + assert global_rows[0]["spend"] == pytest.approx(alice_rows[0]["spend"] + bob_rows[0]["spend"] + 0.01) + assert global_rows[0]["saved_spend"] == pytest.approx(alice_rows[0]["saved_spend"] + bob_rows[0]["saved_spend"] + 0.02) + assert global_rows[0]["tier_turns"] == {"simple": 1, "complex": 1} + assert (key_rows[0]["sessions"], key_rows[0]["turns"]) == (1, 3) + assert key_rows[0]["spend"] == pytest.approx(0.05) + assert (intersection[0]["sessions"], intersection[0]["turns"]) == (1, 1) + assert intersection[0]["spend"] == pytest.approx(0.01) + assert await _scoped_benchmarks(db, router, user_id=bob, key=second_key) == () + assert await _scoped_benchmarks(db, router, user_id=f"u-{uuid.uuid4()}") == () + assert await _scoped_benchmarks(db, router, user_id="") == () + + +async def test_a_failed_user_projection_rolls_back_the_keys_increment(db: Prisma) -> None: + key: Final = f"k-{uuid.uuid4()}" + user_id: Final = "".join(str(uuid.uuid4()) for _ in range(200)) + await _turn(db, key, "A", T0) + before: Final = await _row(db, key) + + with pytest.raises(RawQueryError, match=r"index row (requires|size)"): + await _turn(db, key, "B", T0 + timedelta(seconds=1), user_id=user_id) + + assert await _row(db, key) == before + assert await db.query_raw('SELECT user_id FROM "LiteLLM_AutoRouterUserSession" WHERE user_id = $1', user_id) == [] + + first_user: Final = f"u-{uuid.uuid4()}" + second_user: Final = f"u-{uuid.uuid4()}" + turns: Final = tuple( + AutoRouterTurnTransaction( + api_key=key, + user_id=user, + session_id="s1", + router_name="auto-1", + router_type="complexity", + model=model, + turn_at=T0 + timedelta(seconds=second), + total_tokens=100, + spend=0.01, + saved_spend=0.02, + classifier_cost=0.0, + covered=True, + cache_hit=False, + cache_ttl_seconds=None, + cache_touched=False, + ) + for user, model, second in ( + (first_user, "A", 1), + (user_id, "B", 2), + (first_user, "B", 3), + (second_user, "C", 4), + (first_user, "B", 5), + (second_user, "C", 6), + (user_id, "A", 7), + ) + ) + await flush_autorouter_turn_transactions(SimpleNamespace(db=db), tuple(reversed(turns)), n_retry_times=0) + + key_row: Final = await _row(db, key) + assert (key_row["turns"], key_row["last_model"], key_row["unordered_turns"]) == (2, "A", 0) + assert key_row["spend"] == pytest.approx(0.02) + user_rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterUserSession" WHERE api_key = $1', key) + by_user: Final = {row["user_id"]: row for row in user_rows} + assert set(by_user) == {first_user, second_user} + for user, count, model in ((first_user, 3, "B"), (second_user, 2, "C")): + row: Final = by_user[user] + assert (row["turns"], row["same_model_turns"], row["unordered_turns"], row["last_model"]) == (count, 1, 0, model) + assert row["spend"] == pytest.approx(count * 0.01) + assert row["saved_spend"] == pytest.approx(count * 0.02) + + +async def test_user_session_cleanup_keeps_another_users_recent_keyless_session(db: Prisma) -> None: + router: Final = f"r-{uuid.uuid4()}" + expired_user: Final = f"u-{uuid.uuid4()}" + recent_user: Final = f"u-{uuid.uuid4()}" + await _turn(db, "", "A", T0 - timedelta(days=1), router=router, user_id=expired_user) + await _turn(db, "", "A", T0 + timedelta(days=1), router=router, user_id=recent_user) + cleaner: Final = SpendLogCleanup(general_settings={}) + + await cleaner._delete_old_autorouter_user_session_rows( + SimpleNamespace(db=db), T0.replace(tzinfo=timezone.utc), time.monotonic() + 60 + ) + + assert await db.query_raw( + 'SELECT user_id, turns FROM "LiteLLM_AutoRouterUserSession" WHERE router_name = $1', router + ) == [{"user_id": recent_user, "turns": 1}] + + async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db): key = f"k-{uuid.uuid4()}" router = f"r-{uuid.uuid4()}" @@ -334,6 +495,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) matching = sorted( (row for row in rows if row["router_name"] == router), @@ -418,6 +580,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {"simple": 2, "complex": 1} @@ -446,6 +609,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} @@ -461,6 +625,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {} diff --git a/tests/proxy_behavior/spend/test_baseline_accounting.py b/tests/proxy_behavior/spend/test_baseline_accounting.py index e187a44c29d..3504751d132 100644 --- a/tests/proxy_behavior/spend/test_baseline_accounting.py +++ b/tests/proxy_behavior/spend/test_baseline_accounting.py @@ -56,7 +56,9 @@ def record() -> Callable[..., BaselineAccountingRecord]: }, ) - def create(label: str = "first", started: float = 10000.0, identical: bool = True) -> BaselineAccountingRecord: + def create( + label: str = "first", started: float = 10000.0, identical: bool = True, user_id: str = "" + ) -> BaselineAccountingRecord: return BaselineAccountingRecord( scope="autorouter-baseline:v3:" + run * 2, api_key=run, session_id=run, router_name="test-router", baseline_model="anthropic/claude-opus-5", @@ -76,6 +78,7 @@ def record() -> Callable[..., BaselineAccountingRecord]: total_tokens=6230, spend=0.17, saved_spend=0.0, classifier_cost=0.0, covered=True, cache_hit=False, cache_ttl_seconds=3600, cache_touched=True, baseline_model="anthropic/claude-opus-5", + user_id=user_id, ), daily=DailyBaselineAttribution( date="2026-09-15", api_key=run, model="claude-opus-5", custom_llm_provider="anthropic", @@ -99,21 +102,36 @@ async def _session(db: Prisma, record: BaselineAccountingRecord): return rows[0] +async def _user_sessions(db: Prisma, record: BaselineAccountingRecord) -> dict[str, dict[str, object]]: + rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterUserSession" WHERE api_key=$1', record.api_key) + return {str(row["user_id"]): row for row in rows} + + async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: store: Final = _store(db) - late: Final = record("late", 10001.0) - early: Final = record("early", identical=False) + late: Final = record("late", 10001.0, user_id="late-user") + early: Final = record("early", identical=False, user_id="early-user") await _log(db, late) assert await store.append(late) == "recorded" assert await store.project(late.scope) == "published" before: Final = await _session(db, late) assert before["savings_estimated_actual_spend"] == before["spend"] == 0.17 assert before["saved_spend"] == 0.0 + before_users: Final = await _user_sessions(db, late) + assert set(before_users) == {"late-user"} + assert before_users["late-user"]["savings_estimated_turns"] == 1 + assert before_users["late-user"]["savings_estimated_baseline_models"] == {late.baseline_model: 1} await _log(db, early) assert await store.append(early) == "recorded" pending: Final = await _session(db, late) assert pending["spend"] == 0.34 and pending["savings_estimated_turns"] == 0 assert pending["saved_spend"] == pending["savings_estimated_actual_spend"] == 0.0 + pending_users: Final = await _user_sessions(db, late) + assert set(pending_users) == {"late-user", "early-user"} + for user in pending_users.values(): + assert user["turns"] == 1 and user["spend"] == 0.17 + assert user["savings_estimated_turns"] == user["savings_estimated_actual_spend"] == user["saved_spend"] == 0 + assert user["savings_estimated_baseline_models"] == {} waiting: Final = await db.query_raw('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', late.observation.request_id) assert waiting[0]["metadata"]["autorouter_savings"] is None assert waiting[0]["metadata"]["autorouter_savings_estimate"]["reason"] == "pending_projection" @@ -125,37 +143,69 @@ async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, assert logs[0]["spend"] == 0.17 assert logs[0]["metadata"]["autorouter_savings_estimate"]["provenance"] == "modeled" assert after["saved_spend"] == pytest.approx(logs[0]["metadata"]["autorouter_savings"]) + after_users: Final = await _user_sessions(db, late) + assert after_users["early-user"] == pending_users["early-user"] + for field in ( + "saved_spend", "savings_estimated_turns", "savings_estimated_actual_spend", + "savings_estimated_saved_spend", "savings_estimated_baseline_models", + ): + assert after_users["late-user"][field] == after[field] + assert after_users["late-user"]["turns"] == 1 and after_users["late-user"]["spend"] == 0.17 for table in ("DailyUserSpend", "DailyTeamSpend", "DailyOrganizationSpend", "DailyEndUserSpend", "DailyAgentSpend", "DailyTagSpend"): rows: Final = await db.query_raw(f'SELECT spend,api_requests,autorouter_savings_spend FROM "LiteLLM_{table}" WHERE api_key=$1', late.api_key) assert rows[0]["spend"] == rows[0]["api_requests"] == 0 assert rows[0]["autorouter_savings_spend"] == pytest.approx(after["saved_spend"]) -async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: - event: Final = record() +@pytest.mark.parametrize("attributed", [True, False]) +async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent( + db: Prisma, record: Callable[..., BaselineAccountingRecord], attributed: bool +) -> None: + event: Final = record(user_id="first-user" if attributed else "") + other: Final = record("other", 10001.0, user_id="second-user" if attributed else "") await _log(db, event) assert await _store(db, after_commit=True).append(event) == "unavailable" store: Final = _store(db) assert set(await asyncio.gather(*(store.append(event) for _ in range(4)))) == {"recorded"} + await _log(db, other) + assert await store.append(other) == "recorded" + if not attributed: + await db.execute_raw( + 'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET data=(data::jsonb #- \'{turn,user_id}\')::text WHERE scope=$1', + event.scope, + ) assert await store.project(event.scope) == "published" assert await store.project(event.scope) == "unchanged" session: Final = await _session(db, event) - assert session["turns"] == session["savings_estimated_turns"] == 1 - assert session["spend"] == session["savings_estimated_actual_spend"] == 0.17 + assert session["turns"] == session["savings_estimated_turns"] == 2 + assert session["spend"] == session["savings_estimated_actual_spend"] == 0.34 + users: Final = await _user_sessions(db, event) + assert set(users) == ({"first-user", "second-user"} if attributed else set()) + for user in users.values(): + assert user["turns"] == user["savings_estimated_turns"] == 1 + assert user["spend"] == user["savings_estimated_actual_spend"] == 0.17 + assert user["savings_estimated_baseline_models"] == {event.baseline_model: 1} async def test_publication_rollback_keeps_dirty_revision_for_retry(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: - event: Final = record() + event: Final = record(user_id="rollback-user") await _log(db, event) store: Final = _store(db) assert await store.append(event) == "recorded" assert await _store(db, before_commit=True).project(event.scope) == "unavailable" session: Final = await _session(db, event) assert session["spend"] == 0.17 and session["savings_estimated_turns"] == 0 + before_users: Final = await _user_sessions(db, event) + assert before_users["rollback-user"]["spend"] == 0.17 + assert before_users["rollback-user"]["savings_estimated_turns"] == 0 + assert before_users["rollback-user"]["savings_estimated_baseline_models"] == {} revisions: Final = await db.query_raw('SELECT revision,published_revision FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope=$1', event.scope) assert revisions[0]["revision"] > revisions[0]["published_revision"] assert await store.project(event.scope) == "published" assert (await _session(db, event))["savings_estimated_turns"] == 1 + after_users: Final = await _user_sessions(db, event) + assert after_users["rollback-user"]["turns"] == after_users["rollback-user"]["savings_estimated_turns"] == 1 + assert after_users["rollback-user"]["spend"] == after_users["rollback-user"]["savings_estimated_actual_spend"] == 0.17 async def test_conflicting_duplicate_cannot_restore_an_observed_estimate(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: @@ -196,6 +246,7 @@ async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_a db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch, ) -> None: import os + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index acd3dc18b54..c61a489f894 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -56,6 +56,31 @@ def _build(payload: dict | None = None, metadata: dict | None = None): class TestBuildTransaction: + @pytest.mark.parametrize( + "api_key, user_id, included", + [ + ("hashed-key", "canonical-user", True), + ("hashed-key", None, True), + ("hashed-key", "", True), + ("", "canonical-user", True), + ("", None, False), + ("", "", False), + ], + ) + def test_attribution_uses_the_canonical_user_even_without_a_key( + self, api_key: str, user_id: str | None, included: bool + ) -> None: + transaction: Final = _build( + payload=_payload(api_key=api_key, user=user_id), + metadata=_metadata(user="client-user", user_api_key_user_id="metadata-user"), + ) + if not included: + assert transaction is None + return + assert transaction is not None + assert transaction.api_key == api_key + assert transaction.user_id == (user_id or "") + def test_successful_auto_routed_turn_builds_every_field(self): transaction = _build( metadata=_metadata( @@ -205,23 +230,43 @@ class TestBuildTransaction: class _FakeDB: - def __init__(self, failures: "list[Exception] | None" = None, poison_session: str | None = None): + def __init__( + self, + failures: "list[Exception] | None" = None, + poison_session: str | None = None, + poison_user: str | None = None, + commit_then_error_users: frozenset[str] = frozenset(), + ): self.calls: list[tuple] = [] + self.attempts: list[tuple[str, tuple[object, ...]]] = [] self._failures = list(failures or []) self._poison_session = poison_session + self._poison_user = poison_user + self._commit_then_error_users = commit_then_error_users async def execute_raw(self, sql: str, *params: object) -> int: + self.attempts.append((sql, params)) if self._poison_session is not None and params[1] == self._poison_session: raise RuntimeError("index row size exceeds btree maximum") + if self._poison_user is not None and params[19] == self._poison_user: + raise RuntimeError("index row size exceeds btree maximum") if self._failures: raise self._failures.pop(0) self.calls.append((sql, params)) + if params[19] in self._commit_then_error_users: + raise RuntimeError("commit succeeded but acknowledgement was lost") return 1 class _FakeClient: - def __init__(self, failures: "list[Exception] | None" = None, poison_session: str | None = None): - self.db = _FakeDB(failures, poison_session) + def __init__( + self, + failures: "list[Exception] | None" = None, + poison_session: str | None = None, + poison_user: str | None = None, + commit_then_error_users: frozenset[str] = frozenset(), + ): + self.db = _FakeDB(failures, poison_session, poison_user, commit_then_error_users) def _transaction( @@ -229,9 +274,11 @@ def _transaction( at: datetime = datetime(2026, 8, 1, 12, 0, 0), tier: str | None = "medium", baseline_model: str | None = "anthropic/claude-opus-5", + api_key: str = "k1", + user_id: str = "", ) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( - api_key="k1", + api_key=api_key, session_id=session_id, router_name="live-auto", router_type="complexity", @@ -247,6 +294,7 @@ def _transaction( cache_touched=False, tier=tier, baseline_model=baseline_model, + user_id=user_id, ) @@ -261,7 +309,7 @@ class TestFlush: def test_params_marshal_in_statement_order(self): client = _FakeClient() - asyncio.run(flush_autorouter_turn_transactions(client, [_transaction()])) + asyncio.run(flush_autorouter_turn_transactions(client, [_transaction(user_id="canonical-user")])) sql, params = client.db.calls[0] assert sql == UPSERT_AUTOROUTER_SESSION_SQL assert params == ( @@ -284,8 +332,65 @@ class TestFlush: 0, 0.0, 0.0, + "canonical-user", ) + def test_a_keys_turns_stay_chronological_when_its_canonical_user_changes(self) -> None: + client: Final = _FakeClient() + earlier: Final = _transaction(user_id="z-user", at=datetime(2026, 8, 1, 12, 0, 0)) + later: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 10)) + asyncio.run(flush_autorouter_turn_transactions(client, [later, earlier])) + assert [(params[5], params[19]) for _, params in client.db.calls] == [ + ("2026-08-01T12:00:00", "z-user"), + ("2026-08-01T12:00:10", "a-user"), + ] + + def test_one_keyless_users_failed_session_does_not_drop_another_users_turn(self) -> None: + client: Final = _FakeClient(poison_user="a-user") + failed: Final = _transaction(api_key="", user_id="a-user") + other: Final = _transaction(api_key="", user_id="b-user", at=datetime(2026, 8, 1, 12, 0, 10)) + asyncio.run(flush_autorouter_turn_transactions(client, [other, failed])) + assert [(params[0], params[1], params[19]) for _, params in client.db.calls] == [("", "s1", "b-user")] + + def test_uncertain_commits_quarantine_only_the_key_and_each_failed_user(self) -> None: + client: Final = _FakeClient(commit_then_error_users=frozenset({"a-failed", "c-failed"})) + turns: Final = tuple( + _transaction(user_id=user, at=datetime(2026, 8, 1, 12, 0, second), api_key=key) + for user, second, key in ( + ("b-healthy", 0, "k1"), + ("a-failed", 1, "k1"), + ("b-healthy", 2, "k1"), + ("c-failed", 3, "k1"), + ("b-healthy", 4, "k1"), + ("d-healthy", 5, "k1"), + ("c-failed", 6, "k1"), + ("d-healthy", 7, "k1"), + ("a-failed", 8, "k1"), + ("", 9, "k1"), + ("z-other", 10, "k2"), + ) + ) + asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns)))) + + assert client.db.attempts == client.db.calls + assert [ + (params[0], params[19], params[5]) + for sql, params in client.db.calls + if sql == UPSERT_AUTOROUTER_SESSION_SQL + ] == [ + ("k1", "b-healthy", "2026-08-01T12:00:00"), + ("k1", "a-failed", "2026-08-01T12:00:01"), + ("k2", "z-other", "2026-08-01T12:00:10"), + ] + assert [params[19] for _, params in client.db.attempts].count("a-failed") == 1 + assert [params[19] for _, params in client.db.attempts].count("c-failed") == 1 + for user, seconds in (("b-healthy", (2, 4)), ("c-failed", (3,)), ("d-healthy", (5, 7))): + assert [ + (params[0], params[5]) + for sql, params in client.db.calls + if sql != UPSERT_AUTOROUTER_SESSION_SQL and params[19] == user + ] == [("k1", f"2026-08-01T12:00:{second:02d}") for second in seconds] + def test_a_connect_error_retries_the_same_statement(self): client = _FakeClient(failures=[httpx.ConnectError("boom")]) asyncio.run(flush_autorouter_turn_transactions(client, [_transaction()])) 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 6ac053f4e15..d5ddd5a78c3 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 @@ -4,6 +4,7 @@ Unit tests for auto router management endpoints from collections.abc import Mapping, Sequence from pathlib import Path +from types import SimpleNamespace from typing import Final import pytest @@ -654,17 +655,43 @@ class TestAutoRouterBenchmarks: assert _summed_agg_row([complexity, quality]).tier_turns == {} @pytest.mark.asyncio - async def test_non_admin_roles_cannot_read_benchmarks(self): + @pytest.mark.parametrize("user_id", [None, "own-user", "other-user"]) + async def test_non_admin_roles_cannot_read_benchmarks(self, user_id: str | None): from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks with pytest.raises(HTTPException) as err: await get_auto_router_benchmarks( - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-x"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-x", user_id="own-user" + ), start_date="2026-08-01", end_date="2026-08-02", + user_id=user_id, ) assert err.value.status_code == 403 + @pytest.mark.asyncio + async def test_an_empty_user_filter_is_rejected_before_querying_deployment_data( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import httpx + from fastapi import FastAPI + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks + + query: Final = AsyncMock(return_value=[]) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=SimpleNamespace(query_raw=query))) + app: Final = FastAPI() + app.get("/auto_router/benchmarks")(get_auto_router_benchmarks) + app.dependency_overrides[user_api_key_auth] = lambda: ADMIN + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response: Final = await client.get("/auto_router/benchmarks", params={"user_id": ""}) + + assert response.status_code == 422 + query.assert_not_awaited() + @pytest.mark.asyncio async def test_a_reversed_window_is_rejected(self, monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server @@ -680,7 +707,11 @@ class TestAutoRouterBenchmarks: assert err.value.status_code == 400 @pytest.mark.asyncio - async def test_endpoint_returns_groups_and_totals_from_the_rollup(self, monkeypatch: pytest.MonkeyPatch): + @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) + @pytest.mark.parametrize("user_id", [None, "selected-user"]) + async def test_endpoint_returns_groups_and_totals_from_the_rollup( + self, monkeypatch: pytest.MonkeyPatch, role: LitellmUserRoles, user_id: str | None + ): from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -695,12 +726,13 @@ class TestAutoRouterBenchmarks: monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})()) response = await get_auto_router_benchmarks( - user_api_key_dict=ADMIN, + user_api_key_dict=UserAPIKeyAuth(user_role=role, api_key="sk-admin", user_id="viewer"), start_date="2026-07-01", end_date="2026-08-01", api_key="key-hash", + user_id=user_id, ) - assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash") + assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash", user_id) assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index bf1538183ab..f395c146cf0 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -793,18 +793,20 @@ async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup(): tables = [call[0][0] for call in client.db.execute_raw.call_args_list] assert any('"LiteLLM_SpendLogs"' in sql for sql in tables) assert not any('"LiteLLM_AutoRouterSession"' in sql for sql in tables) + assert not any('"LiteLLM_AutoRouterUserSession"' in sql for sql in tables) assert not any('"LiteLLM_HealthCheckTable"' in sql for sql in tables) @pytest.mark.asyncio -async def test_session_retention_alone_cleans_only_the_session_rollup(): - client = _mock_prisma_for_retention([0]) +async def test_session_retention_alone_cleans_both_session_rollups(): + client = _mock_prisma_for_retention([0, 0]) cleaner = SpendLogCleanup(general_settings={"maximum_autorouter_session_retention_period": "365d"}) cleaner.pod_lock_manager = None await cleaner.cleanup_old_spend_logs(client) tables = [call[0][0] for call in client.db.execute_raw.call_args_list] - assert len(tables) == 1 + assert len(tables) == 2 assert '"LiteLLM_AutoRouterSession"' in tables[0] + assert '"LiteLLM_AutoRouterUserSession"' in tables[1] @pytest.mark.asyncio @@ -825,7 +827,7 @@ async def test_health_check_retention_alone_cleans_only_the_health_check_table() @pytest.mark.asyncio async def test_each_retention_key_cuts_off_at_its_own_horizon(): - client = _mock_prisma_for_retention([0, 0, 0, 0]) + client = _mock_prisma_for_retention([0, 0, 0, 0, 0]) cleaner = SpendLogCleanup( general_settings={ "maximum_spend_logs_retention_period": "7d", @@ -839,6 +841,8 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon(): ( "LiteLLM_AutoRouterSession" if '"LiteLLM_AutoRouterSession"' in call[0][0] + else "LiteLLM_AutoRouterUserSession" + if '"LiteLLM_AutoRouterUserSession"' in call[0][0] else "LiteLLM_HealthCheckTable" if '"LiteLLM_HealthCheckTable"' in call[0][0] else "logs" @@ -848,6 +852,7 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon(): now = datetime.now(timezone.utc) assert (now - cutoffs["logs"]).days == 7 assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365 + assert cutoffs["LiteLLM_AutoRouterUserSession"] == cutoffs["LiteLLM_AutoRouterSession"] assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index 5c7453c1394..a144630cdd0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -434,7 +434,7 @@ describe("AutoRouterBenchmarksTab", () => { mockHook({ data: response([group()]) }); const { dateValue, onDateChange } = renderTab(); - expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, undefined); + expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, undefined, undefined); expect(screen.getByText("Jul 6 – Aug 5 (UTC)")).toBeInTheDocument(); fireEvent.click(screen.getByTestId("date-picker")); @@ -460,7 +460,7 @@ describe("AutoRouterBenchmarksTab", () => { , ); - expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, "key-hash-1"); + expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, "key-hash-1", undefined); expect(screen.getByText("Total estimated savings")).toBeInTheDocument(); expect(screen.queryByRole("tab", { name: "Shadow Evals" })).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index ce55b633b60..063598bd46e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -312,8 +312,7 @@ const BenchmarksBody: React.FC = ({ isPending, error, data, length. Total actual spend includes every turn; savings and baseline spend include only turns with a current estimate, including turns with zero savings. Savings are net of recorded LLM classification cost. Classification cost per 1K turns is averaged over all auto-router turns, including those that skip classification. The range - counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets savings - by UTC day. + counts whole sessions that overlap it, so totals can differ from savings views that group usage by UTC day.

@@ -333,11 +332,17 @@ interface AutoRouterBenchmarksTabProps { accessToken: string | null; activity: Pick; apiKey?: string; + userId?: string; } -export const AutoRouterUsageView: React.FC = ({ accessToken, activity, apiKey }) => { +export const AutoRouterUsageView: React.FC = ({ + accessToken, + activity, + apiKey, + userId, +}) => { const { dateValue, onDateChange } = activity; - const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey); + const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey, userId); const [selectedKey, setSelectedKey] = useState(ALL_ROUTERS); const { data: autoRouters } = useAutoRouters(); @@ -372,6 +377,12 @@ export const AutoRouterUsageView: React.FC = ({ ac
+ {userId && ( +

+ Usage for this user across API keys and JWT-authenticated requests. Older sessions recorded without a user ID + are not included. +

+ )} +export const useAutoRouterBenchmarks = ( + accessToken: string | null, + range: DateRange, + apiKey?: string, + userId?: string, +) => $api.useQuery( "get", "/auto_router/benchmarks", - { params: { query: { ...benchmarksWindow(range, new Date()), api_key: apiKey } } }, + { params: { query: { ...benchmarksWindow(range, new Date()), api_key: apiKey, user_id: userId } } }, { enabled: Boolean(accessToken && range.from && range.to), retry: false }, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx index 4059303d5a5..e501cf00b90 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx @@ -15,6 +15,8 @@ vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", ( isFetchingMore: false, progress: { currentPage: 4, totalPages: 9 }, cancelled: false, + failed: false, + coversRange: true, cancel: mockCancel, }; }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts index 92dd24b8d6d..4eb9f257d30 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts @@ -67,14 +67,16 @@ export const useScopedDailyActivityRange = ( args: [accessToken, startTime, endTime, userId, true, apiKey], enabled: !!accessToken && !!startTime && !!endTime, }; - const { data, loading, isFetchingMore, progress, cancelled, failed, cancel } = + const { data, loading, isFetchingMore, progress, cancelled, failed, coversRange, cancel } = usePaginatedDailyActivity(activityQueryOptions); + const readUnavailable = failed || cancelled; + const waitingForRange = activityQueryOptions.enabled && !coversRange && !readUnavailable; return { dateValue, onDateChange, results: data.results as DailyData[], - loading, + loading: loading || waitingForRange, isFetchingMore, progress, cancelled, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx index 0f1a44851c7..6a0e55a6dda 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx @@ -1,7 +1,17 @@ -import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../../../../tests/test-utils"; +import { + act, + fireEvent, + renderWithProviders as render, + screen, + testQueryClient, + waitFor, +} from "../../../../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { describe, expect, it, vi, beforeEach } from "vitest"; +import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; +import { Profiler } from "react"; import UserInfoView from "./user_info_view"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import type { AutoRouterBenchmarksResponse } from "@/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks"; const mockTeamMemberAddCall = vi.fn(); const mockTeamMemberDeleteCall = vi.fn(); @@ -11,6 +21,8 @@ const mockTeamInfoCall = vi.fn(); const mockUserUpdateUserCall = vi.fn(); const mockFetchMCPServers = vi.fn(); const mockListMCPTools = vi.fn(); +const mockUserDailyActivityCall = vi.fn(); +const mockUserDailyActivityAggregatedCall = vi.fn(); const MCP_SERVER = { server_id: "srv-1", server_name: "GitHub MCP", alias: "GitHub MCP" }; @@ -47,13 +59,18 @@ vi.mock("next/navigation", () => ({ useSearchParams: () => new URLSearchParams(window.location.search), })); -vi.mock("@/components/networking", () => { +vi.mock("@/components/networking", async (importOriginal) => { + const original = await importOriginal(); return { + formatDate: original.formatDate, serverRootPath: "/", userGetInfoV2: (...args: unknown[]) => mockUserGetInfoV2(...args), + userDailyActivityCall: (...args: unknown[]) => mockUserDailyActivityCall(...args), + userDailyActivityAggregatedCall: (...args: unknown[]) => mockUserDailyActivityAggregatedCall(...args), userDeleteCall: vi.fn(), userUpdateUserCall: (...args: unknown[]) => mockUserUpdateUserCall(...args), modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), + modelInfoCall: vi.fn().mockResolvedValue({ data: [], total_pages: 1 }), invitationCreateCall: vi.fn(), teamInfoCall: (...args: unknown[]) => mockTeamInfoCall(...args), teamListCall: (...args: unknown[]) => mockTeamListCall(...args), @@ -291,3 +308,337 @@ describe("UserInfoView add-to-team form", () => { expect(screen.getByText("Add User to Team")).toBeInTheDocument(); }); }); + +const savingsDay = (date: string, metrics: Partial): DailyData => ({ + date, + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 1, + successful_requests: 1, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + ...metrics, + }, + breakdown: { models: {}, model_groups: {}, mcp_servers: {}, providers: {}, api_keys: {}, entities: {} }, +}); + +const savingsResponse = (results: DailyData[]) => ({ + results, + metadata: { total_pages: 1, has_more: false, page: 1 }, +}); + +const routerUsageResponse = (saved: number): AutoRouterBenchmarksResponse => ({ + start_date: "2026-09-01", + end_date: "2026-09-19", + routers_in_scope: 0, + groups: [], + totals: { + sessions: 2, + turns: 2, + avg_turns_per_session: 1, + avg_session_seconds: 0, + avg_tokens_per_session: 100, + spend: 10, + savings_estimated_turns: 2, + savings_estimated_actual_spend: 10, + classifier_cost: 0, + saved_spend: saved, + baseline_spend: 10 + saved, + saved_pct: (100 * saved) / (10 + saved), + saved_per_session: saved / 2, + cache: { + coverage_pct: 100, + hit_rate_pct: 0, + same_model: { turns: 0, hits: 0, hit_rate_pct: 0 }, + first_visit: { turns: 2, hits: 0, hit_rate_pct: 0 }, + return_to_tier: { turns: 0, hits: 0, hit_rate_pct: 0 }, + unordered_turns: 0, + return_misses_expired: 0, + return_misses_within_ttl: 0, + return_misses_unknown: 0, + ttl_5m_turns: 0, + ttl_1h_turns: 0, + }, + }, +}); + +describe("UserInfoView auto-router usage", () => { + const props = { + userId: "user-123", + onClose: vi.fn(), + accessToken: "admin-token", + userRole: "proxy_admin", + possibleUIRoles: null, + }; + const mockFetch = vi.fn(); + + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + mockUserGetInfoV2.mockImplementation((_token: string, userId: string) => + Promise.resolve({ ...MOCK_USER_DATA_NO_TEAMS, user_id: userId }), + ); + mockFetch.mockReset().mockResolvedValue(Response.json(routerUsageResponse(42))); + vi.stubGlobal("fetch", mockFetch); + }); + + afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); + }); + + it.each(["proxy_admin", "proxy_admin_viewer"])( + "loads selected-user usage lazily for %s without a key filter", + async (userRole) => { + const user = userEvent.setup(); + render(); + const tab = await screen.findByRole("tab", { name: "Auto-router usage" }); + expect(mockFetch).not.toHaveBeenCalled(); + await user.click(tab); + + expect(await screen.findByText("$42.00")).toBeInTheDocument(); + const request = mockFetch.mock.calls[0][0] as Request; + const params = new URL(request.url).searchParams; + expect(params.get("user_id")).toBe("user-123"); + expect(params.has("api_key")).toBe(false); + expect(screen.getByText(/Older sessions recorded without a user ID are not included/)).toBeInTheDocument(); + }, + ); + + it("switches query scope without displaying the previous user's usage", async () => { + const nextUser = Promise.withResolvers(); + mockFetch.mockResolvedValueOnce(Response.json(routerUsageResponse(42))).mockReturnValue(nextUser.promise); + const user = userEvent.setup(); + const { rerender } = render(); + await user.click(await screen.findByRole("tab", { name: "Auto-router usage" })); + expect(await screen.findByText("$42.00")).toBeInTheDocument(); + + rerender(); + expect(screen.getByText("Loading auto-router usage...")).toBeInTheDocument(); + expect(screen.queryByText("$42.00")).not.toBeInTheDocument(); + await act(async () => nextUser.resolve(Response.json(routerUsageResponse(-7)))); + expect(await screen.findByText("-$7.00")).toBeInTheDocument(); + expect( + mockFetch.mock.calls.map(([request]) => new URL((request as Request).url).searchParams.get("user_id")), + ).toEqual(["user-123", "user-456"]); + }); + + it.each(["internal_user", "org_admin", null])("keeps the admin-only tab unavailable to %s", async (userRole) => { + render(); + await screen.findByRole("tab", { name: "Overview" }); + expect(screen.queryByRole("tab", { name: "Auto-router usage" })).not.toBeInTheDocument(); + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it("never turns an absent user ID into a deployment-wide request", async () => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Auto-router usage" })); + expect(screen.getByRole("alert")).toHaveTextContent("this user has no ID"); + expect(mockFetch).not.toHaveBeenCalled(); + }); +}); + +describe("UserInfoView savings", () => { + const props = { + userId: "user-123", + onClose: vi.fn(), + accessToken: "admin-token", + userRole: "proxy_admin", + possibleUIRoles: null, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUserGetInfoV2.mockImplementation((_token: string, userId: string) => + Promise.resolve({ ...MOCK_USER_DATA_NO_TEAMS, user_id: userId }), + ); + mockUserDailyActivityAggregatedCall.mockReset().mockResolvedValue(savingsResponse([])); + mockUserDailyActivityCall.mockReset().mockResolvedValue(savingsResponse([])); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it.each(["internal_user", "org_admin", "team_admin"])( + "only offers self savings to %s and stops querying after switching to another user", + async (userRole) => { + const user = userEvent.setup(); + const { rerender } = render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + expect(await screen.findByText("No usage recorded for this user in this range.")).toBeInTheDocument(); + expect(mockUserDailyActivityAggregatedCall.mock.calls[0][3]).toBe("user-1"); + + mockUserDailyActivityAggregatedCall.mockClear(); + mockUserDailyActivityCall.mockClear(); + rerender(); + await screen.findAllByText("another-user"); + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + expect(screen.queryByRole("tab", { name: "Savings" })).not.toBeInTheDocument(); + expect(screen.queryByText("No usage recorded for this user in this range.")).not.toBeInTheDocument(); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalled(); + expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); + }, + ); + + it("loads selected user savings without a key filter, including losses", async () => { + const firstDay: Partial = { + compression_savings_spend: 1.5, + gateway_injected_caching_savings_spend: 0.1, + prompt_caching_savings_spend: 0.25, + autorouter_savings_spend: -1, + }; + const secondDay: Partial = { + compression_savings_spend: 0.5, + gateway_injected_caching_savings_spend: 0.3, + prompt_caching_savings_spend: 0.75, + autorouter_savings_spend: -2, + }; + mockUserDailyActivityAggregatedCall.mockResolvedValue( + savingsResponse([savingsDay("2026-09-18", firstDay), savingsDay("2026-09-19", secondDay)]), + ); + const user = userEvent.setup(); + render(); + const savingsTab = await screen.findByRole("tab", { name: "Savings" }); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalled(); + expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); + + await user.click(savingsTab); + + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$0.6000"); + expect(screen.getByTestId("summary-card-compression-savings")).toHaveTextContent("$2.00"); + expect(screen.getByTestId("summary-card-prompt-caching-savings")).toHaveTextContent("$0.4000"); + expect(screen.getByTestId("summary-card-prompt-caching-savings")).toHaveTextContent("$1.00Total"); + expect(screen.getByTestId("summary-card-auto-router-savings")).toHaveTextContent("-$3.00"); + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledExactlyOnceWith( + "admin-token", + expect.any(Date), + expect.any(Date), + "user-123", + true, + null, + ); + expect(screen.getByTestId("user-savings-scope-note")).toHaveTextContent("JWT-authenticated requests"); + await user.click(screen.getByRole("tab", { name: "Per day" })); + expect(screen.getByRole("tab", { name: "Per day" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$0.6000"); + }); + + it("removes the prior user's savings while the newly selected user's results are loading", async () => { + const nextUser = Promise.withResolvers>(); + mockUserDailyActivityAggregatedCall + .mockResolvedValueOnce(savingsResponse([savingsDay("2026-09-19", { compression_savings_spend: 42 })])) + .mockReturnValueOnce(nextUser.promise); + const user = userEvent.setup(); + const { rerender } = render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("$42.00"); + + rerender(); + + expect(await screen.findByTestId("user-savings-empty")).toHaveTextContent("Loading savings"); + expect(screen.queryByTestId("summary-card-total-recorded-savings")).not.toBeInTheDocument(); + expect(mockUserDailyActivityAggregatedCall).toHaveBeenLastCalledWith( + "admin-token", + expect.any(Date), + expect.any(Date), + "user-456", + true, + null, + ); + await act(async () => { + nextUser.resolve(savingsResponse([savingsDay("2026-09-19", { autorouter_savings_spend: -7 })])); + }); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$7.00"); + expect(screen.queryByText("$42.00")).not.toBeInTheDocument(); + }); + + it("never commits the previous range's savings under the newly selected dates", async () => { + vi.stubGlobal("requestIdleCallback", (callback: IdleRequestCallback) => + window.setTimeout(() => callback({ didTimeout: false, timeRemaining: () => 0 }), 0), + ); + const nextRange = Promise.withResolvers>(); + mockUserDailyActivityAggregatedCall + .mockResolvedValueOnce(savingsResponse([savingsDay("2026-09-19", { compression_savings_spend: 42 })])) + .mockReturnValue(nextRange.promise); + const committedTotals: Array = []; + const captureNewRange = () => { + if (screen.queryByText("Running total saved · Sep 1 – Sep 2 (UTC)")) { + committedTotals.push(screen.queryByTestId("summary-card-total-recorded-savings")?.textContent ?? null); + } + }; + const user = userEvent.setup(); + render( + + + , + ); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("$42.00"); + + await user.click(screen.getByRole("button", { name: / - / })); + const [startDateInput, endDateInput] = screen.getAllByDisplayValue(/^\d{4}-\d{2}-\d{2}$/); + fireEvent.change(startDateInput, { target: { value: "2026-09-01" } }); + fireEvent.change(endDateInput, { target: { value: "2026-09-02" } }); + await user.click(screen.getByRole("button", { name: "Apply" })); + + expect(committedTotals.length).toBeGreaterThan(0); + expect(committedTotals.every((total) => total === null)).toBe(true); + expect(screen.getByTestId("user-savings-empty")).toHaveTextContent("Loading savings"); + await act(async () => { + nextRange.resolve(savingsResponse([savingsDay("2026-09-02", { autorouter_savings_spend: -7 })])); + }); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$7.00"); + }); + + it("reports an incomplete paginated read as unavailable instead of displaying a partial savings total", async () => { + mockUserDailyActivityAggregatedCall.mockRejectedValue(new Error("aggregated unavailable")); + mockUserDailyActivityCall + .mockResolvedValueOnce({ + results: [savingsDay("2026-09-19", { compression_savings_spend: 42 })], + metadata: { total_pages: 2, has_more: true, page: 1 }, + }) + .mockRejectedValueOnce(new Error("next page unavailable")); + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + + expect(await screen.findByRole("alert")).toHaveTextContent("Savings are unavailable for this range"); + expect(mockUserDailyActivityCall).toHaveBeenLastCalledWith( + "admin-token", + expect.any(Date), + expect.any(Date), + 2, + "user-123", + true, + null, + ); + expect(screen.queryByTestId("summary-card-total-recorded-savings")).not.toBeInTheDocument(); + expect(screen.queryByText(/No usage recorded/)).not.toBeInTheDocument(); + }); + + it("distinguishes a user with no usage from an unavailable read", async () => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + + expect(await screen.findByTestId("user-savings-empty")).toHaveTextContent("No usage recorded for this user"); + expect(screen.getByTestId("summary-card-total-recorded-savings")).toHaveTextContent("$0.00"); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it.each(["", " "])("never queries an absent selected user ID (%j)", async (userId) => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + + expect(screen.getByRole("alert")).toHaveTextContent("this user has no ID"); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalled(); + expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index e083e549552..c95badc587a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -28,7 +28,7 @@ import { ComboboxList, } from "@/components/ui/combobox"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { rolesWithWriteAccess } from "@/utils/roles"; +import { hasProxyWideSpendView, rolesWithWriteAccess } from "@/utils/roles"; import { teamDetailHref } from "@/utils/entityLinks"; import { BadgeLink } from "@/components/shared/BadgeLink"; import { UserEditView } from "../user_edit_view"; @@ -44,6 +44,9 @@ import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers" import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; +import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab"; +import { useActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; interface UserInfoViewProps { userId: string; @@ -85,7 +88,10 @@ export default function UserInfoView({ initialTab = 0, startInEditMode = false, }: UserInfoViewProps) { - const { premiumUser } = useAuthorized(); + const { premiumUser, userId: signedInUserId } = useAuthorized(); + const canViewAutoRouterUsage = hasProxyWideSpendView(userRole); + const canViewSavings = canViewAutoRouterUsage || (Boolean(userId.trim()) && userId === signedInUserId); + const activityDateRange = useActivityDateRange(); const [userData, setUserData] = useState(null); const [teamDetails, setTeamDetails] = useState([]); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -97,6 +103,8 @@ export default function UserInfoView({ const [invitationLinkData, setInvitationLinkData] = useState(null); const [baseUrl, setBaseUrl] = useState(null); const [activeTab, setActiveTab] = useState(initialTab === 1 ? "details" : "overview"); + const hiddenSavingsTab = activeTab === "savings" && !canViewSavings; + const hiddenRouterTab = activeTab === "auto-router-usage" && !canViewAutoRouterUsage; const [copiedStates, setCopiedStates] = useState>({}); const [isTeamsExpanded, setIsTeamsExpanded] = useState(false); const [isAddTeamModalOpen, setIsAddTeamModalOpen] = useState(false); @@ -467,7 +475,11 @@ export default function UserInfoView({ confirmLoading={isDeletingUser} /> - setActiveTab(String(v))} className="gap-0"> + setActiveTab(String(v))} + className="gap-0" + > Overview @@ -475,6 +487,16 @@ export default function UserInfoView({ Details + {canViewSavings && ( + + Savings + + )} + {canViewAutoRouterUsage && ( + + Auto-router usage + + )} {/* Overview Panel */} @@ -685,6 +707,38 @@ export default function UserInfoView({ )} + {canViewSavings && ( + + {activeTab === "savings" && + (userId.trim() ? ( + + ) : ( +

Savings are unavailable because this user has no ID.

+ ))} + + )} + {canViewAutoRouterUsage && ( + + {activeTab === "auto-router-usage" && + (userId.trim() ? ( + + ) : ( +

Auto-router usage is unavailable because this user has no ID.

+ ))} + + )} { + const { dateValue, onDateChange, results, loading, isFetchingMore, failed, cancelled } = useScopedDailyActivityRange( + accessToken, + scope, + activity, + ); + const startTime = dateValue.from; + const endTime = dateValue.to; + + const [accumulation, setAccumulation] = useState("cumulative"); + + const perInterval = useMemo(() => savingsSeriesOf(results), [results]); + + const overTime = useMemo(() => { + if (accumulation !== "cumulative") return perInterval; + const startLabel = startTime ? shortDate(localIsoDay(startTime)) : ""; + return withStartAnchor(toCumulative(perInterval), startLabel); + }, [accumulation, perInterval, startTime]); + + const intervalLabel = "Per day"; + const rangeLabel = formatRangeLabel(startTime, endTime); + const savingsSubtitle = [ + accumulation === "cumulative" ? "Running total saved" : `Saved ${intervalLabel.toLowerCase()}`, + rangeLabel && `${rangeLabel} (UTC)`, + ] + .filter(Boolean) + .join(" · "); + + const isLoading = loading || isFetchingMore; + const unavailable = failed || cancelled; + const showResults = !isLoading && !unavailable; + const hasRows = results.length > 0; + const showEmpty = !unavailable && (isLoading || !hasRows); + const showChart = showResults && hasRows; + const chartProps = { + data: overTime, + index: "date", + categories: SAVINGS_SERIES, + colors: SAVINGS_COLORS, + valueFormatter: usd, + showLegend: false, + }; + + return ( +
+
+ Spend is bucketed by UTC day + +
+ + {scopeNote && ( +

+ {scopeNote} +

+ )} + + {unavailable && ( +

+ Savings are unavailable for this range. Try another date range or reopen this tab. +

+ )} + {showResults && } + + + + Savings + {savingsSubtitle} + + + setAccumulation(value as SavingsAccumulation)}> + + Cumulative + {intervalLabel} + + + + + + {showEmpty && ( +

+ {isLoading ? "Loading savings..." : `No usage recorded for this ${entityType} in this range.`} +

+ )} + {showChart && + (accumulation === "cumulative" ? ( + + ) : ( + + ))} + + +
+ ); +}; + +export default ScopedSavingsTab; diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx index c33529eb042..d1395e0be51 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx @@ -1,132 +1,29 @@ "use client"; -import React, { useMemo, useState } from "react"; - -import { AreaChart, BarChart, CustomLegend } from "@/components/shared/charts"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; -import SavingsTiles from "@/components/shared/SavingsTiles"; -import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; -import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; import { hasProxyWideSpendView, spendScopeUserId } from "@/utils/roles"; -import { - formatRangeLabel, - localIsoDay, - MAX_POINTS_WITH_DOTS, - SAVINGS_COLORS, - SAVINGS_SERIES, - SavingsAccumulation, - SavingsPoint, - savingsSeriesOf, - shortDate, - toCumulative, - usd, - withStartAnchor, -} from "@/app/(dashboard)/cost-optimization/_components/costOptimizationUtils"; -import { - useScopedDailyActivityRange, - type ActivityDateRange, -} from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; +import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; interface KeySavingsTabProps { accessToken: string | null; - /** The key's token hash — what spend rows are keyed by, not the one-time plaintext secret. */ keyToken: string; userId: string | null; userRole: string; activity: ActivityDateRange; } -const KeySavingsTab: React.FC = ({ accessToken, keyToken, userId, userRole, activity }) => { - // Proxy admins read the whole key. For anyone else the endpoint applies the caller's own user_id - // alongside the key filter, so the figures cover only that viewer's requests on this key -- said - // plainly in the scope note below rather than left to be misread as the key's total. - const readsWholeKey = hasProxyWideSpendView(userRole); - const { dateValue, onDateChange, results, loading, isFetchingMore } = useScopedDailyActivityRange( - accessToken, - { userId: spendScopeUserId(userRole, userId), apiKey: keyToken }, - activity, - ); - const startTime = dateValue.from ?? null; - const endTime = dateValue.to ?? null; - - const [accumulation, setAccumulation] = useState("cumulative"); - - const perInterval = useMemo(() => savingsSeriesOf(results), [results]); - - const overTime = useMemo(() => { - if (accumulation !== "cumulative") return perInterval; - const startLabel = startTime ? shortDate(localIsoDay(startTime)) : ""; - return withStartAnchor(toCumulative(perInterval), startLabel); - }, [accumulation, perInterval, startTime]); - - const intervalLabel = "Per day"; - const rangeLabel = formatRangeLabel(startTime ?? undefined, endTime ?? undefined); - const savingsSubtitle = [ - accumulation === "cumulative" ? "Running total saved" : `Saved ${intervalLabel.toLowerCase()}`, - rangeLabel && `${rangeLabel} (UTC)`, - ] - .filter(Boolean) - .join(" · "); - - const isLoading = loading || isFetchingMore; - const hasRows = results.length > 0; - const chartProps = { - data: overTime, - index: "date", - categories: SAVINGS_SERIES, - colors: SAVINGS_COLORS, - valueFormatter: usd, - showLegend: false, - }; - - return ( -
-
- Spend is bucketed by UTC day - -
- - {!readsWholeKey && ( -

- Showing your own requests on this key. A key shared across a team will have spend from other members that is - not counted here. -

- )} - - - - - - Savings - {savingsSubtitle} - - - setAccumulation(value as SavingsAccumulation)}> - - Cumulative - {intervalLabel} - - - - - - {/* Distinguishes "still fetching" from "this key genuinely had no traffic": an empty - chart alone reads as a broken panel, and a $0.00 tile reads as a real zero. */} - {!hasRows && ( -

- {isLoading ? "Loading savings..." : "No usage recorded for this key in this range."} -

- )} - {hasRows && accumulation === "cumulative" && ( - - )} - {/* Not stacked: auto-router can go negative on a cold-cache write, and stacking would - draw that below the axis while the rest of the bar still read as the total. */} - {hasRows && accumulation !== "cumulative" && } - - -
- ); -}; +const KeySavingsTab = ({ accessToken, keyToken, userId, userRole, activity }: KeySavingsTabProps) => ( + +); export default KeySavingsTab; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 81580c8bfb1..0055ef4337c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1246,9 +1246,10 @@ export interface paths { * @description Benchmarks for the auto-router dashboard: session shape, savings against the configured * baseline, and prompt-caching behaviour bucketed by what the router did. * - * Reads the LiteLLM_AutoRouterSession rollup, folded once per request at spend-write time, - * so this endpoint never scans LiteLLM_SpendLogs. A session is in the window when it - * overlaps it: its last turn is on or after start_date and its first turn is on or before + * Reads session rollups folded once per request at spend-write time, so this endpoint + * never scans LiteLLM_SpendLogs. A user filter selects only turns attributed to that + * internal user when written; older key-only history remains outside user views. A session + * is in the window when it overlaps it: its last turn is on or after start_date and its first turn is on or before * end_date. Overall hit rate is over telemetry-bearing turns; each bucket's hit rate is * over that bucket's turns. * @@ -43601,6 +43602,8 @@ export interface operations { end_date?: string | null; /** @description Filter to one virtual key token hash */ api_key?: string | null; + /** @description Filter to one canonical internal user recorded on each turn */ + user_id?: string | null; }; header?: never; path?: never; From 39a14f39e5961605ba70deacc0475892f50d3cb7 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:20:47 +0000 Subject: [PATCH 023/109] style(proxy): format cleanup shutdown tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/shutdown/test_scheduled_jobs.py | 4 +- .../proxy/test_spend_log_cleanup.py | 127 +++++------------- 2 files changed, 37 insertions(+), 94 deletions(-) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 7defd6cef6c..87adc464608 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.proxy.shutdown import scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - stop_in_flight_scheduler_jobs, pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 1691b2d174a..b8f28d0c780 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -83,10 +83,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError, match='Wrong number of fields; got'): + with pytest.raises(ValueError, match="Wrong number of fields; got"): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError, match='is higher than the maximum value'): + with pytest.raises(ValueError, match="is higher than the maximum value"): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour @@ -99,6 +99,7 @@ def test_spend_log_cleanup_cron_scheduler_integration(): a real database connection. """ from unittest.mock import MagicMock + from apscheduler.triggers.cron import CronTrigger # Mock scheduler @@ -145,15 +146,11 @@ def test_spend_log_cleanup_cron_scheduler_integration(): # No cron, so it should fall back to interval } - cleanup_cron_fallback = general_settings_interval.get( - "maximum_spend_logs_cleanup_cron" - ) + cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron") assert cleanup_cron_fallback is None # No cron configured # Simulate interval-based scheduling fallback - retention_interval = general_settings_interval.get( - "maximum_spend_logs_retention_interval", "1d" - ) + retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d") from litellm.litellm_core_utils.duration_parser import duration_in_seconds interval_seconds = duration_in_seconds(retention_interval) @@ -181,27 +178,19 @@ async def test_should_delete_spend_logs(): assert cleaner._should_delete_spend_logs() is False # Test case 2: Valid seconds string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "3600s"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"}) assert cleaner._should_delete_spend_logs() is True # Test case 3: Valid days string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "30d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"}) assert cleaner._should_delete_spend_logs() is True # Test case 4: Valid hours string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "24h"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"}) assert cleaner._should_delete_spend_logs() is True # Test case 5: Invalid format - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "invalid"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"}) assert cleaner._should_delete_spend_logs() is False @@ -288,9 +277,7 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Verify the cutoff date is correct cutoff_date = mock_db.execute_raw.call_args[0][1] expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) - assert ( - abs((cutoff_date - expected_cutoff).total_seconds()) < 1 - ) # Allow 1 second difference for test execution time + assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time @pytest.mark.asyncio @@ -310,9 +297,7 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=True) partition_manager.ensure_partitions = AsyncMock(return_value=["p1"]) - partition_manager.drop_partitions_older_than = AsyncMock( - return_value=["LiteLLM_SpendLogs_p20260601"] - ) + partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"]) cleaner = SpendLogCleanup( general_settings={ @@ -450,9 +435,7 @@ async def test_integer_retention_treated_as_days(): An integer value for maximum_spend_logs_retention_period should be treated as days (e.g., 3 → '3d' → 259200 seconds). """ - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": 3} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3}) result = cleaner._should_delete_spend_logs() assert result is True assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds @@ -469,13 +452,11 @@ def test_string_retention_still_works(): ("2w", 2 * 604800), ] for setting, expected_seconds in cases: - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": setting} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting}) assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" - assert ( - cleaner.retention_seconds == expected_seconds - ), f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + assert cleaner.retention_seconds == expected_seconds, ( + f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + ) @pytest.mark.asyncio @@ -489,9 +470,7 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): mock_db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -510,9 +489,7 @@ async def test_delete_old_logs_continues_on_valid_int_return(): mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -559,9 +536,7 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): mock_db.execute_raw = AsyncMock(side_effect=[5, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline()) @@ -581,9 +556,7 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep. - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -591,14 +564,10 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) _wire_tx(mock_db) # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, # batch 5 returns 0 → loop exits naturally. - mock_db.execute_raw = AsyncMock( - side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0] - ) + mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0]) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -615,26 +584,18 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Lower the threshold so the test is fast and deterministic. - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) # Every batch raises — must abort after exactly 3 attempts, not loop forever. - mock_db.execute_raw = AsyncMock( - side_effect=ConnectionError("simulated persistent DB outage") - ) + mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage")) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -649,12 +610,8 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc intermittent timeouts don't trip the abort threshold.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -675,9 +632,7 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -698,9 +653,7 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) # Force the outer try/except to fire by making _should_delete_spend_logs raise. - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = None def boom(): @@ -725,12 +678,8 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch must still be released so the next scheduled run isn't permanently blocked.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -744,9 +693,7 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) mock_pod_lock_manager.release_lock = AsyncMock() - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = mock_pod_lock_manager await cleaner.cleanup_old_spend_logs(mock_prisma_client) @@ -996,9 +943,7 @@ async def test_each_batch_carries_a_statement_and_lock_timeout(): } ) - await cleaner._delete_old_logs( - mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() - ) + await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) assert "SET LOCAL statement_timeout = 12000" in recorded assert "SET LOCAL lock_timeout = 12000" in recorded @@ -1134,9 +1079,7 @@ async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table(): cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) - await cleaner._delete_old_logs( - mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() - ) + await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) count_sql = mock_db.query_raw.call_args[0][0] assert "count(*)" in count_sql From 0b2d52edc29a7098e95314fbf2c45e58b508efc5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:21:35 +0000 Subject: [PATCH 024/109] Revert "style(proxy): format cleanup shutdown tests" This reverts commit 39a14f39e5961605ba70deacc0475892f50d3cb7. --- .../proxy/shutdown/test_scheduled_jobs.py | 4 +- .../proxy/test_spend_log_cleanup.py | 127 +++++++++++++----- 2 files changed, 94 insertions(+), 37 deletions(-) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 87adc464608..7defd6cef6c 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -from litellm.proxy.shutdown import scheduled_jobs +import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - pause_scheduled_jobs, stop_in_flight_scheduler_jobs, + pause_scheduled_jobs, ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index b8f28d0c780..1691b2d174a 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -83,10 +83,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError, match="Wrong number of fields; got"): + with pytest.raises(ValueError, match='Wrong number of fields; got'): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError, match="is higher than the maximum value"): + with pytest.raises(ValueError, match='is higher than the maximum value'): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour @@ -99,7 +99,6 @@ def test_spend_log_cleanup_cron_scheduler_integration(): a real database connection. """ from unittest.mock import MagicMock - from apscheduler.triggers.cron import CronTrigger # Mock scheduler @@ -146,11 +145,15 @@ def test_spend_log_cleanup_cron_scheduler_integration(): # No cron, so it should fall back to interval } - cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron") + cleanup_cron_fallback = general_settings_interval.get( + "maximum_spend_logs_cleanup_cron" + ) assert cleanup_cron_fallback is None # No cron configured # Simulate interval-based scheduling fallback - retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d") + retention_interval = general_settings_interval.get( + "maximum_spend_logs_retention_interval", "1d" + ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds interval_seconds = duration_in_seconds(retention_interval) @@ -178,19 +181,27 @@ async def test_should_delete_spend_logs(): assert cleaner._should_delete_spend_logs() is False # Test case 2: Valid seconds string - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "3600s"} + ) assert cleaner._should_delete_spend_logs() is True # Test case 3: Valid days string - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "30d"} + ) assert cleaner._should_delete_spend_logs() is True # Test case 4: Valid hours string - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "24h"} + ) assert cleaner._should_delete_spend_logs() is True # Test case 5: Invalid format - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "invalid"} + ) assert cleaner._should_delete_spend_logs() is False @@ -277,7 +288,9 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Verify the cutoff date is correct cutoff_date = mock_db.execute_raw.call_args[0][1] expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) - assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time + assert ( + abs((cutoff_date - expected_cutoff).total_seconds()) < 1 + ) # Allow 1 second difference for test execution time @pytest.mark.asyncio @@ -297,7 +310,9 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=True) partition_manager.ensure_partitions = AsyncMock(return_value=["p1"]) - partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"]) + partition_manager.drop_partitions_older_than = AsyncMock( + return_value=["LiteLLM_SpendLogs_p20260601"] + ) cleaner = SpendLogCleanup( general_settings={ @@ -435,7 +450,9 @@ async def test_integer_retention_treated_as_days(): An integer value for maximum_spend_logs_retention_period should be treated as days (e.g., 3 → '3d' → 259200 seconds). """ - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": 3} + ) result = cleaner._should_delete_spend_logs() assert result is True assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds @@ -452,11 +469,13 @@ def test_string_retention_still_works(): ("2w", 2 * 604800), ] for setting, expected_seconds in cases: - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting}) - assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" - assert cleaner.retention_seconds == expected_seconds, ( - f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": setting} ) + assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" + assert ( + cleaner.retention_seconds == expected_seconds + ), f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" @pytest.mark.asyncio @@ -470,7 +489,9 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): mock_db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -489,7 +510,9 @@ async def test_delete_old_logs_continues_on_valid_int_return(): mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -536,7 +559,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): mock_db.execute_raw = AsyncMock(side_effect=[5, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline()) @@ -556,7 +581,9 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep. - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -564,10 +591,14 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) _wire_tx(mock_db) # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, # batch 5 returns 0 → loop exits naturally. - mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0]) + mock_db.execute_raw = AsyncMock( + side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0] + ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -584,18 +615,26 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Lower the threshold so the test is fast and deterministic. - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 + ) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) # Every batch raises — must abort after exactly 3 attempts, not loop forever. - mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage")) + mock_db.execute_raw = AsyncMock( + side_effect=ConnectionError("simulated persistent DB outage") + ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -610,8 +649,12 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc intermittent timeouts don't trip the abort threshold.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 + ) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -632,7 +675,9 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -653,7 +698,9 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) # Force the outer try/except to fire by making _should_delete_spend_logs raise. - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cleaner.pod_lock_manager = None def boom(): @@ -678,8 +725,12 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch must still be released so the next scheduled run isn't permanently blocked.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2 + ) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -693,7 +744,9 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) mock_pod_lock_manager.release_lock = AsyncMock() - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cleaner.pod_lock_manager = mock_pod_lock_manager await cleaner.cleanup_old_spend_logs(mock_prisma_client) @@ -943,7 +996,9 @@ async def test_each_batch_carries_a_statement_and_lock_timeout(): } ) - await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) + await cleaner._delete_old_logs( + mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() + ) assert "SET LOCAL statement_timeout = 12000" in recorded assert "SET LOCAL lock_timeout = 12000" in recorded @@ -1079,7 +1134,9 @@ async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table(): cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) - await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) + await cleaner._delete_old_logs( + mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() + ) count_sql = mock_db.query_raw.call_args[0][0] assert "count(*)" in count_sql From 29bd2ceb2b5c8f2a172d6b08e6c54cb20c41daaa Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:27:06 +0000 Subject: [PATCH 025/109] fix(proxy): avoid mutable shutdown wait set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/shutdown/scheduled_jobs.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index e7625a73b47..5345d380112 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -50,12 +50,13 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: if not scheduler.running: return in_flight: Final = executor.in_flight_jobs() - still_running: set[asyncio.Future[object]] = set() if in_flight: verbose_proxy_logger.info( "Waiting up to %ss for %d in-flight scheduled job(s) to finish", JOB_FINISH_TIMEOUT_SECONDS, len(in_flight) ) - _done, still_running = await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS) + still_running: Final = ( + (await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS))[1] if in_flight else frozenset() + ) scheduler.shutdown(wait=False) if not still_running: return From 6c13e0912b43e54ba912f8257ce436d28e6e8e01 Mon Sep 17 00:00:00 2001 From: ryan Date: Sun, 20 Sep 2026 08:59:30 +0000 Subject: [PATCH 026/109] fix(proxy): apply DB-stored callback redaction settings before logger init Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 6 ++++ tests/test_litellm/proxy/test_proxy_server.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index bbeb4846e27..4bc2e4e0d26 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1853,6 +1853,12 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "max_ui_session_budget", "budget_rollover", "mcp_tool_search", + "turn_off_message_logging", + "datadog_params", + "datadog_llm_observability_params", + "newrelic_params", + "pointfive_params", + "aws_sqs_callback_params", ] SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index f71f9c20f3b..75741b125cb 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11766,6 +11766,42 @@ def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_n assert getattr(litellm, field_name) == db_value +@pytest.mark.asyncio +async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch): + """A DB-only litellm_settings row that pairs success_callback: ["datadog"] with + datadog_params.turn_off_message_logging: true must build the DataDogLogger redacted, the + same as the identical block in YAML. Regression for the redaction keys being absent from + the safe-override allowlist while the callback half of the row was honoured.""" + import litellm.proxy.proxy_server as ps + from litellm.integrations.datadog.datadog import DataDogLogger + from litellm.litellm_core_utils import litellm_logging + + monkeypatch.setenv("DD_API_KEY", "test-key") + monkeypatch.setenv("DD_SITE", "us5.datadoghq.com") + monkeypatch.setattr(litellm, "datadog_params", None) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + + db_row = { + "success_callback": ["datadog"], + "datadog_params": {"turn_off_message_logging": True}, + "turn_off_message_logging": True, + } + pc = ps.ProxyConfig() + pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", db_row)) + pc._add_callbacks_from_db_config({"litellm_settings": db_row}) + + datadog_loggers = [cb for cb in litellm.success_callback if isinstance(cb, DataDogLogger)] + assert len(datadog_loggers) == 1 + assert datadog_loggers[0].turn_off_message_logging is True + assert litellm.turn_off_message_logging is True + + def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): """The flag defaults to False rather than None, so a plain 'is not None' check would report the default as 'In Config' and imply an admin had set it.""" From a90d852d378b46e3541fbb76ec519845e5caa785 Mon Sep 17 00:00:00 2001 From: ryan Date: Sun, 20 Sep 2026 09:26:30 +0000 Subject: [PATCH 027/109] test(proxy): type monkeypatch and cover every DB-overridable callback params key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_server.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 75741b125cb..c56645313c3 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11767,7 +11767,7 @@ def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_n @pytest.mark.asyncio -async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch): +async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch: pytest.MonkeyPatch): """A DB-only litellm_settings row that pairs success_callback: ["datadog"] with datadog_params.turn_off_message_logging: true must build the DataDogLogger redacted, the same as the identical block in YAML. Regression for the redaction keys being absent from @@ -11802,6 +11802,30 @@ async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(mon assert litellm.turn_off_message_logging is True +@pytest.mark.parametrize( + "field_name", + [ + "datadog_params", + "datadog_llm_observability_params", + "newrelic_params", + "pointfive_params", + "aws_sqs_callback_params", + ], +) +def test_db_stored_callback_params_propagate_to_litellm_module(monkeypatch: pytest.MonkeyPatch, field_name: str): + """Every callback init params block stored in the DB litellm_settings row must land on the + litellm module before the matching logger is built, so the DB row behaves like YAML.""" + import litellm.proxy.proxy_server as ps + + monkeypatch.setattr(litellm, field_name, None) + db_value = {"turn_off_message_logging": True} + + pc = ps.ProxyConfig() + pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", {field_name: db_value})) + + assert getattr(litellm, field_name) == db_value + + def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): """The flag defaults to False rather than None, so a plain 'is not None' check would report the default as 'In Config' and imply an admin had set it.""" From bf4fccc937175999d9327051479274bfb8c6d5fd Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 19:25:52 +0000 Subject: [PATCH 028/109] feat(ui): expose remaining complexity router advanced settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../add_model/ClassificationMethodConfig.tsx | 40 ++++++++- .../add_model/ComplexityRouterConfig.tsx | 36 ++++++++ .../add_model/HeuristicKeywordOverrides.tsx | 42 +++++++++ .../add_model/HousekeepingRoutingControls.tsx | 44 ++++++++++ .../add_model/PlanModeOverrideControls.tsx | 18 ++++ .../components/add_model/ReminderMarkers.tsx | 77 +++++++++++++++++ .../add_model/ResponseFormatControls.tsx | 12 +++ .../add_model/add_auto_router_tab.tsx | 14 +++ .../build_complexity_router_config.test.ts | 58 +++++++++++++ .../build_complexity_router_config.ts | 85 +++++++++++++++++++ .../components/add_model/classifier_types.ts | 3 +- ...d_updated_complexity_router_config.test.ts | 17 ++++ .../edit_auto_router_modal.tsx | 54 ++++++++++++ 13 files changed, 498 insertions(+), 2 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/HousekeepingRoutingControls.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index b72b29a29f4..2564137fb02 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -19,7 +19,10 @@ import HeuristicScoringConfig from "./HeuristicScoringConfig"; import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect"; import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig"; import ClassifierVisionConfig from "./ClassifierVisionConfig"; -import { getHeuristicV2SuccessThresholdError } from "./build_complexity_router_config"; +import { + getClassifierPluginTimeoutError, + getHeuristicV2SuccessThresholdError, +} from "./build_complexity_router_config"; import type { ReasoningEffort } from "./complexity_router_tiers"; import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults"; import { @@ -61,6 +64,7 @@ const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin"; const HEURISTIC_V2_SUCCESS_THRESHOLD_ID = "heuristic-v2-success-threshold"; +const CLASSIFIER_PLUGIN_TIMEOUT_ID = "classifier-plugin-timeout-ms"; const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + @@ -467,6 +471,40 @@ const ClassificationMethodConfig: React.FC = ({ <> + {classifierType === "custom" && ( +
+

+ This router uses a custom classifier plugin set in config.yaml. Pick a classifier below to replace it. +

+ + + onChange({ + ...value, + classifier_plugin_timeout_ms: event.target.value.trim() === "" ? undefined : Number(event.target.value), + }) + } + aria-invalid={Boolean( + showValidationErrors && getClassifierPluginTimeoutError(classifierType, value.classifier_plugin_timeout_ms), + )} + /> +

+ Time budget for the plugin call. On expiry the fallback path decides the tier. +

+ {showValidationErrors && getClassifierPluginTimeoutError(classifierType, value.classifier_plugin_timeout_ms) && ( +

+ {getClassifierPluginTimeoutError(classifierType, value.classifier_plugin_timeout_ms)} +

+ )} +
+ )} + {classifierType === "heuristic_v2" && (
)} +
+ Additional plan-mode sentinels + ({ label: pattern, value: pattern }))} + value={value.plan_mode_patterns ?? []} + onValueChange={(patterns) => + onChange({ ...value, plan_mode_patterns: patterns.length > 0 ? patterns : undefined }) + } + placeholder="e.g., enter plan mode" + emptyText="Type to add a sentinel" + allowCustomValues + className="w-full" + /> + + Case-sensitive literal strings added to the built-in Claude Code and Copilot plan-mode markers. + +
); diff --git a/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx new file mode 100644 index 00000000000..452f4dc727f --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx @@ -0,0 +1,77 @@ +import React from "react"; +import { Plus, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { getReminderMarkersError, type ReminderMarkerPair } from "./build_complexity_router_config"; + +const ReminderMarkers: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; + showValidationErrors?: boolean; +}> = ({ value, onChange, showValidationErrors = false }) => { + const markers = value.reminder_markers ?? []; + const update = (index: number, patch: Partial) => + onChange({ + ...value, + reminder_markers: markers.map((marker, markerIndex) => (markerIndex === index ? { ...marker, ...patch } : marker)), + }); + const remove = (index: number) => { + const next = markers.filter((_, markerIndex) => markerIndex !== index); + onChange({ ...value, reminder_markers: next.length > 0 ? next : undefined }); + }; + const error = getReminderMarkersError(value.reminder_markers); + return ( +
+

+ Delimiter pairs that wrap harness-injected reminder blocks, which are stripped before classification. Setting any + pair replaces the built-in pairs, so list every pair your harness emits. Matching is case-insensitive and values + are saved lowercased. +

+
+ {markers.map((marker, index) => ( +
+
+ + update(index, { open: event.target.value })} + /> +
+
+ + update(index, { close: event.target.value })} + /> +
+ +
+ ))} +
+ + {showValidationErrors && error &&

{error}

} +
+ ); +}; + +export default ReminderMarkers; diff --git a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx index 68dd880a684..9edbf204a6d 100644 --- a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx @@ -18,6 +18,18 @@ const ResponseFormatControls: React.FC<{ Return the resolved underlying model name in responses instead of the autorouter alias. +
+ onChange({ ...value, max_tokens_from_tier_model: enabled })} + aria-label="Cap max_tokens at the tier model's output ceiling" + /> + Cap max_tokens at the tier model's output ceiling +
+ + Replace the caller's max_tokens with the routed tier model's output ceiling so one client value fits every + tier. Off forwards the caller's value unchanged. + ); 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 84e44fee9c3..2fa28f761e3 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 @@ -48,6 +48,8 @@ import { getKeywordTierRulesError, getClassifierModelError, getHeuristicV2SuccessThresholdError, + getReminderMarkersError, + getClassifierPluginTimeoutError, getClassifierReasoningEffortError, getMissingTiersError, getPlanModeTierError, @@ -152,6 +154,8 @@ export const getSubmitBlockedReason = ( getKeywordTierRulesError(keywordTierRules, activeTierRows(config)) ?? getClassifierModelError(config) ?? getHeuristicV2SuccessThresholdError(config.heuristic_v2_success_threshold) ?? + getReminderMarkersError(config.reminder_markers) ?? + getClassifierPluginTimeoutError(config.classifier_type, config.classifier_plugin_timeout_ms) ?? (heuristicScoringRole(config) === "decides" ? customDimensionsError(config.custom_dimensions) : null) ?? getClassifierReasoningEffortError(config, modelInfo) ?? getReferencedModelsError(referencedModelsParams, availability) @@ -448,6 +452,16 @@ const AddAutoRouterTab: React.FC = ({ enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation, contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer, sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds, + codeKeywords: complexityRouterConfig.code_keywords, + reasoningKeywords: complexityRouterConfig.reasoning_keywords, + technicalKeywords: complexityRouterConfig.technical_keywords, + simpleKeywords: complexityRouterConfig.simple_keywords, + planModePatterns: complexityRouterConfig.plan_mode_patterns, + routeHousekeepingToCheapestTier: complexityRouterConfig.route_housekeeping_to_cheapest_tier, + housekeepingPatterns: complexityRouterConfig.housekeeping_patterns, + reminderMarkers: complexityRouterConfig.reminder_markers, + maxTokensFromTierModel: complexityRouterConfig.max_tokens_from_tier_model, + classifierPluginTimeoutMs: complexityRouterConfig.classifier_plugin_timeout_ms, }; const submitRecommendedRouter = async (name: string) => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 054aba7a6aa..cf5164d91c4 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 @@ -6,6 +6,8 @@ import { getKeywordTierRulesError, getClassifierModelError, getHeuristicV2SuccessThresholdError, + getReminderMarkersError, + getClassifierPluginTimeoutError, getClassifierReasoningEffortError, getMissingTiersError, hydrateCustomTierSet, @@ -1482,3 +1484,59 @@ describe("classifier vision wire payload", () => { expect(payload.classifier_llm_config).not.toHaveProperty("vision"); }); }); + +describe("advanced complexity router fields", () => { + it("normalizes lists, reminder markers, and explicit false values", () => { + const payload = buildComplexityRouterConfig({ + ...baseParams, + codeKeywords: [" async ", " "], + reasoningKeywords: ["prove"], + technicalKeywords: ["api"], + simpleKeywords: ["hello"], + planModePatterns: [" plan "], + routeHousekeepingToCheapestTier: false, + housekeepingPatterns: [" title "], + reminderMarkers: [{ open: " ", close: " " }], + maxTokensFromTierModel: false, + classifierType: "custom", + classifierPluginTimeoutMs: 3000, + }); + expect(payload).toMatchObject({ + code_keywords: ["async"], + reasoning_keywords: ["prove"], + technical_keywords: ["api"], + simple_keywords: ["hello"], + plan_mode_patterns: ["plan"], + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + classifier_plugin_timeout_ms: 3000, + }); + }); + + it("omits defaults, empty lists, and timeout values for non-custom classifiers", () => { + const payload = buildComplexityRouterConfig({ + ...baseParams, + codeKeywords: [" ", ""], + reminderMarkers: [], + routeHousekeepingToCheapestTier: true, + maxTokensFromTierModel: true, + classifierPluginTimeoutMs: 3000, + }); + expect(payload).not.toHaveProperty("code_keywords"); + expect(payload).not.toHaveProperty("reminder_markers"); + expect(payload).not.toHaveProperty("route_housekeeping_to_cheapest_tier"); + expect(payload).not.toHaveProperty("max_tokens_from_tier_model"); + expect(payload).not.toHaveProperty("classifier_plugin_timeout_ms"); + }); + + it("validates marker pairs and custom classifier timeout", () => { + expect(getReminderMarkersError([{ open: " ", close: " " }])).toContain("different"); + expect(getReminderMarkersError([{ open: "", close: "" }])).toContain("needs both"); + expect(getReminderMarkersError([{ open: "", close: "" }])).toBeNull(); + expect(getClassifierPluginTimeoutError("custom", 0)).toContain("whole number"); + expect(getClassifierPluginTimeoutError("custom", 3000)).toBeNull(); + expect(getClassifierPluginTimeoutError("heuristic", 0)).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 15ae2b4c0b0..96fbe7f2c2e 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 @@ -54,6 +54,10 @@ import { export type ClassifierVisionConfig = { enabled?: boolean; max_images?: number }; export type ClassifierLLMConfigWire = ClassifierLLMConfig & { vision?: ClassifierVisionConfig }; +export interface ReminderMarkerPair { + open: string; + close: string; +} /** * Drop an empty system_prompt so the payload carries an override only when there is one. The @@ -181,6 +185,16 @@ export interface StoredComplexityRouterConfig { stall_escalation_enabled?: unknown; stall_escalation_window?: unknown; stall_escalation_repeat_threshold?: unknown; + code_keywords?: unknown; + reasoning_keywords?: unknown; + technical_keywords?: unknown; + simple_keywords?: unknown; + plan_mode_patterns?: unknown; + route_housekeeping_to_cheapest_tier?: unknown; + housekeeping_patterns?: unknown; + reminder_markers?: unknown; + max_tokens_from_tier_model?: unknown; + classifier_plugin_timeout_ms?: unknown; } export interface BuildComplexityRouterConfigParams { @@ -233,6 +247,16 @@ export interface BuildComplexityRouterConfigParams { enableContextWindowEscalation?: boolean; contextWindowEscalationBuffer?: number; sessionAffinityTtlSeconds?: number; + codeKeywords?: string[]; + reasoningKeywords?: string[]; + technicalKeywords?: string[]; + simpleKeywords?: string[]; + planModePatterns?: string[]; + routeHousekeepingToCheapestTier?: boolean; + housekeepingPatterns?: string[]; + reminderMarkers?: ReminderMarkerPair[]; + maxTokensFromTierModel?: boolean; + classifierPluginTimeoutMs?: number; } /** @@ -302,6 +326,16 @@ export interface ComplexityRouterConfigPayload { enable_context_window_escalation?: boolean; context_window_escalation_buffer?: number; tier_model_configs?: Record; + code_keywords?: string[]; + reasoning_keywords?: string[]; + technical_keywords?: string[]; + simple_keywords?: string[]; + plan_mode_patterns?: string[]; + route_housekeeping_to_cheapest_tier?: boolean; + housekeeping_patterns?: string[]; + reminder_markers?: ReminderMarkerPair[]; + max_tokens_from_tier_model?: boolean; + classifier_plugin_timeout_ms?: number; } export const serializeTierLabels = (tierLabels: ComplexityTierLabels | undefined): ComplexityTierLabels | undefined => { @@ -376,6 +410,26 @@ export const getHeuristicV2SuccessThresholdError = (threshold: number | undefine return validProbability ? null : "Success threshold must be a number between 0 and 1"; }; +export const getReminderMarkersError = (pairs: ReminderMarkerPair[] | undefined): string | null => { + for (const [index, pair] of (pairs ?? []).entries()) { + const open = pair.open.trim().toLowerCase(); + const close = pair.close.trim().toLowerCase(); + if (!open || !close) return `Reminder marker pair ${index + 1} needs both an opening and a closing delimiter`; + if (open === close) return `Reminder marker pair ${index + 1} must use different opening and closing delimiters`; + } + return null; +}; + +export const getClassifierPluginTimeoutError = ( + classifierType: ClassifierType, + timeoutMs: number | undefined, +): string | null => { + if (classifierType !== "custom" || timeoutMs === undefined) return null; + return Number.isInteger(timeoutMs) && timeoutMs > 0 + ? null + : "Classifier plugin timeout must be a whole number of milliseconds greater than 0"; +}; + export const getClassifierModelError = ( config: Pick< ComplexityRouterConfigValue, @@ -640,6 +694,16 @@ export const buildComplexityRouterConfig = ({ enableContextWindowEscalation, contextWindowEscalationBuffer, sessionAffinityTtlSeconds, + codeKeywords, + reasoningKeywords, + technicalKeywords, + simpleKeywords, + planModePatterns, + routeHousekeepingToCheapestTier, + housekeepingPatterns, + reminderMarkers, + maxTokensFromTierModel, + classifierPluginTimeoutMs, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const serializedTierModelConfigs = customTierSet ? serializeTierModelConfigs( @@ -672,6 +736,14 @@ export const buildComplexityRouterConfig = ({ }; const effectiveType = effectiveClassifierType({ custom_tier_set: customTierSet, classifier_type: classifierType }); const forecast = isForecastClassifier(effectiveType); + const cleanList = (items: string[] | undefined): string[] | undefined => { + const cleaned = (items ?? []).map((item) => item.trim()).filter(Boolean); + return cleaned.length > 0 ? cleaned : undefined; + }; + const cleanedReminderMarkers = reminderMarkers?.map(({ open, close }) => ({ + open: open.trim().toLowerCase(), + close: close.trim().toLowerCase(), + })); const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType); const payload: ComplexityRouterConfigPayload = { @@ -740,6 +812,19 @@ export const buildComplexityRouterConfig = ({ ...(sessionAffinityTtlSeconds !== undefined && { session_affinity_ttl_seconds: sessionAffinityTtlSeconds, }), + ...(cleanList(codeKeywords) && { code_keywords: cleanList(codeKeywords) }), + ...(cleanList(reasoningKeywords) && { reasoning_keywords: cleanList(reasoningKeywords) }), + ...(cleanList(technicalKeywords) && { technical_keywords: cleanList(technicalKeywords) }), + ...(cleanList(simpleKeywords) && { simple_keywords: cleanList(simpleKeywords) }), + ...(cleanList(planModePatterns) && { plan_mode_patterns: cleanList(planModePatterns) }), + ...(routeHousekeepingToCheapestTier === false && { route_housekeeping_to_cheapest_tier: false }), + ...(cleanList(housekeepingPatterns) && { housekeeping_patterns: cleanList(housekeepingPatterns) }), + ...(cleanedReminderMarkers && cleanedReminderMarkers.length > 0 && { reminder_markers: cleanedReminderMarkers }), + ...(maxTokensFromTierModel === false && { max_tokens_from_tier_model: false }), + ...(classifierType === "custom" && + classifierPluginTimeoutMs !== undefined && + Number.isInteger(classifierPluginTimeoutMs) && + classifierPluginTimeoutMs > 0 && { classifier_plugin_timeout_ms: classifierPluginTimeoutMs }), ...scorerKnobs, }; if (!customTierSet) return payload; diff --git a/ui/litellm-dashboard/src/components/add_model/classifier_types.ts b/ui/litellm-dashboard/src/components/add_model/classifier_types.ts index ec88166ed2e..aa9d5619052 100644 --- a/ui/litellm-dashboard/src/components/add_model/classifier_types.ts +++ b/ui/litellm-dashboard/src/components/add_model/classifier_types.ts @@ -6,7 +6,8 @@ export type ClassifierType = | "heuristic_first" | "hybrid" | "capability" - | "llm_v2"; + | "llm_v2" + | "custom"; export const usesLlmClassifier = (classifierType: ClassifierType): boolean => (["llm", "heuristic_first", "hybrid", "capability", "llm_v2"] as const).some((type) => type === classifierType); 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 2450f7bce27..eaa6595d8a3 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 @@ -854,6 +854,16 @@ describe("managed keys survive an untouched open-and-save", () => { reasoning_override_min_score: 0.3, enable_context_window_escalation: false, context_window_escalation_buffer: 0.9, + code_keywords: ["async", "await"], + reasoning_keywords: ["prove"], + technical_keywords: ["api"], + simple_keywords: ["hello"], + plan_mode_patterns: ["plan now"], + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + classifier_plugin_timeout_ms: 3000, }; // tier_definitions and fallback_tier cannot sit beside heuristic_first, which this fixture uses, @@ -864,6 +874,7 @@ describe("managed keys survive an untouched open-and-save", () => { "fallback_tier", "hybrid_boundary_margin", "jev_classifier_config", + "classifier_plugin_timeout_ms", ]); // The stall keys are rejected beside the session pinning and user-turn classification this @@ -894,6 +905,12 @@ describe("managed keys survive an untouched open-and-save", () => { expect(dropped).toEqual([]); }); + it("keeps the custom classifier plugin timeout through an untouched save", () => { + const stored = { ...STORED_ALL_MANAGED, classifier_type: "custom", classifier_plugin_timeout_ms: 3000 }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).classifier_plugin_timeout_ms).toBe(3000); + }); + it("carries an enabled non-reasoning tier and its models through their own round trip", () => { // `tiers` is rewritten wholesale on save, so this is the regression that matters: opening an // enabled router and saving an unrelated edit must not delete the tier or its pool. 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 c88bbb101f7..476207553f7 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 @@ -45,6 +45,8 @@ import { buildComplexityRouterConfig, getClassifierModelError, getHeuristicV2SuccessThresholdError, + getReminderMarkersError, + getClassifierPluginTimeoutError, getClassifierReasoningEffortError, getKeywordTierRulesError, getMissingTiersError, @@ -113,6 +115,8 @@ export const hydrateComplexityRouterConfig = ( parsedConfig: StoredComplexityRouterConfig, complexityRouterDefaultModel: string | null | undefined, ): ComplexityRouterConfigValue => { + const stringList = (input: unknown): string[] | undefined => + Array.isArray(input) ? input.filter((item): item is string => typeof item === "string") : undefined; const builtIn = hydrateBuiltInTiers(parsedConfig.tiers, parsedConfig.enable_non_reasoning_tier); const { tiers: hydratedTiers, enable_non_reasoning_tier } = builtIn; const custom_tier_set = hydrateCustomTierSet(parsedConfig); @@ -219,6 +223,31 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.stall_escalation_repeat_threshold === "number" ? parsedConfig.stall_escalation_repeat_threshold : undefined, + code_keywords: stringList(parsedConfig.code_keywords), + reasoning_keywords: stringList(parsedConfig.reasoning_keywords), + technical_keywords: stringList(parsedConfig.technical_keywords), + simple_keywords: stringList(parsedConfig.simple_keywords), + plan_mode_patterns: stringList(parsedConfig.plan_mode_patterns), + route_housekeeping_to_cheapest_tier: + typeof parsedConfig.route_housekeeping_to_cheapest_tier === "boolean" + ? parsedConfig.route_housekeeping_to_cheapest_tier + : undefined, + housekeeping_patterns: stringList(parsedConfig.housekeeping_patterns), + reminder_markers: Array.isArray(parsedConfig.reminder_markers) + ? parsedConfig.reminder_markers.filter( + (pair): pair is { open: string; close: string } => + typeof pair === "object" && + pair !== null && + typeof (pair as { open?: unknown }).open === "string" && + typeof (pair as { close?: unknown }).close === "string", + ) + : undefined, + max_tokens_from_tier_model: + typeof parsedConfig.max_tokens_from_tier_model === "boolean" ? parsedConfig.max_tokens_from_tier_model : undefined, + classifier_plugin_timeout_ms: + typeof parsedConfig.classifier_plugin_timeout_ms === "number" && Number.isFinite(parsedConfig.classifier_plugin_timeout_ms) + ? parsedConfig.classifier_plugin_timeout_ms + : undefined, }; }; @@ -266,6 +295,16 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "stall_escalation_enabled", "stall_escalation_window", "stall_escalation_repeat_threshold", + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "simple_keywords", + "plan_mode_patterns", + "route_housekeeping_to_cheapest_tier", + "housekeeping_patterns", + "reminder_markers", + "max_tokens_from_tier_model", + "classifier_plugin_timeout_ms", ]); // Managed only when the caller passes the corresponding state. A caller that does not render @@ -387,6 +426,16 @@ export const buildUpdatedComplexityRouterConfig = ( stallEscalationEnabled: value.stall_escalation_enabled, stallEscalationWindow: value.stall_escalation_window, stallEscalationRepeatThreshold: value.stall_escalation_repeat_threshold, + codeKeywords: value.code_keywords, + reasoningKeywords: value.reasoning_keywords, + technicalKeywords: value.technical_keywords, + simpleKeywords: value.simple_keywords, + planModePatterns: value.plan_mode_patterns, + routeHousekeepingToCheapestTier: value.route_housekeeping_to_cheapest_tier, + housekeepingPatterns: value.housekeeping_patterns, + reminderMarkers: value.reminder_markers, + maxTokensFromTierModel: value.max_tokens_from_tier_model, + classifierPluginTimeoutMs: value.classifier_plugin_timeout_ms, }; const built = buildComplexityRouterConfig(builderParams); @@ -585,6 +634,11 @@ const EditAutoRouterModal: React.FC = ({ const classifierError = getClassifierModelError(complexityRouterConfig) ?? getHeuristicV2SuccessThresholdError(complexityRouterConfig.heuristic_v2_success_threshold) ?? + getReminderMarkersError(complexityRouterConfig.reminder_markers) ?? + getClassifierPluginTimeoutError( + complexityRouterConfig.classifier_type, + complexityRouterConfig.classifier_plugin_timeout_ms, + ) ?? getForecastConfigError(complexityRouterConfig) ?? (heuristicScoringRole(complexityRouterConfig) === "decides" ? customDimensionsError(complexityRouterConfig.custom_dimensions) From 9c411dd6f2e569ea7ac54c08f03847d5c70cddba Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:35:29 +0000 Subject: [PATCH 029/109] refactor(proxy): make scheduled job shutdown timeouts configurable via env Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 ++ litellm/proxy/shutdown/scheduled_jobs.py | 23 +++++++++++-------- .../proxy/shutdown/test_scheduled_jobs.py | 15 ++++++------ 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index bbeb4846e27..842adf62f6b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1742,6 +1742,8 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300")) SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30")) SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000")) +SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5")) +SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5")) TOOL_SPEND_TOP_TOOLS: Final = 100 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index 5345d380112..cf4937780b8 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -7,9 +7,10 @@ from typing import Final, Protocol from apscheduler.executors.asyncio import AsyncIOExecutor from litellm._logging import verbose_proxy_logger - -JOB_FINISH_TIMEOUT_SECONDS: Final = 5.0 -JOB_CANCEL_TIMEOUT_SECONDS: Final = 5.0 +from litellm.constants import ( + SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, + SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, +) class StoppableScheduler(Protocol): @@ -41,8 +42,8 @@ def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: """ - Let in-flight jobs finish for up to JOB_FINISH_TIMEOUT_SECONDS, then stop the scheduler and - wait, bounded by JOB_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. + Let in-flight jobs finish for up to SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, then stop the scheduler and + wait, bounded by SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. Must run before the database is disconnected: a write job that finishes needs its connection, and a job's cancellation handler is what records the run's outcome. @@ -52,19 +53,23 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: in_flight: Final = executor.in_flight_jobs() if in_flight: verbose_proxy_logger.info( - "Waiting up to %ss for %d in-flight scheduled job(s) to finish", JOB_FINISH_TIMEOUT_SECONDS, len(in_flight) + "Waiting up to %ss for %d in-flight scheduled job(s) to finish", + SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + len(in_flight), ) still_running: Final = ( - (await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS))[1] if in_flight else frozenset() + (await asyncio.wait(in_flight, timeout=SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS))[1] + if in_flight + else frozenset() ) scheduler.shutdown(wait=False) if not still_running: return verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running)) - _done, pending = await asyncio.wait(still_running, timeout=JOB_CANCEL_TIMEOUT_SECONDS) + _done, pending = await asyncio.wait(still_running, timeout=SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS) if pending: verbose_proxy_logger.warning( "%d scheduled job(s) did not finish within %ss of cancellation; giving up on them", len(pending), - JOB_CANCEL_TIMEOUT_SECONDS, + SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, ) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 7defd6cef6c..1f94cee04ed 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.constants import SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - stop_in_flight_scheduler_jobs, pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, ) @@ -75,9 +75,8 @@ async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): @pytest.mark.asyncio -async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(monkeypatch): +async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(): """A spend write cancelled mid-commit drops the rows it popped, so short jobs get to finish first""" - monkeypatch.setattr(scheduled_jobs, "JOB_FINISH_TIMEOUT_SECONDS", 2.0) write = _Job(work_seconds=0.2) stuck = _Job() async with _running_scheduler(write, stuck) as (scheduler, executor): @@ -99,16 +98,18 @@ async def test_every_in_flight_job_is_cancelled_not_only_the_first(): @pytest.mark.asyncio -async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(monkeypatch, caplog): +async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(caplog): """A job that swallows CancelledError must not hold the pod past its termination grace period""" - monkeypatch.setattr(scheduled_jobs, "JOB_CANCEL_TIMEOUT_SECONDS", 0.05) job = _Job(swallow_cancellation=True) async with _running_scheduler(job) as (scheduler, executor): with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await stop_in_flight_scheduler_jobs(scheduler, executor) assert job.events == ["cancelled"] - assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text + assert ( + f"1 scheduled job(s) did not finish within {SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS}s of cancellation" + in caplog.text + ) @pytest.mark.asyncio From 4e388e6aea52ad6c9c2939997f860689e69716aa Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:37:28 +0000 Subject: [PATCH 030/109] refactor(proxy): inject scheduled job shutdown timeouts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/shutdown/scheduled_jobs.py | 20 ++++++++++++------- .../proxy/shutdown/test_scheduled_jobs.py | 10 +++------- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index cf4937780b8..e920ce19eb9 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -40,10 +40,16 @@ def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: scheduler.pause() -async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: +async def stop_in_flight_scheduler_jobs( + scheduler: StoppableScheduler, + executor: AwaitableAsyncIOExecutor, + *, + finish_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + cancel_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, +) -> None: """ - Let in-flight jobs finish for up to SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, then stop the scheduler and - wait, bounded by SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. + Let in-flight jobs finish for up to finish_timeout_seconds, then stop the scheduler and wait, bounded by + cancel_timeout_seconds, for the jobs it cancels. Must run before the database is disconnected: a write job that finishes needs its connection, and a job's cancellation handler is what records the run's outcome. @@ -54,11 +60,11 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: if in_flight: verbose_proxy_logger.info( "Waiting up to %ss for %d in-flight scheduled job(s) to finish", - SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + finish_timeout_seconds, len(in_flight), ) still_running: Final = ( - (await asyncio.wait(in_flight, timeout=SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS))[1] + (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] if in_flight else frozenset() ) @@ -66,10 +72,10 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: if not still_running: return verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running)) - _done, pending = await asyncio.wait(still_running, timeout=SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS) + _done, pending = await asyncio.wait(still_running, timeout=cancel_timeout_seconds) if pending: verbose_proxy_logger.warning( "%d scheduled job(s) did not finish within %ss of cancellation; giving up on them", len(pending), - SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, + cancel_timeout_seconds, ) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 1f94cee04ed..fbce38db39f 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,7 +7,6 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -from litellm.constants import SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, pause_scheduled_jobs, @@ -80,7 +79,7 @@ async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelle write = _Job(work_seconds=0.2) stuck = _Job() async with _running_scheduler(write, stuck) as (scheduler, executor): - await stop_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor, finish_timeout_seconds=2.0) assert write.events == ["committed", "finished"] assert stuck.events == ["cancelled", "finished"] @@ -103,13 +102,10 @@ async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(ca job = _Job(swallow_cancellation=True) async with _running_scheduler(job) as (scheduler, executor): with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - await stop_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor, cancel_timeout_seconds=0.05) assert job.events == ["cancelled"] - assert ( - f"1 scheduled job(s) did not finish within {SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS}s of cancellation" - in caplog.text - ) + assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text @pytest.mark.asyncio From 9644032cb806a8bef55d1bcf4219ddec4886b758 Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 19:37:45 +0000 Subject: [PATCH 031/109] refactor(ui): split complexity router form files and cover advanced fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../add_model/ClassificationMethodConfig.tsx | 118 +-------- .../ClassifierPluginTimeoutField.tsx | 54 ++++ .../add_model/ClassifierTypeRadios.tsx | 89 +++++++ .../ComplexityRouterAdvancedSections.tsx | 240 +++++++++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 58 +++++ .../add_model/ComplexityRouterConfig.tsx | 210 +++------------ .../components/add_model/ReminderMarkers.tsx | 4 +- .../add_model/add_auto_router_tab.tsx | 59 +---- .../build_complexity_router_config.test.ts | 16 ++ .../build_complexity_router_config.ts | 17 +- .../complexity_router_builder_params.ts | 74 ++++++ ...dit_auto_router_modal.integration.test.tsx | 71 +++++ .../edit_auto_router_modal.tsx | 243 +----------------- .../hydrate_complexity_router_config.ts | 183 +++++++++++++ 14 files changed, 839 insertions(+), 597 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/ClassifierPluginTimeoutField.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts create mode 100644 ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 2564137fb02..b7a0fd67443 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -19,10 +19,9 @@ import HeuristicScoringConfig from "./HeuristicScoringConfig"; import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect"; import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig"; import ClassifierVisionConfig from "./ClassifierVisionConfig"; -import { - getClassifierPluginTimeoutError, - getHeuristicV2SuccessThresholdError, -} from "./build_complexity_router_config"; +import { getHeuristicV2SuccessThresholdError } from "./build_complexity_router_config"; +import ClassifierPluginTimeoutField from "./ClassifierPluginTimeoutField"; +import ClassifierTypeRadios from "./ClassifierTypeRadios"; import type { ReasoningEffort } from "./complexity_router_tiers"; import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults"; import { @@ -64,7 +63,6 @@ const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin"; const HEURISTIC_V2_SUCCESS_THRESHOLD_ID = "heuristic-v2-success-threshold"; -const CLASSIFIER_PLUGIN_TIMEOUT_ID = "classifier-plugin-timeout-ms"; const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + @@ -208,84 +206,6 @@ export const InactiveHeuristicV2Threshold: React.FC void; -}> = ({ value, classifierType, onTypeChange }) => { - const scorerLocked = Boolean(value.custom_tier_set); - const scorerLockedReason = restrictedBy(value, "heuristicClassifier")?.reason; - return ( - onTypeChange(classifierType as ClassifierType)} - className="w-full" - > -
- - - - - - - - - - - - - - -
- - ); -}; - const ClassificationMethodConfig: React.FC = ({ value, onChange, @@ -472,37 +392,7 @@ const ClassificationMethodConfig: React.FC = ({ {classifierType === "custom" && ( -
-

- This router uses a custom classifier plugin set in config.yaml. Pick a classifier below to replace it. -

- - - onChange({ - ...value, - classifier_plugin_timeout_ms: event.target.value.trim() === "" ? undefined : Number(event.target.value), - }) - } - aria-invalid={Boolean( - showValidationErrors && getClassifierPluginTimeoutError(classifierType, value.classifier_plugin_timeout_ms), - )} - /> -

- Time budget for the plugin call. On expiry the fallback path decides the tier. -

- {showValidationErrors && getClassifierPluginTimeoutError(classifierType, value.classifier_plugin_timeout_ms) && ( -

- {getClassifierPluginTimeoutError(classifierType, value.classifier_plugin_timeout_ms)} -

- )} -
+ )} {classifierType === "heuristic_v2" && ( diff --git a/ui/litellm-dashboard/src/components/add_model/ClassifierPluginTimeoutField.tsx b/ui/litellm-dashboard/src/components/add_model/ClassifierPluginTimeoutField.tsx new file mode 100644 index 00000000000..45decf09a09 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ClassifierPluginTimeoutField.tsx @@ -0,0 +1,54 @@ +import React from "react"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { getClassifierPluginTimeoutError } from "./build_complexity_router_config"; + +const CLASSIFIER_PLUGIN_TIMEOUT_ID = "classifier-plugin-timeout-ms"; + +interface ClassifierPluginTimeoutFieldProps { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; + showValidationErrors?: boolean; +} + +const ClassifierPluginTimeoutField: React.FC = ({ + value, + onChange, + showValidationErrors = false, +}) => { + const error = getClassifierPluginTimeoutError("custom", value.classifier_plugin_timeout_ms); + return ( +
+

+ This router uses a custom classifier plugin set in config.yaml. Pick a classifier below to replace it. +

+ + + onChange({ + ...value, + classifier_plugin_timeout_ms: event.target.value.trim() === "" ? undefined : Number(event.target.value), + }) + } + aria-invalid={Boolean(showValidationErrors && error)} + /> +

+ Time budget for the plugin call. On expiry the fallback path decides the tier. +

+ {showValidationErrors && error && ( +

+ {error} +

+ )} +
+ ); +}; + +export default ClassifierPluginTimeoutField; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx new file mode 100644 index 00000000000..1602e19069a --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx @@ -0,0 +1,89 @@ +import React from "react"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import type { ClassifierType } from "./classifier_types"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { restrictedBy } from "./TierRestrictions"; + +interface ClassifierTypeRadiosProps { + value: ComplexityRouterConfigValue; + classifierType: ClassifierType; + onTypeChange: (classifierType: ClassifierType) => void; +} + +const ClassifierTypeRadios: React.FC = ({ value, classifierType, onTypeChange }) => { + const scorerLocked = Boolean(value.custom_tier_set); + const scorerLockedReason = restrictedBy(value, "heuristicClassifier")?.reason; + return ( + onTypeChange(nextType as ClassifierType)} + className="w-full" + > +
+ + + + + + + + + + + + + + +
+ + ); +}; + +export default ClassifierTypeRadios; diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx new file mode 100644 index 00000000000..dd822907735 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx @@ -0,0 +1,240 @@ +import React from "react"; +import { ChevronRight } from "lucide-react"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Separator } from "@/components/ui/separator"; +import type { ModelGroup } from "@/components/llm_calls/fetch_models"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; +import ClassificationMethodConfig from "./ClassificationMethodConfig"; +import ContextWindowEscalationConfig from "./ContextWindowEscalationConfig"; +import ResponseFormatControls from "./ResponseFormatControls"; +import StallEscalationConfig from "./StallEscalationConfig"; +import { Restricted, restrictedBy } from "./TierRestrictions"; +import EscalationKeywords from "./EscalationKeywords"; +import KeywordTierRules, { type KeywordTierRule } from "./KeywordTierRules"; +import SemanticKeywordMatching from "./SemanticKeywordMatching"; +import CompressionControls from "./CompressionControls"; +import PlanModeOverrideControls from "./PlanModeOverrideControls"; +import { AffinityControls } from "./AffinityControls"; +import { ModalityRoutingControls } from "./ModalityRoutingControls"; +import HeuristicKeywordOverrides from "./HeuristicKeywordOverrides"; +import HousekeepingRoutingControls from "./HousekeepingRoutingControls"; +import ReminderMarkers from "./ReminderMarkers"; +import type { AutoRouterCompressionState } from "./buildAutoRouterCompression"; +import { activeTierName, type TierRow } from "./tier_rows"; + +interface ComplexityRouterAdvancedSectionsProps { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; + forecast: boolean; + modelOptions: { value: string; label: string }[]; + classifierEffortOptionsByModel: Record; + customTechnicalKeywords?: string[]; + onCustomTechnicalKeywordsChange?: (keywords: string[]) => void; + showValidationErrors: boolean; + defaultModel?: string; + planModeTierOptions: { value: string; label: string }[]; + keywordTierRules: KeywordTierRule[]; + onKeywordTierRulesChange?: (rules: KeywordTierRule[]) => void; + semanticMatchingEnabled: boolean; + onSemanticMatchingEnabledChange?: (enabled: boolean) => void; + embeddingModel?: string; + onEmbeddingModelChange: (model: string) => void; + matchThreshold: number; + onMatchThresholdChange: (threshold: number) => void; + escalationKeywords: string[]; + onEscalationKeywordsChange?: (keywords: string[]) => void; + autoRouterCompression: AutoRouterCompressionState; + onAutoRouterCompressionChange?: (state: AutoRouterCompressionState) => void; + modelInfo: ModelGroup[]; + tierRows: TierRow[]; + customTierSet: ComplexityRouterConfigValue["custom_tier_set"]; +} + +const ComplexityRouterAdvancedSections: React.FC = ({ + value, + onChange, + forecast, + modelOptions, + classifierEffortOptionsByModel, + customTechnicalKeywords, + onCustomTechnicalKeywordsChange, + showValidationErrors, + defaultModel, + planModeTierOptions, + keywordTierRules, + onKeywordTierRulesChange, + semanticMatchingEnabled, + onSemanticMatchingEnabledChange, + embeddingModel, + onEmbeddingModelChange, + matchThreshold, + onMatchThresholdChange, + escalationKeywords, + onEscalationKeywordsChange, + autoRouterCompression, + onAutoRouterCompressionChange, + modelInfo, + tierRows, + customTierSet, +}) => { + const sections = [ + ...(!forecast + ? [ + { + key: "classifier", + label: Advanced: Classification Method, + children: ( + + ), + }, + ] + : []), + ...(!forecast + ? [ + { + key: "keyword-overrides", + label: Advanced: Heuristic Keyword Overrides, + children: , + }, + ] + : []), + { + key: "adaptive", + label: Advanced: Adaptive Routing, + children: ( + + + + ), + }, + { + key: "affinity", + label: Advanced: Affinity, + children: , + }, + { + key: "modality", + label: Advanced: Modality Routing, + children: , + }, + { + key: "plan-mode", + label: Advanced: Plan-Mode Override, + children: , + }, + { + key: "housekeeping", + label: Advanced: Housekeeping Routing, + children: , + }, + { + key: "reminder-markers", + label: Advanced: Reminder Markers, + children: , + }, + { + key: "context-window", + label: Advanced: Context Window Escalation, + children: , + }, + { + key: "stall-escalation", + label: Advanced: Stalled Task Escalation, + children: ( + + + + ), + }, + { + key: "response", + label: Advanced: Response Format, + children: , + }, + ...(onEscalationKeywordsChange + ? [ + { + key: "escalation", + label: Advanced: Escalation Keywords, + children: ( + + + + ), + }, + ] + : []), + ...(onAutoRouterCompressionChange + ? [ + { + key: "compression", + label: Advanced: Compression, + children: , + }, + ] + : []), + ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange + ? [ + { + key: "keyword-semantic", + label: Advanced: Keyword/Semantic Matching, + children: ( + <> + {onKeywordTierRulesChange && ( + + )} + {onKeywordTierRulesChange && onSemanticMatchingEnabledChange && } + {onSemanticMatchingEnabledChange && ( + + )} + + ), + }, + ] + : []), + ]; + + return ( + <> + {sections + .filter(({ key }) => !forecast || !["adaptive", "context-window", "escalation"].includes(key)) + .map(({ key, label, children }) => ( + + + + {label} + + {children} + + ))} + + ); +}; + +export default ComplexityRouterAdvancedSections; 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 70658b787f0..9a8577100df 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -96,6 +96,64 @@ describe("ComplexityRouterConfig", () => { expect(screen.queryByText("Classifier Model")).not.toBeInTheDocument(); }); + it("shows heuristic advanced sections and hides keyword overrides for capability classifiers", () => { + const { rerender } = renderWithProviders(); + + expect(screen.getByText("Advanced: Heuristic Keyword Overrides")).toBeInTheDocument(); + expect(screen.getByText("Advanced: Housekeeping Routing")).toBeInTheDocument(); + expect(screen.getByText("Advanced: Reminder Markers")).toBeInTheDocument(); + + const capabilityValue = { ...defaultValue, classifier_type: "capability" as const }; + rerender(); + expect(screen.queryByText("Advanced: Heuristic Keyword Overrides")).not.toBeInTheDocument(); + + }); + + it.each([ + ["custom", true], + ["heuristic", false], + ] as const)("shows plugin timeout only for %s classifiers", (classifierType, visible) => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + if (visible) { + expect(screen.getByLabelText("Classifier plugin timeout (ms)")).toBeInTheDocument(); + } else { + expect(screen.queryByLabelText("Classifier plugin timeout (ms)")).not.toBeInTheDocument(); + } + }); + + it.each([true, false])("shows reminder marker validation only when requested: %s", (showValidationErrors) => { + const value = { ...defaultValue, reminder_markers: [{ open: "", close: "x" }] }; + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Reminder Markers")); + const validation = screen.queryByText(/needs both/i); + if (showValidationErrors) { + expect(validation).toBeInTheDocument(); + } else { + expect(validation).not.toBeInTheDocument(); + } + }); + + it("disables housekeeping sentinels when cheapest-tier routing is off", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Housekeeping Routing")); + const sentinelInput = screen.getByRole("combobox", { name: "e.g., conversation title" }); + expect(sentinelInput).toBeDisabled(); + }); + it("should toggle returning the raw model name", async () => { const user = userEvent.setup(); const onChange = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index a97e7a77a21..acf6b62a95a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -2,21 +2,17 @@ import RoutingOptions from "./RoutingOptions"; import type { JevClassifierConfig } from "./jev_classifier_config"; import { type ClassifierType } from "./classifier_types"; export { type ClassifierType, usesLlmClassifier, usesClassifierContext } from "./classifier_types"; -import PlanModeOverrideControls from "./PlanModeOverrideControls"; import ForecastClassifierConfig, { ForecastSolverModels } from "./ForecastClassifierConfig"; import { isForecastClassifier, type CapabilitySettings, type FuseSettings } from "./forecast_classifier_config"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { MultiSelect } from "@/components/shared/MultiSelect"; import DefaultModelField from "./DefaultModelField"; -import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react"; +import { Info, Plus, Trash2, X } from "lucide-react"; -import { AffinityControls } from "./AffinityControls"; import NonReasoningTierToggle from "./NonReasoningTierToggle"; import TierConfigIntro from "./TierConfigIntro"; import TierRowSelect from "./TierRowSelect"; -import { ModalityRoutingControls } from "./ModalityRoutingControls"; import { Card, CardContent } from "@/components/ui/card"; -import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { Separator } from "@/components/ui/separator"; import { Button } from "@/components/ui/button"; @@ -39,12 +35,8 @@ import { } from "./tier_rows"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; -import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; -import ClassificationMethodConfig, { InactiveHeuristicV2Threshold } from "./ClassificationMethodConfig"; -import ContextWindowEscalationConfig from "./ContextWindowEscalationConfig"; -import ResponseFormatControls from "./ResponseFormatControls"; -import StallEscalationConfig from "./StallEscalationConfig"; -import { Restricted, restrictedBy } from "./TierRestrictions"; +import { InactiveHeuristicV2Threshold } from "./ClassificationMethodConfig"; +import ComplexityRouterAdvancedSections from "./ComplexityRouterAdvancedSections"; import { type TierSetAction, applyTierSetAction, setFallbackTier } from "./tier_set_actions"; import { ReasoningEffort, @@ -56,16 +48,10 @@ import { tierRowLabel, } from "./complexity_router_tiers"; import TierModelEffortRows from "./TierModelEffortRows"; -import EscalationKeywords from "./EscalationKeywords"; -import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; -import SemanticKeywordMatching from "./SemanticKeywordMatching"; +import { KeywordTierRule } from "./KeywordTierRules"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; import { type CustomDimensionRow } from "./custom_dimensions"; -import CompressionControls from "./CompressionControls"; import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; -import HeuristicKeywordOverrides from "./HeuristicKeywordOverrides"; -import HousekeepingRoutingControls from "./HousekeepingRoutingControls"; -import ReminderMarkers from "./ReminderMarkers"; import { type ReminderMarkerPair } from "./build_complexity_router_config"; export type { DimensionWeights, TierBoundaries, TokenThresholds }; @@ -782,167 +768,33 @@ const ComplexityRouterConfig: React.FC = ({ )}
- {[ - ...(!forecast - ? [ - { - key: "classifier", - label: Advanced: Classification Method, - children: ( - - ), - }, - ] - : []), - ...(!forecast - ? [ - { - key: "keyword-overrides", - label: ( - Advanced: Heuristic Keyword Overrides - ), - children: , - }, - ] - : []), - { - key: "adaptive", - label: Advanced: Adaptive Routing, - children: ( - - - - ), - }, - { - key: "affinity", - label: Advanced: Affinity, - children: , - }, - { - key: "modality", - label: Advanced: Modality Routing, - children: , - }, - { - key: "plan-mode", - label: Advanced: Plan-Mode Override, - children: ( - - ), - }, - { - key: "housekeeping", - label: Advanced: Housekeeping Routing, - children: , - }, - { - key: "reminder-markers", - label: Advanced: Reminder Markers, - children: , - }, - { - key: "context-window", - label: Advanced: Context Window Escalation, - children: , - }, - { - key: "stall-escalation", - label: Advanced: Stalled Task Escalation, - children: ( - - - - ), - }, - { - key: "response", - label: Advanced: Response Format, - children: , - }, - ...(onEscalationKeywordsChange - ? [ - { - key: "escalation", - label: Advanced: Escalation Keywords, - children: ( - - - - ), - }, - ] - : []), - ...(onAutoRouterCompressionChange - ? [ - { - key: "compression", - label: Advanced: Compression, - children: ( - - ), - }, - ] - : []), - ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange - ? [ - { - key: "keyword-semantic", - label: ( - Advanced: Keyword/Semantic Matching - ), - children: ( - <> - {onKeywordTierRulesChange && ( - - )} - {onKeywordTierRulesChange && onSemanticMatchingEnabledChange && } - {onSemanticMatchingEnabledChange && ( - - )} - - ), - }, - ] - : []), - ] - .filter(({ key }) => !forecast || !["adaptive", "context-window", "escalation"].includes(key)) - .map(({ key, label, children }) => ( - - - - {label} - - {children} - - ))} +
diff --git a/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx index 452f4dc727f..970394291d4 100644 --- a/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx @@ -30,14 +30,13 @@ const ReminderMarkers: React.FC<{

{markers.map((marker, index) => ( -
+
update(index, { open: event.target.value })} @@ -49,7 +48,6 @@ const ReminderMarkers: React.FC<{ update(index, { close: event.target.value })} 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 2fa28f761e3..f805a5b5511 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 @@ -28,10 +28,6 @@ import ComplexityRouterConfig, { effectiveClassifierType, usesLlmClassifier, heuristicScoringRole, - DEFAULT_ADAPTIVE_WEIGHTS, - DEFAULT_SESSION_AFFINITY, - DEFAULT_DEPLOYMENT_AFFINITY, - DEFAULT_TIER_DISTANCE_PENALTY, } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { customDimensionsError } from "./custom_dimensions"; @@ -57,6 +53,7 @@ import { getTierLabelsError, dryRunRejection, } from "./build_complexity_router_config"; +import { builderParamsFromValue } from "./complexity_router_builder_params"; import { activeTierName, activeTierRows, getCustomTierRowsError, resolveComplexityDefaultModel } from "./tier_rows"; import { tierRowLabel } from "./complexity_router_tiers"; import { buildAutoRouterTestTargets, AutoRouterTestTarget } from "./build_auto_router_test_targets"; @@ -403,65 +400,13 @@ const AddAutoRouterTab: React.FC = ({ ); const complexityRouterConfigParams: BuildComplexityRouterConfigParams = { - tiers: complexityRouterConfig.tiers, - enableNonReasoningTier: complexityRouterConfig.enable_non_reasoning_tier, - customTierSet: complexityRouterConfig.custom_tier_set, - defaultModel: complexityRouterConfig.default_model, - planModeMinTier: complexityRouterConfig.plan_mode_min_tier, - classificationPrompt: complexityRouterConfig.classification_prompt, - classificationExamples: complexityRouterConfig.classification_examples, - heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, - hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, - classificationMode: complexityRouterConfig.classification_mode, - tierLabels: complexityRouterConfig.tier_labels, - classifierType: complexityRouterConfig.classifier_type, - jevClassifierConfig: complexityRouterConfig.jev_classifier_config, - heuristicV2SuccessThreshold: complexityRouterConfig.heuristic_v2_success_threshold, - capabilityClassifierConfig: complexityRouterConfig.capability_classifier_config, - llmV2Config: complexityRouterConfig.llm_v2_config, - classifierLlmConfig: complexityRouterConfig.classifier_llm_config, - classifierContextWindowSize: complexityRouterConfig.classifier_context_window_size, - classifierContextBudgetChars: complexityRouterConfig.classifier_context_budget_chars, - classifierContextPerTurnChars: complexityRouterConfig.classifier_context_per_turn_chars, - classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns, - classifierFallback: complexityRouterConfig.classifier_fallback, - sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY, - modalityRouting: complexityRouterConfig.modality_routing ?? false, - modalityPinOverride: complexityRouterConfig.modality_pin_override ?? false, - deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + ...builderParamsFromValue(complexityRouterConfig), customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, embeddingModel, matchThreshold, escalationKeywords, - stallEscalationEnabled: complexityRouterConfig.stall_escalation_enabled, - stallEscalationWindow: complexityRouterConfig.stall_escalation_window, - stallEscalationRepeatThreshold: complexityRouterConfig.stall_escalation_repeat_threshold, - adaptive: complexityRouterConfig.adaptive ?? false, - adaptiveWeights: complexityRouterConfig.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, - tierDistancePenalty: complexityRouterConfig.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, - adaptiveEligible: complexityRouterConfig.adaptive_eligible ?? "all", - returnRawModelName: complexityRouterConfig.return_raw_model_name ?? false, - tierModelParams: complexityRouterConfig.tier_model_params, - tierBoundaries: complexityRouterConfig.tier_boundaries, - tokenThresholds: complexityRouterConfig.token_thresholds, - dimensionWeights: complexityRouterConfig.dimension_weights, - customDimensions: complexityRouterConfig.custom_dimensions, - reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score, - enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation, - contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer, - sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds, - codeKeywords: complexityRouterConfig.code_keywords, - reasoningKeywords: complexityRouterConfig.reasoning_keywords, - technicalKeywords: complexityRouterConfig.technical_keywords, - simpleKeywords: complexityRouterConfig.simple_keywords, - planModePatterns: complexityRouterConfig.plan_mode_patterns, - routeHousekeepingToCheapestTier: complexityRouterConfig.route_housekeeping_to_cheapest_tier, - housekeepingPatterns: complexityRouterConfig.housekeeping_patterns, - reminderMarkers: complexityRouterConfig.reminder_markers, - maxTokensFromTierModel: complexityRouterConfig.max_tokens_from_tier_model, - classifierPluginTimeoutMs: complexityRouterConfig.classifier_plugin_timeout_ms, }; const submitRecommendedRouter = async (name: string) => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index cf5164d91c4..6db2b8213d1 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 @@ -1531,6 +1531,22 @@ describe("advanced complexity router fields", () => { expect(payload).not.toHaveProperty("classifier_plugin_timeout_ms"); }); + it.each([ + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "simple_keywords", + "plan_mode_patterns", + "route_housekeeping_to_cheapest_tier", + "housekeeping_patterns", + "reminder_markers", + "max_tokens_from_tier_model", + "classifier_plugin_timeout_ms", + ])("omits unset advanced field %s", (key) => { + const payload = buildComplexityRouterConfig(baseParams); + expect(payload).not.toHaveProperty(key); + }); + it("validates marker pairs and custom classifier timeout", () => { expect(getReminderMarkersError([{ open: " ", close: " " }])).toContain("different"); expect(getReminderMarkersError([{ open: "", close: "" }])).toContain("needs both"); 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 96fbe7f2c2e..ab71b6b3cce 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 @@ -744,6 +744,16 @@ export const buildComplexityRouterConfig = ({ open: open.trim().toLowerCase(), close: close.trim().toLowerCase(), })); + const cleanedLists = Object.fromEntries( + Object.entries({ + code_keywords: cleanList(codeKeywords), + reasoning_keywords: cleanList(reasoningKeywords), + technical_keywords: cleanList(technicalKeywords), + simple_keywords: cleanList(simpleKeywords), + plan_mode_patterns: cleanList(planModePatterns), + housekeeping_patterns: cleanList(housekeepingPatterns), + }).filter(([, list]) => list !== undefined), + ); const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType); const payload: ComplexityRouterConfigPayload = { @@ -812,13 +822,8 @@ export const buildComplexityRouterConfig = ({ ...(sessionAffinityTtlSeconds !== undefined && { session_affinity_ttl_seconds: sessionAffinityTtlSeconds, }), - ...(cleanList(codeKeywords) && { code_keywords: cleanList(codeKeywords) }), - ...(cleanList(reasoningKeywords) && { reasoning_keywords: cleanList(reasoningKeywords) }), - ...(cleanList(technicalKeywords) && { technical_keywords: cleanList(technicalKeywords) }), - ...(cleanList(simpleKeywords) && { simple_keywords: cleanList(simpleKeywords) }), - ...(cleanList(planModePatterns) && { plan_mode_patterns: cleanList(planModePatterns) }), + ...cleanedLists, ...(routeHousekeepingToCheapestTier === false && { route_housekeeping_to_cheapest_tier: false }), - ...(cleanList(housekeepingPatterns) && { housekeeping_patterns: cleanList(housekeepingPatterns) }), ...(cleanedReminderMarkers && cleanedReminderMarkers.length > 0 && { reminder_markers: cleanedReminderMarkers }), ...(maxTokensFromTierModel === false && { max_tokens_from_tier_model: false }), ...(classifierType === "custom" && diff --git a/ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts b/ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts new file mode 100644 index 00000000000..124a85ce9a3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts @@ -0,0 +1,74 @@ +import type { BuildComplexityRouterConfigParams } from "./build_complexity_router_config"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_DEPLOYMENT_AFFINITY, + DEFAULT_SESSION_AFFINITY, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "./ComplexityRouterConfig"; + +export const builderParamsFromValue = ( + value: ComplexityRouterConfigValue, +): Omit< + BuildComplexityRouterConfigParams, + | "customTechnicalKeywords" + | "keywordTierRules" + | "semanticMatchingEnabled" + | "embeddingModel" + | "matchThreshold" + | "escalationKeywords" +> => ({ + tiers: value.tiers, + enableNonReasoningTier: value.enable_non_reasoning_tier, + customTierSet: value.custom_tier_set, + defaultModel: value.default_model, + planModeMinTier: value.plan_mode_min_tier, + classificationPrompt: value.classification_prompt, + classificationExamples: value.classification_examples, + heuristicFirstMaxTier: value.heuristic_first_max_tier, + hybridBoundaryMargin: value.hybrid_boundary_margin, + classificationMode: value.classification_mode, + tierLabels: value.tier_labels, + classifierType: value.classifier_type, + jevClassifierConfig: value.jev_classifier_config, + heuristicV2SuccessThreshold: value.heuristic_v2_success_threshold, + capabilityClassifierConfig: value.capability_classifier_config, + llmV2Config: value.llm_v2_config, + classifierLlmConfig: value.classifier_llm_config, + classifierContextWindowSize: value.classifier_context_window_size, + classifierContextBudgetChars: value.classifier_context_budget_chars, + classifierContextPerTurnChars: value.classifier_context_per_turn_chars, + classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, + classifierFallback: value.classifier_fallback, + sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, + sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds, + modalityRouting: value.modality_routing ?? false, + modalityPinOverride: value.modality_pin_override ?? false, + deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + adaptive: value.adaptive ?? false, + adaptiveWeights: value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, + tierDistancePenalty: value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, + adaptiveEligible: value.adaptive_eligible ?? "all", + returnRawModelName: value.return_raw_model_name ?? false, + tierBoundaries: value.tier_boundaries, + tokenThresholds: value.token_thresholds, + dimensionWeights: value.dimension_weights, + customDimensions: value.custom_dimensions, + reasoningOverrideMinScore: value.reasoning_override_min_score, + tierModelParams: value.tier_model_params, + enableContextWindowEscalation: value.enable_context_window_escalation, + contextWindowEscalationBuffer: value.context_window_escalation_buffer, + stallEscalationEnabled: value.stall_escalation_enabled, + stallEscalationWindow: value.stall_escalation_window, + stallEscalationRepeatThreshold: value.stall_escalation_repeat_threshold, + codeKeywords: value.code_keywords, + reasoningKeywords: value.reasoning_keywords, + technicalKeywords: value.technical_keywords, + simpleKeywords: value.simple_keywords, + planModePatterns: value.plan_mode_patterns, + routeHousekeepingToCheapestTier: value.route_housekeeping_to_cheapest_tier, + housekeepingPatterns: value.housekeeping_patterns, + reminderMarkers: value.reminder_markers, + maxTokensFromTierModel: value.max_tokens_from_tier_model, + classifierPluginTimeoutMs: value.classifier_plugin_timeout_ms, +}); 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 34db61483cf..aa6cf92ceb9 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 @@ -347,6 +347,77 @@ describe("EditAutoRouterModal keyword matching", () => { }); }); +describe("EditAutoRouterModal advanced field round trips", () => { + const storedAdvancedConfig = { + ...STORED_CONFIG, + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + }; + + const renderAdvancedModal = (props: Partial> = {}) => + renderModal({ + modelData: { + ...MODEL_DATA, + litellm_params: { ...MODEL_DATA.litellm_params, complexity_router_config: storedAdvancedConfig }, + }, + ...props, + }); + + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + it("hydrates housekeeping and reminder fields, then omits the default max-token value after editing", async () => { + const user = userEvent.setup(); + renderAdvancedModal(); + + await user.click(await screen.findByText("Advanced: Housekeeping Routing")); + expect(screen.getByRole("switch", { name: "Route housekeeping calls to the cheapest tier" })).not.toBeChecked(); + expect(screen.getByRole("combobox", { name: "e.g., conversation title" })).toHaveValue(""); + + await user.click(screen.getByText("Advanced: Reminder Markers")); + expect(screen.getByLabelText("Opening delimiter")).toHaveValue(""); + expect(screen.getByLabelText("Closing delimiter")).toHaveValue(""); + + await user.click(screen.getByText("Advanced: Response Format")); + const maxTokensSwitch = screen.getByRole("switch", { name: "Cap max_tokens at the tier model's output ceiling" }); + await user.click(maxTokensSwitch); + await user.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + + expect(savedConfig()).not.toHaveProperty("max_tokens_from_tier_model"); + expect(savedConfig()).toMatchObject({ + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + }); + }); + + it("does not PATCH when the edit is cancelled", async () => { + const user = userEvent.setup(); + const onCancel = vi.fn(); + renderAdvancedModal({ onCancel }); + await user.click(screen.getByRole("button", { name: /cancel/i })); + expect(onCancel).toHaveBeenCalledOnce(); + expect(modelPatchUpdateCall).not.toHaveBeenCalled(); + }); + + it("preserves all stored advanced fields through an untouched save", async () => { + const user = userEvent.setup(); + renderAdvancedModal(); + await user.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + expect(savedConfig()).toMatchObject({ + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + }); + }); +}); + describe("EditAutoRouterModal classifier context window", () => { beforeEach(() => { modelPatchUpdateCall.mockClear(); 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 476207553f7..e3a2df39b88 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 @@ -1,13 +1,9 @@ import AutoRouterClassifierTabs from "../add_model/AutoRouterClassifierTabs"; import { usesClassifierContext } from "../add_model/classifier_types"; -import { defaultJevClassifierConfig, jevClassifierConfigSchema } from "../add_model/jev_classifier_config"; -import type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; export type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; import { getForecastConfigError, isForecastClassifier, - capabilitySettingsSchema, - fuseSettingsSchema, } from "../add_model/forecast_classifier_config"; import React, { useEffect, useMemo, useState } from "react"; import { @@ -30,13 +26,10 @@ import ModelChoiceCombobox, { type ModelChoice } from "../add_model/ModelChoiceC import { modelAvailableCall, modelPatchUpdateCall, validateAutoRouterConfig } from "../networking"; import { fetchAutoRouterModels, fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder, { type RouterConfig, serializeRouterConfig } from "../add_model/RouterConfigBuilder"; -import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; import { - type ActiveTierSet, CUSTOM_TIER_OMITTED_KEYS, activeTierRows, getCustomTierRowsError, - tierParamsByRowId, resolveComplexityDefaultModel, } from "../add_model/tier_rows"; import { isComplexityRouter } from "../add_model/auto_router_strategies"; @@ -53,10 +46,6 @@ import { getSemanticConfigError, getPlanModeTierError, getTierLabelsError, - hydrateBuiltInTiers, - hydrateCustomTierSet, - hydratePlanModeMinTier, - hydrateTierLabels, dryRunRejection, } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; @@ -68,22 +57,14 @@ import { hydrateAutoRouterCompression, } from "../add_model/buildAutoRouterCompression"; import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords"; -import { customDimensionsError, hydrateCustomDimensions } from "../add_model/custom_dimensions"; -import { - hydrateDimensionWeights, - hydrateReasoningOverrideMinScore, - hydrateTierBoundaries, - hydrateTokenThresholds, -} from "../add_model/heuristic_scoring_knobs"; +import { customDimensionsError } from "../add_model/custom_dimensions"; import ComplexityRouterConfig, { ComplexityRouterConfigValue, effectiveClassifierType, heuristicScoringRole, - DEFAULT_ADAPTIVE_WEIGHTS, - DEFAULT_SESSION_AFFINITY, - DEFAULT_DEPLOYMENT_AFFINITY, - DEFAULT_TIER_DISTANCE_PENALTY, } from "../add_model/ComplexityRouterConfig"; +import { builderParamsFromValue } from "../add_model/complexity_router_builder_params"; +import { hydrateComplexityRouterConfig, hydratePinnedDefaultModel } from "./hydrate_complexity_router_config"; import { Dialog, DialogContent, @@ -106,151 +87,7 @@ interface EditAutoRouterModalProps { // Keys this modal rewrites from its own form state on save. Anything absent from this set is // carried through untouched from the stored config, so a key only belongs here once the modal // actually renders a control that can set it. - -/** - * The stored complexity_router_config as form state. Every key in MANAGED_COMPLEXITY_ROUTER_KEYS is - * rewritten from this state on save, so a key missing here is silently dropped from the saved config. - */ -export const hydrateComplexityRouterConfig = ( - parsedConfig: StoredComplexityRouterConfig, - complexityRouterDefaultModel: string | null | undefined, -): ComplexityRouterConfigValue => { - const stringList = (input: unknown): string[] | undefined => - Array.isArray(input) ? input.filter((item): item is string => typeof item === "string") : undefined; - const builtIn = hydrateBuiltInTiers(parsedConfig.tiers, parsedConfig.enable_non_reasoning_tier); - const { tiers: hydratedTiers, enable_non_reasoning_tier } = builtIn; - const custom_tier_set = hydrateCustomTierSet(parsedConfig); - const activeTiers = { ...builtIn, custom_tier_set }; - - return { - tiers: hydratedTiers, - enable_non_reasoning_tier, - custom_tier_set, - tier_model_params: tierParamsByRowId( - hydrateTierModelParams(parsedConfig.tiers, parsedConfig.tier_model_configs), - activeTierRows(activeTiers), - ), - default_model: hydratePinnedDefaultModel(parsedConfig.default_model, complexityRouterDefaultModel, activeTiers), - 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_type === "jev" ? undefined : parsedConfig.classifier_llm_config, - jev_classifier_config: - parsedConfig.classifier_type === "jev" - ? jevClassifierConfigSchema.safeParse(parsedConfig.jev_classifier_config ?? {}).data ?? - defaultJevClassifierConfig() - : undefined, - classifier_context_window_size: - typeof parsedConfig.classifier_context_window_size === "number" - ? parsedConfig.classifier_context_window_size - : undefined, - classifier_context_budget_chars: - typeof parsedConfig.classifier_context_budget_chars === "number" - ? parsedConfig.classifier_context_budget_chars - : undefined, - classifier_context_per_turn_chars: - typeof parsedConfig.classifier_context_per_turn_chars === "number" - ? parsedConfig.classifier_context_per_turn_chars - : undefined, - classifier_context_include_assistant_turns: - typeof parsedConfig.classifier_context_include_assistant_turns === "boolean" - ? parsedConfig.classifier_context_include_assistant_turns - : undefined, - classifier_fallback: - parsedConfig.classifier_fallback === "default_model" || parsedConfig.classifier_fallback === "heuristic" - ? parsedConfig.classifier_fallback - : undefined, - classification_prompt: - typeof parsedConfig.classification_prompt === "string" && parsedConfig.classification_prompt.trim() !== "" - ? parsedConfig.classification_prompt - : undefined, - classification_examples: - typeof parsedConfig.classification_examples === "string" && parsedConfig.classification_examples.trim() !== "" - ? parsedConfig.classification_examples - : undefined, - heuristic_first_max_tier: - typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" - ? parsedConfig.heuristic_first_max_tier - : undefined, - hybrid_boundary_margin: - typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, - classification_mode: - parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" - ? parsedConfig.classification_mode - : undefined, - tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), - token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), - dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), - custom_dimensions: hydrateCustomDimensions(parsedConfig.custom_dimensions), - reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), - session_affinity: - typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, - session_affinity_ttl_seconds: - typeof parsedConfig.session_affinity_ttl_seconds === "number" && - Number.isFinite(parsedConfig.session_affinity_ttl_seconds) - ? parsedConfig.session_affinity_ttl_seconds - : undefined, - modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, - modality_pin_override: - typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false, - deployment_affinity: - typeof parsedConfig.deployment_affinity === "boolean" - ? parsedConfig.deployment_affinity - : DEFAULT_DEPLOYMENT_AFFINITY, - adaptive: parsedConfig.adaptive || false, - adaptive_weights: parsedConfig.adaptive_weights, - tier_distance_penalty: parsedConfig.tier_distance_penalty, - adaptive_eligible: parsedConfig.adaptive_eligible || "all", - return_raw_model_name: parsedConfig.return_raw_model_name || false, - enable_context_window_escalation: - typeof parsedConfig.enable_context_window_escalation === "boolean" - ? parsedConfig.enable_context_window_escalation - : undefined, - context_window_escalation_buffer: - typeof parsedConfig.context_window_escalation_buffer === "number" - ? parsedConfig.context_window_escalation_buffer - : undefined, - stall_escalation_enabled: parsedConfig.stall_escalation_enabled === true || undefined, - stall_escalation_window: - typeof parsedConfig.stall_escalation_window === "number" ? parsedConfig.stall_escalation_window : undefined, - stall_escalation_repeat_threshold: - typeof parsedConfig.stall_escalation_repeat_threshold === "number" - ? parsedConfig.stall_escalation_repeat_threshold - : undefined, - code_keywords: stringList(parsedConfig.code_keywords), - reasoning_keywords: stringList(parsedConfig.reasoning_keywords), - technical_keywords: stringList(parsedConfig.technical_keywords), - simple_keywords: stringList(parsedConfig.simple_keywords), - plan_mode_patterns: stringList(parsedConfig.plan_mode_patterns), - route_housekeeping_to_cheapest_tier: - typeof parsedConfig.route_housekeeping_to_cheapest_tier === "boolean" - ? parsedConfig.route_housekeeping_to_cheapest_tier - : undefined, - housekeeping_patterns: stringList(parsedConfig.housekeeping_patterns), - reminder_markers: Array.isArray(parsedConfig.reminder_markers) - ? parsedConfig.reminder_markers.filter( - (pair): pair is { open: string; close: string } => - typeof pair === "object" && - pair !== null && - typeof (pair as { open?: unknown }).open === "string" && - typeof (pair as { close?: unknown }).close === "string", - ) - : undefined, - max_tokens_from_tier_model: - typeof parsedConfig.max_tokens_from_tier_model === "boolean" ? parsedConfig.max_tokens_from_tier_model : undefined, - classifier_plugin_timeout_ms: - typeof parsedConfig.classifier_plugin_timeout_ms === "number" && Number.isFinite(parsedConfig.classifier_plugin_timeout_ms) - ? parsedConfig.classifier_plugin_timeout_ms - : undefined, - }; -}; - +export { hydrateComplexityRouterConfig, hydratePinnedDefaultModel }; export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "tiers", "enable_non_reasoning_tier", @@ -324,24 +161,6 @@ const toRecord = (value: unknown): Record => { : {}; }; -// A pin lives in two places: complexity_router_config.default_model (this UI's own marker, added -// by PR #36615) and litellm_params.complexity_router_default_model (what the backend reads). Only -// the marker proves an operator picked it, because before #36615 every save wrote a tier-derived -// value into litellm_params. So with no marker, a litellm_params value counts as a pin only when -// it diverges from what the tiers alone derive; a match stays unpinned and keeps tracking tiers. -export const hydratePinnedDefaultModel = ( - storedConfigDefaultModel: unknown, - litellmParamsDefaultModel: string | null | undefined, - activeTiers: ActiveTierSet, -): string | undefined => { - if (typeof storedConfigDefaultModel === "string" && storedConfigDefaultModel.trim()) { - return storedConfigDefaultModel; - } - const tierDerived = resolveComplexityDefaultModel(activeTiers); - const externalOverride = litellmParamsDefaultModel?.trim(); - return externalOverride && externalOverride !== tierDerived ? externalOverride : undefined; -}; - export interface KeywordMatchingState { keywordTierRules: KeywordTierRule[]; escalationKeywords: string[]; @@ -377,65 +196,13 @@ export const buildUpdatedComplexityRouterConfig = ( ); const builderParams: BuildComplexityRouterConfigParams = { - tiers: value.tiers, - enableNonReasoningTier: value.enable_non_reasoning_tier, - customTierSet: value.custom_tier_set, - defaultModel: value.default_model, - planModeMinTier: value.plan_mode_min_tier, - classificationPrompt: value.classification_prompt, - classificationExamples: value.classification_examples, - heuristicFirstMaxTier: value.heuristic_first_max_tier, - hybridBoundaryMargin: value.hybrid_boundary_margin, - classificationMode: value.classification_mode, - tierLabels: value.tier_labels, - classifierType: value.classifier_type, - jevClassifierConfig: value.jev_classifier_config, - heuristicV2SuccessThreshold: value.heuristic_v2_success_threshold, - capabilityClassifierConfig: value.capability_classifier_config, - llmV2Config: value.llm_v2_config, - classifierLlmConfig: value.classifier_llm_config, - classifierContextWindowSize: value.classifier_context_window_size, - classifierContextBudgetChars: value.classifier_context_budget_chars, - classifierContextPerTurnChars: value.classifier_context_per_turn_chars, - classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, - classifierFallback: value.classifier_fallback, - sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, - sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds, - modalityRouting: value.modality_routing ?? false, - modalityPinOverride: value.modality_pin_override ?? false, - deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + ...builderParamsFromValue(value), customTechnicalKeywords: customTechnicalKeywords ?? [], keywordTierRules: keywordMatching?.keywordTierRules ?? [], semanticMatchingEnabled: keywordMatching?.semanticMatchingEnabled ?? false, embeddingModel: keywordMatching?.embeddingModel, matchThreshold: keywordMatching?.matchThreshold ?? DEFAULT_MATCH_THRESHOLD, escalationKeywords: keywordMatching?.escalationKeywords ?? [], - adaptive: value.adaptive ?? false, - adaptiveWeights: value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, - tierDistancePenalty: value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, - adaptiveEligible: value.adaptive_eligible ?? "all", - returnRawModelName: value.return_raw_model_name ?? false, - tierBoundaries: value.tier_boundaries, - tokenThresholds: value.token_thresholds, - dimensionWeights: value.dimension_weights, - customDimensions: value.custom_dimensions, - reasoningOverrideMinScore: value.reasoning_override_min_score, - tierModelParams: value.tier_model_params, - enableContextWindowEscalation: value.enable_context_window_escalation, - contextWindowEscalationBuffer: value.context_window_escalation_buffer, - stallEscalationEnabled: value.stall_escalation_enabled, - stallEscalationWindow: value.stall_escalation_window, - stallEscalationRepeatThreshold: value.stall_escalation_repeat_threshold, - codeKeywords: value.code_keywords, - reasoningKeywords: value.reasoning_keywords, - technicalKeywords: value.technical_keywords, - simpleKeywords: value.simple_keywords, - planModePatterns: value.plan_mode_patterns, - routeHousekeepingToCheapestTier: value.route_housekeeping_to_cheapest_tier, - housekeepingPatterns: value.housekeeping_patterns, - reminderMarkers: value.reminder_markers, - maxTokensFromTierModel: value.max_tokens_from_tier_model, - classifierPluginTimeoutMs: value.classifier_plugin_timeout_ms, }; const built = buildComplexityRouterConfig(builderParams); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts new file mode 100644 index 00000000000..c7c7bcb3382 --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts @@ -0,0 +1,183 @@ +import { defaultJevClassifierConfig, jevClassifierConfigSchema } from "../add_model/jev_classifier_config"; +import { capabilitySettingsSchema, fuseSettingsSchema } from "../add_model/forecast_classifier_config"; +import type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; +import { + hydrateBuiltInTiers, + hydrateCustomTierSet, + hydratePlanModeMinTier, + hydrateTierLabels, +} from "../add_model/build_complexity_router_config"; +import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; +import { hydrateCustomDimensions } from "../add_model/custom_dimensions"; +import { + hydrateDimensionWeights, + hydrateReasoningOverrideMinScore, + hydrateTierBoundaries, + hydrateTokenThresholds, +} from "../add_model/heuristic_scoring_knobs"; +import type { ComplexityRouterConfigValue } from "../add_model/ComplexityRouterConfig"; +import { DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_SESSION_AFFINITY } from "../add_model/ComplexityRouterConfig"; +import { + type ActiveTierSet, + activeTierRows, + tierParamsByRowId, + resolveComplexityDefaultModel, +} from "../add_model/tier_rows"; + +const isReminderMarkerPair = ( + input: unknown, +): input is { open: string; close: string } => + typeof input === "object" && + input !== null && + "open" in input && + "close" in input && + typeof input.open === "string" && + typeof input.close === "string"; + +const stringList = (input: unknown): string[] | undefined => + Array.isArray(input) ? input.filter((item): item is string => typeof item === "string") : undefined; + +export const hydratePinnedDefaultModel = ( + storedConfigDefaultModel: unknown, + litellmParamsDefaultModel: string | null | undefined, + activeTiers: ActiveTierSet, +): string | undefined => { + if (typeof storedConfigDefaultModel === "string" && storedConfigDefaultModel.trim()) { + return storedConfigDefaultModel; + } + const tierDerived = resolveComplexityDefaultModel(activeTiers); + const externalOverride = litellmParamsDefaultModel?.trim(); + return externalOverride && externalOverride !== tierDerived ? externalOverride : undefined; +}; + +export const hydrateComplexityRouterConfig = ( + parsedConfig: StoredComplexityRouterConfig, + complexityRouterDefaultModel: string | null | undefined, +): ComplexityRouterConfigValue => { + const builtIn = hydrateBuiltInTiers(parsedConfig.tiers, parsedConfig.enable_non_reasoning_tier); + const { tiers: hydratedTiers, enable_non_reasoning_tier } = builtIn; + const custom_tier_set = hydrateCustomTierSet(parsedConfig); + const activeTiers = { ...builtIn, custom_tier_set }; + + return { + tiers: hydratedTiers, + enable_non_reasoning_tier, + custom_tier_set, + tier_model_params: tierParamsByRowId( + hydrateTierModelParams(parsedConfig.tiers, parsedConfig.tier_model_configs), + activeTierRows(activeTiers), + ), + default_model: hydratePinnedDefaultModel(parsedConfig.default_model, complexityRouterDefaultModel, activeTiers), + 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_type === "jev" ? undefined : parsedConfig.classifier_llm_config, + jev_classifier_config: + parsedConfig.classifier_type === "jev" + ? jevClassifierConfigSchema.safeParse(parsedConfig.jev_classifier_config ?? {}).data ?? + defaultJevClassifierConfig() + : undefined, + classifier_context_window_size: + typeof parsedConfig.classifier_context_window_size === "number" + ? parsedConfig.classifier_context_window_size + : undefined, + classifier_context_budget_chars: + typeof parsedConfig.classifier_context_budget_chars === "number" + ? parsedConfig.classifier_context_budget_chars + : undefined, + classifier_context_per_turn_chars: + typeof parsedConfig.classifier_context_per_turn_chars === "number" + ? parsedConfig.classifier_context_per_turn_chars + : undefined, + classifier_context_include_assistant_turns: + typeof parsedConfig.classifier_context_include_assistant_turns === "boolean" + ? parsedConfig.classifier_context_include_assistant_turns + : undefined, + classifier_fallback: + parsedConfig.classifier_fallback === "default_model" || parsedConfig.classifier_fallback === "heuristic" + ? parsedConfig.classifier_fallback + : undefined, + classification_prompt: + typeof parsedConfig.classification_prompt === "string" && parsedConfig.classification_prompt.trim() !== "" + ? parsedConfig.classification_prompt + : undefined, + classification_examples: + typeof parsedConfig.classification_examples === "string" && parsedConfig.classification_examples.trim() !== "" + ? parsedConfig.classification_examples + : undefined, + heuristic_first_max_tier: + typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" + ? parsedConfig.heuristic_first_max_tier + : undefined, + hybrid_boundary_margin: + typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, + classification_mode: + parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" + ? parsedConfig.classification_mode + : undefined, + tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), + token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), + dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), + custom_dimensions: hydrateCustomDimensions(parsedConfig.custom_dimensions), + reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), + session_affinity: + typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, + session_affinity_ttl_seconds: + typeof parsedConfig.session_affinity_ttl_seconds === "number" && + Number.isFinite(parsedConfig.session_affinity_ttl_seconds) + ? parsedConfig.session_affinity_ttl_seconds + : undefined, + modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, + modality_pin_override: + typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false, + deployment_affinity: + typeof parsedConfig.deployment_affinity === "boolean" + ? parsedConfig.deployment_affinity + : DEFAULT_DEPLOYMENT_AFFINITY, + adaptive: parsedConfig.adaptive || false, + adaptive_weights: parsedConfig.adaptive_weights, + tier_distance_penalty: parsedConfig.tier_distance_penalty, + adaptive_eligible: parsedConfig.adaptive_eligible || "all", + return_raw_model_name: parsedConfig.return_raw_model_name || false, + enable_context_window_escalation: + typeof parsedConfig.enable_context_window_escalation === "boolean" + ? parsedConfig.enable_context_window_escalation + : undefined, + context_window_escalation_buffer: + typeof parsedConfig.context_window_escalation_buffer === "number" + ? parsedConfig.context_window_escalation_buffer + : undefined, + stall_escalation_enabled: parsedConfig.stall_escalation_enabled === true || undefined, + stall_escalation_window: + typeof parsedConfig.stall_escalation_window === "number" ? parsedConfig.stall_escalation_window : undefined, + stall_escalation_repeat_threshold: + typeof parsedConfig.stall_escalation_repeat_threshold === "number" + ? parsedConfig.stall_escalation_repeat_threshold + : undefined, + code_keywords: stringList(parsedConfig.code_keywords), + reasoning_keywords: stringList(parsedConfig.reasoning_keywords), + technical_keywords: stringList(parsedConfig.technical_keywords), + simple_keywords: stringList(parsedConfig.simple_keywords), + plan_mode_patterns: stringList(parsedConfig.plan_mode_patterns), + route_housekeeping_to_cheapest_tier: + typeof parsedConfig.route_housekeeping_to_cheapest_tier === "boolean" + ? parsedConfig.route_housekeeping_to_cheapest_tier + : undefined, + housekeeping_patterns: stringList(parsedConfig.housekeeping_patterns), + reminder_markers: Array.isArray(parsedConfig.reminder_markers) + ? parsedConfig.reminder_markers.filter(isReminderMarkerPair) + : undefined, + max_tokens_from_tier_model: + typeof parsedConfig.max_tokens_from_tier_model === "boolean" ? parsedConfig.max_tokens_from_tier_model : undefined, + classifier_plugin_timeout_ms: + typeof parsedConfig.classifier_plugin_timeout_ms === "number" && Number.isFinite(parsedConfig.classifier_plugin_timeout_ms) + ? parsedConfig.classifier_plugin_timeout_ms + : undefined, + }; +}; From 68074da1d1658e2c1a8337e090ceccf8488994c4 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:39:47 -0700 Subject: [PATCH 032/109] test(mcp): cover stable server ordering and sort priorities --- .../test_mcp_management_endpoints.py | 83 ++++++++++--------- .../_components/mcp_servers.test.tsx | 72 +++++++++++++++- 2 files changed, 112 insertions(+), 43 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index a526bfc90aa..b745c35ba1f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6,7 +6,7 @@ import logging from contextlib import ExitStack from datetime import datetime, timedelta from types import SimpleNamespace -from typing import List, Optional +from typing import Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1539,54 +1539,57 @@ class TestTeamScopedMCPServerAccess: class TestFetchAllMCPServersOrdering: - def test_display_order_is_case_insensitive_name_then_id(self): - servers = [ - generate_mock_mcp_server_db_record(server_id="s-2", alias="github"), - generate_mock_mcp_server_db_record(server_id="s-1", alias="github"), - generate_mock_mcp_server_db_record(server_id="s-0", alias="Slack"), - generate_mock_mcp_server_db_record(server_id="s-3", alias="confluence"), - ] + def test_display_order_is_case_insensitive_name_then_id(self) -> None: + servers: Final = ( + LiteLLM_MCPServerTable(server_id="s-2", server_name="GitHub", alias="aaa", transport=MCPTransport.http), + LiteLLM_MCPServerTable(server_id="s-1", alias="github", transport=MCPTransport.http), + LiteLLM_MCPServerTable(server_id="s-0", server_name="Slack", alias="zzz", transport=MCPTransport.http), + LiteLLM_MCPServerTable(server_id="confluence", server_name="", alias="", transport=MCPTransport.http), + ) - ordered = sorted(servers, key=mgmt_endpoints._mcp_server_display_order) - assert [s.server_id for s in ordered] == ["s-3", "s-1", "s-2", "s-0"] + ordered: Final = sorted(servers, key=mgmt_endpoints._mcp_server_display_order) + assert [s.server_id for s in ordered] == ["confluence", "s-1", "s-2", "s-0"] + @pytest.mark.parametrize("team_id", [None, "team-1"]) + @pytest.mark.parametrize("reverse", [False, True]) @pytest.mark.asyncio - async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(self): - """The registry resolves ids through a set, so the response must impose its own order.""" - mock_user_auth = generate_mock_user_api_key_auth( + async def test_list_is_sorted_by_display_name_regardless_of_resolution_order( + self, team_id: str | None, reverse: bool + ) -> None: + mock_user_auth: Final = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", ) - first_order = [ + servers: Final = ( generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"), generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"), generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"), - ] - second_order = list(reversed(first_order)) - - for resolved in (first_order, second_order): - mock_manager = MagicMock() - mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved) - with ( - patch( # test-quality-ok: the route reads a module-global manager with no injection seam - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), - patch( # test-quality-ok: admin view is derived from module-global proxy settings - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), - patch( # test-quality-ok: auth contexts need a live prisma client - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), - ), - ): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - fetch_all_mcp_servers, - ) - - result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth) - assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"] + ) + resolved: Final = list(reversed(servers) if reverse else servers) + mock_manager: Final = MagicMock() + mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved) + with ( + patch( # test-quality-ok: the route reads a module-global manager with no injection seam + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( # test-quality-ok: admin view is derived from module-global proxy settings + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( # test-quality-ok: auth contexts need a live prisma client + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), + patch( # test-quality-ok: isolate the route's ordering from team database resolution + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list", + AsyncMock(return_value=resolved), + ), + ): + result: Final = await mgmt_endpoints.fetch_all_mcp_servers( + user_api_key_dict=mock_user_auth, team_id=team_id + ) + assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"] @pytest.mark.asyncio async def test_restricted_virtual_key_cannot_use_team_id_filter(self): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 94a2c75016d..61880363234 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -3,7 +3,7 @@ import { render, waitFor, screen, act, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import MCPServers, { compareServers } from "./mcp_servers"; +import MCPServers, { compareServers, type SortKey } from "./mcp_servers"; import type { MCPServer } from "@/components/mcp_tools/types"; import * as networking from "@/components/networking"; @@ -33,8 +33,14 @@ const createQueryClient = () => }); describe("compareServers", () => { - const server = (server_id: string, name: string, created_at = ""): MCPServer => - ({ server_id, server_name: name, created_at, updated_at: created_at }) as MCPServer; + const server = (server_id: string, name: string, created_at = ""): MCPServer => ({ + server_id, + server_name: name, + created_at, + updated_at: created_at, + created_by: "user", + updated_by: "user", + }); const shuffled = [server("c", "github"), server("a", "slack"), server("b", "Jira")]; @@ -56,6 +62,66 @@ describe("compareServers", () => { "old", ]); }); + + it.each(["created_desc", "updated_desc", "name_asc", "health"])( + "breaks equal timestamps and names by ID for %s regardless of input order", + (sort) => { + const servers = [ + server("b", "GitHub", "2026-01-01T00:00:00Z"), + server("c", "Slack", "2026-01-01T00:00:00Z"), + server("a", "github", "2026-01-01T00:00:00Z"), + ]; + for (const input of [servers, [...servers].reverse()]) { + expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual(["a", "b", "c"]); + } + }, + ); + + it("uses the display name before alias, then falls back to alias and ID", () => { + const servers: MCPServer[] = [ + { ...server("s-slack", "Slack"), alias: "aaa" }, + { ...server("s-github", ""), server_name: null, alias: "GitHub" }, + { ...server("confluence", ""), alias: "" }, + ]; + for (const input of [servers, [...servers].reverse()]) { + expect([...input].sort((a, b) => compareServers(a, b, "name_asc")).map((s) => s.server_id)).toEqual([ + "confluence", + "s-github", + "s-slack", + ]); + } + }); + + it.each(["created_desc", "updated_desc", "health"])( + "keeps timestamped servers before missing timestamps for %s", + (sort) => { + const servers = [ + server("config", "aaa"), + server("older", "bbb", "2026-01-01T00:00:00Z"), + server("newer", "zzz", "2026-02-01T00:00:00Z"), + ]; + for (const input of [servers, [...servers].reverse()]) { + expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual([ + "newer", + "older", + "config", + ]); + } + }, + ); + + it("sorts health before recency and display name", () => { + const servers: MCPServer[] = [ + { ...server("healthy", "aaa", "2026-03-01T00:00:00Z"), status: "healthy" }, + { ...server("unknown", "bbb", "2026-02-01T00:00:00Z"), status: "unknown" }, + { ...server("unhealthy", "zzz", "2026-01-01T00:00:00Z"), status: "unhealthy" }, + ]; + expect(servers.sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id)).toEqual([ + "unhealthy", + "unknown", + "healthy", + ]); + }); }); describe("MCPServers", () => { From ee07f710630bdadb7035f08cfb9cd728ec4425db Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:43:11 +0000 Subject: [PATCH 033/109] style(proxy): format scheduled job timeout configuration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 8 ++++++-- litellm/proxy/shutdown/scheduled_jobs.py | 4 +--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 842adf62f6b..1971c336a96 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1742,8 +1742,12 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300")) SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30")) SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000")) -SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5")) -SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5")) +SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float( + os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5") +) +SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float( + os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5") +) TOOL_SPEND_TOP_TOOLS: Final = 100 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index e920ce19eb9..7889c35cf4e 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -64,9 +64,7 @@ async def stop_in_flight_scheduler_jobs( len(in_flight), ) still_running: Final = ( - (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] - if in_flight - else frozenset() + (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] if in_flight else frozenset() ) scheduler.shutdown(wait=False) if not still_running: From f70683ae92748a9e43c5c0faade2e81275e23a8c Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 19:56:30 +0000 Subject: [PATCH 034/109] fix(ui): satisfy complexity router CI lint budgets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../ComplexityRouterAdvancedSections.tsx | 2 +- .../add_model/add_auto_router_tab.tsx | 26 +++++++++--------- .../build_complexity_router_config.ts | 27 ++++++++++--------- .../edit_auto_router_modal.tsx | 13 +++++---- .../hydrate_complexity_router_config.ts | 16 ++++++----- 5 files changed, 45 insertions(+), 39 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx index dd822907735..0412298ccdd 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx @@ -28,7 +28,7 @@ interface ComplexityRouterAdvancedSectionsProps { onChange: (value: ComplexityRouterConfigValue) => void; forecast: boolean; modelOptions: { value: string; label: string }[]; - classifierEffortOptionsByModel: Record; + classifierEffortOptionsByModel: Record; customTechnicalKeywords?: string[]; onCustomTechnicalKeywordsChange?: (keywords: string[]) => void; showValidationErrors: boolean; 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 f805a5b5511..9be49edf08d 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 @@ -408,6 +408,17 @@ const AddAutoRouterTab: React.FC = ({ matchThreshold, escalationKeywords, }; + const jevRequestParams = + effectiveClassifierType(complexityRouterConfig) === "jev" + ? { + prompt: JEV_CONNECTION_TEST_PROMPT, + config: buildComplexityRouterConfig(complexityRouterConfigParams), + defaultModel: resolveComplexityDefaultModel(complexityRouterConfig, complexityRouterConfig.default_model), + routerName: watchedName, + teamId: requiresTeamScope ? watchedTeamId ?? undefined : undefined, + } + : undefined; + const jevRequest = jevRequestParams ? buildAutoRouterRoutingTestRequest(jevRequestParams) : undefined; const submitRecommendedRouter = async (name: string) => { // The one answer the submit button reads, so a disabled button and a refused submit cannot @@ -816,20 +827,7 @@ const AddAutoRouterTab: React.FC = ({ testId={connectionTestId} accessToken={accessToken} targets={testTargets} - jevRequest={ - effectiveClassifierType(complexityRouterConfig) === "jev" - ? buildAutoRouterRoutingTestRequest({ - prompt: JEV_CONNECTION_TEST_PROMPT, - config: buildComplexityRouterConfig(complexityRouterConfigParams), - defaultModel: resolveComplexityDefaultModel( - complexityRouterConfig, - complexityRouterConfig.default_model, - ), - routerName: watchedName, - teamId: requiresTeamScope ? watchedTeamId ?? undefined : undefined, - }) - : undefined - } + jevRequest={jevRequest} onTestComplete={() => setIsTestingConnection(false)} /> 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 ab71b6b3cce..4782a58513d 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 @@ -744,16 +744,22 @@ export const buildComplexityRouterConfig = ({ open: open.trim().toLowerCase(), close: close.trim().toLowerCase(), })); + const cleanedListValues = { + code_keywords: cleanList(codeKeywords), + reasoning_keywords: cleanList(reasoningKeywords), + technical_keywords: cleanList(technicalKeywords), + simple_keywords: cleanList(simpleKeywords), + plan_mode_patterns: cleanList(planModePatterns), + housekeeping_patterns: cleanList(housekeepingPatterns), + }; const cleanedLists = Object.fromEntries( - Object.entries({ - code_keywords: cleanList(codeKeywords), - reasoning_keywords: cleanList(reasoningKeywords), - technical_keywords: cleanList(technicalKeywords), - simple_keywords: cleanList(simpleKeywords), - plan_mode_patterns: cleanList(planModePatterns), - housekeeping_patterns: cleanList(housekeepingPatterns), - }).filter(([, list]) => list !== undefined), + Object.entries(cleanedListValues).filter(([, list]) => list !== undefined), ); + const hasValidCustomClassifierTimeout = + classifierType === "custom" && + classifierPluginTimeoutMs !== undefined && + Number.isInteger(classifierPluginTimeoutMs) && + classifierPluginTimeoutMs > 0; const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType); const payload: ComplexityRouterConfigPayload = { @@ -826,10 +832,7 @@ export const buildComplexityRouterConfig = ({ ...(routeHousekeepingToCheapestTier === false && { route_housekeeping_to_cheapest_tier: false }), ...(cleanedReminderMarkers && cleanedReminderMarkers.length > 0 && { reminder_markers: cleanedReminderMarkers }), ...(maxTokensFromTierModel === false && { max_tokens_from_tier_model: false }), - ...(classifierType === "custom" && - classifierPluginTimeoutMs !== undefined && - Number.isInteger(classifierPluginTimeoutMs) && - classifierPluginTimeoutMs > 0 && { classifier_plugin_timeout_ms: classifierPluginTimeoutMs }), + ...(hasValidCustomClassifierTimeout && { classifier_plugin_timeout_ms: classifierPluginTimeoutMs }), ...scorerKnobs, }; if (!customTierSet) return payload; 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 e3a2df39b88..e6f1ccd0e17 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 @@ -373,12 +373,13 @@ const EditAutoRouterModal: React.FC = ({ setRouterConfig(parsedConfig); // Set form values - form.reset({ + const routerFormValues = { auto_router_name: modelData.model_name, auto_router_default_model: modelData.litellm_params?.auto_router_default_model || null, auto_router_embedding_model: modelData.litellm_params?.auto_router_embedding_model || null, model_access_group: modelData.model_info?.access_groups || [], - }); + }; + form.reset(routerFormValues); } catch (error) { console.error("Error parsing auto router config:", error); toast.fromError("Error loading auto router configuration"); @@ -456,11 +457,12 @@ const EditAutoRouterModal: React.FC = ({ // Dual write: complexity_router_config.default_model (the pin marker hydratePinnedDefaultModel // reads back) and complexity_router_default_model (what the backend routes on) must always be // written together from the same value. Same pairing in add_auto_router_tab.tsx. + const keywordMatching = { keywordTierRules, escalationKeywords, semanticMatchingEnabled, embeddingModel, matchThreshold }; const updatedConfig = buildUpdatedComplexityRouterConfig( modelData.litellm_params?.complexity_router_config, complexityRouterConfig, customTechnicalKeywords, - { keywordTierRules, escalationKeywords, semanticMatchingEnabled, embeddingModel, matchThreshold }, + keywordMatching, ); const serverVerdict = await validateAutoRouterConfig(accessToken, updatedConfig, modelData?.model_info?.team_id); const dryRunError = dryRunRejection(serverVerdict); @@ -497,12 +499,13 @@ const EditAutoRouterModal: React.FC = ({ ); toast.success("Auto router configuration updated successfully"); - onSuccess({ + const updatedModelData = { ...modelData, model_name: values.auto_router_name, litellm_params: updatedLitellmParams, model_info: updatedModelInfo, - }); + }; + onSuccess(updatedModelData); onCancel(); return; } diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts index c7c7bcb3382..b3a9ae0a504 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts @@ -26,13 +26,15 @@ import { const isReminderMarkerPair = ( input: unknown, -): input is { open: string; close: string } => - typeof input === "object" && - input !== null && - "open" in input && - "close" in input && - typeof input.open === "string" && - typeof input.close === "string"; +): input is { open: string; close: string } => { + if (typeof input !== "object" || input === null) { + return false; + } + if (!("open" in input) || !("close" in input)) { + return false; + } + return typeof input.open === "string" && typeof input.close === "string"; +}; const stringList = (input: unknown): string[] | undefined => Array.isArray(input) ? input.filter((item): item is string => typeof item === "string") : undefined; From 4117d9f7860bc7bba6250682654f5cddeca4cd3f Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 20:01:08 +0000 Subject: [PATCH 035/109] style(ui): format complexity router files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../ComplexityRouterAdvancedSections.tsx | 8 ++++---- .../add_model/ComplexityRouterConfig.test.tsx | 12 ++---------- .../add_model/HeuristicKeywordOverrides.tsx | 4 ++-- .../components/add_model/ReminderMarkers.tsx | 17 ++++++++++++----- .../add_model/ResponseFormatControls.tsx | 4 ++-- .../add_model/build_complexity_router_config.ts | 4 +--- .../edit_auto_router/edit_auto_router_modal.tsx | 13 ++++++++----- .../hydrate_complexity_router_config.ts | 11 ++++++----- 8 files changed, 37 insertions(+), 36 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx index 0412298ccdd..71f8bce76b5 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterAdvancedSections.tsx @@ -130,7 +130,9 @@ const ComplexityRouterAdvancedSections: React.FCAdvanced: Plan-Mode Override, - children: , + children: ( + + ), }, { key: "housekeeping", @@ -195,9 +197,7 @@ const ComplexityRouterAdvancedSections: React.FC )} {onKeywordTierRulesChange && onSemanticMatchingEnabledChange && } 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 9a8577100df..56b37b92af3 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -106,7 +106,6 @@ describe("ComplexityRouterConfig", () => { const capabilityValue = { ...defaultValue, classifier_type: "capability" as const }; rerender(); expect(screen.queryByText("Advanced: Heuristic Keyword Overrides")).not.toBeInTheDocument(); - }); it.each([ @@ -127,11 +126,7 @@ describe("ComplexityRouterConfig", () => { it.each([true, false])("shows reminder marker validation only when requested: %s", (showValidationErrors) => { const value = { ...defaultValue, reminder_markers: [{ open: "", close: "x" }] }; renderWithProviders( - , + , ); fireEvent.click(screen.getByText("Advanced: Reminder Markers")); const validation = screen.queryByText(/needs both/i); @@ -144,10 +139,7 @@ describe("ComplexityRouterConfig", () => { it("disables housekeeping sentinels when cheapest-tier routing is off", () => { renderWithProviders( - , + , ); fireEvent.click(screen.getByText("Advanced: Housekeeping Routing")); const sentinelInput = screen.getByRole("combobox", { name: "e.g., conversation title" }); diff --git a/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx b/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx index 696924f4e77..185f187bbfc 100644 --- a/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx +++ b/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx @@ -15,8 +15,8 @@ const HeuristicKeywordOverrides: React.FC<{ }> = ({ value, onChange }) => (

- Each list replaces the built-in keyword list of the same name for the heuristic scorer. Leave a list empty to - keep the built-in one. To add technical terms without replacing the list, use custom technical keywords under + Each list replaces the built-in keyword list of the same name for the heuristic scorer. Leave a list empty to keep + the built-in one. To add technical terms without replacing the list, use custom technical keywords under Classification Method.

{fields.map(([key, label]) => { diff --git a/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx index 970394291d4..c7f9ee48e0b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx @@ -14,7 +14,9 @@ const ReminderMarkers: React.FC<{ const update = (index: number, patch: Partial) => onChange({ ...value, - reminder_markers: markers.map((marker, markerIndex) => (markerIndex === index ? { ...marker, ...patch } : marker)), + reminder_markers: markers.map((marker, markerIndex) => + markerIndex === index ? { ...marker, ...patch } : marker, + ), }); const remove = (index: number) => { const next = markers.filter((_, markerIndex) => markerIndex !== index); @@ -24,9 +26,9 @@ const ReminderMarkers: React.FC<{ return (

- Delimiter pairs that wrap harness-injected reminder blocks, which are stripped before classification. Setting any - pair replaces the built-in pairs, so list every pair your harness emits. Matching is case-insensitive and values - are saved lowercased. + Delimiter pairs that wrap harness-injected reminder blocks, which are stripped before classification. Setting + any pair replaces the built-in pairs, so list every pair your harness emits. Matching is case-insensitive and + values are saved lowercased.

{markers.map((marker, index) => ( @@ -53,7 +55,12 @@ const ReminderMarkers: React.FC<{ onChange={(event) => update(index, { close: event.target.value })} />
-
diff --git a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx index 9edbf204a6d..6e7b1489a4c 100644 --- a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx @@ -27,8 +27,8 @@ const ResponseFormatControls: React.FC<{ Cap max_tokens at the tier model's output ceiling
- Replace the caller's max_tokens with the routed tier model's output ceiling so one client value fits every - tier. Off forwards the caller's value unchanged. + Replace the caller's max_tokens with the routed tier model's output ceiling so one client value fits + every tier. Off forwards the caller's value unchanged. ); 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 4782a58513d..ca46e171970 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 @@ -752,9 +752,7 @@ export const buildComplexityRouterConfig = ({ plan_mode_patterns: cleanList(planModePatterns), housekeeping_patterns: cleanList(housekeepingPatterns), }; - const cleanedLists = Object.fromEntries( - Object.entries(cleanedListValues).filter(([, list]) => list !== undefined), - ); + const cleanedLists = Object.fromEntries(Object.entries(cleanedListValues).filter(([, list]) => list !== undefined)); const hasValidCustomClassifierTimeout = classifierType === "custom" && classifierPluginTimeoutMs !== undefined && 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 e6f1ccd0e17..0dafa9b330a 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 @@ -1,10 +1,7 @@ import AutoRouterClassifierTabs from "../add_model/AutoRouterClassifierTabs"; import { usesClassifierContext } from "../add_model/classifier_types"; export type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; -import { - getForecastConfigError, - isForecastClassifier, -} from "../add_model/forecast_classifier_config"; +import { getForecastConfigError, isForecastClassifier } from "../add_model/forecast_classifier_config"; import React, { useEffect, useMemo, useState } from "react"; import { complexityRouterSchema, @@ -457,7 +454,13 @@ const EditAutoRouterModal: React.FC = ({ // Dual write: complexity_router_config.default_model (the pin marker hydratePinnedDefaultModel // reads back) and complexity_router_default_model (what the backend routes on) must always be // written together from the same value. Same pairing in add_auto_router_tab.tsx. - const keywordMatching = { keywordTierRules, escalationKeywords, semanticMatchingEnabled, embeddingModel, matchThreshold }; + const keywordMatching = { + keywordTierRules, + escalationKeywords, + semanticMatchingEnabled, + embeddingModel, + matchThreshold, + }; const updatedConfig = buildUpdatedComplexityRouterConfig( modelData.litellm_params?.complexity_router_config, complexityRouterConfig, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts index b3a9ae0a504..6dbd2b19b52 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts @@ -24,9 +24,7 @@ import { resolveComplexityDefaultModel, } from "../add_model/tier_rows"; -const isReminderMarkerPair = ( - input: unknown, -): input is { open: string; close: string } => { +const isReminderMarkerPair = (input: unknown): input is { open: string; close: string } => { if (typeof input !== "object" || input === null) { return false; } @@ -176,9 +174,12 @@ export const hydrateComplexityRouterConfig = ( ? parsedConfig.reminder_markers.filter(isReminderMarkerPair) : undefined, max_tokens_from_tier_model: - typeof parsedConfig.max_tokens_from_tier_model === "boolean" ? parsedConfig.max_tokens_from_tier_model : undefined, + typeof parsedConfig.max_tokens_from_tier_model === "boolean" + ? parsedConfig.max_tokens_from_tier_model + : undefined, classifier_plugin_timeout_ms: - typeof parsedConfig.classifier_plugin_timeout_ms === "number" && Number.isFinite(parsedConfig.classifier_plugin_timeout_ms) + typeof parsedConfig.classifier_plugin_timeout_ms === "number" && + Number.isFinite(parsedConfig.classifier_plugin_timeout_ms) ? parsedConfig.classifier_plugin_timeout_ms : undefined, }; From d2f457f144a430ad2848b33839d298d9312c8d4e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:20:31 +0000 Subject: [PATCH 036/109] feat(cache): add semantic cache context and unsupported operation error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/base_cache.rs | 50 +++++++++++++++++++++ litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/cache/src/lib.rs | 2 +- 3 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 8bd69ba5ad6..5c10e7fd5c3 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -32,6 +32,28 @@ impl CacheContext for ExactCacheContext { } } +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Option, + pub metadata: Option, + pub scope: Option, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + ttl, + ..self.clone() + } + } +} + #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum CacheConnectionStatus { @@ -105,3 +127,31 @@ pub trait BaseCache: Send + Sync { fn test_connection(&self) -> impl Future> + Send; } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use serde_json::json; + + use super::{CacheContext, SemanticCacheContext}; + + #[test] + fn semantic_context_with_ttl_only_replaces_ttl() { + let context = SemanticCacheContext { + input: Some(json!({"input": "hello"})), + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"tenant": "team"})), + scope: Some("scope".into()), + ttl: Some(Duration::from_secs(10)), + }; + + let updated = context.with_ttl(Some(Duration::from_secs(20))); + + assert_eq!(updated.ttl, Some(Duration::from_secs(20))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + } +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..51e4fe2d66a 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("cache operation is not supported by this backend")] + UnsupportedOperation, } diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ce9f93b6dc4..8364c635e3a 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -8,7 +8,7 @@ mod error; pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, - ExactCacheContext, + ExactCacheContext, SemanticCacheContext, }; pub use cache_type::CacheType; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; From df97b274fc78ab261064f0a691016488c6c709dc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:21:19 +0000 Subject: [PATCH 037/109] refactor(cache-response): generalize ResponseCache over the backend context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-response/README.md | 2 +- .../crates/cache-response/src/response.rs | 43 +++++++++++-------- .../crates/python-bridge/src/cache/request.rs | 3 +- 3 files changed, 28 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 56c1646d343..46e561ddad1 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..e70e07a5d26 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,21 +1,21 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache, }; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; #[derive(Clone)] -pub struct ResponseCacheRequest { +pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, - pub context: ExactCacheContext, + pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { pub fn new(key: CacheKeyInput) -> Self { Self { key, @@ -26,17 +26,24 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - context: ExactCacheContext::default(), + context: C::default(), max_age: None, } } } -pub struct ResponseCache> { +pub struct ResponseCache> +where + B::Context: Default + PartialEq, +{ backend: Arc, } -impl> ResponseCache { +impl ResponseCache +where + B: BaseCache, + B::Context: Default + PartialEq, +{ pub fn new(backend: Arc) -> Self { Self { backend } } @@ -46,7 +53,7 @@ impl> ResponseCach } pub fn default_ttl(&self) -> Option { - self.backend.get_ttl(&ExactCacheContext::default()) + self.backend.get_ttl(&B::Context::default()) } pub async fn async_flush(&self) -> Result<(), Error> @@ -62,7 +69,7 @@ impl> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +88,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +108,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +133,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +160,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +179,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,7 +200,7 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { self.async_store_entries( @@ -209,7 +216,7 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() @@ -249,8 +256,8 @@ impl> ResponseCach } fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0c5343a63d0..52a5f7d9055 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,5 +1,6 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use litellm_cache::ExactCacheContext; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; @@ -20,7 +21,7 @@ pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult PyResult { - let mut request = ResponseCacheRequest::new(input.key); + let mut request = ResponseCacheRequest::::new(input.key); if let Some(controls) = input.controls { request.controls = controls; } From 8dc960c928614c2276c8f5e7cc8d1658009ba796 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:22:13 +0000 Subject: [PATCH 038/109] feat(cache): add SemanticCacheContext Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/base_cache.rs | 25 +++++++++++++++++++++ litellm-rust/crates/cache/src/lib.rs | 2 +- litellm-rust/crates/cache/tests/caching.rs | 24 +++++++++++++++++++- 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 8bd69ba5ad6..6f961798795 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -32,6 +32,31 @@ impl CacheContext for ExactCacheContext { } } +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Vec, + pub metadata: serde_json::Map, + pub scope: Option, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + input: self.input.clone(), + messages: self.messages.clone(), + metadata: self.metadata.clone(), + scope: self.scope.clone(), + ttl, + } + } +} + #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum CacheConnectionStatus { diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ce9f93b6dc4..8364c635e3a 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -8,7 +8,7 @@ mod error; pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, - ExactCacheContext, + ExactCacheContext, SemanticCacheContext, }; pub use cache_type::CacheType; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 9180ee9d0dc..2e65b4eeae5 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,7 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, + SemanticCacheContext, get_cache, }; struct TestCache { @@ -126,6 +127,27 @@ fn associated_context_preserves_backend_specific_lookup_inputs() { ); } +#[test] +fn semantic_context_with_ttl_preserves_lookup_inputs() { + let context = SemanticCacheContext { + input: Some(serde_json::json!("text")), + messages: vec![serde_json::json!({"role": "user", "content": "hi"})], + metadata: serde_json::Map::from_iter([( + "key".into(), + serde_json::json!("value"), + )]), + scope: Some("scope".into()), + ttl: None, + }; + let updated = context.with_ttl(Some(Duration::from_secs(30))); + assert_eq!(updated.ttl(), Some(Duration::from_secs(30))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + assert_eq!(context.with_ttl(None).ttl(), None); +} + #[tokio::test] async fn default_batch_operations_use_async_writes_and_stop_on_failure() { let cache = TestCache { From 1320eeeb41fbcd2a5842889e39f768d446fa6a05 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:23:29 +0000 Subject: [PATCH 039/109] refactor(cache-response): generalize ResponseCache over the backend context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-response/src/response.rs | 50 +++++++++++-------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..a27cc1967d5 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,21 +1,22 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, + ExactCacheContext, FlushCache, }; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; #[derive(Clone)] -pub struct ResponseCacheRequest { +pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, - pub context: ExactCacheContext, + pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { pub fn new(key: CacheKeyInput) -> Self { Self { key, @@ -32,11 +33,11 @@ impl ResponseCacheRequest { } } -pub struct ResponseCache> { +pub struct ResponseCache> { backend: Arc, } -impl> ResponseCache { +impl> ResponseCache { pub fn new(backend: Arc) -> Self { Self { backend } } @@ -45,8 +46,11 @@ impl> ResponseCach &self.backend } - pub fn default_ttl(&self) -> Option { - self.backend.get_ttl(&ExactCacheContext::default()) + pub fn default_ttl(&self) -> Option + where + B::Context: Default, + { + self.backend.get_ttl(&B::Context::default()) } pub async fn async_flush(&self) -> Result<(), Error> @@ -62,7 +66,7 @@ impl> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +85,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +105,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +130,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +157,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +176,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,9 +197,12 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, - ) -> Result<(), Error> { + ) -> Result<(), Error> + where + B::Context: PartialEq, + { self.async_store_entries( entries .into_iter() @@ -209,8 +216,11 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, - ) -> Result<(), Error> { + entries: Vec<(ResponseCacheRequest, Value, Duration)>, + ) -> Result<(), Error> + where + B::Context: PartialEq, + { let writable = entries .into_iter() .filter(|(request, _, _)| request.controls.writes()) @@ -249,8 +259,8 @@ impl> ResponseCach } fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { From 8d9ab9eeaafe88efdf19fe67809e5fd21b2adb03 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:24:23 +0000 Subject: [PATCH 040/109] feat(cache-redis): expose the pooled connection handling for reuse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-redis/src/cache.rs | 93 +++++++++++-------- .../cache-redis/src/cache/operations.rs | 28 +++--- litellm-rust/crates/cache-redis/src/lib.rs | 4 + 3 files changed, 72 insertions(+), 53 deletions(-) diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index a960c383bf4..6388448accc 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -19,7 +19,7 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; -struct PooledConnection { +pub struct PooledConnection { connection: redis::Connection, failed: bool, } @@ -27,16 +27,19 @@ struct PooledConnection { /// Pools connections without a checkout PING, which would double every operation's round trips. /// A timed-out command leaves its reply on the socket while redis still reports the connection /// open, so any connection whose operation failed is discarded instead of being reused. -struct ConnectionManager(redis::Client); +pub struct ConnectionManager { + client: redis::Client, + timeout: Duration, +} impl r2d2::ManageConnection for ConnectionManager { type Connection = PooledConnection; type Error = redis::RedisError; fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - connection.set_read_timeout(Some(REDIS_TIMEOUT))?; - connection.set_write_timeout(Some(REDIS_TIMEOUT))?; + let connection = self.client.get_connection()?; + connection.set_read_timeout(Some(self.timeout))?; + connection.set_write_timeout(Some(self.timeout))?; Ok(PooledConnection { connection, failed: false, @@ -68,12 +71,12 @@ const CLAIM_SCRIPT: &str = concat!( ); const CLAIM_ATTEMPTS: usize = 8; -enum Connections { +pub enum Connections { Pool(r2d2::Pool), Fixed(Mutex), } -struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); +pub struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); impl redis::ConnectionLike for ConnectionRef<'_> { fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { @@ -110,7 +113,23 @@ impl Connections where C: redis::ConnectionLike + Send + 'static, { - fn execute( + pub fn pooled(url: &str, timeout: Duration, pool_size: u32) -> Result { + let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; + let pool = r2d2::Pool::builder() + .max_size(pool_size) + .min_idle(Some(0)) + .connection_timeout(timeout) + .test_on_check_out(false) + .build(ConnectionManager { client, timeout }) + .map_err(|_| Error::Unavailable)?; + Ok(Self::Pool(pool)) + } + + pub fn fixed(connection: C) -> Self { + Self::Fixed(Mutex::new(connection)) + } + + pub fn execute( &self, operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, ) -> Result { @@ -127,6 +146,16 @@ where } } } + + pub async fn run_blocking(connections: Arc, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, + { + tokio::task::spawn_blocking(move || connections.execute(operation)) + .await + .map_err(|_| Error::Unavailable)? + } } pub struct RedisCache { @@ -138,16 +167,8 @@ pub struct RedisCache { impl RedisCache { pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { - let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; - let pool = r2d2::Pool::builder() - .max_size(REDIS_POOL_SIZE) - .min_idle(Some(0)) - .connection_timeout(REDIS_TIMEOUT) - .test_on_check_out(false) - .build(ConnectionManager(client)) - .map_err(|_| Error::Unavailable)?; Ok(Self { - connections: Arc::new(Connections::Pool(pool)), + connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, @@ -162,7 +183,7 @@ where { pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { - connections: Arc::new(Connections::Fixed(Mutex::new(connection))), + connections: Arc::new(Connections::fixed(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, @@ -241,20 +262,14 @@ where } fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs() - .saturating_add(u64::from(ttl.subsec_nanos() > 0)) - .max(1) + ttl_seconds(ttl) } +} - async fn run_blocking(connections: Arc>, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - tokio::task::spawn_blocking(move || connections.execute(operation)) - .await - .map_err(|_| Error::Unavailable)? - } +pub fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) } fn namespaced_key(namespace: Option<&str>, key: &str) -> String { @@ -313,7 +328,7 @@ where let payload = self.codec.encode(&value)?; let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) @@ -327,7 +342,7 @@ where _: &ExactCacheContext, ) -> Result, Error> { let key = self.namespaced_key(key); - let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) @@ -350,7 +365,7 @@ where }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); for (key, payload) in entries { pipeline @@ -372,7 +387,7 @@ where } async fn test_connection(&self) -> Result { - match Self::run_blocking(Arc::clone(&self.connections), |connection| { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { Ok(match redis::cmd("PING").query::(connection) { Ok(_) => CacheConnectionResult { status: CacheConnectionStatus::Success, @@ -433,7 +448,7 @@ where .iter() .map(|key| self.namespaced_key(key)) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("MGET") .arg(keys) .query::>(connection) @@ -460,7 +475,7 @@ where async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) .await @@ -480,7 +495,7 @@ where async fn async_flush_cache(&self) -> Result<(), Error> { let pattern = self.namespaced_pattern()?; - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { Self::flush_matching(connection, &pattern) }) .await @@ -512,7 +527,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { increment(connection, key, amount, ttl) }) .await @@ -623,7 +638,7 @@ where let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); let codec = self.codec.clone(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { claim(connection, &codec, &key, candidate, &eligible, ttl) }) .await diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index d8d9ae24c4c..f27a7802bab 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -144,7 +144,7 @@ where .into_iter() .map(|key| self.namespaced_key(&key)) .collect::>(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del(keys).map_err(|_| Error::Unavailable) }) .await @@ -172,7 +172,7 @@ where .iter() .map(|key| self.namespaced_key(key)) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("MGET") .arg(keys) .query::>(connection) @@ -192,7 +192,7 @@ where } pub async fn ping(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connections), |connection| { + Connections::run_blocking(Arc::clone(&self.connections), |connection| { redis::cmd("PING") .query::(connection) .map(|response| response == "PONG") @@ -203,7 +203,7 @@ where pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { let key = self.namespaced_key(key); - let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("TTL") .arg(key) .query::(connection) @@ -215,7 +215,7 @@ where pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { let pattern = format!("{}*", self.namespaced_key(pattern)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut cursor = 0u64; let mut matches = Vec::new(); loop { @@ -249,7 +249,7 @@ where } let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); pipeline.cmd("SADD").arg(&key).arg(values); pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore(); @@ -266,7 +266,7 @@ where return Err(Error::InvalidEntry); } let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("RPUSH") .arg(key) .arg(values) @@ -292,7 +292,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); for (key, values) in operations { pipeline.cmd("RPUSH").arg(key).arg(values); @@ -309,7 +309,7 @@ where ) -> Result { let key = self.namespaced_key(key); let multiple = count.is_some(); - let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut command = redis::cmd("LPOP"); command.arg(key); if let Some(count) = count { @@ -338,7 +338,7 @@ where .iter() .map(|(_, count)| count.is_some()) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); for (key, count) in operations { let command = pipeline.cmd("LPOP").arg(key); @@ -368,7 +368,7 @@ where .into_iter() .map(|key| self.namespaced_key(&key)) .collect::>(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("EVAL") .arg(script) .arg(keys.len()) @@ -440,7 +440,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); for (key, amount, ttl) in operations { pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount); @@ -461,7 +461,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { increment_with_floor(connection, key, amount, ttl) }) .await @@ -475,7 +475,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("EVAL") .arg(SET_MAX_SCRIPT) .arg(1) diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 98f6bfd8ce5..ea75906e9c9 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -1,6 +1,10 @@ mod cache; mod topology; +pub mod connection { + pub use crate::cache::{ConnectionRef, Connections, ttl_seconds}; +} + pub use cache::{ RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, }; From 1f86bb8e4640fd7e758e10f66106fd7c1bda01de Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:24:34 +0000 Subject: [PATCH 041/109] feat(cache-valkey-semantic): add native Valkey semantic cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 16 + .../crates/cache-valkey-semantic/Cargo.toml | 20 + .../crates/cache-valkey-semantic/src/lib.rs | 844 ++++++++++++++++++ 3 files changed, 880 insertions(+) create mode 100644 litellm-rust/crates/cache-valkey-semantic/Cargo.toml create mode 100644 litellm-rust/crates/cache-valkey-semantic/src/lib.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..5daf691d4c3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2502,6 +2502,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-valkey-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "rstest", + "serde_json", + "sha2 0.10.9", + "tokio", + "uuid", +] + [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml new file mode 100644 index 00000000000..9a0a566ca3b --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-cache-valkey-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-response.workspace = true +r2d2 = "0.8.10" +redis = { version = "1.7.0", features = ["tls-rustls"] } +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true +uuid = { version = "1", features = ["v4"] } + +[dev-dependencies] +redis-test = "1.0.4" +rstest.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs new file mode 100644 index 00000000000..85c4c9af15c --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -0,0 +1,844 @@ +use std::{ + future::Future, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; +use litellm_cache_response::CacheEntry; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug, PartialEq)] +pub struct ValkeySemanticConfig { + pub similarity_threshold: f64, + pub index_name: String, +} + +pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; + +struct PooledConnection { + connection: redis::Connection, + failed: bool, +} + +struct ConnectionManager(redis::Client); + +impl r2d2::ManageConnection for ConnectionManager { + type Connection = PooledConnection; + type Error = redis::RedisError; + + fn connect(&self) -> Result { + let connection = self.0.get_connection()?; + Ok(PooledConnection { + connection, + failed: false, + }) + } + + fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), Self::Error> { + redis::cmd("PING").query::(&mut connection.connection)?; + Ok(()) + } + + fn has_broken(&self, connection: &mut Self::Connection) -> bool { + connection.failed || !redis::ConnectionLike::is_open(&connection.connection) + } +} + +enum Connections { + Pool(r2d2::Pool), + Fixed(Mutex), +} + +struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); + +impl redis::ConnectionLike for ConnectionRef<'_> { + fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { + self.0.req_packed_command(cmd) + } + + fn req_packed_commands( + &mut self, + cmd: &[u8], + offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.0.req_packed_commands(cmd, offset, count) + } + + fn get_db(&self) -> i64 { + self.0.get_db() + } + + fn supports_pipelining(&self) -> bool { + self.0.supports_pipelining() + } + + fn check_connection(&mut self) -> bool { + self.0.check_connection() + } + + fn is_open(&self) -> bool { + self.0.is_open() + } +} + +impl Connections +where + C: redis::ConnectionLike + Send + 'static, +{ + fn execute( + &self, + operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, + ) -> Result { + match self { + Self::Pool(pool) => { + let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let result = operation(&mut ConnectionRef(&mut pooled.connection)); + pooled.failed = matches!(result, Err(Error::Unavailable)); + result + } + Self::Fixed(connection) => { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + } + } +} + +pub struct ValkeySemanticCache< + E: Embedder, + S: CacheCodec, + C = redis::Connection, +> { + connections: Arc>, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + index_dimension: Arc>>, +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, +{ + pub fn new( + url: &str, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + ) -> Result { + let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; + let pool = r2d2::Pool::builder() + .max_size(16) + .min_idle(Some(0)) + .test_on_check_out(false) + .build(ConnectionManager(client)) + .map_err(|_| Error::Unavailable)?; + Ok(Self { + connections: Arc::new(Connections::Pool(pool)), + embedder, + codec, + config, + index_dimension: Arc::new(Mutex::new(None)), + }) + } +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub fn with_connection( + connection: C, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + ) -> Self { + Self { + connections: Arc::new(Connections::Fixed(Mutex::new(connection))), + embedder, + codec, + config, + index_dimension: Arc::new(Mutex::new(None)), + } + } + + pub fn similarity_threshold(&self) -> f64 { + self.config.similarity_threshold + } + + pub fn index_name(&self) -> &str { + &self.config.index_name + } + + fn key_prefix(&self) -> String { + format!("{}:", self.config.index_name) + } + + fn ensure_index(&self, dimension: usize) -> Result<(), Error> { + ensure_index( + &self.connections, + &self.config.index_name, + &self.key_prefix(), + &self.index_dimension, + dimension, + ) + } +} + +impl BaseCache for ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; + self.ensure_index(embedding.len())?; + let scope = scope_tag(key); + let document = format!("{}{}:{}", self.key_prefix(), scope, Uuid::new_v4()); + let response = self.codec.encode(&value)?; + let vector = embedding_bytes(&embedding); + let ttl = self.get_ttl(context); + self.connections.execute(|connection| { + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(&scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { + pipeline + .cmd("EXPIRE") + .arg(&document) + .arg(ttl.as_secs()) + .ignore(); + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; + self.ensure_index(embedding.len())?; + let scope = scope_tag(key); + let query = + format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); + let vector = embedding_bytes(&embedding); + let response = self.connections.execute(|connection| { + redis::cmd("FT.SEARCH") + .arg(&self.config.index_name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let Some(fields) = search_fields(response)? else { + return Ok(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < self.config.similarity_threshold { + return Ok(None); + } + self.codec.decode(&response).map(Some) + } + + fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> impl Future> + Send { + let key = key.to_owned(); + let prompt = prompt_from_context(&context); + let metadata = context.metadata.clone(); + async move { + let Some(prompt) = prompt else { + return Ok(()); + }; + let embedding = self + .embedder + .async_embed(&prompt, metadata.as_ref()) + .await?; + let connections = Arc::clone(&self.connections); + let config = self.config.clone(); + let index_dimension = Arc::clone(&self.index_dimension); + let response = self.codec.encode(&value)?; + let vector = embedding_bytes(&embedding); + let prefix = format!("{}:", config.index_name); + let scope = scope_tag(&key); + let document = format!("{prefix}{scope}:{}", Uuid::new_v4()); + let ttl = context.ttl; + tokio::task::spawn_blocking(move || { + ensure_index( + &connections, + &config.index_name, + &prefix, + &index_dimension, + embedding.len(), + )?; + connections.execute(|connection| { + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(&scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { + pipeline + .cmd("EXPIRE") + .arg(&document) + .arg(ttl.as_secs()) + .ignore(); + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) + }) + }) + .await + .map_err(|_| Error::Unavailable)? + } + } + + fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> impl Future, Error>> + Send { + let key = key.to_owned(); + let prompt = prompt_from_context(context); + let metadata = context.metadata.clone(); + async move { + let Some(prompt) = prompt else { + return Ok(None); + }; + let embedding = self + .embedder + .async_embed(&prompt, metadata.as_ref()) + .await?; + let connections = Arc::clone(&self.connections); + let config = self.config.clone(); + let index_dimension = Arc::clone(&self.index_dimension); + let threshold = config.similarity_threshold; + tokio::task::spawn_blocking(move || { + let prefix = format!("{}:", config.index_name); + ensure_index( + &connections, + &config.index_name, + &prefix, + &index_dimension, + embedding.len(), + )?; + let scope = scope_tag(&key); + let query = format!( + "(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]" + ); + let vector = embedding_bytes(&embedding); + let response = connections.execute(|connection| { + redis::cmd("FT.SEARCH") + .arg(&config.index_name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let Some(fields) = search_fields(response)? else { + return Ok(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < threshold { + return Ok(None); + } + Ok(Some(response)) + }) + .await + .map_err(|_| Error::Unavailable)? + .and_then(|response| response.map(|bytes| self.codec.decode(&bytes)).transpose()) + } + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(Value::Array(messages)) = context.messages.as_ref() + && !messages.is_empty() + { + return Some( + messages + .iter() + .filter_map(Value::as_object) + .map(message_text) + .collect(), + ); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_owned(); + (!prompt.is_empty()).then_some(prompt) +} + +fn message_text(message: &serde_json::Map) -> String { + let content = match message.get("content") { + Some(Value::String(value)) => value.clone(), + Some(Value::Array(parts)) => parts + .iter() + .filter_map(Value::as_object) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .filter(|text| !text.is_empty()) + .collect(), + _ => String::new(), + }; + format!( + "{content}{}", + search_results_text(message.get("search_results")) + ) +} + +fn search_results_text(value: Option<&Value>) -> String { + let Some(Value::Array(results)) = value else { + return String::new(); + }; + results + .iter() + .filter_map(Value::as_object) + .map(|result| { + let source = result.get("source").and_then(Value::as_str).unwrap_or(""); + let title = result.get("title").and_then(Value::as_str).unwrap_or(""); + let content = result + .get("content") + .and_then(Value::as_array) + .map(|blocks| { + blocks + .iter() + .filter_map(Value::as_object) + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect::() + }) + .unwrap_or_default(); + let citations = result + .get("citations") + .filter(|value| !value.is_null()) + .and_then(|value| serde_json::to_string(value).ok()) + .unwrap_or_default(); + format!("{source}{title}{content}{citations}") + }) + .collect() +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(value) => { + let value = value.trim(); + if !value.is_empty() { + parts.push(value.to_owned()); + } + } + Value::Array(values) => values + .iter() + .for_each(|value| collect_input_text(value, parts)), + Value::Object(object) => { + if let Some(content) = object.get("content").filter(|value| !value.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(value)) = object.get(key) { + let value = value.trim(); + if !value.is_empty() { + parts.push(value.to_owned()); + return; + } + } + } + } + _ => {} + } +} + +fn scope_tag(key: &str) -> String { + let digest = Sha256::digest(key.as_bytes()); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn embedding_bytes(embedding: &[f32]) -> Vec { + embedding + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn ensure_index( + connections: &Connections, + index_name: &str, + prefix: &str, + index_dimension: &Mutex>, + dimension: usize, +) -> Result<(), Error> +where + C: redis::ConnectionLike + Send + 'static, +{ + if index_dimension + .lock() + .map_err(|_| Error::Unavailable)? + .is_some_and(|existing| existing == dimension) + { + return Ok(()); + } + let create = connections.execute(|connection| { + Ok(redis::cmd("FT.CREATE") + .arg(index_name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(prefix) + .arg("SCHEMA") + .arg("litellm_cache_key") + .arg("TAG") + .arg("embedding") + .arg("VECTOR") + .arg("HNSW") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dimension) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .query::(connection) + .map(|_| ()) + .map_err(|error| error.to_string())) + })?; + if let Err(message) = create { + if !message.to_ascii_lowercase().contains("already exists") { + return Err(Error::Unavailable); + } + let info = connections.execute(|connection| { + redis::cmd("FT.INFO") + .arg(index_name) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let existing = index_dimension_from_info(&info).ok_or(Error::Unavailable)?; + if existing != dimension { + return Err(Error::Unavailable); + } + } + *index_dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension); + Ok(()) +} + +fn index_dimension_from_info(value: &redis::Value) -> Option { + let redis::Value::Array(values) = value else { + return None; + }; + let attributes = values.windows(2).find_map(|pair| { + (value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1]) + })?; + let redis::Value::Array(fields) = attributes else { + return None; + }; + fields.iter().find_map(|field| { + let redis::Value::Array(values) = field else { + return None; + }; + let flattened = values.iter().flat_map(|value| match value { + redis::Value::Array(values) => values.as_slice(), + _ => std::slice::from_ref(value), + }); + let values = flattened.collect::>(); + values.windows(2).find_map(|pair| { + if value_text(pair[0]).as_deref() == Some("dimensions") { + return value_text(pair[1]).and_then(|value| value.parse().ok()); + } + None + }) + }) +} + +type SearchFields = Vec<(String, Vec)>; + +fn search_fields(value: redis::Value) -> Result, Error> { + let redis::Value::Array(values) = value else { + return Err(Error::InvalidEntry); + }; + let total = parse_i64(values.first().ok_or(Error::InvalidEntry)?)?; + if total <= 0 || values.len() < 3 { + return Ok(None); + } + let redis::Value::Array(fields) = &values[2] else { + return Err(Error::InvalidEntry); + }; + let (pairs, remainder) = fields.as_chunks::<2>(); + if !remainder.is_empty() { + return Err(Error::InvalidEntry); + } + let pairs = pairs + .iter() + .map(|pair| { + Ok(( + value_text(&pair[0]).ok_or(Error::InvalidEntry)?, + value_bytes(&pair[1])?, + )) + }) + .collect::, Error>>()?; + Ok(Some(pairs)) +} + +fn parse_i64(value: &redis::Value) -> Result { + value_text(value) + .ok_or(Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn parse_f64(value: &[u8]) -> Result { + std::str::from_utf8(value) + .map_err(|_| Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn value_text(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(value) => Some(value.clone()), + redis::Value::Int(value) => Some(value.to_string()), + _ => None, + } +} + +fn value_bytes(value: &redis::Value) -> Result, Error> { + match value { + redis::Value::BulkString(bytes) => Ok(bytes.clone()), + redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), + redis::Value::Int(value) => Ok(value.to_string().into_bytes()), + _ => Err(Error::InvalidEntry), + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use litellm_cache::BaseCache; + use litellm_cache_response::ResponseCacheCodec; + use redis_test::MockRedisConnection; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::{ + Embedder, ValkeySemanticCache, ValkeySemanticConfig, index_dimension_from_info, + prompt_from_context, scope_tag, + }; + + #[derive(Clone)] + struct FixedEmbedder { + vector: Vec, + calls: EmbedderCalls, + } + + type EmbedderCalls = Arc)>>>; + + impl Embedder for FixedEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, super::Error> { + self.calls + .lock() + .unwrap() + .push((prompt.into(), metadata.cloned())); + Ok(self.vector.clone()) + } + + async fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> Result, super::Error> { + self.embed(prompt, metadata) + } + } + + fn context( + messages: Option, + input: Option, + ) -> litellm_cache::SemanticCacheContext { + litellm_cache::SemanticCacheContext { + messages, + input, + ..Default::default() + } + } + + #[rstest] + #[case(json!([{"content": "hello"}]), None, Some("hello"))] + #[case(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] + #[case(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] + #[case(Value::Array(vec![]), Some(json!(" hello ")), Some("hello"))] + #[case(Value::Array(vec![]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] + #[case(Value::Array(vec![]), Some(json!(" ")), None)] + fn prompt_shapes( + #[case] messages: Value, + #[case] input: Option, + #[case] expected: Option<&str>, + ) { + assert_eq!( + prompt_from_context(&context(Some(messages), input)), + expected.map(str::to_owned) + ); + } + + #[test] + fn scope_tags_are_lowercase_sha256() { + assert_eq!( + scope_tag("key"), + "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" + ); + } + + #[test] + fn existing_index_dimension_is_read_from_attributes() { + let info = redis::Value::Array(vec![ + redis::Value::SimpleString("attributes".into()), + redis::Value::Array(vec![redis::Value::Array(vec![ + redis::Value::SimpleString("identifier".into()), + redis::Value::SimpleString("embedding".into()), + redis::Value::Array(vec![ + redis::Value::SimpleString("dimensions".into()), + redis::Value::SimpleString("2".into()), + ]), + ])]), + ]); + assert_eq!(index_dimension_from_info(&info), Some(2)); + } + + #[tokio::test] + async fn unsupported_connection_test_is_reported() { + let cache = ValkeySemanticCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + assert_eq!( + cache.test_connection().await, + Err(super::Error::UnsupportedOperation) + ); + } + + #[test] + fn missing_prompt_does_not_touch_redis() { + let cache = ValkeySemanticCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None); + assert_eq!(cache.get_ttl(&context(None, None)), None); + } +} From be2f0d081b6c7ac41090ad9300b18402f226ef5f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 21 Sep 2026 13:25:39 -0700 Subject: [PATCH 042/109] fix(proxy): report sources only on the read endpoints main does not cover /config/field/info and /config/list already report per-key source on main, so this drops the branch's versions of those and keeps /alerting/settings, /get/ui_settings and /router/settings. Read endpoints no longer write the freshly read database row back into the shared settings store; the reload path already keeps it current, and a GET that mutates global state leaks across callers. Regenerates the lazy OpenAPI snapshot on Python 3.12, matching CI, and the dashboard API types for the two new response fields. --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../router_settings_endpoints.py | 31 +++++++------- litellm/proxy/proxy_server.py | 9 ++-- .../proxy_setting_endpoints.py | 41 +++++++++++-------- .../test_router_settings_endpoints.py | 41 ++++++++++--------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 11 +++++ 6 files changed, 80 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 391f0042ed0..06e157498aa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19632,7 +19632,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 019ac68ae23..d6d74ada35a 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -8,7 +8,9 @@ GET /router/fields - Get router settings field definitions without values (for U """ import inspect -from typing import Any, Final, cast, get_args +from collections.abc import Mapping +from types import MappingProxyType +from typing import Any, Final, get_args from fastapi import APIRouter, Depends from pydantic import BaseModel, Field @@ -127,19 +129,20 @@ async def get_router_settings( if field.field_name in current_values: field.field_value = current_values[field.field_name] - field_defaults: Final[dict[str, object]] = { - field.field_name: cast(object, field.field_default) # cast-ok: Pydantic field defaults are untyped - for field in router_fields - } - source: Final[dict[str, FieldSource]] = { - key: _router_setting_source( - proxy_config.router_settings, - key, - cast(object, current_values[key]), # cast-ok: current values are stored in a typed response map - field_defaults.get(key), - ) - for key in current_values - } + field_defaults: Final[Mapping[str, object]] = MappingProxyType( + {field.field_name: field.field_default for field in router_fields} + ) + source: Final[Mapping[str, FieldSource]] = MappingProxyType( + { + key: _router_setting_source( + proxy_config.router_settings, + key, + current_values[key], + field_defaults.get(key), + ) + for key in current_values + } + ) return RouterSettingsResponse( fields=router_fields, current_values=current_values, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4b2781e031b..24af6d7f7d3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15991,7 +15991,7 @@ def _nested_setting_source( field_default: JsonValue, ) -> FieldSource: db_value: Final = db_values.get(field_name) - if db_value is not None and db_value != []: + if db_value is not None and not (isinstance(db_value, list) and len(db_value) == 0): return "db" parent_value: Final = settings.config_value(parent_key) if isinstance(parent_value, Mapping) and field_name in parent_value: @@ -16036,13 +16036,13 @@ async def alerting_settings( where={"param_name": "general_settings"} ) - db_general_settings_dict: Final[Mapping[str, JsonValue]] = ( - dict(db_general_settings.param_value) + db_general_settings_dict: Final[Mapping[str, JsonValue]] = MappingProxyType( + dict(db_general_settings.param_value) # mutable-ok: Prisma returns the JSON column as a plain dict if db_general_settings is not None and db_general_settings.param_value is not None else {} ) alerting_args_value: Final = db_general_settings_dict.get("alerting_args") - alerting_args_dict: Final[Mapping[str, JsonValue]] = ( + alerting_args_dict: Final[Mapping[str, JsonValue]] = MappingProxyType( alerting_args_value if isinstance(alerting_args_value, dict) else {} ) alerting_values: Final = cast( # cast-ok: alerting is stored as a JSON list when present @@ -16050,7 +16050,6 @@ async def alerting_settings( ) settings: Final = proxy_config.settings - settings.apply_db_row("general_settings", db_general_settings_dict) allowed_args: Final = MappingProxyType( { diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 4505f3144ee..ed626bdb624 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -1755,18 +1755,21 @@ async def get_ui_settings(): ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS} apply_runtime_general_settings_flags(ui_settings) - proxy_config.settings.apply_db_row("ui_settings", ui_settings) # Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values from litellm.proxy.proxy_server import user_api_key_cache await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) - effective_ui_settings: Final = { - **{key: proxy_config.settings[key] for key in ALLOWED_UI_SETTINGS_FIELDS if key in proxy_config.settings}, - **ui_settings, - } - config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": effective_ui_settings}} + effective_ui_settings: Final[Mapping[str, object]] = MappingProxyType( + { + **{key: proxy_config.settings[key] for key in ALLOWED_UI_SETTINGS_FIELDS if key in proxy_config.settings}, + **ui_settings, + } + ) + config: Final[Mapping[str, object]] = MappingProxyType( + {"litellm_settings": MappingProxyType({"ui_settings": effective_ui_settings})} + ) settings_class: Final = _get_effective_ui_settings_class() resolved_settings: Final = _SettingsWithSchema.model_validate( await _get_settings_with_schema( @@ -1775,16 +1778,22 @@ async def get_ui_settings(): config=config, ) ) - values: Final = { - **resolved_settings.values, - ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), - } - source: Final[dict[str, FieldSource]] = { - key: ( - "db" if key in ui_settings else _ui_setting_source(key, values[key], proxy_config.settings, settings_class) - ) - for key in values - } + values: Final[Mapping[str, object]] = MappingProxyType( + { + **resolved_settings.values, + ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), + } + ) + source: Final[Mapping[str, FieldSource]] = MappingProxyType( + { + key: ( + "db" + if key in ui_settings + else _ui_setting_source(key, values[key], proxy_config.settings, settings_class) + ) + for key in values + } + ) return UISettingsResponse( values=values, field_schema=resolved_settings.field_schema, diff --git a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py index 3af7de62abe..51c8679e89e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py @@ -15,12 +15,24 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.router_settings_endpoints import ( get_router_settings, ) +from litellm.proxy.config_resolvers import SettingsStore from litellm.proxy.proxy_server import app from litellm.router import Router client = TestClient(app) +def _stub_proxy_config(router_settings, config_router_settings): + class _StubProxyConfig: + def __init__(self): + self.router_settings = router_settings + + async def get_config(self, config_file_path=None): + return {"router_settings": dict(config_router_settings)} + + return _StubProxyConfig() + + class TestRouterSettingsEndpoints: """Test suite for router settings endpoints""" @@ -77,25 +89,18 @@ class TestRouterSettingsEndpoints: @pytest.mark.asyncio async def test_get_router_settings_reports_sources(self, monkeypatch): - from litellm.proxy.config_resolvers import SettingsStore - store = SettingsStore("router_settings") store.load_yaml({"routing_strategy": "simple-shuffle"}) store.apply_db_row("router_settings", {"num_retries": 3}) - monkeypatch.setattr(proxy_server.proxy_config, "router_settings", store) - monkeypatch.setattr(proxy_server, "llm_router", None) - - async def fake_get_config(self, config_file_path=None): - return { - "router_settings": { - "routing_strategy": "simple-shuffle", - "num_retries": 3, - } - } - monkeypatch.setattr( - proxy_server.ProxyConfig, "get_config", fake_get_config, raising=True + proxy_server, + "proxy_config", + _stub_proxy_config( + store, + {"routing_strategy": "simple-shuffle", "num_retries": 3}, + ), ) + monkeypatch.setattr(proxy_server, "llm_router", None) admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-x" @@ -132,12 +137,10 @@ class TestRouterSettingsEndpoints: ) monkeypatch.setattr(proxy_server, "llm_router", llm_router) - - async def fake_get_config(self, config_file_path=None): - return {} - monkeypatch.setattr( - proxy_server.ProxyConfig, "get_config", fake_get_config, raising=True + proxy_server, + "proxy_config", + _stub_proxy_config(SettingsStore("router_settings"), {}), ) admin_user = UserAPIKeyAuth( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6211eeaf962..ac89d676921 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37069,6 +37069,13 @@ export interface components { routing_strategy_descriptions: { [key: string]: string; }; + /** + * Source + * @description Source of each current router setting + */ + source: { + [key: string]: "config" | "db" | "env" | "default" | "unset"; + }; }; /** * RoutingGroup @@ -39532,6 +39539,10 @@ export interface components { field_schema: { [key: string]: unknown; }; + /** Source */ + source: { + [key: string]: "config" | "db" | "env" | "default" | "unset"; + }; /** Values */ values: { [key: string]: unknown; From 0a88658227f8e0d7e2d4928df7c5ee63bd83fd1d Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 20:25:56 +0000 Subject: [PATCH 043/109] chore: retrigger ci after docs merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 3ba4a60d5ed5eeeb217b63ac7898746927f0db12 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:26:28 +0000 Subject: [PATCH 044/109] feat(rust): add HashiCorp Vault secret manager crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 2 +- litellm-rust/Cargo.lock | 19 + litellm-rust/Cargo.toml | 1 + .../crates/secrets-hashicorp/Cargo.toml | 23 + .../crates/secrets-hashicorp/src/config.rs | 169 +++++++ .../crates/secrets-hashicorp/src/error.rs | 32 ++ .../crates/secrets-hashicorp/src/lib.rs | 9 + .../secrets-hashicorp/src/secret_manager.rs | 333 ++++++++++++++ .../secrets-hashicorp/tests/secret_manager.rs | 424 ++++++++++++++++++ litellm-rust/crates/secrets/Cargo.toml | 2 + litellm-rust/crates/secrets/README.md | 2 + litellm-rust/crates/secrets/src/error.rs | 3 + litellm-rust/crates/secrets/src/handler.rs | 10 + litellm-rust/crates/secrets/src/lib.rs | 2 + litellm-rust/crates/secrets/tests/handler.rs | 128 ++++++ .../hashicorp_secret_manager.py | 16 +- .../hashicorp_vault_parity.json | 97 ++++ .../test_hashicorp_secret_manager.py | 50 +++ 18 files changed, 1314 insertions(+), 8 deletions(-) create mode 100644 litellm-rust/crates/secrets-hashicorp/Cargo.toml create mode 100644 litellm-rust/crates/secrets-hashicorp/src/config.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/src/error.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/src/lib.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs create mode 100644 tests/test_litellm/secret_managers/hashicorp_vault_parity.json diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 278fa7c425f..56bb9a568fc 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -130,7 +130,7 @@ jobs: - 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 + for features in '' aws google hashicorp aws,google,hashicorp; do cargo test -p litellm-secrets --locked --no-default-features --features "$features" done diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..6f05bba9417 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2698,6 +2698,7 @@ dependencies = [ "litellm-core-utils", "litellm-secrets-aws", "litellm-secrets-google", + "litellm-secrets-hashicorp", "litellm-secrets-types", "moka", "reqwest 0.12.28", @@ -2755,6 +2756,24 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-secrets-hashicorp" +version = "0.1.0" +dependencies = [ + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + [[package]] name = "litellm-secrets-types" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..5030a94b140 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -22,6 +22,7 @@ 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-secrets-hashicorp = { path = "crates/secrets-hashicorp" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml new file mode 100644 index 00000000000..0656980a033 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "litellm-secrets-hashicorp" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-core-utils.workspace = true +litellm-secrets-types.workspace = true +moka.workspace = true +reqwest.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true +tempfile = "3" +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs new file mode 100644 index 00000000000..f9491f71afb --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -0,0 +1,169 @@ +use std::{path::PathBuf, time::Duration}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; +use litellm_secrets_types::SecretValue; + +use crate::Error; + +const DEFAULT_ADDRESS: &str = "http://127.0.0.1:8200"; +const DEFAULT_MOUNT: &str = "secret"; +const DEFAULT_APPROLE_MOUNT_PATH: &str = "approle"; +const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400); +const HCP_VAULT_ADDR: &str = "HCP_VAULT_ADDR"; +const HCP_VAULT_TOKEN: &str = "HCP_VAULT_TOKEN"; +const HCP_VAULT_NAMESPACE: &str = "HCP_VAULT_NAMESPACE"; +const HCP_VAULT_LOGIN_NAMESPACE: &str = "HCP_VAULT_LOGIN_NAMESPACE"; +const HCP_VAULT_SECRET_NAMESPACE: &str = "HCP_VAULT_SECRET_NAMESPACE"; +const HCP_VAULT_MOUNT_NAME: &str = "HCP_VAULT_MOUNT_NAME"; +const HCP_VAULT_PATH_PREFIX: &str = "HCP_VAULT_PATH_PREFIX"; +const HCP_VAULT_APPROLE_ROLE_ID: &str = "HCP_VAULT_APPROLE_ROLE_ID"; +const HCP_VAULT_APPROLE_SECRET_ID: &str = "HCP_VAULT_APPROLE_SECRET_ID"; +const HCP_VAULT_APPROLE_MOUNT_PATH: &str = "HCP_VAULT_APPROLE_MOUNT_PATH"; +const HCP_VAULT_CLIENT_CERT: &str = "HCP_VAULT_CLIENT_CERT"; +const HCP_VAULT_CLIENT_KEY: &str = "HCP_VAULT_CLIENT_KEY"; +const HCP_VAULT_CERT_ROLE: &str = "HCP_VAULT_CERT_ROLE"; +const HCP_VAULT_REFRESH_INTERVAL: &str = "HCP_VAULT_REFRESH_INTERVAL"; +const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL"; + +#[derive(Clone, Debug)] +pub struct AppRoleAuth { + pub role_id: String, + pub secret_id: SecretValue, + pub mount_path: String, +} + +#[derive(Clone, Debug)] +pub struct TlsCertAuth { + pub cert_path: PathBuf, + pub key_path: PathBuf, + pub role: Option, +} + +#[derive(Clone, Debug)] +pub struct HashicorpVaultConfig { + pub address: String, + pub token: Option, + pub namespace: Option, + pub login_namespace: Option, + pub secret_namespace: Option, + pub mount: String, + pub path_prefix: Option, + pub approle: Option, + pub tls_cert: Option, + pub refresh_interval: Duration, +} + +impl HashicorpVaultConfig { + pub fn from_environment(environment: &dyn Lookup) -> Result { + let address: String = environment + .get(HCP_VAULT_ADDR) + .and_then(|value| nonempty(value.trim())) + .map(|value| value.trim_end_matches('/').to_owned()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| DEFAULT_ADDRESS.to_owned()); + let token: Option = environment + .get(HCP_VAULT_TOKEN) + .and_then(nonempty) + .map(SecretValue::new); + let namespace: Option = path_component(environment.get(HCP_VAULT_NAMESPACE)); + let login_namespace: Option = + path_component(environment.get(HCP_VAULT_LOGIN_NAMESPACE)); + let secret_namespace: Option = + path_component(environment.get(HCP_VAULT_SECRET_NAMESPACE)); + let mount: String = path_component(environment.get(HCP_VAULT_MOUNT_NAME)) + .unwrap_or_else(|| DEFAULT_MOUNT.to_owned()); + let path_prefix: Option = path_component(environment.get(HCP_VAULT_PATH_PREFIX)); + let approle: Option = match ( + environment + .get(HCP_VAULT_APPROLE_ROLE_ID) + .and_then(nonempty), + environment + .get(HCP_VAULT_APPROLE_SECRET_ID) + .and_then(nonempty) + .map(SecretValue::new), + ) { + (Some(role_id), Some(secret_id)) => Some(AppRoleAuth { + role_id, + secret_id, + mount_path: path_component(environment.get(HCP_VAULT_APPROLE_MOUNT_PATH)) + .unwrap_or_else(|| DEFAULT_APPROLE_MOUNT_PATH.to_owned()), + }), + _ => None, + }; + let tls_cert: Option = match ( + environment.get(HCP_VAULT_CLIENT_CERT).and_then(nonempty), + environment.get(HCP_VAULT_CLIENT_KEY).and_then(nonempty), + ) { + (Some(cert_path), Some(key_path)) => Some(TlsCertAuth { + cert_path: PathBuf::from(cert_path), + key_path: PathBuf::from(key_path), + role: environment.get(HCP_VAULT_CERT_ROLE).and_then(nonempty), + }), + _ => None, + }; + let refresh_interval: Duration = refresh_interval(environment)?; + Ok(Self { + address, + token, + namespace, + login_namespace, + secret_namespace, + mount, + path_prefix, + approle, + tls_cert, + refresh_interval, + }) + } + + pub fn from_settings( + _settings: &KeyManagementSettings, + environment: &dyn Lookup, + ) -> Result { + Self::from_environment(environment) + } + + pub fn login_namespace(&self) -> Option<&str> { + self.login_namespace + .as_deref() + .or(self.namespace.as_deref()) + } + + pub fn secret_namespace(&self) -> Option<&str> { + self.secret_namespace + .as_deref() + .or(self.namespace.as_deref()) + } +} + +fn nonempty(value: impl AsRef) -> Option { + let value: &str = value.as_ref(); + (!value.is_empty()).then(|| value.to_owned()) +} + +fn path_component(value: Option) -> Option { + value + .and_then(|value| nonempty(value.trim())) + .map(|value| value.trim_matches('/').to_owned()) + .filter(|value| !value.is_empty()) +} + +fn refresh_interval(environment: &dyn Lookup) -> Result { + let value: Option = environment + .get(HCP_VAULT_REFRESH_INTERVAL) + .and_then(nonempty) + .or_else(|| { + environment + .get(SECRET_MANAGER_REFRESH_INTERVAL) + .and_then(nonempty) + }); + let Some(value) = value else { + return Ok(DEFAULT_REFRESH_INTERVAL); + }; + let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?; + if seconds < 0 { + return Ok(Duration::from_nanos(1)); + } + Ok(Duration::from_secs(seconds as u64)) +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs new file mode 100644 index 00000000000..3f6a5b66475 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -0,0 +1,32 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("HashiCorp Vault requires an enterprise license")] + EnterpriseRequired, + #[error("invalid secret name")] + InvalidSecretName(#[from] litellm_secrets_types::Error), + #[error("HashiCorp Vault request failed")] + Request( + #[from] + #[redact] + reqwest::Error, + ), + #[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")] + TlsIdentity { + path: std::path::PathBuf, + message: String, + }, + #[error("HashiCorp Vault login returned HTTP {status}")] + LoginStatus { status: u16 }, + #[error("HashiCorp Vault login response is malformed")] + MalformedLogin, + #[error("HashiCorp Vault authentication is not configured")] + NoAuthConfigured, + #[error("HashiCorp Vault returned HTTP {status}")] + Status { status: u16 }, + #[error("HashiCorp Vault response payload is malformed")] + MalformedPayload, + #[error("HashiCorp Vault secret value is not a string")] + NonStringValue, + #[error("invalid HashiCorp Vault refresh interval")] + RefreshInterval, +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs new file mode 100644 index 00000000000..a561a58b59e --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -0,0 +1,9 @@ +#![forbid(unsafe_code)] + +mod config; +mod error; +pub mod secret_manager; + +pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth}; +pub use error::Error; +pub use secret_manager::HashicorpVault; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs new file mode 100644 index 00000000000..6b887289443 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -0,0 +1,333 @@ +use std::{ + fmt, + sync::Arc, + time::{Duration, Instant}, +}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{ + BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name, +}; +use moka::future::Cache; +use reqwest::Client; +use serde_json::{Value, json}; +use tokio::sync::Mutex; + +use crate::{Error, HashicorpVaultConfig}; + +const CACHE_CAPACITY: u64 = 200; + +#[derive(Clone)] +struct CachedToken { + token: SecretValue, + expires_at: Option, +} + +#[derive(Clone)] +pub struct HashicorpVault { + client: Arc, + config: HashicorpVaultConfig, + cache: Cache, + auth_token: Arc>>, +} + +impl fmt::Debug for HashicorpVault { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("HashicorpVault") + .field("config", &self.config) + .finish_non_exhaustive() + } +} + +impl HashicorpVault { + pub fn new( + environment: Arc, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let config: HashicorpVaultConfig = + HashicorpVaultConfig::from_environment(environment.as_ref())?; + let client: Client = client_for_config(&config)?; + Self::with_client(client, config, enterprise_enabled) + } + + pub fn with_client( + client: Client, + config: HashicorpVaultConfig, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let cache: Cache = Cache::builder() + .max_capacity(CACHE_CAPACITY) + .time_to_live(config.refresh_interval) + .build(); + Ok(Self { + client: Arc::new(client), + config, + cache, + auth_token: Arc::new(Mutex::new(None)), + }) + } + + pub fn secret_url(&self, secret_name: &str) -> Result { + validate_secret_name(secret_name).map_err(Error::InvalidSecretName)?; + let namespace: String = self + .config + .secret_namespace() + .map(|value| format!("{value}/")) + .unwrap_or_default(); + let path_prefix: String = self + .config + .path_prefix + .as_deref() + .map(|value| format!("{value}/")) + .unwrap_or_default(); + Ok(format!( + "{}/v1/{}{}/data/{}{}", + self.config.address, namespace, self.config.mount, path_prefix, secret_name + )) + } + + pub fn login_url(&self) -> Option { + self.config.approle.as_ref().map_or_else( + || { + self.config + .tls_cert + .as_ref() + .map(|_| format!("{}/v1/auth/cert/login", self.config.address)) + }, + |approle| { + Some(format!( + "{}/v1/auth/{}/login", + self.config.address, approle.mount_path + )) + }, + ) + } + + pub fn config(&self) -> &HashicorpVaultConfig { + &self.config + } + + pub async fn async_read_secret(&self, secret_name: &str) -> Result, Error> { + let url: String = self.secret_url(secret_name)?; + if let Some(value) = self.cache.get(&url).await { + return Ok(Some(value)); + } + let token: SecretValue = self.vault_token().await?; + let mut request = self + .client + .get(&url) + .header("X-Vault-Token", token.expose()); + if let Some(namespace) = self.config.secret_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.send().await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if !response.status().is_success() { + return Err(Error::Status { + status: response.status().as_u16(), + }); + } + let body = response.bytes().await?; + let body: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload)?; + let data: &Value = body + .get("data") + .and_then(Value::as_object) + .and_then(|value| value.get("data")) + .ok_or(Error::MalformedPayload)?; + let data: &serde_json::Map = + data.as_object().ok_or(Error::MalformedPayload)?; + let Some(value) = data.get("key") else { + return Ok(None); + }; + let value: &str = value.as_str().ok_or(Error::NonStringValue)?; + let value: SecretValue = SecretValue::new(value); + self.cache.insert(url, value.clone()).await; + Ok(Some(value)) + } + + pub async fn async_write_secret( + &self, + secret_name: &str, + value: SecretValue, + description: Option<&str>, + ) -> Result { + let url: String = self.secret_url(secret_name)?; + let data: Value = match description { + Some(description) => json!({"key": value.expose(), "description": description}), + None => json!({"key": value.expose()}), + }; + let token: SecretValue = self.vault_token().await?; + let mut request = self + .client + .post(&url) + .header("X-Vault-Token", token.expose()); + if let Some(namespace) = self.config.secret_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.json(&json!({"data": data})).send().await?; + if !response.status().is_success() { + return Err(Error::Status { + status: response.status().as_u16(), + }); + } + self.cache.invalidate(&url).await; + let body = response.bytes().await?; + if body.is_empty() { + return Ok(Value::Null); + } + serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload) + } + + pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> { + let url: String = self.secret_url(secret_name)?; + let token: SecretValue = self.vault_token().await?; + let mut request = self + .client + .delete(&url) + .header("X-Vault-Token", token.expose()); + if let Some(namespace) = self.config.secret_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.send().await?; + if !response.status().is_success() { + return Err(Error::Status { + status: response.status().as_u16(), + }); + } + self.cache.invalidate(&url).await; + Ok(()) + } + + pub async fn async_rotate_secret( + &self, + current_name: &str, + new_name: &str, + value: &SecretValue, + ) -> Result { + async_rotate_secret(self, current_name, new_name, value).await + } + + async fn vault_token(&self) -> Result { + let mut cached: tokio::sync::MutexGuard<'_, Option> = + self.auth_token.lock().await; + if let Some(entry) = cached.as_ref() + && entry + .expires_at + .is_none_or(|expires_at| expires_at > Instant::now()) + { + return Ok(entry.token.clone()); + } + let Some(login_url) = self.login_url() else { + let Some(token) = self.config.token.clone() else { + return Err(Error::NoAuthConfigured); + }; + return Ok(token); + }; + let body: Value = match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) { + (Some(approle), _) => { + json!({"role_id": approle.role_id, "secret_id": approle.secret_id.expose()}) + } + (None, Some(tls)) => tls + .role + .as_deref() + .map_or_else(|| json!({}), |role| json!({"name": role})), + (None, None) => { + let Some(token) = self.config.token.clone() else { + return Err(Error::NoAuthConfigured); + }; + return Ok(token); + } + }; + let mut request = self.client.post(login_url).json(&body); + if let Some(namespace) = self.config.login_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.send().await?; + if !response.status().is_success() { + return Err(Error::LoginStatus { + status: response.status().as_u16(), + }); + } + let body = response.bytes().await?; + let payload: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedLogin)?; + let auth: &serde_json::Map = payload + .get("auth") + .and_then(Value::as_object) + .ok_or(Error::MalformedLogin)?; + let token: SecretValue = SecretValue::new( + auth.get("client_token") + .and_then(Value::as_str) + .ok_or(Error::MalformedLogin)?, + ); + let lease_duration: u64 = auth + .get("lease_duration") + .and_then(Value::as_u64) + .ok_or(Error::MalformedLogin)?; + let expires_at: Option = + (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)); + *cached = Some(CachedToken { + token: token.clone(), + expires_at, + }); + Ok(token) + } +} + +impl BaseSecretManager for HashicorpVault { + type Error = Error; + type WriteResponse = Value; + type DeleteResponse = (); + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + HashicorpVault::async_read_secret(self, name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + HashicorpVault::async_write_secret(self, name, value.clone(), description).await + } + + async fn async_delete_secret( + &self, + name: &str, + _recovery_window_in_days: i64, + ) -> Result<(), Error> { + HashicorpVault::async_delete_secret(self, name).await + } +} + +fn client_for_config(config: &HashicorpVaultConfig) -> Result { + let mut builder: reqwest::ClientBuilder = Client::builder(); + if let Some(tls) = config.tls_cert.as_ref() { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { + path: tls.key_path.clone(), + message: source.to_string(), + })?; + let identity: reqwest::Identity = reqwest::Identity::from_pem( + &[cert.as_slice(), key.as_slice()].concat(), + ) + .map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + builder = builder.identity(identity); + } + builder.build().map_err(Error::Request) +} diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs new file mode 100644 index 00000000000..1cf13a79c77 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -0,0 +1,424 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig}; +use litellm_secrets_types::SecretValue; +use serde::Deserialize; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header, method, path}, +}; + +fn config(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVaultConfig { + let mut environment_values: HashMap = values + .iter() + .map(|(name, value)| ((*name).to_owned(), (*value).to_owned())) + .collect(); + environment_values.insert("HCP_VAULT_ADDR".to_owned(), server.uri()); + let environment: Arc = + Arc::new(move |name: &str| environment_values.get(name).cloned()); + HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap() +} + +fn manager(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVault { + HashicorpVault::with_client(reqwest::Client::new(), config(server, values), true).unwrap() +} + +#[tokio::test] +async fn token_reads_use_vault_headers_and_cache_values() { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name")) + .and(header("X-Vault-Token", "token")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]); + + assert_eq!( + manager + .async_read_secret("name") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); + assert_eq!( + manager + .async_read_secret("name") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[tokio::test] +async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/team-a/kv-prod/data/virtual-keys/name")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager( + &server, + &[ + ("HCP_VAULT_TOKEN", "token"), + ("HCP_VAULT_SECRET_NAMESPACE", " /team-a/ "), + ("HCP_VAULT_MOUNT_NAME", " /kv-prod/ "), + ("HCP_VAULT_PATH_PREFIX", " /virtual-keys/ "), + ], + ); + + assert_eq!( + manager.secret_url("name").unwrap(), + format!("{}/v1/team-a/kv-prod/data/virtual-keys/name", server.uri()) + ); + assert!(manager.async_read_secret("name").await.unwrap().is_some()); +} + +#[tokio::test] +async fn approle_login_uses_namespace_and_reuses_the_token() { + let server: MockServer = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/auth/custom-approle/login")) + .and(header("X-Vault-Namespace", "login-root")) + .and(body_json(json!({"role_id": "role", "secret_id": "secret"}))) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"auth": {"client_token": "login-token", "lease_duration": 3600}}), + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/secret-root/secret/data/name")) + .and(header("X-Vault-Token", "login-token")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager( + &server, + &[ + ("HCP_VAULT_APPROLE_ROLE_ID", "role"), + ("HCP_VAULT_APPROLE_SECRET_ID", "secret"), + ("HCP_VAULT_APPROLE_MOUNT_PATH", "custom-approle"), + ("HCP_VAULT_NAMESPACE", "secret-root"), + ("HCP_VAULT_LOGIN_NAMESPACE", "login-root"), + ], + ); + + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + assert!(manager.async_read_secret("name-2").await.unwrap().is_none()); +} + +#[tokio::test] +async fn approle_tokens_expire_after_the_vault_lease() { + let server: MockServer = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/auth/approle/login")) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"auth": {"client_token": "login-token", "lease_duration": 1}}), + ), + ) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(2) + .mount(&server) + .await; + let manager: HashicorpVault = manager( + &server, + &[ + ("HCP_VAULT_APPROLE_ROLE_ID", "role"), + ("HCP_VAULT_APPROLE_SECRET_ID", "secret"), + ("HCP_VAULT_REFRESH_INTERVAL", "0"), + ], + ); + + assert!(manager.async_read_secret("first").await.unwrap().is_some()); + tokio::time::sleep(Duration::from_secs(1) + Duration::from_millis(50)).await; + assert!(manager.async_read_secret("second").await.unwrap().is_some()); +} + +#[tokio::test] +async fn tls_login_posts_the_role_and_uses_the_client_identity() { + let server: MockServer = MockServer::start().await; + let directory: tempfile::TempDir = tempfile::tempdir().unwrap(); + let cert_path = directory.path().join("client.crt"); + let key_path = directory.path().join("client.key"); + std::fs::write(&cert_path, TEST_CERTIFICATE).unwrap(); + std::fs::write(&key_path, TEST_PRIVATE_KEY).unwrap(); + Mock::given(method("POST")) + .and(path("/v1/auth/cert/login")) + .and(body_json(json!({"name": "vault-role"}))) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"auth": {"client_token": "cert-token", "lease_duration": 0}}), + ), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .mount(&server) + .await; + let manager: HashicorpVault = HashicorpVault::new( + { + let environment_values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + key_path.to_str().unwrap().to_owned(), + ), + ("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()), + ]); + Arc::new(move |name: &str| environment_values.get(name).cloned()) + }, + true, + ) + .unwrap(); + + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + assert_eq!( + manager.login_url().as_deref(), + Some(format!("{}/v1/auth/cert/login", server.uri()).as_str()) + ); +} + +#[rstest::rstest] +#[case::missing(404, json!({}), 0)] +#[case::malformed(200, json!({}), 1)] +#[case::missing_key(200, json!({"data": {"data": {}}}), 0)] +#[case::non_string(200, json!({"data": {"data": {"key": 1}}}), 2)] +#[tokio::test] +async fn read_responses_distinguish_absence_and_malformed_payloads( + #[case] status: u16, + #[case] body: serde_json::Value, + #[case] expected: u8, +) { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .expect(1) + .mount(&server) + .await; + let result: Result, Error> = + manager(&server, &[("HCP_VAULT_TOKEN", "token")]) + .async_read_secret("name") + .await; + match expected { + 0 => assert!(result.unwrap().is_none()), + 1 => assert!(matches!(result, Err(Error::MalformedPayload))), + 2 => assert!(matches!(result, Err(Error::NonStringValue))), + _ => unreachable!(), + } +} + +#[tokio::test] +async fn write_and_delete_invalidate_the_read_cache() { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/secret/data/name")) + .and(body_json( + json!({"data": {"key": "updated", "description": "description"}}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": {"version": 2}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/v1/secret/data/name")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]); + + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + assert!( + manager + .async_write_secret("name", SecretValue::new("updated"), Some("description")) + .await + .is_ok() + ); + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + manager.async_delete_secret("name").await.unwrap(); +} + +#[tokio::test] +async fn no_auth_and_invalid_names_fail_without_requests() { + let server: MockServer = MockServer::start().await; + let manager: HashicorpVault = manager(&server, &[]); + + assert!(matches!( + manager.async_read_secret("name").await, + Err(Error::NoAuthConfigured) + )); + assert!(matches!( + manager.async_read_secret("../name").await, + Err(Error::InvalidSecretName(_)) + )); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn debug_output_redacts_authentication_values() { + let server: MockServer = MockServer::start().await; + let manager: HashicorpVault = HashicorpVault::with_client( + reqwest::Client::new(), + config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), + true, + ) + .unwrap(); + let debug: String = format!("{manager:?}"); + assert!(!debug.contains("token-value")); + assert!(!debug.contains("secret-id")); +} + +#[derive(Deserialize)] +struct ParityCase { + env: HashMap, + expected_secret_url: String, + expected_login_url: Option, + expected_login_namespace: Option, + expected_secret_namespace: Option, + secret_name: String, +} + +#[test] +fn configuration_matches_python_parity_fixture() { + let cases: Vec = serde_json::from_str(include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json" + ))) + .unwrap(); + for case in cases { + let values: HashMap = case.env.clone(); + let environment: Arc = + Arc::new(move |name: &str| values.get(name).cloned()); + let config: HashicorpVaultConfig = + HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); + let manager: HashicorpVault = + HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + assert_eq!( + manager.secret_url(&case.secret_name).unwrap(), + case.expected_secret_url + ); + assert_eq!(manager.login_url(), case.expected_login_url); + assert_eq!( + manager.config().login_namespace(), + case.expected_login_namespace.as_deref() + ); + assert_eq!( + manager.config().secret_namespace(), + case.expected_secret_namespace.as_deref() + ); + } +} + +#[tokio::test] +#[ignore] +async fn live_vault_round_trip() { + let environment: Arc = + Arc::new(litellm_core_utils::settings::ProcessEnvironment); + let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap(); + let name: String = std::env::var("LITELLM_VAULT_LIVE_SECRET_NAME").unwrap(); + let value: SecretValue = SecretValue::new("native-live-value"); + let url: String = manager.secret_url(&name).unwrap(); + println!("native provenance: {} {}", module_path!(), url); + manager + .async_write_secret(&name, value.clone(), None) + .await + .unwrap(); + assert_eq!( + manager.async_read_secret(&name).await.unwrap().unwrap(), + value + ); + manager.async_delete_secret(&name).await.unwrap(); + assert!(manager.async_read_secret(&name).await.unwrap().is_none()); +} + +const TEST_CERTIFICATE: &str = "-----BEGIN CERTIFICATE----- +MIIDDzCCAfegAwIBAgIUeMzLFLM/mRbPGbNAew5N2UTscocwDQYJKoZIhvcNAQEL +BQAwFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MB4XDTI2MDkyMTIwMjA1OVoXDTI2 +MDkyMjIwMjA1OVowFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MIIBIjANBgkqhkiG +9w0BAQEFAAOCAQ8AMIIBCgKCAQEAveYoSUJXybmkHmQsBfhBcv2Ob5Oy8ejZu+B3 +vTnrPumW4ANi1XXKBSazRGB3fEtAgr+3KhKeHaSKEQeBwJkAEBfdmQv0tpXICwHs +1kFNtU0owy54HVW5/ia+LMszsFcPzVIoMnbUOuiKr9RaV7P+IEFzILPBVuV4DoYH +yocjD3+9QNqokWgNL8LK37JijmNEFVaKFz0X6SyL2VRDlfPWTEBK52Gp/pvDgA6G +eTSfyI+kCm9h5ECTYUAtmatk9WPVS8sWOqV1EXVanFyYBU+mDxoywAS1/6CHeIPh +bNmCOZjPoO9qWBJ7ZyGhOconBigXY8qnlXymev+44IPHrx4urwIDAQABo1MwUTAd +BgNVHQ4EFgQUvaZrZ6HKtbr3ekeZmgy4b5Pq95QwHwYDVR0jBBgwFoAUvaZrZ6HK +tbr3ekeZmgy4b5Pq95QwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOC +AQEAEejrD8d1qDxW55XxQ4IC31rufoEvDV955jyvh2kALPaN/i5oWsBGI+UAQZna +aaoQXwzlmHrtDUBWl0LztVTUamIleUep2+PLLauqqt43vxppxMX8Jn2mnPO20YE/ +hIzGx0jN/LBG8PDyLSvHdlgjP9ofA4Vg4rTQugdXRgOvlCE/epnH/MADcg9KYJtJ +C1RObCIkL3LcdUbjStJRCY/U/FeWcgyncEPz95OFDkbrlNDajb6o6CkYfouqvhTc +8XlgjjAVKIbAbRgbVu3elsquuFM97x2DzWDjkrMNmDt1FJ9ubK36gL6B3o0UMaoQ +00R7x/eqvH+EkWa/2ekW9lpleQ== +-----END CERTIFICATE----- +"; + +const TEST_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC95ihJQlfJuaQe +ZCwF+EFy/Y5vk7Lx6Nm74He9Oes+6ZbgA2LVdcoFJrNEYHd8S0CCv7cqEp4dpIoR +B4HAmQAQF92ZC/S2lcgLAezWQU21TSjDLngdVbn+Jr4syzOwVw/NUigydtQ66Iqv +1FpXs/4gQXMgs8FW5XgOhgfKhyMPf71A2qiRaA0vwsrfsmKOY0QVVooXPRfpLIvZ +VEOV89ZMQErnYan+m8OADoZ5NJ/Ij6QKb2HkQJNhQC2Zq2T1Y9VLyxY6pXURdVqc +XJgFT6YPGjLABLX/oId4g+Fs2YI5mM+g72pYEntnIaE5yicGKBdjyqeVfKZ6/7jg +g8evHi6vAgMBAAECggEAGdJjlP6b8Fa5bdaCM/ebcrbuuNZVJVbb0JPHxGfNSLs7 +pE9hj5QaOdQW2Uviw3h6F61ZCzQH4xD+Iy2po5ZKb2XHYKnDB1bboj+LRGER337T +9aJqe9at2VTMVEv3Rdm40NsEk0QcPLxlK16NQFK90gYEUSSQPDAswJDSG2R/zHn+ +vADI907mW/goEJHeLn8PWGlNlSiR6x+5JJtq+GXCzUzVvJYQSCLGxCSl2x2H+0g7 +NhFI0zPpdzNmO/h+yhzaFb6Rp5U8+ZsnZ3qYjQ/03gw1myTDKJt1YaO9JvArnNYX +hcJQQ8Rt0bHhcrZA16bBOpqZlo5pKCicwI/netgN8QKBgQDcFz7AzdJ26sMSV32V +rwrMgIoggt8qDjO1ARwqW35A1TIge0FoW4M4KpsXQGGfT341uU1esXEcyZ/1L/5X +3ql2gX4DbOYLZLWYzZGR2hq33oi8HkhN98QrEwL9emSH8NqYX3Xxja3PrmCrSYJe +Zbnd9TIm2XkxyMoyXJu6M/QvnwKBgQDc4dzqTbxoGEGa5MuJoGmMwPnqgdG9UM5J +eExVnh7osxc2sOdsiPeRjjQTxs9v2kJwctC359OJoo9yGaaJeSghU4LEWJo1sqnA +fzSCLammYvtVAtniyNv5Mxk/6Uimi4NNDKaAKB+m4K2uSn3U9AmY7KPYMGaSbS9W +XSnobjxm8QKBgC8bPpAvvWs8ZhIn7bY659nLbUT2HeO3dHO6UBf0yzn/J6JyHxbB +93zvCZDZc8uQTRgcmCW7XtVlhjoJUqvl+Wlm39zF0xr/LCsPXKfWAb/2/lcdOCaP +8Emz4QD10EyUTYUtcWYJB/mafhBLRH8F0Nlj4J8WDu2L51MOJTqeYhZLAoGAWffN +icocAbJPlo22sdoa4+/+W5yBF8GAJMDRJtZ+9H1t6SLpQHYRkMIBSETkXUTjZvX9 +Ocs9iIQkNW9pO/mTdO+VBfCo71JUfknR02xR+6m5gYjlws/ZeYlssXGN2/hbhNiw +QOcW7Vv6olFJK6Iy/oz0t6wPO3kpnN3Zogi0paECgYEAwo44M1DdYCtV0snhmYM9 +5u0mPfYt5P2SVLXyUbr+vFTfrTL/WKnXIJgbsnj3Gvf+GIZv9tKcXhSNmEHQCYX4 +X3w9iTPddCHuvZ1fpufi2TyArJh0OkoNtLXJHTKrHjf2N+61AQzFiv5WieJrdE+H +qr32PTUuVGPyO9LyTY4/RL0= +-----END PRIVATE KEY----- +"; diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index a7e7ec80636..7369f9e1db6 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -9,11 +9,13 @@ repository.workspace = true default = [] aws = ["dep:litellm-secrets-aws"] google = ["dep:litellm-secrets-google"] +hashicorp = ["dep:litellm-secrets-hashicorp"] [dependencies] litellm-secrets-types.workspace = true litellm-secrets-aws = { workspace = true, optional = true } litellm-secrets-google = { workspace = true, optional = true } +litellm-secrets-hashicorp = { workspace = true, optional = true } litellm-core-utils.workspace = true base64.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md index 183a39e15bb..0d333c7d116 100644 --- a/litellm-rust/crates/secrets/README.md +++ b/litellm-rust/crates/secrets/README.md @@ -9,3 +9,5 @@ Backend failures propagate by default. To allow fallback during a backend failur `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 + +The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV v2 values from `HCP_VAULT_*` environment variables. It supports static tokens, AppRole authentication, and TLS certificate authentication diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 0c6e681b8aa..e5b82580cfe 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -30,4 +30,7 @@ pub enum Error { #[cfg(feature = "google")] #[error(transparent)] Google(#[from] litellm_secrets_google::Error), + #[cfg(feature = "hashicorp")] + #[error(transparent)] + Hashicorp(#[from] litellm_secrets_hashicorp::Error), } diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 943ffdf6158..c035d4cb586 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -13,6 +13,8 @@ pub enum SecretManager { GoogleKms(crate::google::GoogleKms), #[cfg(feature = "google")] GoogleSecretManager(crate::google::GoogleSecretManager), + #[cfg(feature = "hashicorp")] + HashicorpVault(crate::hashicorp::HashicorpVault), } impl SecretManager { @@ -27,6 +29,8 @@ impl SecretManager { Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, #[cfg(feature = "google")] Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + #[cfg(feature = "hashicorp")] + Self::HashicorpVault(_) => KeyManagementSystem::HashicorpVault, } } } @@ -78,6 +82,12 @@ pub async fn get_secret_from_manager( .get_secret_from_google_secret_manager(secret_name) .await .map_err(Error::from), + #[cfg(feature = "hashicorp")] + SecretManager::HashicorpVault(client) => client + .async_read_secret(secret_name) + .await + .map(|value| value.map(Secret::String)) + .map_err(Error::from), } } diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index ff2e95f7b2f..6cb391ff844 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -19,3 +19,5 @@ pub use state::{SecretManagerState, secret_manager_would_be_consulted}; pub use litellm_secrets_aws as aws; #[cfg(feature = "google")] pub use litellm_secrets_google as google; +#[cfg(feature = "hashicorp")] +pub use litellm_secrets_hashicorp as hashicorp; diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index a2cbbd843e1..58941171fd9 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -105,3 +105,131 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites Err(Error::MissingCiphertext) )); } + +#[cfg(feature = "hashicorp")] +#[tokio::test] +async fn hashicorp_handler_resolves_found_missing_and_failed_values() { + use std::sync::Arc; + + use litellm_core_utils::settings::Lookup; + use litellm_secrets::{ + Error, FailurePolicy, KeyManagementSettings, SecretManager, SecretManagerState, + SecretResolver, hashicorp::HashicorpVault, hashicorp::HashicorpVaultConfig, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, + }; + + let found_server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/KEY")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"data": {"data": {"key": "remote"}}})), + ) + .mount(&found_server) + .await; + let found_environment: Arc = Arc::new({ + let address = found_server.uri(); + move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(address.clone()), + "HCP_VAULT_TOKEN" => Some("token".into()), + _ => None, + } + }); + let found_config = HashicorpVaultConfig::from_environment(found_environment.as_ref()).unwrap(); + let found_manager = + HashicorpVault::with_client(reqwest::Client::new(), found_config, true).unwrap(); + let found_resolver = SecretResolver::new( + Arc::new(SecretManagerState::new( + SecretManager::HashicorpVault(found_manager), + KeyManagementSettings { + hosted_keys: Some(vec!["KEY".into()]), + ..Default::default() + }, + )), + Arc::new(|_: &str| None), + litellm_secrets::OidcResolver::default(), + ); + assert_eq!( + found_resolver + .get_secret_str("KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "remote" + ); + + let missing_server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(404)) + .mount(&missing_server) + .await; + let missing_environment: Arc = Arc::new({ + let address = missing_server.uri(); + move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(address.clone()), + "HCP_VAULT_TOKEN" => Some("token".into()), + _ => None, + } + }); + let missing_config = + HashicorpVaultConfig::from_environment(missing_environment.as_ref()).unwrap(); + let missing_manager = + HashicorpVault::with_client(reqwest::Client::new(), missing_config, true).unwrap(); + let missing_state = SecretManagerState::new( + SecretManager::HashicorpVault(missing_manager), + KeyManagementSettings { + hosted_keys: Some(vec!["KEY".into()]), + ..Default::default() + }, + ); + let missing = litellm_secrets::get_secret_from_manager( + missing_state.backend().unwrap(), + "KEY", + missing_state.settings().unwrap(), + &|_: &str| None, + ) + .await + .unwrap(); + assert!(missing.is_none()); + + let failed_server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(500)) + .mount(&failed_server) + .await; + let failed_environment: Arc = Arc::new({ + let address = failed_server.uri(); + move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(address.clone()), + "HCP_VAULT_TOKEN" => Some("token".into()), + _ => None, + } + }); + let failed_config = + HashicorpVaultConfig::from_environment(failed_environment.as_ref()).unwrap(); + let failed_manager = + HashicorpVault::with_client(reqwest::Client::new(), failed_config, true).unwrap(); + let failed_state = SecretManagerState::new( + SecretManager::HashicorpVault(failed_manager), + KeyManagementSettings { + hosted_keys: Some(vec!["KEY".into()]), + ..Default::default() + }, + ); + let failed_resolver = SecretResolver::new( + Arc::new(failed_state), + Arc::new(|_: &str| None), + litellm_secrets::OidcResolver::default(), + ) + .with_failure_policy(FailurePolicy::Propagate); + assert!(matches!( + failed_resolver.get_secret_str("KEY", None).await, + Err(Error::Hashicorp( + litellm_secrets::hashicorp::Error::Status { status: 500 } + )) + )); +} diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index e37a912c7e1..259b2416d7f 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -95,16 +95,16 @@ class HashicorpSecretManager(BaseSecretManager): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user # Vault-specific config - self.vault_addr = os.getenv("HCP_VAULT_ADDR", "http://127.0.0.1:8200") + self.vault_addr = (os.getenv("HCP_VAULT_ADDR") or "http://127.0.0.1:8200").rstrip("/") self.vault_token = os.getenv("HCP_VAULT_TOKEN", "") - self.vault_namespace = os.getenv("HCP_VAULT_NAMESPACE", None) - self.login_namespace_override = os.getenv("HCP_VAULT_LOGIN_NAMESPACE", None) - self.secret_namespace_override = os.getenv("HCP_VAULT_SECRET_NAMESPACE", None) + self.vault_namespace = self._sanitize_path_component(os.getenv("HCP_VAULT_NAMESPACE")) + self.login_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_LOGIN_NAMESPACE")) + self.secret_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_SECRET_NAMESPACE")) # KV engine mount name (default: "secret") # If your KV engine is mounted somewhere other than "secret", set HCP_VAULT_MOUNT_NAME - self.vault_mount_name = os.getenv("HCP_VAULT_MOUNT_NAME", "secret") + self.vault_mount_name = self._sanitize_path_component(os.getenv("HCP_VAULT_MOUNT_NAME")) or "secret" # Optional path prefix for secrets (e.g., "myapp" -> secret/data/myapp/{secret_name}) - self.vault_path_prefix = os.getenv("HCP_VAULT_PATH_PREFIX", None) + self.vault_path_prefix = self._sanitize_path_component(os.getenv("HCP_VAULT_PATH_PREFIX")) # Optional config for TLS cert auth self.tls_cert_path = os.getenv("HCP_VAULT_CLIENT_CERT", "") @@ -114,7 +114,9 @@ class HashicorpSecretManager(BaseSecretManager): # Optional config for AppRole auth self.approle_role_id = os.getenv("HCP_VAULT_APPROLE_ROLE_ID", "") self.approle_secret_id = os.getenv("HCP_VAULT_APPROLE_SECRET_ID", "") - self.approle_mount_path = os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH", "approle") + self.approle_mount_path = self._sanitize_path_component( + os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH") + ) or "approle" self._verify_required_credentials_exist() diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json new file mode 100644 index 00000000000..5f28d9602a0 --- /dev/null +++ b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json @@ -0,0 +1,97 @@ +[ + { + "name": "defaults", + "env": { + "HCP_VAULT_TOKEN": "token" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://127.0.0.1:8200/v1/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": null, + "expected_secret_namespace": null + }, + { + "name": "global_namespace", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_TOKEN": "token", + "HCP_VAULT_NAMESPACE": " admin/ " + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": "admin", + "expected_secret_namespace": "admin" + }, + { + "name": "namespace_overrides", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_TOKEN": "token", + "HCP_VAULT_NAMESPACE": "admin", + "HCP_VAULT_LOGIN_NAMESPACE": " /root/ ", + "HCP_VAULT_SECRET_NAMESPACE": " /teams/team-a/ " + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/teams/team-a/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": "root", + "expected_secret_namespace": "teams/team-a" + }, + { + "name": "custom_mount_and_prefix", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_TOKEN": "token", + "HCP_VAULT_MOUNT_NAME": " /kv-prod/ ", + "HCP_VAULT_PATH_PREFIX": " /virtual-keys/ " + }, + "secret_name": "DB_PASSWORD", + "expected_secret_url": "http://vault.test:8200/v1/kv-prod/data/virtual-keys/DB_PASSWORD", + "expected_login_url": null, + "expected_login_namespace": null, + "expected_secret_namespace": null + }, + { + "name": "approle_custom_mount", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_APPROLE_ROLE_ID": "role-id", + "HCP_VAULT_APPROLE_SECRET_ID": "secret-id", + "HCP_VAULT_APPROLE_MOUNT_PATH": " /custom-approle/ ", + "HCP_VAULT_NAMESPACE": "admin" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", + "expected_login_url": "http://vault.test:8200/v1/auth/custom-approle/login", + "expected_login_namespace": "admin", + "expected_secret_namespace": "admin" + }, + { + "name": "tls_cert", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_CLIENT_CERT": "/tmp/client.crt", + "HCP_VAULT_CLIENT_KEY": "/tmp/client.key", + "HCP_VAULT_CERT_ROLE": "vault-role", + "HCP_VAULT_NAMESPACE": "admin" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", + "expected_login_url": "http://vault.test:8200/v1/auth/cert/login", + "expected_login_namespace": "admin", + "expected_secret_namespace": "admin" + }, + { + "name": "trailing_address_slash", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200///", + "HCP_VAULT_TOKEN": "token" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": null, + "expected_secret_namespace": null + } +] diff --git a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py b/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py index 1676540e4ec..fc18cb8b8f7 100644 --- a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py @@ -1,4 +1,5 @@ import datetime +import json from collections.abc import Mapping from pathlib import Path from typing import Final @@ -18,6 +19,23 @@ LOGIN_RESPONSE: Final = {"auth": {"client_token": "hvs.login-token", "lease_dura SECRET_RESPONSE: Final = {"data": {"data": {"key": "sk-from-vault", "password": "pw-from-vault"}}} NAMESPACE_ENV_VARS: Final = ("HCP_VAULT_NAMESPACE", "HCP_VAULT_LOGIN_NAMESPACE", "HCP_VAULT_SECRET_NAMESPACE") +PARITY_ENV_VARS: Final = ( + "HCP_VAULT_ADDR", + "HCP_VAULT_TOKEN", + "HCP_VAULT_NAMESPACE", + "HCP_VAULT_LOGIN_NAMESPACE", + "HCP_VAULT_SECRET_NAMESPACE", + "HCP_VAULT_MOUNT_NAME", + "HCP_VAULT_PATH_PREFIX", + "HCP_VAULT_APPROLE_ROLE_ID", + "HCP_VAULT_APPROLE_SECRET_ID", + "HCP_VAULT_APPROLE_MOUNT_PATH", + "HCP_VAULT_CLIENT_CERT", + "HCP_VAULT_CLIENT_KEY", + "HCP_VAULT_CERT_ROLE", + "HCP_VAULT_REFRESH_INTERVAL", + "SECRET_MANAGER_REFRESH_INTERVAL", +) def _build_manager(monkeypatch: pytest.MonkeyPatch, env: Mapping[str, str]) -> HashicorpSecretManager: @@ -236,3 +254,35 @@ def test_tls_login_uses_login_namespace(monkeypatch: pytest.MonkeyPatch, tmp_pat assert manager._auth_via_tls_cert() == "hvs.login-token" assert login_route.calls.last.request.headers["X-Vault-Namespace"] == "root" + + +with Path(__file__).with_name("hashicorp_vault_parity.json").open() as parity_file: + PARITY_CASES: Final = json.load(parity_file) + + +@pytest.mark.parametrize("case", PARITY_CASES, ids=lambda case: case["name"]) +def test_configuration_matches_native_parity_fixture( + monkeypatch: pytest.MonkeyPatch, case: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + for name in PARITY_ENV_VARS: + monkeypatch.delenv(name, raising=False) + for name, value in case["env"].items(): + monkeypatch.setenv(name, value) + + manager: Final = HashicorpSecretManager() + env: Final = case["env"] + expected_login_url: Final = case["expected_login_url"] + if env.get("HCP_VAULT_APPROLE_ROLE_ID") and env.get("HCP_VAULT_APPROLE_SECRET_ID"): + login_url: str | None = ( + f"{manager.vault_addr}/v1/auth/{manager.approle_mount_path}/login" + ) + elif env.get("HCP_VAULT_CLIENT_CERT") and env.get("HCP_VAULT_CLIENT_KEY"): + login_url = f"{manager.vault_addr}/v1/auth/cert/login" + else: + login_url = None + + assert manager.get_url(case["secret_name"]) == case["expected_secret_url"] + assert manager.vault_login_namespace == case["expected_login_namespace"] + assert manager.vault_secret_namespace == case["expected_secret_namespace"] + assert login_url == expected_login_url From beba2576be4b015574dccb3f916af5dfaf5d189f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:29:06 +0000 Subject: [PATCH 045/109] fix(rust): align vault namespace handling with python and drop unused settings hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/secrets-hashicorp/src/config.rs | 10 +----- .../secrets-hashicorp/src/secret_manager.rs | 31 ++++++++--------- .../secrets-hashicorp/tests/secret_manager.rs | 33 +++++++++++++++++++ .../hashicorp_secret_manager.py | 16 ++++----- .../hashicorp_vault_parity.json | 20 +++-------- 5 files changed, 58 insertions(+), 52 deletions(-) diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index f9491f71afb..d32491199c4 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -1,7 +1,6 @@ use std::{path::PathBuf, time::Duration}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::KeyManagementSettings; use litellm_secrets_types::SecretValue; use crate::Error; @@ -117,13 +116,6 @@ impl HashicorpVaultConfig { }) } - pub fn from_settings( - _settings: &KeyManagementSettings, - environment: &dyn Lookup, - ) -> Result { - Self::from_environment(environment) - } - pub fn login_namespace(&self) -> Option<&str> { self.login_namespace .as_deref() @@ -163,7 +155,7 @@ fn refresh_interval(environment: &dyn Lookup) -> Result { }; let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?; if seconds < 0 { - return Ok(Duration::from_nanos(1)); + return Err(Error::RefreshInterval); } Ok(Duration::from_secs(seconds as u64)) } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 6b887289443..e4ca8c387c3 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -120,14 +120,12 @@ impl HashicorpVault { return Ok(Some(value)); } let token: SecretValue = self.vault_token().await?; - let mut request = self + let response: reqwest::Response = self .client .get(&url) - .header("X-Vault-Token", token.expose()); - if let Some(namespace) = self.config.secret_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } - let response: reqwest::Response = request.send().await?; + .header("X-Vault-Token", token.expose()) + .send() + .await?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } @@ -166,14 +164,13 @@ impl HashicorpVault { None => json!({"key": value.expose()}), }; let token: SecretValue = self.vault_token().await?; - let mut request = self + let response: reqwest::Response = self .client .post(&url) - .header("X-Vault-Token", token.expose()); - if let Some(namespace) = self.config.secret_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } - let response: reqwest::Response = request.json(&json!({"data": data})).send().await?; + .header("X-Vault-Token", token.expose()) + .json(&json!({"data": data})) + .send() + .await?; if !response.status().is_success() { return Err(Error::Status { status: response.status().as_u16(), @@ -190,14 +187,12 @@ impl HashicorpVault { pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> { let url: String = self.secret_url(secret_name)?; let token: SecretValue = self.vault_token().await?; - let mut request = self + let response: reqwest::Response = self .client .delete(&url) - .header("X-Vault-Token", token.expose()); - if let Some(namespace) = self.config.secret_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } - let response: reqwest::Response = request.send().await?; + .header("X-Vault-Token", token.expose()) + .send() + .await?; if !response.status().is_success() { return Err(Error::Status { status: response.status().as_u16(), diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index 1cf13a79c77..b7a81f49ad4 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -87,6 +87,39 @@ async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { assert!(manager.async_read_secret("name").await.unwrap().is_some()); } +#[test] +fn trailing_address_slashes_are_removed() { + let environment: Arc = Arc::new(|name: &str| match name { + "HCP_VAULT_ADDR" => Some("http://vault.test:8200///".to_owned()), + "HCP_VAULT_TOKEN" => Some("token".to_owned()), + _ => None, + }); + let config: HashicorpVaultConfig = + HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); + let manager: HashicorpVault = + HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + + assert_eq!( + manager.secret_url("name").unwrap(), + "http://vault.test:8200/v1/secret/data/name" + ); +} + +#[rstest::rstest] +#[case("-1")] +#[case("not-a-number")] +fn invalid_refresh_intervals_are_rejected(#[case] value: &str) { + let environment: Arc = Arc::new(move |name: &str| match name { + "HCP_VAULT_REFRESH_INTERVAL" => Some(value.to_owned()), + _ => None, + }); + + assert!(matches!( + HashicorpVaultConfig::from_environment(environment.as_ref()), + Err(Error::RefreshInterval) + )); +} + #[tokio::test] async fn approle_login_uses_namespace_and_reuses_the_token() { let server: MockServer = MockServer::start().await; diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index 259b2416d7f..e37a912c7e1 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -95,16 +95,16 @@ class HashicorpSecretManager(BaseSecretManager): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user # Vault-specific config - self.vault_addr = (os.getenv("HCP_VAULT_ADDR") or "http://127.0.0.1:8200").rstrip("/") + self.vault_addr = os.getenv("HCP_VAULT_ADDR", "http://127.0.0.1:8200") self.vault_token = os.getenv("HCP_VAULT_TOKEN", "") - self.vault_namespace = self._sanitize_path_component(os.getenv("HCP_VAULT_NAMESPACE")) - self.login_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_LOGIN_NAMESPACE")) - self.secret_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_SECRET_NAMESPACE")) + self.vault_namespace = os.getenv("HCP_VAULT_NAMESPACE", None) + self.login_namespace_override = os.getenv("HCP_VAULT_LOGIN_NAMESPACE", None) + self.secret_namespace_override = os.getenv("HCP_VAULT_SECRET_NAMESPACE", None) # KV engine mount name (default: "secret") # If your KV engine is mounted somewhere other than "secret", set HCP_VAULT_MOUNT_NAME - self.vault_mount_name = self._sanitize_path_component(os.getenv("HCP_VAULT_MOUNT_NAME")) or "secret" + self.vault_mount_name = os.getenv("HCP_VAULT_MOUNT_NAME", "secret") # Optional path prefix for secrets (e.g., "myapp" -> secret/data/myapp/{secret_name}) - self.vault_path_prefix = self._sanitize_path_component(os.getenv("HCP_VAULT_PATH_PREFIX")) + self.vault_path_prefix = os.getenv("HCP_VAULT_PATH_PREFIX", None) # Optional config for TLS cert auth self.tls_cert_path = os.getenv("HCP_VAULT_CLIENT_CERT", "") @@ -114,9 +114,7 @@ class HashicorpSecretManager(BaseSecretManager): # Optional config for AppRole auth self.approle_role_id = os.getenv("HCP_VAULT_APPROLE_ROLE_ID", "") self.approle_secret_id = os.getenv("HCP_VAULT_APPROLE_SECRET_ID", "") - self.approle_mount_path = self._sanitize_path_component( - os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH") - ) or "approle" + self.approle_mount_path = os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH", "approle") self._verify_required_credentials_exist() diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json index 5f28d9602a0..f1faefd48e8 100644 --- a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json +++ b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json @@ -15,7 +15,7 @@ "env": { "HCP_VAULT_ADDR": "http://vault.test:8200", "HCP_VAULT_TOKEN": "token", - "HCP_VAULT_NAMESPACE": " admin/ " + "HCP_VAULT_NAMESPACE": "admin" }, "secret_name": "OPENAI_API_KEY", "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", @@ -29,8 +29,8 @@ "HCP_VAULT_ADDR": "http://vault.test:8200", "HCP_VAULT_TOKEN": "token", "HCP_VAULT_NAMESPACE": "admin", - "HCP_VAULT_LOGIN_NAMESPACE": " /root/ ", - "HCP_VAULT_SECRET_NAMESPACE": " /teams/team-a/ " + "HCP_VAULT_LOGIN_NAMESPACE": "root", + "HCP_VAULT_SECRET_NAMESPACE": "teams/team-a" }, "secret_name": "OPENAI_API_KEY", "expected_secret_url": "http://vault.test:8200/v1/teams/team-a/secret/data/OPENAI_API_KEY", @@ -58,7 +58,7 @@ "HCP_VAULT_ADDR": "http://vault.test:8200", "HCP_VAULT_APPROLE_ROLE_ID": "role-id", "HCP_VAULT_APPROLE_SECRET_ID": "secret-id", - "HCP_VAULT_APPROLE_MOUNT_PATH": " /custom-approle/ ", + "HCP_VAULT_APPROLE_MOUNT_PATH": "custom-approle", "HCP_VAULT_NAMESPACE": "admin" }, "secret_name": "OPENAI_API_KEY", @@ -81,17 +81,5 @@ "expected_login_url": "http://vault.test:8200/v1/auth/cert/login", "expected_login_namespace": "admin", "expected_secret_namespace": "admin" - }, - { - "name": "trailing_address_slash", - "env": { - "HCP_VAULT_ADDR": "http://vault.test:8200///", - "HCP_VAULT_TOKEN": "token" - }, - "secret_name": "OPENAI_API_KEY", - "expected_secret_url": "http://vault.test:8200/v1/secret/data/OPENAI_API_KEY", - "expected_login_url": null, - "expected_login_namespace": null, - "expected_secret_namespace": null } ] From 6ab121a3e7364178544bcbee529f6c57b3115a0a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:29:49 +0000 Subject: [PATCH 046/109] feat(cache-redis-semantic): add native Redis Semantic cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 15 + litellm-rust/Cargo.toml | 1 + .../crates/cache-redis-semantic/Cargo.toml | 21 + .../crates/cache-redis-semantic/src/cache.rs | 599 ++++++++++++++ .../crates/cache-redis-semantic/src/lib.rs | 4 + .../crates/cache-redis-semantic/src/prompt.rs | 95 +++ .../cache-redis-semantic/tests/cache.rs | 751 ++++++++++++++++++ litellm-rust/crates/cache/tests/caching.rs | 9 +- 8 files changed, 1489 insertions(+), 6 deletions(-) create mode 100644 litellm-rust/crates/cache-redis-semantic/Cargo.toml create mode 100644 litellm-rust/crates/cache-redis-semantic/src/cache.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/lib.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/prompt.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..e911d0d9c45 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2486,6 +2486,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "serde_json", + "sha2 0.10.9", + "tokio", +] + [[package]] name = "litellm-cache-response" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..05eea6bc299 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml new file mode 100644 index 00000000000..9a8755a189e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +r2d2 = "0.8.10" +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs new file mode 100644 index 00000000000..26d19d34670 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -0,0 +1,599 @@ +use std::{ + future::Future, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; +use litellm_cache_redis::connection::{ConnectionRef, Connections, ttl_seconds}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::prompt::prompt_from_context; + +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; +const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; +const CACHE_KEY_FIELD: &str = "litellm_cache_key"; +const VECTOR_FIELD: &str = "prompt_vector"; + +pub trait Embedder: Send + Sync + 'static { + fn embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug)] +pub struct RedisSemanticConfig { + pub index_name: String, + pub similarity_threshold: f32, +} + +impl Default for RedisSemanticConfig { + fn default() -> Self { + Self { + index_name: DEFAULT_INDEX_NAME.into(), + similarity_threshold: 0.9, + } + } +} + +struct Inner { + index_name: String, + distance_threshold: f64, + resolved_index: OnceLock, + codec: ResponseCacheCodec, + clock: fn() -> f64, +} + +impl Inner { + fn new(config: RedisSemanticConfig) -> Self { + Self { + index_name: config.index_name, + distance_threshold: 1.0 - f64::from(config.similarity_threshold), + resolved_index: OnceLock::new(), + codec: ResponseCacheCodec, + clock: timestamp, + } + } + + fn ensure_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + if let Some(name) = self.resolved_index.get() { + return Ok(name.clone()); + } + let name = match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => { + create_index(connection, &self.index_name, dims)?; + self.index_name.clone() + } + }; + let _ = self.resolved_index.set(name.clone()); + Ok(name) + } + + fn isolated_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + let name = format!("{}_isolated", self.index_name); + match index_compatible(connection, &name, dims)? { + Some(true) => Ok(name), + Some(false) => { + redis::cmd("FT.DROPINDEX") + .arg(&name) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + create_index(connection, &name, dims)?; + Ok(name) + } + None => { + create_index(connection, &name, dims)?; + Ok(name) + } + } + } + + fn store( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + value: &CacheEntry, + prompt: &str, + vector: &[f32], + ttl: Option, + ) -> Result<(), Error> { + let index = self.ensure_index(connection, vector.len())?; + let entry_id = entry_id(prompt, tag); + let hash_key = format!("{index}:{entry_id}"); + let response = self.codec.encode(value)?; + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(&entry_id) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg(VECTOR_FIELD) + .arg(vector_buffer(vector)) + .arg("inserted_at") + .arg(format!("{}", (self.clock)())) + .arg("updated_at") + .arg(format!("{}", (self.clock)())) + .arg(CACHE_KEY_FIELD) + .arg(tag) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + if let Some(ttl) = ttl { + redis::cmd("EXPIRE") + .arg(&hash_key) + .arg(ttl_seconds(ttl)) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + } + + fn lookup( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + vector: &[f32], + ) -> Result, Error> { + let index = self.ensure_index(connection, vector.len())?; + let query = format!( + "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", + escape_tag(tag) + ); + let result = redis::cmd("FT.SEARCH") + .arg(&index) + .arg(query) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg(CACHE_KEY_FIELD) + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_buffer(vector)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let Some(fields) = first_document(&result) else { + return Ok(None); + }; + if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { + return Ok(None); + } + if number_field(fields, "vector_distance") + .is_none_or(|distance| distance > self.distance_threshold) + { + return Ok(None); + } + let Some(response) = bytes_field(fields, "response") else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } +} + +pub struct RedisSemanticCache { + connections: Arc>, + embedder: E, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { + Ok(Self { + connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), + embedder, + inner: Arc::new(Inner::new(config)), + }) + } +} + +impl RedisSemanticCache { + pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self { + Self { + connections: Arc::new(Connections::fixed(connection)), + embedder, + inner: Arc::new(Inner::new(config)), + } + } + + pub fn with_clock(self, clock: fn() -> f64) -> Self { + Self { + inner: Arc::new(Inner { + index_name: self.inner.index_name.clone(), + distance_threshold: self.inner.distance_threshold, + resolved_index: OnceLock::new(), + codec: self.inner.codec, + clock, + }), + ..self + } + } + + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { + context.scope.as_deref().unwrap_or(key) + } +} + +impl BaseCache + for RedisSemanticCache +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let vector = self.embedder.embed(&prompt, &context.metadata)?; + let tag = Self::tag(key, context).to_string(); + self.connections.execute(|connection| { + self.inner + .store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self.embedder.embed(&prompt, &context.metadata)?; + let tag = Self::tag(key, context).to_string(); + self.connections + .execute(|connection| self.inner.lookup(connection, &tag, &vector)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(&context) else { + return Ok(()); + }; + let vector = self + .embedder + .async_embed(&prompt, &context.metadata) + .await?; + let tag = Self::tag(key, &context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self + .embedder + .async_embed(&prompt, &context.metadata) + .await?; + let tag = Self::tag(key, context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.lookup(connection, &tag, &vector) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +fn timestamp() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(CACHE_KEY_FIELD.as_bytes()); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn vector_buffer(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn escape_tag(value: &str) -> String { + value + .chars() + .flat_map(|ch| { + if matches!( + ch, + ',' | '.' + | '<' + | '>' + | '{' + | '}' + | '[' + | ']' + | '\\' + | '"' + | '\'' + | ':' + | ';' + | '!' + | '@' + | '#' + | '$' + | '%' + | '^' + | '&' + | '*' + | '(' + | ')' + | '-' + | '+' + | '=' + | '~' + | '|' + | '/' + | ' ' + | '?' + ) { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect() +} + +fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { + redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg(VECTOR_FIELD) + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg(CACHE_KEY_FIELD) + .arg("TAG") + .arg("SEPARATOR") + .arg(",") + .query::<()>(connection) + .map_err(|_| Error::Unavailable) +} + +fn index_compatible( + connection: &mut ConnectionRef<'_>, + name: &str, + dims: usize, +) -> Result, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(connection) + { + Ok(info) => info, + Err(error) if unknown_index(&error) => return Ok(None), + Err(_) => return Err(Error::Unavailable), + }; + Ok(Some(schema_compatible(&info, dims))) +} + +fn unknown_index(error: &redis::RedisError) -> bool { + let message = error.to_string().to_lowercase(); + message.contains("unknown") && message.contains("index") +} + +fn schema_compatible(info: &redis::Value, dims: usize) -> bool { + let redis::Value::Array(entries) = info else { + return false; + }; + let attributes = entries + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) + .map(|pair| &pair[1]); + let Some(redis::Value::Array(attributes)) = attributes else { + return false; + }; + let fields = attributes + .iter() + .map(|attribute| { + let redis::Value::Array(attribute) = attribute else { + return (None, None, None); + }; + let mut name = None; + let mut field_type = None; + let mut dim = None; + for pair in attribute.as_chunks::<2>().0 { + match string_value(&pair[0]).as_deref() { + Some("identifier") => name = string_value(&pair[1]), + Some("type") => field_type = string_value(&pair[1]), + Some("dim") => dim = number_value(&pair[1]), + _ => {} + } + } + (name, field_type, dim) + }) + .collect::>(); + let has_field = |name: &str, field_type: &str| { + fields + .iter() + .any(|(n, t, _)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + }; + has_field("prompt", "TEXT") + && has_field("response", "TEXT") + && has_field("inserted_at", "NUMERIC") + && has_field("updated_at", "NUMERIC") + && has_field(CACHE_KEY_FIELD, "TAG") + && fields.iter().any(|(n, t, d)| { + n.as_deref() == Some(VECTOR_FIELD) + && t.as_deref() == Some("VECTOR") + && *d == Some(dims as f64) + }) +} + +fn string_value(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::VerbatimString { text, .. } => Some(text.clone()), + _ => None, + } +} + +fn number_value(value: &redis::Value) -> Option { + match value { + redis::Value::Int(number) => Some(*number as f64), + redis::Value::Double(number) => Some(*number), + _ => string_value(value).and_then(|text| text.parse().ok()), + } +} + +fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { + let redis::Value::Array(items) = result else { + return None; + }; + let [count, _document_id, fields, ..] = items.as_slice() else { + return None; + }; + if !matches!(count, redis::Value::Int(count) if *count > 0) { + return None; + } + match fields { + redis::Value::Array(fields) => Some(fields.as_slice()), + _ => None, + } +} + +fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { + fields + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) + .map(|pair| &pair[1]) +} + +fn string_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(string_value) +} + +fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs new file mode 100644 index 00000000000..a34603cd18f --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -0,0 +1,4 @@ +mod cache; +mod prompt; + +pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs new file mode 100644 index 00000000000..fc99898d847 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -0,0 +1,95 @@ +use litellm_cache::SemanticCacheContext; +use serde_json::Value; + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if !context.messages.is_empty() { + return Some(messages_text(&context.messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_string(); + (!prompt.is_empty()).then_some(prompt) +} + +fn messages_text(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages { + let Some(message) = message.as_object() else { + continue; + }; + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(text_content) = part.get("text").and_then(Value::as_str) { + text.push_str(text_content); + } + } + } + _ => {} + } + text.push_str(&search_results_text(message.get("search_results"))); + } + text +} + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + let mut text = String::new(); + for result in results { + let Some(result) = result.as_object() else { + continue; + }; + for key in ["source", "title"] { + if let Some(value) = result.get(key).and_then(Value::as_str) { + text.push_str(value); + } + } + if let Some(Value::Array(content)) = result.get("content") { + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations") { + text.push_str(&citations.to_string()); + } + } + text +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + Value::Array(items) => { + for item in items { + collect_input_text(item, parts); + } + } + Value::Object(map) => { + if let Some(content) = map.get("content").filter(|content| !content.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(text)) = map.get(key) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + return; + } + } + } + } + _ => {} + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs new file mode 100644 index 00000000000..77b057ae3b9 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -0,0 +1,751 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; +use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +const INDEX: &str = "litellm_semantic_cache_index"; + +struct FakeEmbedder { + vectors: HashMap>, + calls: Arc>>, +} + +impl FakeEmbedder { + fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + vectors: vectors + .iter() + .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) + .collect(), + calls: Arc::clone(&calls), + }, + calls, + ) + } +} + +impl Embedder for FakeEmbedder { + fn embed(&self, prompt: &str, _: &serde_json::Map) -> Result, Error> { + self.calls.lock().unwrap().push(prompt.to_string()); + + Ok(self + .vectors + .get(prompt) + .cloned() + .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) + } + + async fn async_embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> Result, Error> { + self.embed(prompt, metadata) + } +} + +fn config() -> RedisSemanticConfig { + RedisSemanticConfig { + index_name: INDEX.into(), + similarity_threshold: 0.9, + } +} + +fn messages_context(messages: Vec) -> SemanticCacheContext { + SemanticCacheContext { + messages, + ..Default::default() + } +} + +fn entry() -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "yes"}), + } +} + +fn encoded(entry: &CacheEntry) -> Vec { + ResponseCacheCodec.encode(entry).unwrap() +} + +fn vector_bytes(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(b"litellm_cache_key"); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn s(value: &str) -> redis::Value { + redis::Value::BulkString(value.as_bytes().to_vec()) +} + +fn unknown_index_error() -> redis::RedisError { + redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) +} + +fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { + let mut parts = vec![ + s("identifier"), + s(name), + s("attribute"), + s(name), + s("type"), + s(field_type), + ]; + parts.extend(extra); + redis::Value::Array(parts) +} + +fn index_info(attributes: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("index_name"), + s(INDEX), + s("attributes"), + redis::Value::Array(attributes), + ]) +} + +fn vector_attribute(dims: i64) -> redis::Value { + attribute( + "prompt_vector", + "VECTOR", + vec![ + s("algorithm"), + s("FLAT"), + s("data_type"), + s("FLOAT32"), + s("dim"), + redis::Value::Int(dims), + s("distance_metric"), + s("COSINE"), + ], + ) +} + +fn compatible_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + attribute("litellm_cache_key", "TAG", vec![]), + ]) +} + +fn unscoped_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + ]) +} + +fn create_index_command(name: &str, dims: usize) -> redis::Cmd { + let mut command = redis::cmd("FT.CREATE"); + command + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg("prompt_vector") + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg("litellm_cache_key") + .arg("TAG") + .arg("SEPARATOR") + .arg(","); + command +} + +fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { + let mut command = redis::cmd("FT.SEARCH"); + command + .arg(index) + .arg(format!( + "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" + )) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg("litellm_cache_key") + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_bytes(vector)); + command +} + +fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s("hello prompt"), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]) +} + +fn search_result(fields: redis::Value) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + s("litellm_semantic_cache_index:stored-id"), + fields, + ]) +} + +fn empty_result() -> redis::Value { + redis::Value::Array(vec![redis::Value::Int(0)]) +} + +#[test] +fn store_creates_index_and_writes_hash_with_expire() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(5)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + cache.set_cache(tag, value, &context).unwrap(); +} + +#[test] +fn store_without_ttl_skips_expire() { + let prompt = "hello prompt"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) + .arg("entry_id") + .arg(entry_id(prompt, "key1")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("key1"), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + "key1", + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn lookup_returns_hit_below_distance_threshold() { + let vector = vec![0.1f32, 0.2, 0.3]; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let hit = cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), + ) + .unwrap(); + assert_eq!(hit, Some(value)); +} + +#[test] +fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), + ), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "other", + "0.05", + encoded(&entry()), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); + + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); +} + +#[test] +fn lookup_returns_invalid_entry_on_malformed_response() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "key1", + "0.05", + b"not json!".to_vec(), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn missing_prompt_is_noop_and_never_embeds() { + let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let context = SemanticCacheContext::default(); + cache.set_cache("key1", entry(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert!(calls.lock().unwrap().is_empty()); +} + +#[test] +fn scope_overrides_key_as_filter_tag() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) + .arg("entry_id") + .arg(entry_id(prompt, "scope-a")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("scope-a"), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, "scope\\-a", &vector), + Ok(search_result(hit_fields( + "scope-a", + "0.05", + encoded(&value), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = SemanticCacheContext { + scope: Some("scope-a".into()), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + + cache.set_cache("key1", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); +} + +#[test] +fn incompatible_schema_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn tag_special_characters_are_escaped_in_search_filter() { + let vector = vec![0.1f32, 0.2, 0.3]; + let tag = "a:b, c|d"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), + Ok(empty_result()), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + tag, + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap(), + None + ); +} + +#[test] +fn prompt_extraction_matches_python_message_and_input_shapes() { + let vector = vec![0.1f32, 0.2, 0.3]; + let lookups = 5; + let mut commands = vec![MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(compatible_info(3)), + )]; + for _ in 0..lookups { + commands.push(MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(empty_result()), + )); + } + let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + cache + .get_cache( + "key1", + &messages_context(vec![ + json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), + json!({"role": "assistant", "content": "reply"}), + ]), + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!(" plain input ")), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some( + json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), + ), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!({"output_text": " result text "})), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + ) + .unwrap(); + + assert_eq!( + *calls.lock().unwrap(), + vec![ + "firstsecondreply", + "plain input", + "nested\ntail", + "result text", + "questionsrctfound{\"a\":1}", + ] + ); +} + +#[test] +fn ttl_passes_through_context_only() { + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection( + MockRedisConnection::new(Vec::::new()), + embedder, + config(), + ); + assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&SemanticCacheContext { + ttl: Some(Duration::from_secs(9)), + ..Default::default() + }), + Some(Duration::from_secs(9)) + ); +} + +#[tokio::test] +async fn async_paths_embed_then_run_blocking_redis_work() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, tag, &vector), + Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = messages_context(vec![json!({"role": "user", "content": prompt})]); + + cache + .async_set_cache(tag, value.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache(tag, &context).await.unwrap(), + Some(value) + ); +} + +#[test] +fn live_store_lookup_and_ttl_against_redis_stack() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + let vector = vec![0.1f32, 0.2, 0.3, 0.4]; + let prompt = "rust semantic cache live prompt"; + let tag = "live-key"; + let index_name = format!("rust_semantic_test_{}", std::process::id()); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: index_name.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap(); + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(120)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + let value = entry(); + + cache.set_cache(tag, value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + let ttl: i64 = redis::Commands::ttl( + &mut connection, + format!("{index_name}:{}", entry_id(prompt, tag)), + ) + .unwrap(); + assert!( + ttl > 0, + "expected stored hash to carry an expiry, got {ttl}" + ); +} diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 2e65b4eeae5..33171f36f46 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,8 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, - SemanticCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, + get_cache, }; struct TestCache { @@ -132,10 +132,7 @@ fn semantic_context_with_ttl_preserves_lookup_inputs() { let context = SemanticCacheContext { input: Some(serde_json::json!("text")), messages: vec![serde_json::json!({"role": "user", "content": "hi"})], - metadata: serde_json::Map::from_iter([( - "key".into(), - serde_json::json!("value"), - )]), + metadata: serde_json::Map::from_iter([("key".into(), serde_json::json!("value"))]), scope: Some("scope".into()), ttl: None, }; From 4db34812449038fed2c724a3ef099fefb3198ac2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:33:32 +0000 Subject: [PATCH 047/109] feat(python-bridge): serve ValkeySemanticCache natively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/config.rs | 84 +++++- .../python-bridge/src/cache/embedder.rs | 58 ++++ .../crates/python-bridge/src/cache/facade.rs | 58 +++- .../crates/python-bridge/src/cache/handle.rs | 44 ++- .../crates/python-bridge/src/cache/mod.rs | 4 +- .../crates/python-bridge/src/cache/native.rs | 253 +++++++++++++----- .../crates/python-bridge/src/cache/request.rs | 39 ++- .../test_valkey_semantic_cache.py | 149 +++++++++++ 10 files changed, 598 insertions(+), 95 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/embedder.rs create mode 100644 tests/test_litellm_rust/test_valkey_semantic_cache.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5daf691d4c3..e4c9c385f1f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2685,6 +2685,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", + "litellm-cache-valkey-semantic", "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", @@ -2695,6 +2696,7 @@ dependencies = [ "litellm-types", "pyo3", "pyo3-async-runtimes", + "redis", "rstest", "serde", "serde_json", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..ce2405f33c4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,6 +24,7 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true +litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" } serde.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy-python.workspace = true @@ -37,6 +38,7 @@ litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..7218805b8df 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,18 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] +pub(super) struct ValkeySemanticCacheConfig { + pub(super) similarity_threshold: f64, + pub(super) index_name: String, + pub(super) embedding_model: String, + pub(super) connection: RedisConnectionConfig, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + ValkeySemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -142,9 +151,15 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some( CacheType::RedisSemantic - | CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic @@ -158,11 +173,13 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) + if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) + && service.default_ttl() + != Some(match &self.backend { + CacheBackendConfig::Memory(config) => config.default_ttl, + CacheBackendConfig::Redis(config) => config.default_ttl, + CacheBackendConfig::ValkeySemantic(_) => Duration::ZERO, + }) { return Some("facade and native backend default TTLs must match"); } @@ -185,6 +202,16 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::ValkeySemantic(config) => { + if service.kind() != "valkey-semantic" { + return Some("facade and native backend types must match"); + } + let Some((threshold, index_name)) = service.semantic_config() else { + return Some("facade and native backend types must match"); + }; + (threshold != config.similarity_threshold || index_name != config.index_name) + .then_some("facade and native semantic settings must match") + } } } } @@ -299,6 +326,51 @@ fn project_redis( })) } +#[inline(never)] +fn project_valkey_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let client = backend.getattr("sync_client")?; + let pool = client.getattr("connection_pool")?; + if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + if !class_is(&connection_class, "redis.connection", "Connection")? + && !class_is(&connection_class, "redis.connection", "SSLConnection")? + { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let connection = RedisConnectionConfig { + host: required_string(&resolved, "host")?, + port: u16::try_from(required_i64(&resolved, "port")?) + .map_err(|_| PyValueError::new_err("invalid Redis port"))?, + database: optional_i64(&resolved, "db")?.unwrap_or(0), + username: optional_dict_string(&resolved, "username")?, + password: optional_dict_string(&resolved, "password")?, + protocol: RedisProtocol::Resp2, + pool_size: pool.getattr("max_connections")?.extract::()?, + read_timeout: None, + connect_timeout: None, + socket_keepalive: None, + health_check_interval: Duration::ZERO, + client_name: None, + tls: None, + }; + if connection.host.is_empty() { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + Ok(Ok(ValkeySemanticCacheConfig { + similarity_threshold: backend.getattr("similarity_threshold")?.extract()?, + index_name: backend.getattr("index_name")?.extract()?, + embedding_model: backend.getattr("embedding_model")?.extract()?, + connection, + })) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs new file mode 100644 index 00000000000..d240d9d019e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -0,0 +1,58 @@ +use std::{future::Future, sync::Arc}; + +use litellm_cache::Error; +use litellm_cache_valkey_semantic::Embedder; +use litellm_host_python::to_py; +use pyo3::prelude::*; +use serde_json::Value; + +#[derive(Clone)] +pub(super) struct PythonEmbedder { + sync_embed: Arc>, + async_embed: Arc>, +} + +impl PythonEmbedder { + pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self { + sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()), + async_embed: Arc::new(backend.getattr("_get_async_embedding")?.unbind()), + }) + } +} + +impl Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + let result = Python::attach(|py| -> PyResult> { + let metadata = to_py(py, &metadata)?; + self.sync_embed + .bind(py) + .call1((prompt, metadata))? + .extract() + }) + .map_err(|_| Error::Unavailable)?; + Ok(result.into_iter().map(|value| value as f32).collect()) + } + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + let callable = Arc::clone(&self.async_embed); + let prompt = prompt.to_owned(); + let metadata = metadata.cloned(); + async move { + let future = Python::attach(|py| -> PyResult<_> { + let metadata = to_py(py, &metadata)?; + let awaitable = callable.bind(py).call1((prompt, metadata))?; + pyo3_async_runtimes::tokio::into_future(awaitable) + }) + .map_err(|_| Error::Unavailable)?; + let result = future.await.map_err(|_| Error::Unavailable)?; + let result = Python::attach(|py| result.bind(py).extract::>()) + .map_err(|_| Error::Unavailable)?; + Ok(result.into_iter().map(|value| value as f32).collect()) + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..eb76d22097d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -36,6 +36,7 @@ pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, redis_pool: Option, + redis_client_name: Option<&'static str>, } impl ObjectGuard { @@ -117,7 +118,9 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + if (instance.contains(name)? && !self.config_names.contains(&name.as_str())) + || !attributes.get_item(name)?.is(value.bind(py)) + { return Ok(false); } } @@ -138,10 +141,8 @@ impl ObjectGuard { } impl RedisPoolGuard { - fn capture(backend: &Bound<'_, PyAny>) -> PyResult { - let pool = backend - .getattr("redis_client")? - .getattr("connection_pool")?; + fn capture(backend: &Bound<'_, PyAny>, client_name: &str) -> PyResult { + let pool = backend.getattr(client_name)?.getattr("connection_pool")?; Ok(Self { reference: pool.clone().unbind(), connection_class: pool.getattr("connection_class")?.unbind(), @@ -153,10 +154,13 @@ impl RedisPoolGuard { }) } - fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { - let pool = backend - .getattr("redis_client")? - .getattr("connection_pool")?; + fn matches( + &self, + py: Python<'_>, + backend: &Bound<'_, PyAny>, + client_name: &str, + ) -> PyResult { + let pool = backend.getattr(client_name)?.getattr("connection_pool")?; Ok(self.reference.bind(py).is(&pool) && self .connection_class @@ -192,6 +196,11 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "valkey-semantic" => ( + "litellm.caching.valkey_semantic_cache", + "ValkeySemanticCache", + "valkey-semantic", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -235,11 +244,32 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "similarity_threshold", + "embedding_model", + "index_name", + "embedding_max_input_tokens", + "embedding_timeout", ], )?, - redis_pool: (kind == "redis") - .then(|| RedisPoolGuard::capture(&backend)) + redis_pool: (kind == "redis" || kind == "valkey-semantic") + .then(|| { + RedisPoolGuard::capture( + &backend, + if kind == "redis" { + "redis_client" + } else { + "sync_client" + }, + ) + }) .transpose()?, + redis_client_name: (kind == "redis" || kind == "valkey-semantic").then_some( + if kind == "redis" { + "redis_client" + } else { + "sync_client" + }, + ), }) } @@ -252,7 +282,11 @@ impl FacadeGuard { return Ok(false); } match &self.redis_pool { - Some(guard) => guard.matches(py, &backend), + Some(guard) => guard.matches( + py, + &backend, + self.redis_client_name.unwrap_or("redis_client"), + ), None => Ok(true), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..119bd35cd25 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,7 +1,10 @@ use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; -use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use super::{ + cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache, + request::duration, +}; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -51,6 +54,29 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (url, similarity_threshold, index_name, embedder))] + fn valkey_semantic( + url: String, + similarity_threshold: f64, + index_name: String, + embedder: &Bound<'_, PyAny>, + ) -> PyResult { + let python_embedder = PythonEmbedder::from_backend(embedder)?; + let service = NativeResponseCache::valkey_semantic( + &url, + similarity_threshold, + index_name, + python_embedder, + ) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -59,11 +85,17 @@ impl CacheTestHandle { fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; let guard = FacadeGuard::capture(py, facade, &service)?; - let service = service.with_redis_flush_size( - facade - .getattr("redis_flush_size")? - .extract::>()?, - ); + let service = service + .with_scope( + facade + .getattr("semantic_cache_scope")? + .extract::()?, + ) + .with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); let handle = Py::new( py, Self { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..4cc87367d91 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,6 +1,7 @@ mod binding; mod callback; mod config; +mod embedder; mod facade; mod future; mod handle; @@ -10,7 +11,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +22,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..d314cd41ac5 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,13 +1,18 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache::{ + CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext, +}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; +use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; use serde_json::Value; +use super::{embedder::PythonEmbedder, request::NativeRequest}; + #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), @@ -15,6 +20,10 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + ValkeySemantic { + cache: Arc>>, + scope: String, + }, } impl NativeResponseCache { @@ -43,41 +52,52 @@ impl NativeResponseCache { buffer: None, }) } -} -impl NativeResponseCache { - pub fn kind(&self) -> &'static str { - match self { - Self::Memory(_) => "memory", - Self::Redis { .. } => "redis", + pub fn valkey_semantic( + url: &str, + similarity_threshold: f64, + index_name: String, + embedder: PythonEmbedder, + ) -> Result { + let backend = ValkeySemanticCache::new( + url, + embedder, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold, + index_name, + }, + )?; + Ok(Self::ValkeySemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + scope: String::from("key"), + }) + } + + fn exact(request: &NativeRequest) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: ExactCacheContext { ttl: request.ttl }, + max_age: request.max_age, } } - pub fn default_ttl(&self) -> Option { - match self { - Self::Memory(cache) => cache.default_ttl(), - Self::Redis { cache, .. } => cache.default_ttl(), - } - } - - pub fn namespace(&self) -> Option<&str> { - match self { - Self::Memory(_) => None, - Self::Redis { cache, .. } => cache.backend().namespace(), - } - } - - pub fn capacity(&self) -> Option { - match self { - Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, - } - } - - pub fn max_entry_bytes(&self) -> Option { - match self { - Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } => None, + fn semantic( + request: &NativeRequest, + scope: &str, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: SemanticCacheContext { + input: request.input.clone(), + messages: request.messages.clone(), + metadata: request.metadata.clone(), + scope: Some(scope.to_owned()), + ttl: request.ttl, + }, + max_age: request.max_age, } } @@ -85,95 +105,206 @@ impl NativeResponseCache { match self { Self::Redis { cache, .. } => Self::Redis { cache, - buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), + buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))), }, - memory => memory, + value => value, } } - pub fn lookup( - &self, - request: &ResponseCacheRequest, - now: Duration, - ) -> Result, Error> { + pub fn with_scope(self, scope: String) -> Self { match self { - Self::Memory(cache) => cache.lookup(request, now), - Self::Redis { cache, .. } => cache.lookup(request, now), + Self::ValkeySemantic { cache, .. } => Self::ValkeySemantic { cache, scope }, + value => value, + } + } + + pub fn kind(&self) -> &'static str { + match self { + Self::Memory(_) => "memory", + Self::Redis { .. } => "redis", + Self::ValkeySemantic { .. } => "valkey-semantic", + } + } + + pub fn default_ttl(&self) -> Option { + match self { + Self::Memory(cache) => cache.default_ttl(), + Self::Redis { cache, .. } => cache.default_ttl(), + Self::ValkeySemantic { cache, .. } => cache.default_ttl(), + } + } + + pub fn namespace(&self) -> Option<&str> { + match self { + Self::Memory(_) | Self::ValkeySemantic { .. } => None, + Self::Redis { cache, .. } => cache.backend().namespace(), + } + } + + pub fn capacity(&self) -> Option { + match self { + Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), + Self::Redis { .. } | Self::ValkeySemantic { .. } => None, + } + } + + pub fn max_entry_bytes(&self) -> Option { + match self { + Self::Memory(cache) => cache.backend().max_entry_bytes(), + Self::Redis { .. } | Self::ValkeySemantic { .. } => None, + } + } + + pub fn semantic_config(&self) -> Option<(f64, &str)> { + match self { + Self::ValkeySemantic { cache, .. } => Some(( + cache.backend().similarity_threshold(), + cache.backend().index_name(), + )), + _ => None, + } + } + + pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result, Error> { + match self { + Self::Memory(cache) => cache.lookup(&Self::exact(request), now), + Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now), + Self::ValkeySemantic { cache, scope } => { + cache.lookup(&Self::semantic(request, scope), now) + } } } pub fn store( &self, - request: &ResponseCacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(request, response, now), - Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Memory(cache) => cache.store(&Self::exact(request), response, now), + Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), + Self::ValkeySemantic { cache, scope } => { + cache.store(&Self::semantic(request, scope), response, now) + } } } pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(requests, now), - Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Memory(cache) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), } } pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &NativeRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, + Self::ValkeySemantic { cache, scope } => { + cache + .async_lookup(&Self::semantic(request, scope), now) + .await + } } } pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Memory(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: None, - } => cache.async_store(request, response, now).await, + } => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: Some(buffer), - } => buffer.async_store(cache, request, response, now).await, + } => { + buffer + .async_store(cache, &Self::exact(request), response, now) + .await + } + Self::ValkeySemantic { cache, scope } => { + cache + .async_store(&Self::semantic(request, scope), response, now) + .await + } } } pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, - Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Memory(cache) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), } } pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(NativeRequest, Value)>, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store_batch(entries, now).await, - Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Memory(cache) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::Redis { cache, .. } => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::ValkeySemantic { cache, scope } => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::semantic(&request, scope), value)) + .collect(); + cache.async_store_batch(entries, now).await + } } } @@ -186,6 +317,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), } } @@ -193,6 +325,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::ValkeySemantic { cache, .. } => cache.test_connection().await, } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 52a5f7d9055..3e19e7fdc22 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -5,6 +5,7 @@ use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest} use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; +use serde_json::Value; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -13,24 +14,42 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + messages: Option, + input: Option, + metadata: Option, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) struct NativeRequest { + pub(super) key: CacheKeyInput, + pub(super) controls: CacheControls, + pub(super) ttl: Option, + pub(super) max_age: Option, + pub(super) messages: Option, + pub(super) input: Option, + pub(super) metadata: Option, +} + +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::::new(input.key); - if let Some(controls) = input.controls { - request.controls = controls; - } - request.context.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) +fn request_input(input: RequestInput) -> PyResult { + let controls = input.controls.unwrap_or_else(|| { + ResponseCacheRequest::::new(input.key.clone()).controls + }); + Ok(NativeRequest { + key: input.key, + controls, + ttl: input.ttl_seconds.map(duration).transpose()?, + max_age: input.max_age_seconds.map(duration).transpose()?, + messages: input.messages, + input: input.input, + metadata: input.metadata, + }) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { from_py::>(value)? .into_iter() .map(request_input) diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache.py b/tests/test_litellm_rust/test_valkey_semantic_cache.py new file mode 100644 index 00000000000..00037e6e29f --- /dev/null +++ b/tests/test_litellm_rust/test_valkey_semantic_cache.py @@ -0,0 +1,149 @@ +import os +from collections.abc import Generator, Mapping +from types import SimpleNamespace +from typing import Final, cast +from uuid import uuid4 + +import pytest +import redis + +from litellm.caching.caching import Cache +from litellm.caching.valkey_semantic_cache import ValkeySemanticCache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def valkey_url() -> str: + url: Final = os.environ.get("LITELLM_TEST_VALKEY_URL") + if url is None: + pytest.skip("LITELLM_TEST_VALKEY_URL is not set") + return url + + +@pytest.fixture +def index_name(valkey_url: str) -> Generator[str]: + index: Final = f"litellm_test_{uuid4().hex}" + yield index + client: Final = redis.Redis.from_url(valkey_url) + try: + client.ft(index).dropindex(delete_documents=True) + except redis.ResponseError: + pass + finally: + client.close() + + +def _request() -> dict[str, object]: + return { + "key": {"preset": "key"}, + "messages": [{"role": "user", "content": "semantic cache prompt"}], + } + + +def _backend(url: str, index_name: str) -> ValkeySemanticCache: + backend: Final = ValkeySemanticCache( + redis_url=url, + similarity_threshold=0.8, + index_name=index_name, + ) + backend._get_embedding = lambda prompt, metadata=None: [1.0, 0.0] + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + return [1.0, 0.0] + + backend._get_async_embedding = async_embedding + return backend + + +def test_python_write_native_read( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + response: Final = {"answer": "python"} + backend.set_cache("key", response, messages=_request()["messages"]) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + assert binding.lookup(_request()) == response + + +def test_native_write_python_read( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + response: Final = {"answer": "native"} + binding.store({**_request(), "ttl_seconds": 2.0}, response) + cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"])) + assert cached["response"] == response + + +async def test_async_lookup_and_store( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "ttl_seconds": 2.0} + await binding.async_store(request, {"answer": "async"}) + assert await binding.async_lookup(request) == {"answer": "async"} + + +def test_facade_activation_and_mutation_fallback( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=valkey_url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + facade.cache.similarity_threshold = 0.7 + assert resolver.resolve().kind == "python_callback" + + +def test_batch_lookup_is_unsupported( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + with pytest.raises(NotImplementedError): + binding.lookup_batch([_request()]) From f6db876a3d72e143fd6638f226ac4c01a34ca088 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:21 +0000 Subject: [PATCH 048/109] test(cache): disambiguate Valkey semantic test module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...key_semantic_cache.py => test_valkey_semantic_cache_native.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/test_litellm_rust/{test_valkey_semantic_cache.py => test_valkey_semantic_cache_native.py} (100%) diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py similarity index 100% rename from tests/test_litellm_rust/test_valkey_semantic_cache.py rename to tests/test_litellm_rust/test_valkey_semantic_cache_native.py From a69ebb7ac4584e3a91497e44fd65b4bf86d52815 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:39 +0000 Subject: [PATCH 049/109] feat(cache): add the unsupported operation error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/error.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..2418ab978dc 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("cache backend does not support this operation")] + UnsupportedOperation, } From c311073a178d295a5133d2e68300c9f5f83664bc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:51 +0000 Subject: [PATCH 050/109] feat(cache-redis-semantic): expose backend accessors for bridge binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-redis-semantic/src/cache.rs | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 26d19d34670..5aa484356d7 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -41,15 +41,6 @@ pub struct RedisSemanticConfig { pub similarity_threshold: f32, } -impl Default for RedisSemanticConfig { - fn default() -> Self { - Self { - index_name: DEFAULT_INDEX_NAME.into(), - similarity_threshold: 0.9, - } - } -} - struct Inner { index_name: String, distance_threshold: f64, @@ -247,6 +238,18 @@ impl RedisSemanticCache< } } + pub fn embedder(&self) -> &E { + &self.embedder + } + + pub fn index_name(&self) -> &str { + &self.inner.index_name + } + + pub fn similarity_threshold(&self) -> f32 { + (1.0 - self.inner.distance_threshold) as f32 + } + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { context.scope.as_deref().unwrap_or(key) } From ac1c4a399ff2dcdf82ce82b93ef95ac748d1d01c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:57 +0000 Subject: [PATCH 051/109] refactor(cache-redis-semantic): drop the unused default index name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-redis-semantic/src/cache.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 5aa484356d7..b1440e80de5 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -17,7 +17,6 @@ use crate::prompt::prompt_from_context; const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; -const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; const CACHE_KEY_FIELD: &str = "litellm_cache_key"; const VECTOR_FIELD: &str = "prompt_vector"; From 788e24655a0e253f553af95da8ac69a51552a6cd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:35:08 +0000 Subject: [PATCH 052/109] feat(rust): add native S3 cache backend Mirror the Redis vertical slice for S3Cache: a litellm-cache-s3 crate built on aws-sdk-s3 with path-style custom endpoints, python-identical put_object metadata (cache-control, expires, content headers), expires-aware get_object, and no-op flush/unsupported test_connection. Wire it through python-bridge config projection, facade guards, test handle, and binding dispatch, plus an in-process S3 stub and parity tests. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 165 +++++++++- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/cache-s3/Cargo.toml | 20 ++ litellm-rust/crates/cache-s3/src/auth.rs | 32 ++ litellm-rust/crates/cache-s3/src/cache.rs | 204 +++++++++++++ litellm-rust/crates/cache-s3/src/lib.rs | 4 + litellm-rust/crates/cache-s3/tests/cache.rs | 264 ++++++++++++++++ litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/framer/Cargo.toml | 2 +- litellm-rust/crates/llms/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/config.rs | 281 +++++++++++++++++- .../crates/python-bridge/src/cache/facade.rs | 39 ++- .../crates/python-bridge/src/cache/handle.rs | 38 ++- .../crates/python-bridge/src/cache/mod.rs | 3 +- .../crates/python-bridge/src/cache/native.rs | 43 ++- litellm/rust_bridge/_native.pyi | 90 ++++++ tests/test_litellm_rust/support/s3_stub.py | 112 +++++++ tests/test_litellm_rust/test_cache.py | 182 ++++++++++++ 19 files changed, 1465 insertions(+), 21 deletions(-) create mode 100644 litellm-rust/crates/cache-s3/Cargo.toml create mode 100644 litellm-rust/crates/cache-s3/src/auth.rs create mode 100644 litellm-rust/crates/cache-s3/src/cache.rs create mode 100644 litellm-rust/crates/cache-s3/src/lib.rs create mode 100644 litellm-rust/crates/cache-s3/tests/cache.rs create mode 100644 tests/test_litellm_rust/support/s3_stub.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..4018ee2875b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -40,6 +40,12 @@ dependencies = [ "cc", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "android_system_properties" version = "0.1.6" @@ -208,6 +214,7 @@ dependencies = [ "aws-credential-types", "aws-sigv4", "aws-smithy-async", + "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -216,7 +223,9 @@ dependencies = [ "bytes", "bytes-utils", "fastrand", + "http 0.2.12", "http 1.4.2", + "http-body 0.4.6", "http-body 1.1.0", "percent-encoding", "pin-project-lite", @@ -250,6 +259,43 @@ dependencies = [ "tracing", ] +[[package]] +name = "aws-sdk-s3" +version = "1.148.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a61c1db3987ab6c8740fb87f248864687d50549401115244a84f02c3047b2d5" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-checksums", + "aws-smithy-eventstream", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-smithy-xml 0.62.1", + "aws-types", + "bytes", + "fastrand", + "hex", + "hmac", + "http 0.2.12", + "http 1.4.2", + "http-body 1.1.0", + "lru", + "percent-encoding", + "regex-lite", + "sha2 0.11.0", + "tracing", + "url", +] + [[package]] name = "aws-sdk-secretsmanager" version = "1.117.0" @@ -294,7 +340,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-smithy-schema", "aws-smithy-types", - "aws-smithy-xml", + "aws-smithy-xml 0.61.1", "aws-types", "fastrand", "http 0.2.12", @@ -310,6 +356,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70" dependencies = [ "aws-credential-types", + "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-runtime-api", "aws-smithy-types", @@ -337,10 +384,31 @@ dependencies = [ ] [[package]] -name = "aws-smithy-eventstream" -version = "0.61.1" +name = "aws-smithy-checksums" +version = "0.65.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944" +checksum = "b67ecd999972b58e67cab052f5129906c08c25883bd0788ceefc55ef97d61307" +dependencies = [ + "aws-smithy-http", + "aws-smithy-types", + "bytes", + "crc-fast", + "hex", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "md-5", + "pin-project-lite", + "sha1 0.11.0", + "sha2 0.11.0", + "tracing", +] + +[[package]] +name = "aws-smithy-eventstream" +version = "0.61.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80c2051c2f1016fb8e6548dd07b8bc2ac9c3fe583721444b92f515e856d31609" dependencies = [ "aws-smithy-types", "bytes", @@ -353,6 +421,7 @@ version = "0.64.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d" dependencies = [ + "aws-smithy-eventstream", "aws-smithy-runtime-api", "aws-smithy-types", "bytes", @@ -532,6 +601,18 @@ dependencies = [ "xmlparser", ] +[[package]] +name = "aws-smithy-xml" +version = "0.62.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b932c8d6dc127fc980eecd78f8694ae9b9551b69a93a7def2a199c1c0033daf" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "xmlparser", +] + [[package]] name = "aws-types" version = "1.6.0" @@ -927,6 +1008,16 @@ dependencies = [ "libc", ] +[[package]] +name = "crc-fast" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5" +dependencies = [ + "digest 0.10.7", + "spin", +] + [[package]] name = "crc32fast" version = "1.5.1" @@ -1369,6 +1460,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1838,6 +1935,11 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] [[package]] name = "heck" @@ -2502,6 +2604,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-s3" +version = "0.1.0" +dependencies = [ + "aws-credential-types", + "aws-sdk-s3", + "aws-smithy-types", + "aws-types", + "litellm-auth-aws", + "litellm-cache", + "serde_json", + "tokio", + "wiremock", +] + [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" @@ -2664,11 +2781,13 @@ dependencies = [ "criterion", "futures-util", "litellm-auth", + "litellm-auth-aws", "litellm-auth-gcp", "litellm-cache", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", + "litellm-cache-s3", "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", @@ -2846,6 +2965,15 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "lru" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25" +dependencies = [ + "hashbrown 0.17.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -2868,6 +2996,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + [[package]] name = "memchr" version = "2.8.3" @@ -4297,6 +4435,17 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + [[package]] name = "sha1_smol" version = "1.0.1" @@ -4413,6 +4562,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spin" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" + [[package]] name = "spm_precompiled" version = "0.1.4" @@ -5010,7 +5165,7 @@ dependencies = [ "rand 0.8.7", "rustls 0.23.42", "rustls-pki-types", - "sha1", + "sha1 0.10.7", "thiserror 1.0.69", "utf-8", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..031cfcbbb21 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-s3 = { path = "crates/cache-s3" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml new file mode 100644 index 00000000000..c5d084207c8 --- /dev/null +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-cache-s3" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-auth-aws.workspace = true +aws-sdk-s3 = { version = "1.141.0", default-features = false, features = ["rustls", "rt-tokio"] } +aws-credential-types = "1.3.0" +aws-smithy-types = "1.6.0" +aws-types = "1.6.0" +tokio.workspace = true + +[dev-dependencies] +wiremock = "0.6.5" +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs new file mode 100644 index 00000000000..8f7eb66059d --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -0,0 +1,32 @@ +use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; +use litellm_auth_aws::{AwsAuthConfig, resolve_credentials}; + +#[derive(Clone)] +pub(crate) struct Credentials { + config: AwsAuthConfig, +} + +impl Credentials { + pub(crate) fn new(config: AwsAuthConfig) -> Self { + Self { config } + } +} + +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| std::env::var(name).ok()) + .await + .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed")) + }) + } +} + +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/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs new file mode 100644 index 00000000000..1aa16ecaaa0 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -0,0 +1,204 @@ +use std::{ + future::Future, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use aws_sdk_s3::{ + config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation}, + error::SdkError, + primitives::ByteStream, +}; +use aws_smithy_types::{DateTime, date_time::Format}; +use litellm_auth_aws::AwsAuthConfig; +use litellm_cache::{ + BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache, +}; +use tokio::runtime::Handle; + +use crate::auth::Credentials; + +pub struct S3Endpoint { + pub url: String, +} + +pub struct S3CacheConfig { + pub bucket: String, + pub key_prefix: String, + pub region: String, + pub endpoint: Option, + pub auth: AwsAuthConfig, +} + +pub struct S3Cache { + client: aws_sdk_s3::Client, + codec: C, + runtime: Handle, + bucket: Arc, + key_prefix: Arc, +} + +impl S3Cache { + pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self { + let mut builder = aws_sdk_s3::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(config.region)) + .credentials_provider(Credentials::new(config.auth)) + .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) + .response_checksum_validation(ResponseChecksumValidation::WhenRequired); + if let Some(endpoint) = config.endpoint { + builder = builder.endpoint_url(endpoint.url).force_path_style(true); + } + Self { + client: aws_sdk_s3::Client::from_conf(builder.build()), + codec, + runtime, + bucket: config.bucket.into(), + key_prefix: config.key_prefix.into(), + } + } + + pub fn bucket(&self) -> &str { + &self.bucket + } + + pub fn key_prefix(&self) -> &str { + &self.key_prefix + } + + pub fn to_s3_key(&self, key: &str) -> String { + format!("{}{}", self.key_prefix, key.replace(':', "/")) + } + + fn block_on(&self, future: F) -> F::Output { + if Handle::try_current().is_ok() { + tokio::task::block_in_place(|| self.runtime.block_on(future)) + } else { + self.runtime.block_on(future) + } + } + + async fn put( + &self, + key: &str, + value: C::Value, + context: &ExactCacheContext, + ) -> Result<(), Error> { + let s3_key = self.to_s3_key(key); + let body = self.codec.encode(&value)?; + let request = self + .client + .put_object() + .bucket(self.bucket.as_ref()) + .key(&s3_key) + .body(ByteStream::from(body)) + .content_type("application/json") + .content_language("en") + .content_disposition(format!("inline; filename=\"{s3_key}.json\"")); + let request = match context.ttl { + Some(ttl) => { + let seconds = ttl.as_secs_f64(); + request + .cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}")) + .expires(DateTime::from(SystemTime::now() + ttl)) + } + None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"), + }; + request.send().await.map_err(|_| Error::Unavailable)?; + Ok(()) + } + + async fn get(&self, key: &str) -> Result, Error> { + let output = match self + .client + .get_object() + .bucket(self.bucket.as_ref()) + .key(self.to_s3_key(key)) + .send() + .await + { + Ok(output) => output, + Err(error) => { + if let SdkError::ServiceError(service) = &error { + let not_found = service.err().is_no_such_key() + || error + .raw_response() + .map(|response| response.status().as_u16()) + == Some(404); + if not_found { + return Ok(None); + } + } + return Err(Error::Unavailable); + } + }; + if let Some(expires) = output.expires_string() + && let Ok(expires) = DateTime::from_str(expires, Format::HttpDate) + && expires < DateTime::from(SystemTime::now()) + { + return Ok(None); + } + let bytes = output + .body + .collect() + .await + .map_err(|_| Error::Unavailable)? + .into_bytes(); + self.codec.decode(&bytes).map(Some) + } +} + +impl BaseCache for S3Cache { + type Value = C::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _context: &Self::Context) -> Option { + None + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + self.block_on(self.put(key, value, context)) + } + + fn get_cache(&self, key: &str, _context: &Self::Context) -> Result, Error> { + self.block_on(self.get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + self.put(key, value, &context).await + } + + async fn async_get_cache( + &self, + key: &str, + _context: &Self::Context, + ) -> Result, Error> { + self.get(key).await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for S3Cache {} + +impl FlushCache for S3Cache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs new file mode 100644 index 00000000000..f6126dfa908 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -0,0 +1,4 @@ +mod auth; +mod cache; + +pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs new file mode 100644 index 00000000000..93e88cf7cbd --- /dev/null +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -0,0 +1,264 @@ +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_auth_aws::AwsAuthConfig; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec, +}; +use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint}; +use serde_json::{Value, json}; +use tokio::runtime::Handle; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, +}; + +fn config(endpoint: String) -> S3CacheConfig { + S3CacheConfig { + bucket: "cache-bucket".to_string(), + key_prefix: "team/".to_string(), + region: "us-east-1".to_string(), + endpoint: Some(S3Endpoint { url: endpoint }), + auth: AwsAuthConfig { + access_key_id: Some("key".to_string()), + secret_access_key: Some("secret".to_string()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + } +} + +fn cache(endpoint: &str) -> S3Cache> { + S3Cache::new( + config(endpoint.to_string()), + JsonCodec::::new(), + Handle::current(), + ) +} + +async fn mock_server() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\"")) + .mount(&server) + .await; + server +} + +fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option { + use aws_smithy_types::{DateTime, date_time::Format}; + headers + .get(name) + .and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok()) + .map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos())) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn set_writes_python_metadata_with_and_without_ttl() { + let server = mock_server().await; + let cache = cache(&server.uri()); + let context = ExactCacheContext { + ttl: Some(Duration::from_secs(90)), + }; + cache + .set_cache("alpha:beta", json!({"answer": 1}), &context) + .unwrap(); + cache + .set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default()) + .unwrap(); + + let requests = server.received_requests().await.unwrap(); + let ttl_request = requests + .iter() + .find(|request| request.url.path() == "/cache-bucket/team/alpha/beta") + .expect("ttl write should hit the converted S3 key"); + assert_eq!( + ttl_request.headers["cache-control"].to_str().unwrap(), + "immutable, max-age=90, s-maxage=90" + ); + assert_eq!( + ttl_request.headers["content-type"].to_str().unwrap(), + "application/json" + ); + assert_eq!( + ttl_request.headers["content-language"].to_str().unwrap(), + "en" + ); + assert_eq!( + ttl_request.headers["content-disposition"].to_str().unwrap(), + "inline; filename=\"team/alpha/beta.json\"" + ); + let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires"); + let remaining = expires.duration_since(SystemTime::now()).unwrap(); + assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91)); + assert_eq!( + serde_json::from_slice::(&ttl_request.body).unwrap(), + json!({"answer": 1}) + ); + + let plain = requests + .iter() + .find(|request| request.url.path() == "/cache-bucket/team/plain") + .expect("no-ttl write should hit the converted S3 key"); + assert_eq!( + plain.headers["cache-control"].to_str().unwrap(), + "immutable, max-age=31536000, s-maxage=31536000" + ); + assert!(plain.headers.get("expires").is_none()); + assert_eq!( + plain.headers["content-disposition"].to_str().unwrap(), + "inline; filename=\"team/plain.json\"" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_hit_miss_expired_and_invalid_entries() { + let server = mock_server().await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/hit")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/missing")) + .respond_with( + ResponseTemplate::new(404).set_body_string("NoSuchKey"), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/expired")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT") + .set_body_json(json!({"answer": 4})), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/malformed")) + .respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry")) + .mount(&server) + .await; + let cache = cache(&server.uri()); + let context = ExactCacheContext::default(); + + assert_eq!( + cache.get_cache("hit", &context).unwrap(), + Some(json!({"answer": 3})) + ); + assert_eq!(cache.get_cache("missing", &context).unwrap(), None); + assert_eq!(cache.get_cache("expired", &context).unwrap(), None); + assert_eq!( + cache.get_cache("malformed", &context), + Err(Error::InvalidEntry) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn batch_get_preserves_order_with_hits_misses_and_invalid() { + let server = mock_server().await; + for (key, status, body) in [ + ("first", 200, "{\"answer\": 1}"), + ("invalid", 200, "garbage"), + ] { + Mock::given(method("GET")) + .and(path(format!("/cache-bucket/team/{key}"))) + .respond_with(ResponseTemplate::new(status).set_body_string(body)) + .mount(&server) + .await; + } + Mock::given(method("GET")) + .and(path("/cache-bucket/team/miss")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + let cache = cache(&server.uri()); + let context = ExactCacheContext::default(); + let keys = vec![ + "first".to_string(), + "miss".to_string(), + "invalid".to_string(), + ]; + + let entries = cache.batch_get_cache(&keys, &context).unwrap(); + + assert_eq!( + entries, + vec![ + BatchEntry::Hit(json!({"answer": 1})), + BatchEntry::Miss, + BatchEntry::Invalid, + ] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unsupported_and_noop_capabilities_match_python() { + let server = mock_server().await; + let cache = cache(&server.uri()); + + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); + cache.flush_cache().unwrap(); + cache.disconnect().await.unwrap(); + assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[test] +fn key_conversion_prefixes_and_splits_colons() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap(); + let _guard = runtime.enter(); + let cache = S3Cache::new( + S3CacheConfig { + key_prefix: "team/".to_string(), + ..config("http://localhost".to_string()) + }, + JsonCodec::::new(), + runtime.handle().clone(), + ); + + assert_eq!(cache.bucket(), "cache-bucket"); + assert_eq!(cache.key_prefix(), "team/"); + assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c"); + assert_eq!(cache.to_s3_key("plain"), "team/plain"); + + let unprefixed = S3Cache::new( + S3CacheConfig { + key_prefix: String::new(), + ..config("http://localhost".to_string()) + }, + JsonCodec::::new(), + runtime.handle().clone(), + ); + assert_eq!(unprefixed.to_s3_key("a:b"), "a/b"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sync_methods_block_inside_and_outside_the_runtime() { + let server = mock_server().await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9}))) + .mount(&server) + .await; + let uri = server.uri(); + let cache = tokio::task::spawn_blocking(move || { + let cache = cache(&uri); + let context = ExactCacheContext::default(); + cache + .set_cache("key", json!({"answer": 9}), &context) + .unwrap(); + cache.get_cache("key", &context).unwrap() + }) + .await + .unwrap(); + + assert_eq!(cache, Some(json!({"answer": 9}))); +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..e373853e70a 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("cache operation is not supported")] + UnsupportedOperation, } diff --git a/litellm-rust/crates/framer/Cargo.toml b/litellm-rust/crates/framer/Cargo.toml index d22502f871a..60dfa9469d2 100644 --- a/litellm-rust/crates/framer/Cargo.toml +++ b/litellm-rust/crates/framer/Cargo.toml @@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"] sse = ["dep:sse-stream"] [dependencies] -aws-smithy-eventstream = { version = "=0.61.1", optional = true } +aws-smithy-eventstream = { version = "0.61.1", optional = true } aws-smithy-types = { version = "1.6.1", optional = true } bytes = "1" futures-util.workspace = true diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 0cc7af1836f..8bdbb7c3b77 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] } url.workspace = true [dev-dependencies] -aws-smithy-eventstream = "=0.61.1" +aws-smithy-eventstream = "0.61.1" aws-smithy-types = "1.6.1" rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..2aea8122fa6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,9 +23,11 @@ bytes.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-s3.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true +litellm-auth-aws.workspace = true litellm-callbacks-legacy-python.workspace = true litellm-core.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..5d37ed27083 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,10 +1,12 @@ use std::time::Duration; +use litellm_auth_aws::AwsAuthConfig; use litellm_cache::CacheType; +use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use pyo3::{ - exceptions::{PyTypeError, PyValueError}, + exceptions::{PyAttributeError, PyTypeError, PyValueError}, prelude::*, - types::{PyAny, PyDict, PyString}, + types::{PyAny, PyBool, PyDict, PyString}, }; use super::{native::NativeResponseCache, request::duration}; @@ -76,6 +78,7 @@ pub(super) struct RedisCacheConfig { pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + S3(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -90,6 +93,9 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + S3Client, + S3Credentials, + S3Option, } impl UnsupportedCacheConfig { @@ -100,6 +106,9 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::S3Client => "native S3 client type is not implemented", + Self::S3Credentials => "native S3 credentials require Python", + Self::S3Option => "native S3 configuration requires Python", } } } @@ -142,10 +151,16 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::S3) => match project_s3(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::S3(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some( CacheType::RedisSemantic | CacheType::ValkeySemantic - | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic | CacheType::AzureBlob @@ -159,10 +174,11 @@ impl NativeCacheConfig { pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) + != match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::S3(_) => None, + } { return Some("facade and native backend default TTLs must match"); } @@ -185,6 +201,18 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::S3(_) if service.kind() != "s3" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => { + Some("facade and native backend buckets must match") + } + CacheBackendConfig::S3(config) + if service.key_prefix() != Some(config.key_prefix.as_str()) => + { + Some("facade and native backend key prefixes must match") + } + CacheBackendConfig::S3(_) => None, } } } @@ -299,6 +327,77 @@ fn project_redis( })) } +#[inline(never)] +fn project_s3( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let client = backend.getattr("s3_client")?; + if !instance_class_is(&client, "botocore.client", "S3")? { + return Ok(Err(UnsupportedCacheConfig::S3Client)); + } + let meta = client.getattr("meta")?; + let Some(region) = optional_string(meta.getattr("region_name")?)? else { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + }; + let Some(endpoint_url) = optional_string(meta.getattr("endpoint_url")?)? else { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + }; + let client_config = meta.getattr("config")?; + for name in ["s3", "proxies", "client_cert"] { + if optional_attribute(&client_config, name)?.is_some_and(|value| !value.is_none()) { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + } + } + let signature = match optional_attribute(&client_config, "signature_version")? { + Some(value) => value.extract::>()?, + None => None, + }; + if signature.as_deref() != Some("s3v4") { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + } + let insecure = endpoint_url.starts_with("http://"); + let verify = optional_attribute_chain(&client, &["_endpoint", "http_session", "_verify"])?; + let verified = verify + .and_then(|value| value.cast::().ok().map(|value| value.is_true())) + .unwrap_or(false); + if !verified && !insecure { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + } + let credentials = optional_attribute_chain(&client, &["_request_signer", "_credentials"])? + .ok_or(UnsupportedCacheConfig::S3Credentials); + let credentials = match credentials { + Ok(credentials) if !credentials.is_none() => credentials, + _ => return Ok(Err(UnsupportedCacheConfig::S3Credentials)), + }; + let auth = if credentials.getattr("method")?.extract::()?.as_str() == "explicit" { + AwsAuthConfig { + access_key_id: credentials + .getattr("access_key")? + .extract::>()?, + secret_access_key: credentials + .getattr("secret_key")? + .extract::>()?, + session_token: credentials.getattr("token")?.extract::>()?, + region_name: Some(region.clone()), + ..Default::default() + } + } else { + AwsAuthConfig { + region_name: Some(region.clone()), + ..Default::default() + } + }; + let default_endpoint = endpoint_url == format!("https://s3.{region}.amazonaws.com") + || (region == "us-east-1" && endpoint_url == "https://s3.amazonaws.com"); + Ok(Ok(S3CacheConfig { + bucket: backend.getattr("bucket_name")?.extract::()?, + key_prefix: backend.getattr("key_prefix")?.extract::()?, + region, + endpoint: (!default_endpoint).then_some(S3Endpoint { url: endpoint_url }), + auth, + })) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { @@ -380,6 +479,33 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult( + value: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + match value.getattr(name) { + Ok(value) => Ok(Some(value)), + Err(error) if error.is_instance_of::(value.py()) => Ok(None), + Err(error) => Err(error), + } +} + +#[inline(never)] +fn optional_attribute_chain<'py>( + value: &Bound<'py, PyAny>, + names: &[&str], +) -> PyResult>> { + let mut current = value.clone(); + for name in names { + match optional_attribute(¤t, name)? { + Some(next) => current = next, + None => return Ok(None), + } + } + Ok(Some(current)) +} + #[inline(never)] fn optional_string(value: Bound<'_, PyAny>) -> PyResult> { Ok(value @@ -591,4 +717,145 @@ mod tests { assert_eq!(reason.message(), "native Redis credentials require Python"); }); } + + fn s3_facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> { + let locals = PyDict::new(py); + py.run( + &CString::new(format!( + "from types import SimpleNamespace\n\ + S3Client = type('S3', (), {{'__module__': 'botocore.client'}})\n\ + client = S3Client()\n\ + client.meta = SimpleNamespace(region_name='us-east-1', endpoint_url='https://example.test', config=SimpleNamespace(s3=None, proxies=None, client_cert=None, signature_version='s3v4'))\n\ + client._endpoint = SimpleNamespace(http_session=SimpleNamespace(_verify=True))\n\ + client._request_signer = SimpleNamespace(_credentials=SimpleNamespace(method='explicit', access_key='key', secret_key='secret', token='token'))\n\ + backend = SimpleNamespace(bucket_name='bucket', key_prefix='team/', s3_client=client)\n\ + facade = SimpleNamespace(type='s3', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\ + {body}" + )) + .unwrap(), + None, + Some(&locals), + ) + .unwrap(); + locals.get_item("facade").unwrap().unwrap() + } + + #[test] + fn projects_s3_configuration_with_explicit_credentials_and_custom_endpoint() { + Python::initialize(); + Python::attach(|py| { + let facade = s3_facade(py, ""); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("S3 cache should be supported"); + }; + let CacheBackendConfig::S3(s3) = config.backend else { + panic!("expected S3 configuration"); + }; + assert_eq!(s3.bucket, "bucket"); + assert_eq!(s3.key_prefix, "team/"); + assert_eq!(s3.region, "us-east-1"); + assert_eq!( + s3.endpoint.map(|endpoint| endpoint.url).as_deref(), + Some("https://example.test") + ); + assert_eq!(s3.auth.access_key_id.as_deref(), Some("key")); + assert_eq!(s3.auth.secret_access_key.as_deref(), Some("secret")); + assert_eq!(s3.auth.session_token.as_deref(), Some("token")); + assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1")); + }); + } + + #[test] + fn default_s3_endpoint_projects_no_custom_endpoint() { + Python::initialize(); + Python::attach(|py| { + let facade = s3_facade( + py, + "facade.cache.s3_client.meta.endpoint_url = 'https://s3.us-east-1.amazonaws.com'", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("S3 cache should be supported"); + }; + let CacheBackendConfig::S3(s3) = config.backend else { + panic!("expected S3 configuration"); + }; + assert!(s3.endpoint.is_none()); + }); + } + + #[test] + fn non_sigv4_proxies_and_disabled_verification_stay_on_python() { + Python::initialize(); + Python::attach(|py| { + for (body, message) in [ + ( + "facade.cache.s3_client.meta.config.signature_version = 's3'", + "native S3 configuration requires Python", + ), + ( + "facade.cache.s3_client.meta.config.proxies = {'https': 'proxy'}", + "native S3 configuration requires Python", + ), + ( + "facade.cache.s3_client._endpoint.http_session._verify = False", + "native S3 configuration requires Python", + ), + ( + "del facade.cache.s3_client._endpoint.http_session._verify", + "native S3 configuration requires Python", + ), + ] { + let facade = s3_facade(py, body); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("{body} must stay on Python"); + }; + assert_eq!(reason.message(), message); + } + let facade = s3_facade(py, "facade.cache.s3_client = SimpleNamespace()"); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("non-botocore client must stay on Python"); + }; + assert_eq!(reason.message(), "native S3 client type is not implemented"); + let facade = s3_facade( + py, + "facade.cache.s3_client._request_signer._credentials = None", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("missing credentials must stay on Python"); + }; + assert_eq!(reason.message(), "native S3 credentials require Python"); + }); + } + + #[test] + fn non_explicit_s3_credentials_use_the_default_chain() { + Python::initialize(); + Python::attach(|py| { + let facade = s3_facade( + py, + "facade.cache.s3_client._request_signer._credentials = SimpleNamespace(method='sso', access_key=None, secret_key=None, token=None)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("default-chain credentials should be supported"); + }; + let CacheBackendConfig::S3(s3) = config.backend else { + panic!("expected S3 configuration"); + }; + assert_eq!(s3.auth.access_key_id, None); + assert_eq!(s3.auth.secret_access_key, None); + assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1")); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..0f7b9778d3e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -32,10 +32,15 @@ struct RedisPoolGuard { max_connections: usize, } +struct S3ClientGuard { + reference: Py, +} + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, redis_pool: Option, + s3_client: Option, } impl ObjectGuard { @@ -176,6 +181,22 @@ impl RedisPoolGuard { } } +impl S3ClientGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self { + reference: backend.getattr("s3_client")?.unbind(), + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference) + } +} + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -192,6 +213,7 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "s3" => ("litellm.caching.s3_cache", "S3Cache", "s3"), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -235,11 +257,16 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "bucket_name", + "key_prefix", ], )?, redis_pool: (kind == "redis") .then(|| RedisPoolGuard::capture(&backend)) .transpose()?, + s3_client: (kind == "s3") + .then(|| S3ClientGuard::capture(&backend)) + .transpose()?, }) } @@ -252,9 +279,14 @@ impl FacadeGuard { return Ok(false); } match &self.redis_pool { - Some(guard) => guard.matches(py, &backend), - None => Ok(true), + Some(guard) if !guard.matches(py, &backend)? => return Ok(false), + _ => {} } + match &self.s3_client { + Some(guard) if !guard.matches(py, &backend)? => return Ok(false), + _ => {} + } + Ok(true) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -263,6 +295,9 @@ impl FacadeGuard { if let Some(guard) = &self.redis_pool { guard.traverse(&visit)?; } + if let Some(guard) = &self.s3_client { + guard.traverse(&visit)?; + } Ok(()) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..0d5a2ac9d8c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,4 +1,6 @@ -use litellm_host_python::release_gil; +use litellm_auth_aws::AwsAuthConfig; +use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; +use litellm_host_python::{release_gil, run_sync_value}; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; @@ -51,6 +53,40 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (bucket, *, region, endpoint_url=None, key_prefix="", access_key_id=None, secret_access_key=None, session_token=None))] + fn s3( + py: Python<'_>, + bucket: String, + region: String, + endpoint_url: Option, + key_prefix: &str, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + ) -> PyResult { + let config = S3CacheConfig { + bucket, + key_prefix: key_prefix.to_string(), + region: region.clone(), + endpoint: endpoint_url.map(|url| S3Endpoint { url }), + auth: AwsAuthConfig { + access_key_id, + secret_access_key, + session_token, + region_name: Some(region), + ..Default::default() + }, + }; + let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..ac4494150d9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -10,7 +10,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +21,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..b9ea443119b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -6,6 +6,7 @@ use litellm_cache_redis::RedisCache; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; +use litellm_cache_s3::{S3Cache, S3CacheConfig}; use serde_json::Value; #[derive(Clone)] @@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + S3(Arc>>), } impl NativeResponseCache { @@ -43,6 +45,15 @@ impl NativeResponseCache { buffer: None, }) } + + pub async fn s3(config: S3CacheConfig) -> Self { + let runtime = tokio::runtime::Handle::current(); + Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new( + config, + ResponseCacheCodec, + runtime, + ))))) + } } impl NativeResponseCache { @@ -50,6 +61,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::S3(_) => "s3", } } @@ -57,6 +69,21 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::S3(cache) => cache.default_ttl(), + } + } + + pub fn bucket(&self) -> Option<&str> { + match self { + Self::S3(cache) => Some(cache.backend().bucket()), + _ => None, + } + } + + pub fn key_prefix(&self) -> Option<&str> { + match self { + Self::S3(cache) => Some(cache.backend().key_prefix()), + _ => None, } } @@ -64,20 +91,21 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), + Self::S3(_) => None, } } pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::S3(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::S3(_) => None, } } @@ -87,7 +115,7 @@ impl NativeResponseCache { cache, buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, - memory => memory, + cache => cache, } } @@ -99,6 +127,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(request, now), Self::Redis { cache, .. } => cache.lookup(request, now), + Self::S3(cache) => cache.lookup(request, now), } } @@ -111,6 +140,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(request, response, now), Self::Redis { cache, .. } => cache.store(request, response, now), + Self::S3(cache) => cache.store(request, response, now), } } @@ -122,6 +152,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup_batch(requests, now), Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::S3(cache) => cache.lookup_batch(requests, now), } } @@ -133,6 +164,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::S3(cache) => cache.async_lookup(request, now).await, } } @@ -152,6 +184,7 @@ impl NativeResponseCache { cache, buffer: Some(buffer), } => buffer.async_store(cache, request, response, now).await, + Self::S3(cache) => cache.async_store(request, response, now).await, } } @@ -163,6 +196,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::S3(cache) => cache.async_lookup_batch(requests, now).await, } } @@ -174,6 +208,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_store_batch(entries, now).await, Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::S3(cache) => cache.async_store_batch(entries, now).await, } } @@ -186,6 +221,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::S3(cache) => cache.async_flush().await, } } @@ -193,6 +229,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::S3(cache) => cache.test_connection().await, } } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..f1b010fb26e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -93,6 +93,96 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... +@final +class _CacheTestBinding: + @property + def kind(self) -> str: ... + def lookup( + self, + request: object, + *, + callback_kwargs: Mapping[str, object] | Sequence[object] | None = None, + ) -> object: ... + def store( + self, + request: object, + response: object, + *, + callback_kwargs: Mapping[str, object] | None = None, + ) -> None: ... + def lookup_batch( + self, + requests: Sequence[object], + *, + callback_kwargs: Sequence[object] | None = None, + ) -> object: ... + def async_lookup( + self, + request: object, + *, + callback_kwargs: Mapping[str, object] | None = None, + ) -> Future[object]: ... + def async_store( + self, + request: object, + response: object, + *, + callback_kwargs: Mapping[str, object] | None = None, + ) -> Future[None]: ... + def async_lookup_batch( + self, + requests: Sequence[object], + *, + callback_kwargs: Sequence[object] | None = None, + ) -> Future[object]: ... + def async_store_batch( + self, + requests: Sequence[object], + responses: Sequence[object], + *, + callback_result: object = None, + callback_kwargs: Mapping[str, object] | None = None, + ) -> Future[object]: ... + def async_flush(self) -> Future[None]: ... + def ping(self) -> Future[object]: ... + +@final +class _CacheTestHandle: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @staticmethod + def memory( + *, + capacity: int = 200, + ttl_seconds: float = 600.0, + max_entry_bytes: int = 1048576, + ) -> _CacheTestHandle: ... + @staticmethod + def redis( + url: str, + *, + ttl_seconds: float = 60.0, + namespace: str | None = None, + ) -> _CacheTestHandle: ... + @staticmethod + def s3( + bucket: str, + *, + region: str, + endpoint_url: str | None = None, + key_prefix: str = "", + access_key_id: str | None = None, + secret_access_key: str | None = None, + session_token: str | None = None, + ) -> _CacheTestHandle: ... + @property + def backend(self) -> str: ... + def _bind_facade(self, facade: object) -> None: ... + +@final +class _CacheTestResolver: + def __new__(cls, namespace: object) -> _CacheTestResolver: ... + def resolve(self) -> _CacheTestBinding: ... + @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... diff --git a/tests/test_litellm_rust/support/s3_stub.py b/tests/test_litellm_rust/support/s3_stub.py new file mode 100644 index 00000000000..18863ad8083 --- /dev/null +++ b/tests/test_litellm_rust/support/s3_stub.py @@ -0,0 +1,112 @@ +"""In-process path-style S3 stub for native cache parity tests.""" + +import threading +from dataclasses import dataclass, field +from email.utils import parsedate_to_datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final +from urllib.parse import unquote, urlsplit + +_STORED_HEADERS: Final = ( + "cache-control", + "content-type", + "content-language", + "content-disposition", + "expires", +) + + +@dataclass +class S3Object: + body: bytes + headers: dict[str, str] = field(default_factory=dict) + + +class S3Stub: + """Minimal path-style S3 endpoint serving PUT and GET object operations.""" + + def __init__(self) -> None: + self._objects: dict[str, S3Object] = {} + stub: Final = self + + class Handler(BaseHTTPRequestHandler): + def _key(self) -> str: + parts: Final = urlsplit(self.path).path.lstrip("/").split("/", 1) + return unquote(parts[1]) if len(parts) == 2 else "" + + def _read_body(self) -> bytes: + transfer: Final = self.headers.get("transfer-encoding", "") + if "chunked" not in transfer: + return self.rfile.read(int(self.headers.get("content-length", 0))) + chunks: Final = bytearray() + while True: + size = int(self.rfile.readline().split(b";")[0].strip(), 16) + if size == 0: + while self.rfile.readline().strip(): + pass + return bytes(chunks) + chunks.extend(self.rfile.read(size)) + self.rfile.readline() + + def do_PUT(self) -> None: + body: Final = self._read_body() + headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers} + stub._objects[self._key()] = S3Object(body=body, headers=headers) + self.send_response(200) + self.send_header("ETag", '"stub"') + self.send_header("Content-Length", "0") + self.end_headers() + + def do_HEAD(self) -> None: + self._object(send_body=False) + + def do_GET(self) -> None: + self._object(send_body=True) + + def _object(self, send_body: bool) -> None: + entry: Final = stub._objects.get(self._key()) + if entry is None: + self.send_response(404) + self.send_header("Content-Type", "application/xml") + body: Final = b'NoSuchKey' + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if send_body: + self.wfile.write(body) + return + self.send_response(200) + for name, value in entry.headers.items(): + self.send_header(name, value) + self.send_header("ETag", '"stub"') + self.send_header("Content-Length", str(len(entry.body))) + self.end_headers() + if send_body: + self.wfile.write(entry.body) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 + pass + + self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self._worker: Final = threading.Thread(target=self._server.serve_forever, daemon=True) + self._worker.start() + + @property + def url(self) -> str: + host, port = self._server.server_address[:2] + return f"http://{host}:{port}" + + @property + def objects(self) -> dict[str, S3Object]: + return self._objects + + def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None: + self._objects[key] = S3Object(body=body, headers=headers or {}) + + def expires(self, key: str) -> object: + header: Final = self._objects[key].headers.get("expires") + return parsedate_to_datetime(header) if header else None + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._worker.join(timeout=5) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..f18454ea8d1 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -6,10 +6,13 @@ import threading import time import weakref from collections.abc import Generator +from datetime import datetime from types import SimpleNamespace from typing import Final, Protocol, cast from urllib.parse import urlparse +import boto3 +import botocore.config import fakeredis import pytest import redis @@ -17,9 +20,11 @@ import redis import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.s3_stub import S3Stub pytestmark: Final = pytest.mark.requires_rust_extension @@ -393,3 +398,180 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + +@pytest.fixture +def s3_stub() -> Generator[S3Stub]: + stub: Final = S3Stub() + try: + yield stub + finally: + stub.close() + + +def python_s3(url: str) -> S3Cache: + return S3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + + +async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None: + python_cache: Final = python_s3(s3_stub.url) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90) + python_cache.set_cache("plain", {"timestamp": time.time(), "response": response}) + s3_stub.put_object("team/malformed", b"not a cache entry") + s3_stub.put_object( + "team/expired", + json.dumps({"timestamp": time.time(), "response": response}).encode(), + {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"}, + ) + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + ) + ).resolve() + + assert binding.lookup(request("sync:key")) == response + assert await binding.async_lookup(request("plain")) == response + assert binding.lookup(request("malformed")) is None + assert binding.lookup(request("expired")) is None + assert binding.lookup(request("absent")) is None + + binding.store({**request("native:key"), "ttl_seconds": 90.0}, response) + await binding.async_store(request("no_ttl"), response) + stored: Final = s3_stub.objects["team/native/key"] + assert stored.headers["content-type"] == "application/json" + assert stored.headers["content-language"] == "en" + assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"' + assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90" + expires: Final = cast(datetime, s3_stub.expires("team/native/key")) + remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds() + assert 60 < remaining <= 91 + no_ttl: Final = s3_stub.objects["team/no_ttl"] + assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000" + assert "expires" not in no_ttl.headers + assert python_cache.get_cache("native:key")["response"] == response + + partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")]) + assert partial == {"values": [response, None, None], "missing_indices": [1, 2]} + + +def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + handle: Final = _native._CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + with pytest.raises(TypeError, match="buckets must match"): + _native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade) + with pytest.raises(TypeError, match="key prefixes must match"): + _native._CacheTestHandle.s3( + "cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/" + )._bind_facade(facade) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + + calls: Final = [] + facade.cache.s3_client.meta.events.register("before-call.s3.*", lambda **_kwargs: calls.append(1)) + binding.store(request("native"), {"answer": 1}) + assert binding.lookup(request("native")) == {"answer": 1} + assert calls == [] + assert "team/native" in s3_stub.objects + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + other_client: Final = boto3.client( + "s3", + region_name="us-east-1", + endpoint_url=s3_stub.url, + aws_access_key_id="key", + aws_secret_access_key="secret", + ) + with rebound(facade.cache, "s3_client", other_client): + assert resolver.resolve().kind == "python_callback" + + class CustomS3Cache(S3Cache): + pass + + subclassed: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + subclassed.cache = CustomS3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + with pytest.raises(TypeError): + handle._bind_facade(subclassed) + assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback" + + +def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None: + handle: Final = _native._CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + unverified: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url="https://s3.example.test", + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_verify=False, + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unverified) + proxied: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}), + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(proxied) From 10b977fe29caccc1a2730568d33c21aa751cddab Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:38:33 +0000 Subject: [PATCH 053/109] feat(python-bridge): serve redis-semantic caches natively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 67 +++++++- .../python-bridge/src/cache/embedder.rs | 85 +++++++++ .../crates/python-bridge/src/cache/facade.rs | 21 +++ .../crates/python-bridge/src/cache/handle.rs | 43 ++++- .../crates/python-bridge/src/cache/mod.rs | 4 +- .../crates/python-bridge/src/cache/native.rs | 161 ++++++++++++++---- .../crates/python-bridge/src/cache/request.rs | 66 +++++-- 9 files changed, 399 insertions(+), 50 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/embedder.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e911d0d9c45..0030018df34 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2683,6 +2683,7 @@ dependencies = [ "litellm-cache", "litellm-cache-memory", "litellm-cache-redis", + "litellm-cache-redis-semantic", "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..635c0942ceb 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,6 +23,7 @@ bytes.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-redis-semantic.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..85e400d89a3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,23 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[allow( + dead_code, + reason = "embedding settings are projected so drift falls back to Python" +)] +pub(super) struct RedisSemanticCacheConfig { + pub(super) redis_url: String, + pub(super) index_name: String, + pub(super) similarity_threshold: f64, + pub(super) embedding_model: String, + pub(super) embedding_max_input_tokens: Option, + pub(super) embedding_timeout: Option, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + RedisSemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -142,9 +156,14 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), + })) + }), Some( - CacheType::RedisSemantic - | CacheType::ValkeySemantic + CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic @@ -159,10 +178,11 @@ impl NativeCacheConfig { pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) + != match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::RedisSemantic(_) => None, + } { return Some("facade and native backend default TTLs must match"); } @@ -185,10 +205,45 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.index_name() != Some(config.index_name.as_str()) => + { + Some("facade and native backend index names must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.similarity_threshold() != Some(config.similarity_threshold as f32) => + { + Some("facade and native backend similarity thresholds must match") + } + CacheBackendConfig::RedisSemantic(_) => None, } } } +#[inline(never)] +pub(super) fn project_redis_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult { + Ok(RedisSemanticCacheConfig { + redis_url: backend.getattr("_redis_url")?.extract::()?, + index_name: backend + .getattr("_index_name")? + .extract::>()? + .unwrap_or_else(|| "litellm_semantic_cache_index".into()), + similarity_threshold: backend.getattr("similarity_threshold")?.extract::()?, + embedding_model: backend.getattr("embedding_model")?.extract::()?, + embedding_max_input_tokens: backend + .getattr("embedding_max_input_tokens")? + .extract::>()?, + embedding_timeout: backend + .getattr("embedding_timeout")? + .extract::>()?, + }) +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs new file mode 100644 index 00000000000..63e078cd815 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -0,0 +1,85 @@ +use std::future::Future; + +use litellm_cache::Error; +use litellm_cache_redis_semantic::Embedder; +use litellm_host_python::to_py; +use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; +use serde_json::{Map, Value}; + +pub(super) struct PythonEmbedder(Py); + +impl PythonEmbedder { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + + pub(super) fn object(&self) -> &Py { + &self.0 + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + + fn metadata_kwargs<'py>( + py: Python<'py>, + metadata: &Map, + ) -> PyResult> { + let kwargs = PyDict::new(py); + if metadata.is_empty() { + kwargs.set_item("metadata", py.None())?; + } else { + kwargs.set_item("metadata", to_py(py, metadata)?)?; + } + Ok(kwargs) + } + + fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + Ok(vector + .extract::>()? + .into_iter() + .map(|value| value as f32) + .collect()) + } +} + +impl Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: &Map) -> Result, Error> { + Python::attach(|py| { + let kwargs = Self::metadata_kwargs(py, metadata)?; + Self::extract(self.0.bind(py).call_method( + "_get_embedding", + (prompt,), + Some(&kwargs), + )?) + }) + .map_err(|_| Error::Unavailable) + } + + fn async_embed( + &self, + prompt: &str, + metadata: &Map, + ) -> impl Future, Error>> + Send { + let coroutine = Python::attach(|py| { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + }) + .map_err(|_| Error::Unavailable); + async move { + let coroutine = coroutine?; + let awaited = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py)) + }) + .map_err(|_| Error::Unavailable)? + .await + .map_err(|_| Error::Unavailable)?; + let vector = Python::attach(|py| awaited.extract::>(py)) + .map_err(|_| Error::Unavailable)?; + Ok(vector.into_iter().map(|value| value as f32).collect()) + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..58730857d60 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -192,6 +192,11 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "redis_semantic" => ( + "litellm.caching.redis_semantic_cache", + "RedisSemanticCache", + "redis-semantic", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -211,6 +216,15 @@ impl FacadeGuard { if let Some(message) = config.service_mismatch(service) { return Err(PyTypeError::new_err(message)); } + if kind == "redis_semantic" + && service + .embedder_object() + .is_none_or(|embedder| !backend.is(embedder.bind(py))) + { + return Err(PyTypeError::new_err( + "facade backend must be the native embedder", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -235,6 +249,13 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "similarity_threshold", + "distance_threshold", + "embedding_model", + "embedding_max_input_tokens", + "embedding_timeout", + "_index_name", + "_redis_url", ], )?, redis_pool: (kind == "redis") diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..b61ae59bb58 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,7 +1,15 @@ +use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_host_python::release_gil; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, +}; -use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use super::{ + cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard, + native::NativeResponseCache, request::duration, +}; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -51,6 +59,36 @@ impl CacheTestHandle { }) } + #[staticmethod] + fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult { + let class = py + .import("litellm.caching.redis_semantic_cache")? + .getattr("RedisSemanticCache")?; + if !backend.get_type().is(&class) { + return Err(PyTypeError::new_err( + "native redis-semantic handles require the built-in RedisSemanticCache", + )); + } + let config = project_redis_semantic(&backend)?; + let embedder = PythonEmbedder::new(backend.unbind()); + let service = release_gil(py, move || { + NativeResponseCache::redis_semantic( + &config.redis_url, + embedder, + RedisSemanticConfig { + index_name: config.index_name, + similarity_threshold: config.similarity_threshold as f32, + }, + ) + }) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -76,6 +114,7 @@ impl CacheTestHandle { } fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(&visit)?; if let Some(guard) = &self.guard { guard.traverse(visit)?; } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..4cc87367d91 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,6 +1,7 @@ mod binding; mod callback; mod config; +mod embedder; mod facade; mod future; mod handle; @@ -10,7 +11,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +22,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..8cd77fa8eb0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,13 +1,17 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error, ExactCacheContext}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; +use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; +use pyo3::{Py, PyAny, PyTraverseError, PyVisit}; use serde_json::Value; +use super::{embedder::PythonEmbedder, request::CacheRequest}; + #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), @@ -15,6 +19,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + RedisSemantic(Arc>>), } impl NativeResponseCache { @@ -43,6 +48,17 @@ impl NativeResponseCache { buffer: None, }) } + + pub fn redis_semantic( + url: &str, + embedder: PythonEmbedder, + config: RedisSemanticConfig, + ) -> Result { + let backend = RedisSemanticCache::new(url, embedder, config)?; + Ok(Self::RedisSemantic(Arc::new(ResponseCache::new(Arc::new( + backend, + ))))) + } } impl NativeResponseCache { @@ -50,6 +66,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::RedisSemantic(_) => "redis_semantic", } } @@ -57,12 +74,13 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::RedisSemantic(cache) => cache.default_ttl(), } } pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) => None, + Self::Memory(_) | Self::RedisSemantic(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), } } @@ -70,110 +88,189 @@ impl NativeResponseCache { pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::RedisSemantic(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::RedisSemantic(_) => None, } } + pub fn index_name(&self) -> Option<&str> { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().index_name()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + + pub fn similarity_threshold(&self) -> Option { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().similarity_threshold()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + + pub fn embedder_object(&self) -> Option<&Py> { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Self::RedisSemantic(cache) = self { + cache.backend().embedder().traverse(visit)?; + } + Ok(()) + } + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { match self { Self::Redis { cache, .. } => Self::Redis { cache, buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, - memory => memory, + other => other, } } - pub fn lookup( - &self, - request: &ResponseCacheRequest, - now: Duration, - ) -> Result, Error> { + fn exact_requests(requests: &[CacheRequest]) -> Vec> { + requests.iter().map(CacheRequest::exact).collect() + } + + pub fn lookup(&self, request: &CacheRequest, now: Duration) -> Result, Error> { match self { - Self::Memory(cache) => cache.lookup(request, now), - Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Memory(cache) => cache.lookup(&request.exact(), now), + Self::Redis { cache, .. } => cache.lookup(&request.exact(), now), + Self::RedisSemantic(cache) => cache.lookup(&request.semantic(), now), } } pub fn store( &self, - request: &ResponseCacheRequest, + request: &CacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(request, response, now), - Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Memory(cache) => cache.store(&request.exact(), response, now), + Self::Redis { cache, .. } => cache.store(&request.exact(), response, now), + Self::RedisSemantic(cache) => cache.store(&request.semantic(), response, now), } } pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[CacheRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(requests, now), - Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Memory(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), + Self::Redis { cache, .. } => cache.lookup_batch(&Self::exact_requests(requests), now), + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &CacheRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Memory(cache) => cache.async_lookup(&request.exact(), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&request.exact(), now).await, + Self::RedisSemantic(cache) => cache.async_lookup(&request.semantic(), now).await, } } pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &CacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Memory(cache) => cache.async_store(&request.exact(), response, now).await, Self::Redis { cache, buffer: None, - } => cache.async_store(request, response, now).await, + } => cache.async_store(&request.exact(), response, now).await, Self::Redis { cache, buffer: Some(buffer), - } => buffer.async_store(cache, request, response, now).await, + } => { + buffer + .async_store(cache, &request.exact(), response, now) + .await + } + Self::RedisSemantic(cache) => { + cache.async_store(&request.semantic(), response, now).await + } } } pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[CacheRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, - Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Memory(cache) => { + cache + .async_lookup_batch(&Self::exact_requests(requests), now) + .await + } + Self::Redis { cache, .. } => { + cache + .async_lookup_batch(&Self::exact_requests(requests), now) + .await + } + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(CacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store_batch(entries, now).await, - Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Memory(cache) => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.exact(), value)) + .collect(), + now, + ) + .await + } + Self::Redis { cache, .. } => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.exact(), value)) + .collect(), + now, + ) + .await + } + Self::RedisSemantic(cache) => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.semantic(), value)) + .collect(), + now, + ) + .await + } } } @@ -186,6 +283,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } @@ -193,6 +291,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0c5343a63d0..26e0fe4e62c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,9 +1,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use litellm_cache::{ExactCacheContext, SemanticCacheContext}; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; +use serde_json::{Map, Value}; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -12,24 +14,68 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + input: Option, + messages: Option>, + metadata: Option>, + scope: Option, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) struct CacheRequest { + key: CacheKeyInput, + controls: CacheControls, + ttl: Option, + max_age: Option, + input: Option, + messages: Vec, + metadata: Map, + scope: Option, +} + +impl CacheRequest { + pub(super) fn exact(&self) -> ResponseCacheRequest { + let mut request = ResponseCacheRequest::new(self.key.clone()); + request.controls = self.controls; + request.context.ttl = self.ttl; + request.max_age = self.max_age; + request + } + + pub(super) fn semantic(&self) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key.clone(), + controls: self.controls, + context: SemanticCacheContext { + input: self.input.clone(), + messages: self.messages.clone(), + metadata: self.metadata.clone(), + scope: self.scope.clone(), + ttl: self.ttl, + }, + max_age: self.max_age, + } + } +} + +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::new(input.key); - if let Some(controls) = input.controls { - request.controls = controls; - } - request.context.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) +fn request_input(input: RequestInput) -> PyResult { + let defaults = ResponseCacheRequest::::new(input.key.clone()); + Ok(CacheRequest { + key: input.key, + controls: input.controls.unwrap_or(defaults.controls), + ttl: input.ttl_seconds.map(duration).transpose()?, + max_age: input.max_age_seconds.map(duration).transpose()?, + input: input.input, + messages: input.messages.unwrap_or_default(), + metadata: input.metadata.unwrap_or_default(), + scope: input.scope, + }) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { from_py::>(value)? .into_iter() .map(request_input) From 605729a194c997b8f90a5edb5494465be9cc4d95 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:39:17 +0000 Subject: [PATCH 054/109] fix(rust): pin aws-sdk-s3 without relaxing eventstream pins Keep the exact eventstream pins but move them to =0.61.4 so aws-runtime 1.9.4's eventstream requirement resolves, downgrade aws-sdk-s3 to 1.146.1, and echo the context ttl from S3Cache::get_ttl to match Python Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 24 ++++++++++----------- litellm-rust/crates/cache-s3/Cargo.toml | 2 +- litellm-rust/crates/cache-s3/src/cache.rs | 4 ++-- litellm-rust/crates/cache-s3/tests/cache.rs | 6 ++++++ litellm-rust/crates/framer/Cargo.toml | 2 +- litellm-rust/crates/llms/Cargo.toml | 2 +- 6 files changed, 23 insertions(+), 17 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4018ee2875b..e2337f41ebc 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -261,9 +261,9 @@ dependencies = [ [[package]] name = "aws-sdk-s3" -version = "1.148.0" +version = "1.146.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a61c1db3987ab6c8740fb87f248864687d50549401115244a84f02c3047b2d5" +checksum = "2cd651b4400d4011b8927b83a9552bf90ff11e6e5da0b9f0a7583247aceec971" dependencies = [ "arc-swap", "aws-credential-types", @@ -1380,7 +1380,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2152,7 +2152,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.5", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -3638,7 +3638,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.42", - "socket2 0.6.5", + "socket2 0.5.10", "thiserror 2.0.19", "tokio", "tracing", @@ -3677,9 +3677,9 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.5", + "socket2 0.5.10", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4104,7 +4104,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4175,7 +4175,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4699,10 +4699,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5542,7 +5542,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index c5d084207c8..cdc17e732cb 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -8,7 +8,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true litellm-auth-aws.workspace = true -aws-sdk-s3 = { version = "1.141.0", default-features = false, features = ["rustls", "rt-tokio"] } +aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } aws-credential-types = "1.3.0" aws-smithy-types = "1.6.0" aws-types = "1.6.0" diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 1aa16ecaaa0..a2ea33c33e4 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -152,8 +152,8 @@ impl BaseCache for S3Cache { type Value = C::Value; type Context = ExactCacheContext; - fn get_ttl(&self, _context: &Self::Context) -> Option { - None + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl } fn set_cache( diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs index 93e88cf7cbd..6ca3eb9a758 100644 --- a/litellm-rust/crates/cache-s3/tests/cache.rs +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -204,6 +204,12 @@ async fn unsupported_and_noop_capabilities_match_python() { cache.flush_cache().unwrap(); cache.disconnect().await.unwrap(); assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&ExactCacheContext { + ttl: Some(Duration::from_secs(45)), + }), + Some(Duration::from_secs(45)) + ); assert!(server.received_requests().await.unwrap().is_empty()); } diff --git a/litellm-rust/crates/framer/Cargo.toml b/litellm-rust/crates/framer/Cargo.toml index 60dfa9469d2..62bfcc7da3d 100644 --- a/litellm-rust/crates/framer/Cargo.toml +++ b/litellm-rust/crates/framer/Cargo.toml @@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"] sse = ["dep:sse-stream"] [dependencies] -aws-smithy-eventstream = { version = "0.61.1", optional = true } +aws-smithy-eventstream = { version = "=0.61.4", optional = true } aws-smithy-types = { version = "1.6.1", optional = true } bytes = "1" futures-util.workspace = true diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 8bdbb7c3b77..f04b78feee1 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] } url.workspace = true [dev-dependencies] -aws-smithy-eventstream = "0.61.1" +aws-smithy-eventstream = "=0.61.4" aws-smithy-types = "1.6.1" rstest.workspace = true tokio.workspace = true From 56237af7a9dbfcdb20b41b3959620128c8fbccd2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:42:41 +0000 Subject: [PATCH 055/109] test(cache): add Valkey semantic contract coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-valkey-semantic/src/lib.rs | 659 +++++++++++++----- .../crates/python-bridge/src/cache/config.rs | 84 ++- .../test_valkey_semantic_cache_native.py | 140 +++- 3 files changed, 695 insertions(+), 188 deletions(-) diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index 85c4c9af15c..ef028f0a7b2 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -62,6 +62,14 @@ enum Connections { Fixed(Mutex), } +#[derive(Clone)] +struct IndexState { + name: String, + prefix: String, + dimension: Arc>>, + similarity_threshold: f64, +} + struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); impl redis::ConnectionLike for ConnectionRef<'_> { @@ -187,18 +195,13 @@ where &self.config.index_name } - fn key_prefix(&self) -> String { - format!("{}:", self.config.index_name) - } - - fn ensure_index(&self, dimension: usize) -> Result<(), Error> { - ensure_index( - &self.connections, - &self.config.index_name, - &self.key_prefix(), - &self.index_dimension, - dimension, - ) + fn index_state(&self) -> IndexState { + IndexState { + name: self.config.index_name.clone(), + prefix: format!("{}:", self.config.index_name), + dimension: Arc::clone(&self.index_dimension), + similarity_threshold: self.config.similarity_threshold, + } } } @@ -225,37 +228,19 @@ where return Ok(()); }; let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - self.ensure_index(embedding.len())?; let scope = scope_tag(key); - let document = format!("{}{}:{}", self.key_prefix(), scope, Uuid::new_v4()); let response = self.codec.encode(&value)?; let vector = embedding_bytes(&embedding); - let ttl = self.get_ttl(context); - self.connections.execute(|connection| { - let mut pipeline = redis::pipe(); - pipeline - .cmd("HSET") - .arg(&document) - .arg("litellm_cache_key") - .arg(&scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) - .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) - }) + let index = self.index_state(); + write_document( + &self.connections, + &index, + &scope, + &prompt, + response, + vector, + self.get_ttl(context), + ) } fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { @@ -263,43 +248,14 @@ where return Ok(None); }; let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - self.ensure_index(embedding.len())?; let scope = scope_tag(key); - let query = - format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); let vector = embedding_bytes(&embedding); - let response = self.connections.execute(|connection| { - redis::cmd("FT.SEARCH") - .arg(&self.config.index_name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; - let Some(fields) = search_fields(response)? else { + let index = self.index_state(); + let Some(response) = + search_document(&self.connections, &index, &scope, vector, embedding.len())? + else { return Ok(None); }; - let response = fields - .iter() - .find_map(|(name, value)| (name == "response").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = fields - .iter() - .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = parse_f64(&distance)?; - if 1.0 - distance < self.config.similarity_threshold { - return Ok(None); - } self.codec.decode(&response).map(Some) } @@ -321,47 +277,13 @@ where .async_embed(&prompt, metadata.as_ref()) .await?; let connections = Arc::clone(&self.connections); - let config = self.config.clone(); - let index_dimension = Arc::clone(&self.index_dimension); + let index = self.index_state(); let response = self.codec.encode(&value)?; let vector = embedding_bytes(&embedding); - let prefix = format!("{}:", config.index_name); let scope = scope_tag(&key); - let document = format!("{prefix}{scope}:{}", Uuid::new_v4()); let ttl = context.ttl; tokio::task::spawn_blocking(move || { - ensure_index( - &connections, - &config.index_name, - &prefix, - &index_dimension, - embedding.len(), - )?; - connections.execute(|connection| { - let mut pipeline = redis::pipe(); - pipeline - .cmd("HSET") - .arg(&document) - .arg("litellm_cache_key") - .arg(&scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) - .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) - }) + write_document(&connections, &index, &scope, &prompt, response, vector, ttl) }) .await .map_err(|_| Error::Unavailable)? @@ -385,56 +307,11 @@ where .async_embed(&prompt, metadata.as_ref()) .await?; let connections = Arc::clone(&self.connections); - let config = self.config.clone(); - let index_dimension = Arc::clone(&self.index_dimension); - let threshold = config.similarity_threshold; + let index = self.index_state(); tokio::task::spawn_blocking(move || { - let prefix = format!("{}:", config.index_name); - ensure_index( - &connections, - &config.index_name, - &prefix, - &index_dimension, - embedding.len(), - )?; let scope = scope_tag(&key); - let query = format!( - "(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]" - ); let vector = embedding_bytes(&embedding); - let response = connections.execute(|connection| { - redis::cmd("FT.SEARCH") - .arg(&config.index_name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; - let Some(fields) = search_fields(response)? else { - return Ok(None); - }; - let response = fields - .iter() - .find_map(|(name, value)| (name == "response").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = fields - .iter() - .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = parse_f64(&distance)?; - if 1.0 - distance < threshold { - return Ok(None); - } - Ok(Some(response)) + search_document(&connections, &index, &scope, vector, embedding.len()) }) .await .map_err(|_| Error::Unavailable)? @@ -560,6 +437,108 @@ fn embedding_bytes(embedding: &[f32]) -> Vec { .collect() } +fn write_document( + connections: &Connections, + index: &IndexState, + scope: &str, + prompt: &str, + response: Vec, + vector: Vec, + ttl: Option, +) -> Result<(), Error> +where + C: redis::ConnectionLike + Send + 'static, +{ + let dimension = vector.len() / std::mem::size_of::(); + ensure_index( + connections, + &index.name, + &index.prefix, + &index.dimension, + dimension, + )?; + let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4()); + connections.execute(|connection| { + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { + pipeline + .cmd("EXPIRE") + .arg(&document) + .arg(ttl.as_secs()) + .ignore(); + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) + }) +} + +fn search_document( + connections: &Connections, + index: &IndexState, + scope: &str, + vector: Vec, + dimension: usize, +) -> Result>, Error> +where + C: redis::ConnectionLike + Send + 'static, +{ + ensure_index( + connections, + &index.name, + &index.prefix, + &index.dimension, + dimension, + )?; + let query = + format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); + let response = connections.execute(|connection| { + redis::cmd("FT.SEARCH") + .arg(&index.name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let Some(fields) = search_fields(response)? else { + return Ok(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < index.similarity_threshold { + return Ok(None); + } + Ok(Some(response)) +} + fn ensure_index( connections: &Connections, index_name: &str, @@ -712,10 +691,16 @@ fn value_bytes(value: &redis::Value) -> Result, Error> { #[cfg(test)] mod tests { - use std::sync::{Arc, Mutex}; + use std::{ + collections::VecDeque, + sync::{Arc, Mutex}, + time::Duration, + }; - use litellm_cache::BaseCache; - use litellm_cache_response::ResponseCacheCodec; + use litellm_cache::{BaseCache, CacheCodec}; + use litellm_cache_response::{ + CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + }; use redis_test::MockRedisConnection; use rstest::rstest; use serde_json::{Value, json}; @@ -732,6 +717,9 @@ mod tests { } type EmbedderCalls = Arc)>>>; + type RecordingCache = + ValkeySemanticCache; + type RecordingSetup = (RecordingCache, Arc>>>, EmbedderCalls); impl Embedder for FixedEmbedder { fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, super::Error> { @@ -751,6 +739,61 @@ mod tests { } } + struct RecordingConnection { + requests: Arc>>>, + replies: Mutex>>, + } + + impl RecordingConnection { + fn new(replies: impl IntoIterator>) -> Self { + Self { + requests: Arc::default(), + replies: Mutex::new(replies.into_iter().collect()), + } + } + + fn requests(&self) -> Arc>>> { + Arc::clone(&self.requests) + } + + fn reply(&self) -> redis::RedisResult { + self.replies + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into()))) + } + } + + impl redis::ConnectionLike for RecordingConnection { + fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { + self.requests.lock().unwrap().push(command.to_vec()); + self.reply() + } + + fn req_packed_commands( + &mut self, + command: &[u8], + _offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.requests.lock().unwrap().push(command.to_vec()); + (0..count).map(|_| self.reply()).collect() + } + + fn get_db(&self) -> i64 { + 0 + } + + fn check_connection(&mut self) -> bool { + true + } + + fn is_open(&self) -> bool { + true + } + } + fn context( messages: Option, input: Option, @@ -841,4 +884,302 @@ mod tests { assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None); assert_eq!(cache.get_ttl(&context(None, None)), None); } + + fn semantic_context(ttl: Option) -> litellm_cache::SemanticCacheContext { + litellm_cache::SemanticCacheContext { + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"source": "test"})), + ttl, + ..Default::default() + } + } + + fn cache_with_recording( + replies: impl IntoIterator>, + vector: Vec, + threshold: f64, + ) -> RecordingSetup { + let connection = RecordingConnection::new(replies); + let requests = connection.requests(); + let calls: EmbedderCalls = Arc::default(); + let cache = ValkeySemanticCache::with_connection( + connection, + FixedEmbedder { + vector, + calls: Arc::clone(&calls), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: threshold, + index_name: "test".into(), + }, + ); + (cache, requests, calls) + } + + fn ok() -> redis::RedisResult { + Ok(redis::Value::SimpleString("OK".into())) + } + + fn already_exists() -> redis::RedisResult { + Err(redis::RedisError::from(( + redis::ErrorKind::Io, + "already exists", + ))) + } + + fn info_dimension(dimension: usize) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::SimpleString("attributes".into()), + redis::Value::Array(vec![redis::Value::Array(vec![ + redis::Value::SimpleString("embedding".into()), + redis::Value::Array(vec![ + redis::Value::SimpleString("dimensions".into()), + redis::Value::Int(dimension as i64), + ]), + ])]), + ]) + } + + fn search_hit(response: Vec, distance: &str) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"test:document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(response), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(distance.as_bytes().to_vec()), + ]), + ]) + } + + fn requests_text(requests: &Arc>>>) -> String { + requests + .lock() + .unwrap() + .iter() + .map(|request| String::from_utf8_lossy(request)) + .collect::>() + .join("\n") + } + + #[test] + fn set_without_ttl_writes_hset_without_expire() { + let (cache, requests, calls) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ) + .unwrap(); + let text = requests_text(&requests); + assert!(text.contains("FT.CREATE")); + assert!(text.contains("HSET")); + assert!( + text.contains("test:2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683:") + ); + assert!(!text.contains("EXPIRE")); + assert_eq!( + *calls.lock().unwrap(), + vec![("hello".into(), Some(json!({"source": "test"})))] + ); + } + + #[test] + fn set_with_ttl_truncates_expire_seconds() { + let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(Some(Duration::from_millis(1900))), + ) + .unwrap(); + let text = requests_text(&requests); + assert!(text.contains("EXPIRE")); + assert!(text.contains("\r\n$1\r\n1\r\n")); + } + + #[test] + fn second_set_skips_create_after_dimension_is_cached() { + let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + let context = semantic_context(None); + let entry = CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }; + cache.set_cache("key", entry.clone(), &context).unwrap(); + cache.set_cache("key", entry, &context).unwrap(); + let text = requests_text(&requests); + assert_eq!(text.matches("FT.CREATE").count(), 1); + assert_eq!(text.matches("HSET").count(), 2); + } + + #[test] + fn existing_index_dimension_must_match_embedding() { + let (cache, _, _) = cache_with_recording( + [already_exists(), Ok(info_dimension(2))], + vec![1.0, 0.0], + 0.8, + ); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ) + .unwrap(); + + let (cache, _, _) = cache_with_recording( + [already_exists(), Ok(info_dimension(3))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ), + Err(super::Error::Unavailable) + ); + } + + #[test] + fn get_applies_threshold_and_decodes_entry() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let (cache, _, _) = cache_with_recording( + [ok(), Ok(search_hit(encoded.clone(), "0.1"))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + Some(entry) + ); + + let (cache, _, _) = + cache_with_recording([ok(), Ok(search_hit(encoded, "0.5"))], vec![1.0, 0.0], 0.8); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + None + ); + } + + #[test] + fn get_zero_docs_is_a_miss() { + let (cache, _, _) = cache_with_recording( + [ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + None + ); + } + + #[rstest] + #[case(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"0.1".to_vec()), + ]), + ]))] + #[case(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(b"not-json".to_vec()), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"abc".to_vec()), + ]), + ]))] + fn malformed_entries_are_invalid(#[case] search: redis::Value) { + let (cache, _, _) = cache_with_recording([ok(), Ok(search)], vec![1.0, 0.0], 0.8); + assert_eq!( + cache.get_cache("key", &semantic_context(None)), + Err(super::Error::InvalidEntry) + ); + } + + #[test] + fn response_cache_turns_invalid_entries_into_misses() { + let (cache, _, _) = cache_with_recording( + [ + ok(), + Ok(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(b"not-json".to_vec()), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"0.1".to_vec()), + ]), + ])), + ], + vec![1.0, 0.0], + 0.8, + ); + let service = ResponseCache::new(Arc::new(cache)); + let request = ResponseCacheRequest { + key: CacheKeyInput { + preset: Some("key".into()), + ..Default::default() + }, + context: semantic_context(None), + ..ResponseCacheRequest::new(CacheKeyInput::default()) + }; + assert_eq!(service.lookup(&request, Duration::ZERO).unwrap(), None); + } + + #[tokio::test] + async fn async_set_and_get_use_shared_document_helpers() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let (cache, requests, calls) = cache_with_recording( + [ok(), ok(), ok(), Ok(search_hit(encoded, "0.1"))], + vec![1.0, 0.0], + 0.8, + ); + let context = semantic_context(Some(Duration::from_millis(1900))); + cache + .async_set_cache("key", entry.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("key", &context).await.unwrap(), + Some(entry) + ); + let text = requests_text(&requests); + assert!(text.contains("FT.CREATE")); + assert!(text.contains("HSET")); + assert!(text.contains("EXPIRE")); + assert_eq!(calls.lock().unwrap().len(), 2); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 7218805b8df..e074c5e2f5d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -276,25 +276,15 @@ fn project_redis( let client = backend.getattr("redis_client")?; let pool = client.getattr("connection_pool")?; - if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } - let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + }; for key in ["credential_provider", "redis_connect_func"] { if has_value(&resolved, key)? { return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } } - let connection_class = resolved - .get_item("connection_class")? - .unwrap_or(pool.getattr("connection_class")?); - let tls = if class_is(&connection_class, "redis.connection", "Connection")? { - None - } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { - Some(project_tls(&resolved)?) - } else { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - }; + let tls = is_tls.then(|| project_tls(&resolved)).transpose()?; let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, @@ -332,18 +322,9 @@ fn project_valkey_semantic( ) -> PyResult> { let client = backend.getattr("sync_client")?; let pool = client.getattr("connection_pool")?; - if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + let Ok((resolved, _is_tls)) = project_connection_pool(&pool)? else { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } - let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; - let connection_class = resolved - .get_item("connection_class")? - .unwrap_or(pool.getattr("connection_class")?); - if !class_is(&connection_class, "redis.connection", "Connection")? - && !class_is(&connection_class, "redis.connection", "SSLConnection")? - { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } + }; let connection = RedisConnectionConfig { host: required_string(&resolved, "host")?, port: u16::try_from(required_i64(&resolved, "port")?) @@ -371,6 +352,27 @@ fn project_valkey_semantic( })) } +#[inline(never)] +fn project_connection_pool<'py>( + pool: &Bound<'py, PyAny>, +) -> PyResult, bool), UnsupportedCacheConfig>> { + if !instance_class_is(pool, "redis.connection", "ConnectionPool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? { + false + } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { + true + } else { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + }; + Ok(Ok((resolved, is_tls))) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { @@ -646,6 +648,40 @@ mod tests { }); } + #[test] + fn projects_valkey_semantic_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = Connection\n\ + pool.max_connections = 12\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Valkey semantic cache should be supported"); + }; + let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else { + panic!("expected Valkey semantic configuration"); + }; + assert_eq!(valkey.similarity_threshold, 0.85); + assert_eq!(valkey.index_name, "semantic_idx"); + assert_eq!(valkey.embedding_model, "text-embedding-3-small"); + assert_eq!(valkey.connection.host, "cache.internal"); + assert_eq!(valkey.connection.port, 6390); + assert_eq!(valkey.connection.database, 2); + assert_eq!(valkey.connection.pool_size, 12); + assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2); + assert!(valkey.connection.tls.is_none()); + }); + } + #[test] fn dynamic_redis_auth_stays_on_python() { Python::initialize(); diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index 00037e6e29f..dc4ede8ea30 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -1,4 +1,7 @@ +import hashlib import os +import struct +import time from collections.abc import Generator, Mapping from types import SimpleNamespace from typing import Final, cast @@ -36,24 +39,32 @@ def index_name(valkey_url: str) -> Generator[str]: client.close() -def _request() -> dict[str, object]: +def _request(prompt: str = "semantic cache prompt") -> dict[str, object]: return { "key": {"preset": "key"}, - "messages": [{"role": "user", "content": "semantic cache prompt"}], + "messages": [{"role": "user", "content": prompt}], } -def _backend(url: str, index_name: str) -> ValkeySemanticCache: +def _backend( + url: str, + index_name: str, + embeddings: Mapping[str, list[float]] | None = None, +) -> ValkeySemanticCache: + vectors: Final = embeddings or {"semantic cache prompt": [1.0, 0.0]} backend: Final = ValkeySemanticCache( redis_url=url, similarity_threshold=0.8, index_name=index_name, ) - backend._get_embedding = lambda prompt, metadata=None: [1.0, 0.0] + + def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + return vectors[prompt] async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: - return [1.0, 0.0] + return vectors[prompt] + backend._get_embedding = embed backend._get_async_embedding = async_embedding return backend @@ -147,3 +158,122 @@ def test_batch_lookup_is_unsupported( binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() with pytest.raises(NotImplementedError): binding.lookup_batch([_request()]) + + +def test_ttl_expiry( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store({**_request(), "ttl_seconds": 1.0}, {"answer": "expires"}) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + assert client.ttl(documents[0]) > 0 + time.sleep(1.5) + assert binding.lookup(_request()) is None + + +def test_no_ttl_is_persistent_and_python_reads_native_value( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + response: Final = {"answer": "persistent"} + binding.store(_request(), response) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + assert client.ttl(documents[0]) == -1 + cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"])) + assert cached["response"] == response + + +def test_below_threshold_misses_on_native_and_python( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend( + valkey_url, + index_name, + {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store(_request("prompt A"), {"answer": "A"}) + assert binding.lookup(_request("prompt B")) is None + assert backend.get_cache("key", messages=_request("prompt B")["messages"]) is None + + +def test_malformed_entry_is_a_miss_on_native_and_python( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + client: Final = redis.Redis.from_url(valkey_url) + scope: Final = hashlib.sha256(b"key").hexdigest() + document: Final = f"{index_name}:{scope}:{uuid4().hex}" + client.hset( + document, + mapping={ + "litellm_cache_key": scope, + "prompt": "semantic cache prompt", + "response": "not json", + "embedding": struct.pack("<2f", 1.0, 0.0), + }, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + assert binding.lookup(_request()) is None + assert backend.get_cache("key", messages=_request()["messages"]) is None + + +async def test_async_store_batch_and_lookup( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend( + valkey_url, + index_name, + {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + requests: Final = [_request("prompt A"), _request("prompt B")] + responses: Final = [{"answer": "A"}, {"answer": "B"}] + await binding.async_store_batch(requests, responses) + assert await binding.async_lookup(requests[0]) == responses[0] + assert await binding.async_lookup(requests[1]) == responses[1] + + +def test_subclass_backend_falls_back_to_python( + valkey_url: str, + index_name: str, +) -> None: + class Custom(ValkeySemanticCache): + pass + + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=valkey_url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + facade.cache = Custom(redis_url=valkey_url, similarity_threshold=0.8, index_name=index_name) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "python_callback" + + +async def test_ping_maps_unsupported_native_operation_to_not_implemented( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + with pytest.raises(NotImplementedError): + await binding.ping() From c7458bf5fe20bcbb7902fd259fd3b113b2b9e168 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:42:49 +0000 Subject: [PATCH 056/109] refactor(rust): build vault requests without local reassignment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../secrets-hashicorp/src/secret_manager.rs | 59 +++++++++++-------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index e4ca8c387c3..52cc143c1e3 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -9,11 +9,11 @@ use litellm_secrets_types::{ BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name, }; use moka::future::Cache; -use reqwest::Client; +use reqwest::{Client, Identity, RequestBuilder}; use serde_json::{Value, json}; use tokio::sync::Mutex; -use crate::{Error, HashicorpVaultConfig}; +use crate::{Error, HashicorpVaultConfig, TlsCertAuth}; const CACHE_CAPACITY: u64 = 200; @@ -242,10 +242,10 @@ impl HashicorpVault { return Ok(token); } }; - let mut request = self.client.post(login_url).json(&body); - if let Some(namespace) = self.config.login_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } + let request: RequestBuilder = with_namespace( + self.client.post(login_url).json(&body), + self.config.login_namespace(), + ); let response: reqwest::Response = request.send().await?; if !response.status().is_success() { return Err(Error::LoginStatus { @@ -305,24 +305,33 @@ impl BaseSecretManager for HashicorpVault { } fn client_for_config(config: &HashicorpVaultConfig) -> Result { - let mut builder: reqwest::ClientBuilder = Client::builder(); - if let Some(tls) = config.tls_cert.as_ref() { - let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - })?; - let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { - path: tls.key_path.clone(), - message: source.to_string(), - })?; - let identity: reqwest::Identity = reqwest::Identity::from_pem( - &[cert.as_slice(), key.as_slice()].concat(), - ) - .map_err(|source| Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - })?; - builder = builder.identity(identity); - } + let builder: reqwest::ClientBuilder = match config.tls_cert.as_ref() { + Some(tls) => Client::builder().identity(identity_for(tls)?), + None => Client::builder(), + }; builder.build().map_err(Error::Request) } + +fn with_namespace(request: RequestBuilder, namespace: Option<&str>) -> RequestBuilder { + match namespace { + Some(namespace) => request.header("X-Vault-Namespace", namespace), + None => request, + } +} + +fn identity_for(tls: &TlsCertAuth) -> Result { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { + path: tls.key_path.clone(), + message: source.to_string(), + })?; + Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { + Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + } + }) +} From a45be4f276e2628248026731dc4f8b7b014fef1a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 21 Sep 2026 13:48:37 -0700 Subject: [PATCH 057/109] fix(proxy): let the config file win when reporting nested alerting sources _nested_setting_source returned "db" whenever the stored row held a value, without first asking whether the config file declares the same key. For a config-owned alerting_args field that disagrees with the database, the endpoint reported source "db" while the proxy actually serves the file's value and rejects any write to it. Config ownership is now checked first, matching SettingsStore.source and the precedence the rest of the resolver applies. The source test set grows a field that only the database sets, a field only the file sets, and a stored empty list, so each reported source is discriminating. --- litellm/proxy/proxy_server.py | 6 +- .../proxy_server/test_routes_model_metrics.py | 69 +++++++++++++------ 2 files changed, 51 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5f0295f8208..7b5086f3e8b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15993,12 +15993,12 @@ def _nested_setting_source( field_name: str, field_default: JsonValue, ) -> FieldSource: - db_value: Final = db_values.get(field_name) - if db_value is not None and not (isinstance(db_value, list) and len(db_value) == 0): - return "db" parent_value: Final = settings.config_value(parent_key) if isinstance(parent_value, Mapping) and field_name in parent_value: return "config" + db_value: Final = db_values.get(field_name) + if db_value is not None and not (isinstance(db_value, list) and len(db_value) == 0): + return "db" return "default" if field_default is not None else "unset" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py index f65db69ce89..c5287db5027 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py @@ -186,20 +186,21 @@ def test_model_settings_method_not_allowed(client, auth_as): def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): from litellm.proxy.config_resolvers import SettingsStore + db_alerting_args = { + "daily_report_frequency": 7, + "outage_alert_ttl": 99, + "region_outage_alert_ttl": [], + } + pc = MagicMock() row = MagicMock() - row.param_value = { - "alerting_args": { - "daily_report_frequency": 7, - "report_check_interval": None, - } - } + row.param_value = {"alerting_args": db_alerting_args} pc.db.litellm_config.find_first = AsyncMock(return_value=row) monkeypatch.setattr(proxy_server, "prisma_client", pc) logging_obj = MagicMock() args_model = MagicMock() - args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 7}) + args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 3}) logging_obj.slack_alerting_instance.alerting_args = args_model monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) @@ -207,21 +208,10 @@ def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): store.load_yaml( { "alerting": ["slack"], - "alerting_args": { - "daily_report_frequency": 3, - "report_check_interval": 300, - }, + "alerting_args": {"daily_report_frequency": 3, "report_check_interval": 300}, } ) - store.apply_db_row( - "general_settings", - { - "alerting_args": { - "daily_report_frequency": 7, - "report_check_interval": None, - } - }, - ) + store.apply_db_row("general_settings", {"alerting_args": db_alerting_args}) monkeypatch.setattr(proxy_server.proxy_config, "settings", store) monkeypatch.setattr(proxy_server, "general_settings", store) @@ -230,12 +220,48 @@ def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): assert response.status_code == 200 by_name = {entry["field_name"]: entry for entry in response.json()} + assert by_name["slack_alerting"]["source"] == "config" - assert by_name["daily_report_frequency"]["source"] == "db" + assert by_name["daily_report_frequency"]["source"] == "config" assert by_name["report_check_interval"]["source"] == "config" + assert by_name["outage_alert_ttl"]["source"] == "db" + assert by_name["region_outage_alert_ttl"]["source"] == "default" assert by_name["budget_alert_ttl"]["source"] == "default" +def test_alerting_settings_reports_config_source_when_db_disagrees(client, auth_as, monkeypatch): + from litellm.proxy.config_resolvers import SettingsStore + + db_alerting_args = {"daily_report_frequency": 7} + + pc = MagicMock() + row = MagicMock() + row.param_value = {"alerting_args": db_alerting_args} + pc.db.litellm_config.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", pc) + + logging_obj = MagicMock() + args_model = MagicMock() + args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 3}) + logging_obj.slack_alerting_instance.alerting_args = args_model + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) + + store = SettingsStore("general_settings") + store.load_yaml({"alerting_args": {"daily_report_frequency": 3}}) + store.apply_db_row("general_settings", {"alerting_args": db_alerting_args}) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + assert store.source("alerting_args") == "config" + assert by_name["daily_report_frequency"]["field_value"] == 3 + assert by_name["daily_report_frequency"]["source"] == "config" + + @pytest.mark.parametrize("db_alerting_args", [None, []]) def test_alerting_settings_handles_empty_db_args( client: TestClient, @@ -268,6 +294,7 @@ def test_alerting_settings_handles_empty_db_args( assert response.status_code == 200 by_name = {entry["field_name"]: entry for entry in response.json()} assert by_name["report_check_interval"]["source"] == "config" + assert by_name["budget_alert_ttl"]["source"] == "default" def test_alerting_settings_no_db_error(client, auth_as, no_prisma): From f998ab53d5e68d46c05238bf8ff05d711ddf7085 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 21 Sep 2026 13:52:16 -0700 Subject: [PATCH 058/109] fix(proxy): treat a config-owned alerting_args as shadowing the stored row When the config file declares alerting_args at all, the resolver hands the file's dict to every reader and the stored row never reaches one. Reporting a nested field as "db" because the row happens to carry it told the admin a value was in effect that the proxy does not serve: a live proxy answered source "db" for outage_alert_ttl while serving the default. A config-owned parent now reports the field's own default, and the DB is consulted only when the file leaves the parent alone. --- litellm/proxy/proxy_server.py | 5 +- .../proxy_server/test_routes_model_metrics.py | 53 +++++++++++++------ 2 files changed, 41 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7b5086f3e8b..d162ce2914e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15996,10 +15996,13 @@ def _nested_setting_source( parent_value: Final = settings.config_value(parent_key) if isinstance(parent_value, Mapping) and field_name in parent_value: return "config" + unset_source: Final[FieldSource] = "default" if field_default is not None else "unset" + if settings.owned_by_config(parent_key): + return unset_source db_value: Final = db_values.get(field_name) if db_value is not None and not (isinstance(db_value, list) and len(db_value) == 0): return "db" - return "default" if field_default is not None else "unset" + return unset_source @router.get( diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py index c5287db5027..3fb6e6fcb45 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py @@ -183,37 +183,39 @@ def test_model_settings_method_not_allowed(client, auth_as): # --------------------------------------------------------------------------- -def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): +def _alerting_client(monkeypatch, *, yaml_values, db_row, live_args): from litellm.proxy.config_resolvers import SettingsStore - db_alerting_args = { - "daily_report_frequency": 7, - "outage_alert_ttl": 99, - "region_outage_alert_ttl": [], - } - pc = MagicMock() row = MagicMock() - row.param_value = {"alerting_args": db_alerting_args} + row.param_value = db_row pc.db.litellm_config.find_first = AsyncMock(return_value=row) monkeypatch.setattr(proxy_server, "prisma_client", pc) logging_obj = MagicMock() args_model = MagicMock() - args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 3}) + args_model.model_dump = MagicMock(return_value=live_args) logging_obj.slack_alerting_instance.alerting_args = args_model monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) store = SettingsStore("general_settings") - store.load_yaml( - { - "alerting": ["slack"], - "alerting_args": {"daily_report_frequency": 3, "report_check_interval": 300}, - } - ) - store.apply_db_row("general_settings", {"alerting_args": db_alerting_args}) + store.load_yaml(yaml_values) + store.apply_db_row("general_settings", db_row) monkeypatch.setattr(proxy_server.proxy_config, "settings", store) monkeypatch.setattr(proxy_server, "general_settings", store) + return store + + +def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): + _alerting_client( + monkeypatch, + yaml_values={ + "alerting": ["slack"], + "alerting_args": {"daily_report_frequency": 3, "report_check_interval": 300}, + }, + db_row={"alerting_args": {"daily_report_frequency": 7, "outage_alert_ttl": 4242}}, + live_args={"daily_report_frequency": 3}, + ) with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/alerting/settings") @@ -224,6 +226,25 @@ def test_alerting_settings_reports_sources(client, auth_as, monkeypatch): assert by_name["slack_alerting"]["source"] == "config" assert by_name["daily_report_frequency"]["source"] == "config" assert by_name["report_check_interval"]["source"] == "config" + assert by_name["outage_alert_ttl"]["source"] == "default" + assert by_name["budget_alert_ttl"]["source"] == "default" + + +def test_alerting_settings_reports_db_source_when_the_file_omits_alerting_args(client, auth_as, monkeypatch): + store = _alerting_client( + monkeypatch, + yaml_values={"alerting": ["slack"]}, + db_row={"alerting_args": {"outage_alert_ttl": 4242, "region_outage_alert_ttl": []}}, + live_args={"outage_alert_ttl": 4242}, + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + + assert store.owned_by_config("alerting_args") is False assert by_name["outage_alert_ttl"]["source"] == "db" assert by_name["region_outage_alert_ttl"]["source"] == "default" assert by_name["budget_alert_ttl"]["source"] == "default" From 25af094e27ba4b40ceabaccbae944528a807c3cf Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:52:28 +0000 Subject: [PATCH 059/109] fix(python-bridge): allow instance attributes to shadow class defaults Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/facade.rs | 5 +- litellm/rust_bridge/_native.pyi | 64 ++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 58730857d60..d7ec2052dd0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -117,7 +117,10 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + if !attributes.get_item(name)?.is(value.bind(py)) { + return Ok(false); + } + if instance.contains(name)? && value.bind(py).is_callable() { return Ok(false); } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..7eb266d5a09 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence -from typing import Never, final +from typing import Literal, Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest @@ -93,6 +93,68 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... +@final +class _CacheTestHandle: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @staticmethod + def memory( + *, + capacity: int = 200, + ttl_seconds: float = 600.0, + max_entry_bytes: int = 1048576, + ) -> _CacheTestHandle: ... + @staticmethod + def redis( + url: str, + *, + ttl_seconds: float = 60.0, + namespace: str | None = None, + ) -> _CacheTestHandle: ... + @staticmethod + def redis_semantic(backend: object) -> _CacheTestHandle: ... + @property + def backend(self) -> Literal["memory", "redis", "redis_semantic"]: ... + def _bind_facade(self, facade: object) -> None: ... + +@final +class _CacheTestBinding: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @property + def kind(self) -> Literal["disabled", "native", "python_callback"]: ... + def lookup( + self, request: object, *, callback_kwargs: object = None + ) -> object: ... + def store( + self, request: object, response: object, *, callback_kwargs: object = None + ) -> None: ... + def lookup_batch( + self, requests: object, *, callback_kwargs: object = None + ) -> object: ... + def async_lookup( + self, request: object, *, callback_kwargs: object = None + ) -> Future[object]: ... + def async_store( + self, request: object, response: object, *, callback_kwargs: object = None + ) -> Future[object]: ... + def async_lookup_batch( + self, requests: object, *, callback_kwargs: object = None + ) -> Future[object]: ... + def async_store_batch( + self, + requests: object, + responses: object, + *, + callback_result: object = None, + callback_kwargs: object = None, + ) -> Future[object]: ... + def async_flush(self) -> Future[object]: ... + def ping(self) -> Future[object]: ... + +@final +class _CacheTestResolver: + def __new__(cls, namespace: object) -> _CacheTestResolver: ... + def resolve(self) -> _CacheTestBinding: ... + @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... From 4a0151f77fc3d8b68606a59171412fbe55b4015e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:52:28 +0000 Subject: [PATCH 060/109] test(rust): add redis-semantic native parity fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm_rust/test_cache.py | 549 ++++++++++++++++++++++++-- 1 file changed, 518 insertions(+), 31 deletions(-) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..e9f4b99a3b7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -1,14 +1,19 @@ import asyncio import contextvars import gc +import hashlib import json +import math +import os import threading import time import weakref -from collections.abc import Generator +from collections.abc import Callable, Generator +from contextlib import ExitStack from types import SimpleNamespace from typing import Final, Protocol, cast from urllib.parse import urlparse +from uuid import uuid4 import fakeredis import pytest @@ -17,10 +22,16 @@ import redis import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.isolation import rebound +_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name +_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + pytestmark: Final = pytest.mark.requires_rust_extension @@ -50,14 +61,14 @@ def test_existing_constructor_and_global_are_unchanged() -> None: assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) with rebound(litellm, "cache", facade): - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) enabled: Final = litellm.cache @@ -80,13 +91,13 @@ def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> Non async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) - resolver: Final = _native._CacheTestResolver(namespace) + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) + resolver: Final = _CacheTestResolver(namespace) selected: Final = resolver.resolve() assert selected.kind == "native" selected.store(request(), {"answer": 1}) assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _native._CacheTestHandle.memory()): + with rebound(namespace, "cache", _CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -119,7 +130,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -140,7 +151,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -155,9 +166,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -188,12 +199,12 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() with pytest.raises(TypeError): handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) facade: Final = Cache(type=LiteLLMCacheType.LOCAL) handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): @@ -218,7 +229,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: def cyclic_reference() -> weakref.ReferenceType[CustomCache]: callback: Final = CustomCache() namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -229,8 +240,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _native._CacheTestResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _CacheTestResolver(namespace).resolve() response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) @@ -252,33 +263,33 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _native._CacheTestHandle.memory(ttl_seconds=-1) + _CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve() small: Final = {"answer": "ok"} binding.store(request("small"), small) assert await binding.async_lookup(request("small")) == small await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) + disabled: Final = _CacheTestResolver( + SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0)) ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -316,7 +327,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: ) -> object: return result, kwargs - binding: Final = _native._CacheTestResolver( + binding: Final = _CacheTestResolver( SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) ).resolve() assert binding.kind == "python_callback" @@ -346,7 +357,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) cache.cache.set_cache("key", "value") - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() assert binding.kind == "python_callback" setattr(cache.cache, "ping", ping) @@ -358,7 +369,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: def test_facade_registration_rejects_mismatched_capacity() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) with pytest.raises(TypeError, match="capacities must match"): - _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) + _CacheTestHandle.memory(capacity=7)._bind_facade(facade) async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: @@ -371,19 +382,19 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: redis_flush_size=2, ) with pytest.raises(TypeError, match="default TTLs must match"): - _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) with pytest.raises(TypeError, match="namespaces must match"): - _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" pool: Final = facade.cache.redis_client.connection_pool with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None @@ -393,3 +404,479 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + +PARAPHRASE_MARKER: Final = " (paraphrase)" +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [ + (1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] + for index in range(8) + ] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context( + rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"]) + ) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, _CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store( + semantic_request("async", "name a primary color"), {"answer": "blue"} + ) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads( + cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response")) + ) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) == expected[key], key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup( + semantic_request("async-python", "python written prompt") + ) == {"answer": "python"} + client.close() + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store( + {**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1} + ) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + _CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + _CacheTestHandle.redis_semantic( + subclassed_facade.cache + )._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) From d3f2ddba050bad5af5ec662ac83bac263919d907 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:58:02 +0000 Subject: [PATCH 061/109] fix(python-bridge): allow instance shadowing only for validated config attributes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/facade.rs | 2 +- tests/test_litellm_rust/test_cache.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index d7ec2052dd0..a6395d086cb 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -120,7 +120,7 @@ impl ObjectGuard { if !attributes.get_item(name)?.is(value.bind(py)) { return Ok(false); } - if instance.contains(name)? && value.bind(py).is_callable() { + if instance.contains(name)? && !self.config_names.contains(&name.as_str()) { return Ok(false); } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index e9f4b99a3b7..73a6321011b 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -820,6 +820,8 @@ def test_redis_semantic_configuration_drift_falls_back_to_python( assert resolver.resolve().kind == "python_callback" with rebound(facade.cache, "_index_name", "other-index"): assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: return _semantic_embedding(prompt) From 0672fcbafeabb33cc0deeb325afbe53d14e431f1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:59:13 +0000 Subject: [PATCH 062/109] fix(rust): shrink release wheel under the 25 MB native limit Optimize the aws-sdk dependency tree and cache-s3 for size in the release profile and switch to fat LTO so the native extension stays under the wheel verification gate (27.87 -> 21.88 MB) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 74 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 73 insertions(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 031cfcbbb21..9cded32b7e0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -75,9 +75,81 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "thin" +lto = "fat" codegen-units = 1 panic = "unwind" debug = false incremental = false strip = "symbols" + +[profile.release.package."aws-config"] +opt-level = "s" + +[profile.release.package."aws-credential-types"] +opt-level = "s" + +[profile.release.package."aws-runtime"] +opt-level = "s" + +[profile.release.package."aws-sdk-kms"] +opt-level = "s" + +[profile.release.package."aws-sdk-s3"] +opt-level = "s" + +[profile.release.package."aws-sdk-secretsmanager"] +opt-level = "s" + +[profile.release.package."aws-sdk-sts"] +opt-level = "s" + +[profile.release.package."aws-sigv4"] +opt-level = "s" + +[profile.release.package."aws-smithy-async"] +opt-level = "s" + +[profile.release.package."aws-smithy-checksums"] +opt-level = "s" + +[profile.release.package."aws-smithy-eventstream"] +opt-level = "s" + +[profile.release.package."aws-smithy-http"] +opt-level = "s" + +[profile.release.package."aws-smithy-http-client"] +opt-level = "s" + +[profile.release.package."aws-smithy-json"] +opt-level = "s" + +[profile.release.package."aws-smithy-observability"] +opt-level = "s" + +[profile.release.package."aws-smithy-query"] +opt-level = "s" + +[profile.release.package."aws-smithy-runtime"] +opt-level = "s" + +[profile.release.package."aws-smithy-runtime-api"] +opt-level = "s" + +[profile.release.package."aws-smithy-runtime-api-macros"] +opt-level = "s" + +[profile.release.package."aws-smithy-schema"] +opt-level = "s" + +[profile.release.package."aws-smithy-types"] +opt-level = "s" + +[profile.release.package."aws-smithy-xml"] +opt-level = "s" + +[profile.release.package."aws-types"] +opt-level = "s" + +[profile.release.package."litellm-cache-s3"] +opt-level = "s" From 0d09e9d8929f0bac68932ec59dfbe465a805f2c1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:04:57 +0000 Subject: [PATCH 063/109] fix(rust-wheel): reduce native extension size Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..8725b25cfc7 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -73,7 +73,7 @@ fancy-regex = "0.19.2" veil = "0.3.0" [profile.release] -opt-level = 3 +opt-level = 2 lto = "thin" codegen-units = 1 panic = "unwind" From a9ff1de42bfdf85c6ae327113a968bacb914e99d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:09:36 +0000 Subject: [PATCH 064/109] build(rust): switch release LTO to fat for wheel size headroom Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 05eea6bc299..7fa8de05f60 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -75,7 +75,7 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "thin" +lto = "fat" codegen-units = 1 panic = "unwind" debug = false From 793efe3eb43c6e73778f3c5c846f5c6df91eb593 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:15:42 +0000 Subject: [PATCH 065/109] fix(rust): restore hot-path opt levels and validate s3 binding destination Keep sigv4 signing, eventstream decoding, smithy runtime api and types at opt-level 3 since they serve Bedrock request and streaming hot paths, and make the S3 facade binding reject region and endpoint mismatches between the projected configuration and the native handle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 8 --- litellm-rust/crates/cache-s3/src/cache.rs | 19 +++++- .../crates/python-bridge/src/cache/config.rs | 58 +++++++++++++++++++ .../crates/python-bridge/src/cache/native.rs | 14 +++++ 4 files changed, 88 insertions(+), 11 deletions(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 9cded32b7e0..5a72437d767 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -103,8 +103,6 @@ opt-level = "s" [profile.release.package."aws-sdk-sts"] opt-level = "s" -[profile.release.package."aws-sigv4"] -opt-level = "s" [profile.release.package."aws-smithy-async"] opt-level = "s" @@ -112,8 +110,6 @@ opt-level = "s" [profile.release.package."aws-smithy-checksums"] opt-level = "s" -[profile.release.package."aws-smithy-eventstream"] -opt-level = "s" [profile.release.package."aws-smithy-http"] opt-level = "s" @@ -133,8 +129,6 @@ opt-level = "s" [profile.release.package."aws-smithy-runtime"] opt-level = "s" -[profile.release.package."aws-smithy-runtime-api"] -opt-level = "s" [profile.release.package."aws-smithy-runtime-api-macros"] opt-level = "s" @@ -142,8 +136,6 @@ opt-level = "s" [profile.release.package."aws-smithy-schema"] opt-level = "s" -[profile.release.package."aws-smithy-types"] -opt-level = "s" [profile.release.package."aws-smithy-xml"] opt-level = "s" diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index a2ea33c33e4..88575f36fce 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -36,18 +36,21 @@ pub struct S3Cache { runtime: Handle, bucket: Arc, key_prefix: Arc, + region: Arc, + endpoint: Option>, } impl S3Cache { pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self { let mut builder = aws_sdk_s3::Config::builder() .behavior_version(BehaviorVersion::latest()) - .region(Region::new(config.region)) + .region(Region::new(config.region.clone())) .credentials_provider(Credentials::new(config.auth)) .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) .response_checksum_validation(ResponseChecksumValidation::WhenRequired); - if let Some(endpoint) = config.endpoint { - builder = builder.endpoint_url(endpoint.url).force_path_style(true); + let endpoint_url: Option = config.endpoint.map(|endpoint| endpoint.url); + if let Some(url) = &endpoint_url { + builder = builder.endpoint_url(url).force_path_style(true); } Self { client: aws_sdk_s3::Client::from_conf(builder.build()), @@ -55,6 +58,8 @@ impl S3Cache { runtime, bucket: config.bucket.into(), key_prefix: config.key_prefix.into(), + region: config.region.into(), + endpoint: endpoint_url.map(Into::into), } } @@ -66,6 +71,14 @@ impl S3Cache { &self.key_prefix } + pub fn region(&self) -> &str { + &self.region + } + + pub fn endpoint(&self) -> Option<&str> { + self.endpoint.as_deref() + } + pub fn to_s3_key(&self, key: &str) -> String { format!("{}{}", self.key_prefix, key.replace(':', "/")) } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 5d37ed27083..003a2687e88 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -212,6 +212,18 @@ impl NativeCacheConfig { { Some("facade and native backend key prefixes must match") } + CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => { + Some("facade and native backend regions must match") + } + CacheBackendConfig::S3(config) + if service.endpoint() + != config + .endpoint + .as_ref() + .map(|endpoint| endpoint.url.as_str()) => + { + Some("facade and native backend endpoints must match") + } CacheBackendConfig::S3(_) => None, } } @@ -593,6 +605,10 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; + use litellm_auth_aws::AwsAuthConfig; + use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; + use litellm_host_python::run_sync_value; + use super::{ CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, RedisProtocol, @@ -858,4 +874,46 @@ mod tests { assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1")); }); } + + fn s3_service(py: Python<'_>, region: &str, endpoint: Option<&str>) -> NativeResponseCache { + let config = S3CacheConfig { + bucket: "bucket".to_string(), + key_prefix: "team/".to_string(), + region: region.to_string(), + endpoint: endpoint.map(|url| S3Endpoint { + url: url.to_string(), + }), + auth: AwsAuthConfig::default(), + }; + run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) }).unwrap() + } + + #[test] + fn s3_binding_rejects_region_and_endpoint_mismatches() { + Python::initialize(); + Python::attach(|py| { + let facade = s3_facade(py, ""); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("S3 cache should be supported"); + }; + assert_eq!( + config.service_mismatch(&s3_service(py, "us-east-1", Some("https://example.test"))), + None + ); + assert_eq!( + config.service_mismatch(&s3_service(py, "us-west-2", Some("https://example.test"))), + Some("facade and native backend regions must match") + ); + assert_eq!( + config.service_mismatch(&s3_service(py, "us-east-1", Some("https://other.test"))), + Some("facade and native backend endpoints must match") + ); + assert_eq!( + config.service_mismatch(&s3_service(py, "us-east-1", None)), + Some("facade and native backend endpoints must match") + ); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index b9ea443119b..a6eef420ace 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -87,6 +87,20 @@ impl NativeResponseCache { } } + pub fn region(&self) -> Option<&str> { + match self { + Self::S3(cache) => Some(cache.backend().region()), + _ => None, + } + } + + pub fn endpoint(&self) -> Option<&str> { + match self { + Self::S3(cache) => cache.backend().endpoint(), + _ => None, + } + } + pub fn namespace(&self) -> Option<&str> { match self { Self::Memory(_) => None, From 64078689d303fdd85d2c9c25ca496cb7bebd834d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:18:01 -0700 Subject: [PATCH 066/109] refactor(mcp): clear server list lint warnings --- .../_components/mcp_servers.test.tsx | 128 +++++++++++- .../mcp-servers/_components/mcp_servers.tsx | 195 ++++++++++-------- 2 files changed, 232 insertions(+), 91 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 61880363234..7198cb3cce9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -167,6 +167,130 @@ describe("MCPServers", () => { const myConnections = await screen.findByRole("link", { name: "My Connections" }); expect(myConnections).toBeVisible(); expect(myConnections).toHaveAttribute("href", "/ui/connect"); + for (const name of ["Semantic Filter", "Tool Search", "Network Settings", "Submitted MCPs"]) { + const tab = screen.queryByRole("tab", { name }); + if (userRole === "Admin") { + expect(tab).toBeVisible(); + } else { + expect(tab).not.toBeInTheDocument(); + } + } + expect( + screen.getByRole("button", { + name: userRole === "Admin" ? "+ Add New MCP Server" : "+ Submit MCP Server", + }), + ).toBeVisible(); + }); + + it.each(["cancel", "success", "failure", "unnamed"])("preserves delete confirmation on %s", async (outcome) => { + const server: MCPServer = { + server_id: "delete-server", + server_name: outcome === "unnamed" ? null : "Delete fixture", + alias: "delete-alias", + url: outcome === "unnamed" ? null : "https://example.com/mcp", + created_by: "user", + updated_by: "user", + }; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + let finishDelete: () => void = () => {}; + vi.mocked(networking.deleteMCPServer).mockImplementation( + () => + new Promise((resolve, reject) => { + finishDelete = () => (outcome === "failure" ? reject(new Error("Delete failed")) : resolve(undefined)); + }), + ); + render( + + + , + ); + await userEvent.click(await screen.findByRole("button", { name: "Server actions" })); + await userEvent.click(await screen.findByRole("menuitem", { name: "Delete" })); + const dialog = await screen.findByRole("alertdialog", { name: "Delete MCP Server?" }); + expect(within(dialog).getByText("delete-server")).toBeVisible(); + if (outcome === "unnamed") { + expect(within(dialog).queryByText("Name")).not.toBeInTheDocument(); + expect(within(dialog).queryByText("URL")).not.toBeInTheDocument(); + } else { + expect(within(dialog).getByText("Delete fixture")).toBeVisible(); + expect(within(dialog).getByText("https://example.com/mcp")).toBeVisible(); + } + if (outcome === "cancel") { + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + expect(networking.deleteMCPServer).not.toHaveBeenCalled(); + } else { + await userEvent.click(within(dialog).getByRole("button", { name: "Delete", exact: true })); + expect(within(dialog).getByRole("button", { name: "Deleting..." })).toBeDisabled(); + expect(within(dialog).getByRole("button", { name: "Cancel" })).toBeDisabled(); + expect(networking.deleteMCPServer).toHaveBeenCalledWith("123", "delete-server"); + await act(async () => finishDelete()); + } + await waitFor(() => expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument()); + }); + + it("filters servers by access group", async () => { + const server = { created_by: "user", updated_by: "user" }; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { + ...server, + server_id: "string-group", + server_name: "String group", + alias: "string-alias", + mcp_access_groups: ["shared"], + }, + { + ...server, + server_id: "legacy-group", + server_name: "Legacy group", + alias: "legacy-alias", + mcp_access_groups: ["shared"], + }, + { + ...server, + server_id: "other-group", + server_name: "Other group", + alias: "other-alias", + mcp_access_groups: ["different"], + }, + ]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + render( + + + , + ); + await screen.findByText("String group"); + await userEvent.click(screen.getByRole("combobox", { name: "Access Group" })); + await userEvent.click(await screen.findByRole("option", { name: "shared", exact: true })); + expect(screen.getByText("String group")).toBeVisible(); + expect(screen.getByText("Legacy group")).toBeVisible(); + expect(screen.queryByText("Other group")).not.toBeInTheDocument(); + }); + + it.each(["server_name", "alias", "url", "server_id"] as const)("searches by %s case-insensitively", async (field) => { + const server: MCPServer = { + server_id: "search-server", + server_name: "Search fixture", + created_by: "user", + updated_by: "user", + [field]: "Needle", + }; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + render( + + + , + ); + await screen.findByTestId("mcp-servers-grid"); + const search = screen.getByPlaceholderText("Search by name, alias, URL, or ID"); + await userEvent.type(search, " NEEDLE "); + expect(screen.getByTestId("mcp-servers-grid")).toBeVisible(); + await userEvent.clear(search); + await userEvent.type(search, "no-match"); + expect(screen.queryByTestId("mcp-servers-grid")).not.toBeInTheDocument(); + expect(screen.getByText("No servers match the current filters or search.")).toBeVisible(); }); it("should render mocked MCP servers data in the table", async () => { @@ -409,9 +533,7 @@ describe("MCPServers", () => { expect(screen.getByText("Team B Server")).toBeInTheDocument(); expect(screen.getByText("Team A Server 2")).toBeInTheDocument(); - // Find the team select by its "Team" label, then the combobox it labels - const teamLabel = screen.getByText("Team"); - const teamSelect = within(teamLabel.parentElement!).getByRole("combobox"); + const teamSelect = screen.getByRole("combobox", { name: "Team" }); await userEvent.click(teamSelect); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index 4a5ab379969..818b5150650 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -113,6 +113,62 @@ const readToolsOAuthServerId = (): string | null => { } }; +function DeleteServerDialog({ + open, + onOpenChange, + server, + isDeleting, + onConfirm, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + server: MCPServer | undefined; + isDeleting: boolean; + onConfirm: () => Promise; +}) { + return ( + + + + Delete MCP Server? + +
+

+ This action is permanent and cannot be undone. All associated configurations will be removed. +

+ + {server && ( +
+ {server.server_name && ( +
+
Name
+
{server.server_name}
+
+ )} +
+
ID
+
{server.server_id}
+
+ {server.url && ( +
+
URL
+
{server.url}
+
+ )} +
+ )} +
+ + Cancel + + + + + ); +} + const MCPServers: React.FC = ({ accessToken, userRole, userID, isViewOnly = false }) => { const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers(); @@ -299,7 +355,9 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i } if (group !== "all") { filtered = filtered.filter((server) => - server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)), + server.mcp_access_groups?.some((g: string | { name?: string } | null) => + typeof g === "string" ? g === group : g?.name === group, + ), ); } setFilteredServers(filtered); @@ -333,7 +391,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i const alias = (s.alias || "").toLowerCase(); const url = (s.url || "").toLowerCase(); const id = s.server_id.toLowerCase(); - return name.includes(q) || alias.includes(q) || url.includes(q) || id.includes(q); + return [name, alias, url, id].some((value) => value.includes(q)); }) : filteredServers; return [...matches].sort((a, b) => compareServers(a, b, sortKey)); @@ -376,9 +434,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i }; // Find the server to delete from the servers list - const serverToDelete = serverIdToDelete - ? (mcpServers || []).find((server) => server.server_id === serverIdToDelete) - : null; + const serverToDelete = mcpServers?.find((server) => server.server_id === serverIdToDelete); const handleCreateSuccess = (newMcpServer: MCPServer) => { setFilteredServers((prev) => [...prev, newMcpServer]); @@ -420,45 +476,13 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i return (
- !open && cancelDelete()}> - - - Delete MCP Server? - -
-

- This action is permanent and cannot be undone. All associated configurations will be removed. -

- - {serverToDelete && ( -
- {serverToDelete.server_name && ( -
-
Name
-
{serverToDelete.server_name}
-
- )} -
-
ID
-
{serverToDelete.server_id}
-
- {serverToDelete.url && ( -
-
URL
-
{serverToDelete.url}
-
- )} -
- )} -
- - Cancel - - - - + !open && cancelDelete()} + server={serverToDelete} + isDeleting={isDeletingServer} + onConfirm={confirmDelete} + /> = ({ accessToken, userRole, userID, i My Connections - {isAdminRole(userRole) && ( + {isAdminRole(userRole) ? ( <> - )} - {!isAdminRole(userRole) && ( + ) : (