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 01/41] 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 02/41] 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 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 03/41] 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 d2f457f144a430ad2848b33839d298d9312c8d4e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:20:31 +0000 Subject: [PATCH 04/41] 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 05/41] 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 1f86bb8e4640fd7e758e10f66106fd7c1bda01de Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:24:34 +0000 Subject: [PATCH 06/41] 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 4db34812449038fed2c724a3ef099fefb3198ac2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:33:32 +0000 Subject: [PATCH 07/41] 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 08/41] 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 788e24655a0e253f553af95da8ac69a51552a6cd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:35:08 +0000 Subject: [PATCH 09/41] 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 605729a194c997b8f90a5edb5494465be9cc4d95 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:39:17 +0000 Subject: [PATCH 10/41] 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 11/41] 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 0672fcbafeabb33cc0deeb325afbe53d14e431f1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:59:13 +0000 Subject: [PATCH 12/41] 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 13/41] 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 793efe3eb43c6e73778f3c5c846f5c6df91eb593 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:15:42 +0000 Subject: [PATCH 14/41] 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 15/41] 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) && ( + ) : (