From 6bbcf1dbbf3d533ee1948b9b3183d8962ace4263 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:08:08 +0000 Subject: [PATCH 01/56] fix(mcp): keep the streamable-HTTP routing peek on a UTF-8 boundary Fixes https://github.com/BerriAI/litellm/issues/34917 --- .../proxy/_experimental/mcp_server/server.py | 24 +++- .../mcp_server/test_mcp_server.py | 104 ++++++++++++++++++ 2 files changed, 125 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 14673cf12c1..af5da275961 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -237,6 +237,24 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: return False +def _utf8_boundary_prefix(data: bytes) -> bytes: + """``data`` with any trailing incomplete UTF-8 sequence removed. + + Cutting a body at a fixed byte budget can land in the middle of a multibyte + character, and ``json.loads`` on such bytes raises ``UnicodeDecodeError`` + rather than ``JSONDecodeError``. Trimming to a character boundary keeps the + truncated peek decodable so callers only have to handle malformed JSON. + """ + for trailing in range(0, min(3, len(data)) + 1): + candidate = data[: len(data) - trailing] + try: + candidate.decode("utf-8") + except UnicodeDecodeError: + continue + return candidate + return data + + def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. @@ -3411,7 +3429,7 @@ if MCP_AVAILABLE: try: data = json.loads(body) return isinstance(data, dict) and data.get("method") == "initialize" - except (json.JSONDecodeError, TypeError): + except (json.JSONDecodeError, UnicodeDecodeError, TypeError): return False async def _read_request_body_for_routing( @@ -3462,7 +3480,7 @@ if MCP_AVAILABLE: # directly from the original `receive` via wrapped_receive. break - return consumed_messages, b"".join(body_chunks) + return consumed_messages, _utf8_boundary_prefix(b"".join(body_chunks)) async def _handle_stale_mcp_session( scope: Scope, @@ -4227,7 +4245,7 @@ if MCP_AVAILABLE: "MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock", _peeked.get("id"), ) - except (json.JSONDecodeError, TypeError): + except (json.JSONDecodeError, UnicodeDecodeError, TypeError): # Peek cap truncated the body, so it can't be fully parsed. # Scan the top-level keys (depth-aware) instead of a flat # substring search: a response's result payload may nest a diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 1753b0d92a8..7f79e5aebda 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,5 +1,6 @@ import asyncio import contextvars +import json from datetime import datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch @@ -1689,6 +1690,109 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): assert total_streamed == len(first_chunk) + sum(len(b) for b in oversized_tail) +@pytest.mark.asyncio +async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap(): + """ + A tool-call POST whose UTF-8 body is larger than the routing peek cap, with a + multibyte character straddling the cap boundary, must still be forwarded + intact instead of blowing up with a UnicodeDecodeError 500. + + Regression test for https://github.com/BerriAI/litellm/issues/34917 + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + peek_cap = mcp_server._MCP_ROUTING_PEEK_MAX_BYTES + + def _splits_multibyte_at_cap(candidate: bytes) -> bool: + try: + candidate[:peek_cap].decode("utf-8") + except UnicodeDecodeError: + return True + return False + + def _build_body() -> bytes: + for pad in range(4): + candidate = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "update_full_document" + "x" * pad, + "arguments": {"markdown": "щ" * 3000}, + }, + }, + ensure_ascii=False, + ).encode("utf-8") + if len(candidate) > peek_cap and _splits_multibyte_at_cap(candidate): + return candidate + raise AssertionError("could not build a body splitting a multibyte char at the peek cap") + + body = _build_body() + + messages = [{"type": "http.request", "body": body, "more_body": False}] + receive_calls = {"count": 0} + + async def receive(): + idx = receive_calls["count"] + receive_calls["count"] += 1 + return messages[idx] + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/progress_test", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + send = AsyncMock() + + streamed_chunks = [] + + async def stateless_handle(s, r, se): + while True: + msg = await r() + if msg.get("type") != "http.request": + break + streamed_chunks.append(msg.get("body", b"") or b"") + if not msg.get("more_body", False): + break + + async def stateful_handle(s, r, se): + raise AssertionError("non-initialize POST should not reach stateful manager") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object(session_manager_stateless, "handle_request", side_effect=stateless_handle), + patch.object(session_manager_stateful, "handle_request", side_effect=stateful_handle), + patch.object(session_manager_stateless, "_server_instances", {}), + patch.object(session_manager_stateful, "_server_instances", {}), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert send.await_count == 0, f"unexpected response emitted by the proxy: {send.await_args_list}" + assert b"".join(streamed_chunks) == body + + @pytest.mark.asyncio async def test_enforce_stateful_session_cap_evicts_oldest_idle_then_rejects(): """ From 4a6a387ca1e511e35858fee0c92fe3e3415d03ee Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:48:22 +0000 Subject: [PATCH 02/56] fix(mcp): follow nextCursor on paginated tools/prompts/resources list operations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/experimental_mcp_client/client.py | 56 ++++++----- litellm/experimental_mcp_client/pagination.py | 92 +++++++++++++++++++ litellm/experimental_mcp_client/tools.py | 7 +- .../mcp_server/rest_endpoints.py | 6 +- .../test_mcp_client.py | 52 ++++++++++- .../test_pagination.py | 80 ++++++++++++++++ 7 files changed, 258 insertions(+), 36 deletions(-) create mode 100644 litellm/experimental_mcp_client/pagination.py create mode 100644 tests/test_litellm/experimental_mcp_client/test_pagination.py diff --git a/litellm/constants.py b/litellm/constants.py index 9a50797f517..11f35177636 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -136,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_LIST_MAX_PAGES: Final = int(os.getenv("LITELLM_MCP_LIST_MAX_PAGES", "100")) # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index f0a1bff8fdc..814074d35a4 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -48,6 +48,12 @@ from pydantic import AnyUrl from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR +from litellm.experimental_mcp_client.pagination import ( + list_all_prompts, + list_all_resource_templates, + list_all_resources, + list_all_tools, +) from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -603,17 +609,17 @@ class MCPClient: """ verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_tools_operation(session: ClientSession): - return await session.list_tools() + async def _list_tools_operation(session: ClientSession) -> tuple[MCPTool, ...]: + return await list_all_tools(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) - tool_count: Final = len(result.tools) - tool_names: Final = [tool.name for tool in result.tools] + tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) + tool_count: Final = len(tools) + tool_names: Final = [tool.name for tool in tools] verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) - return result.tools + return list(tools) except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") raise @@ -734,17 +740,17 @@ class MCPClient: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_prompts_operation(session: ClientSession): - return await session.list_prompts() + async def _list_prompts_operation(session: ClientSession) -> tuple[Prompt, ...]: + return await list_all_prompts(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_prompts_operation) - prompt_count: Final = len(result.prompts) - prompt_names: Final = [prompt.name for prompt in result.prompts] + prompts: Final = await self.run_with_session(_list_prompts_operation) + prompt_count: Final = len(prompts) + prompt_names: Final = [prompt.name for prompt in prompts] verbose_logger.info( - "MCP client listed %s tools from %s: %s", prompt_count, self.server_url or "stdio", prompt_names + "MCP client listed %s prompts from %s: %s", prompt_count, self.server_url or "stdio", prompt_names ) - return result.prompts + return list(prompts) except asyncio.CancelledError: verbose_logger.warning("MCP client list_prompts was cancelled") raise @@ -811,17 +817,17 @@ class MCPClient: """List available resources from the server.""" verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") - async def _list_resources_operation(session: ClientSession): - return await session.list_resources() + async def _list_resources_operation(session: ClientSession) -> tuple[Resource, ...]: + return await list_all_resources(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_resources_operation) - resource_count: Final = len(result.resources) - resource_names: Final = [resource.name for resource in result.resources] + resources: Final = await self.run_with_session(_list_resources_operation) + resource_count: Final = len(resources) + resource_names: Final = [resource.name for resource in resources] verbose_logger.info( "MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names ) - return result.resources + return list(resources) except asyncio.CancelledError: verbose_logger.warning("MCP client list_resources was cancelled") raise @@ -847,20 +853,20 @@ class MCPClient: """List available resource templates from the server.""" verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") - async def _list_resource_templates_operation(session: ClientSession): - return await session.list_resource_templates() + async def _list_resource_templates_operation(session: ClientSession) -> tuple[ResourceTemplate, ...]: + return await list_all_resource_templates(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_resource_templates_operation) - resource_template_count: Final = len(result.resourceTemplates) - resource_template_names: Final = [resourceTemplate.name for resourceTemplate in result.resourceTemplates] + resource_templates: Final = await self.run_with_session(_list_resource_templates_operation) + resource_template_count: Final = len(resource_templates) + resource_template_names: Final = [resource_template.name for resource_template in resource_templates] verbose_logger.info( "MCP client listed %s resource templates from %s: %s", resource_template_count, self.server_url or "stdio", resource_template_names, ) - return result.resourceTemplates + return list(resource_templates) except asyncio.CancelledError: verbose_logger.warning("MCP client list_resource_templates was cancelled") raise diff --git a/litellm/experimental_mcp_client/pagination.py b/litellm/experimental_mcp_client/pagination.py new file mode 100644 index 00000000000..8852715aba5 --- /dev/null +++ b/litellm/experimental_mcp_client/pagination.py @@ -0,0 +1,92 @@ +""" +Follows ``nextCursor`` on the paginated MCP list operations so a multi-page catalog is read in full. +""" + +from collections.abc import Awaitable, Callable, Sequence +from typing import Final, TypeVar + +from mcp import ClientSession, Resource +from mcp.types import PaginatedRequestParams, PaginatedResult, Prompt, ResourceTemplate +from mcp.types import Tool as MCPTool + +from litellm._logging import verbose_logger +from litellm.constants import MCP_LIST_MAX_PAGES + +TPage = TypeVar("TPage", bound=PaginatedResult) +TItem = TypeVar("TItem") + + +async def collect_pages( + fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[TPage]], + items_of: Callable[[TPage], Sequence[TItem]], + *, + method: str, + server: str, + cursor: str | None = None, + seen_cursors: frozenset[str] = frozenset(), +) -> tuple[TItem, ...]: + page: Final = await fetch_page(None if cursor is None else PaginatedRequestParams(cursor=cursor)) + items: Final = tuple(items_of(page)) + next_cursor: Final = page.nextCursor + pages_read: Final = len(seen_cursors) + 1 + if next_cursor is None: + return items + if next_cursor in seen_cursors: + verbose_logger.warning( + "MCP %s from %s repeated cursor %r; returning the %s page(s) read so far", + method, + server, + next_cursor, + pages_read, + ) + return items + if pages_read >= MCP_LIST_MAX_PAGES: + verbose_logger.warning( + "MCP %s from %s still paginating after %s pages (LITELLM_MCP_LIST_MAX_PAGES); returning what was read", + method, + server, + pages_read, + ) + return items + rest: Final = await collect_pages( + fetch_page, + items_of, + method=method, + server=server, + cursor=next_cursor, + seen_cursors=seen_cursors | frozenset((next_cursor,)), + ) + return items + rest + + +async def list_all_tools(session: ClientSession, server: str) -> tuple[MCPTool, ...]: + return await collect_pages( + lambda params: session.list_tools(params=params), lambda page: page.tools, method="tools/list", server=server + ) + + +async def list_all_prompts(session: ClientSession, server: str) -> tuple[Prompt, ...]: + return await collect_pages( + lambda params: session.list_prompts(params=params), + lambda page: page.prompts, + method="prompts/list", + server=server, + ) + + +async def list_all_resources(session: ClientSession, server: str) -> tuple[Resource, ...]: + return await collect_pages( + lambda params: session.list_resources(params=params), + lambda page: page.resources, + method="resources/list", + server=server, + ) + + +async def list_all_resource_templates(session: ClientSession, server: str) -> tuple[ResourceTemplate, ...]: + return await collect_pages( + lambda params: session.list_resource_templates(params=params), + lambda page: page.resourceTemplates, + method="resources/templates/list", + server=server, + ) diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 30d50e2a74b..adaca0888aa 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -9,6 +9,7 @@ from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from litellm.experimental_mcp_client.pagination import list_all_tools from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall @@ -103,10 +104,10 @@ async def load_mcp_tools( If format is set to "openai", the tools are converted to OpenAI API compatible tools. """ - tools: Final = await session.list_tools() + tools: Final = await list_all_tools(session, "upstream") if format == "openai": - return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools] - return tools.tools + return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools] + return list(tools) ######################################################## diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..3ca6b6c5f90 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1402,11 +1402,7 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - async def _list_tools_session_operation(session): - return await session.list_tools() - - list_tools_response: Final = await client.run_with_session(_list_tools_session_operation) - list_tools_result: Final[list[MCPTool]] = list_tools_response.tools + list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True) model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index fd7ab3afdab..3f501d3859a 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -20,8 +20,10 @@ from mcp.types import ( JSONRPCError, JSONRPCMessage, JSONRPCResponse, + ListToolsResult, ServerCapabilities, ) +from mcp.types import Tool as MCPTool # Add the parent directory to the path so we can import litellm @@ -740,8 +742,14 @@ class _ScriptedUpstream: error, the shape an upstream application uses to report its own failure. """ - def __init__(self, tools_list_error: ErrorData | None = None): + def __init__( + self, + tools_list_error: ErrorData | None = None, + tool_pages: tuple[tuple[MCPTool, ...], ...] = (), + ): self._tools_list_error = tools_list_error + self._tool_pages = tool_pages + self.tools_list_cursors: list[str | None] = [] self._to_client_tx, self._to_client_rx = anyio.create_memory_object_stream(10) self._from_client_tx, self._from_client_rx = anyio.create_memory_object_stream(10) self._task_group = None @@ -778,15 +786,37 @@ class _ScriptedUpstream: ) elif method == "tools/list" and self._tools_list_error is not None: await self._send(JSONRPCError(jsonrpc="2.0", id=request.id, error=self._tools_list_error)) + elif method == "tools/list" and self._tool_pages: + cursor = (request.params or {}).get("cursor") + self.tools_list_cursors.append(cursor) + page_index = int(cursor) if cursor else 0 + has_more = page_index + 1 < len(self._tool_pages) + page = ListToolsResult( + tools=list(self._tool_pages[page_index]), + nextCursor=str(page_index + 1) if has_more else None, + ) + await self._send( + JSONRPCResponse( + jsonrpc="2.0", + id=request.id, + result=page.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) class _ScriptedClient(MCPClient): """An MCPClient whose transport is a scripted in-memory upstream instead of a real connection, so the real ``ClientSession`` and its real timeout machinery are what run.""" - def __init__(self, *, timeout: float, tools_list_error: ErrorData | None = None): + def __init__( + self, + *, + timeout: float, + tools_list_error: ErrorData | None = None, + tool_pages: tuple[tuple[MCPTool, ...], ...] = (), + ): super().__init__(server_url="http://upstream.local/mcp", timeout=timeout) - self._upstream = _ScriptedUpstream(tools_list_error=tools_list_error) + self._upstream = _ScriptedUpstream(tools_list_error=tools_list_error, tool_pages=tool_pages) def _create_transport_context(self): return self._upstream, None @@ -821,6 +851,22 @@ async def test_list_tools_fails_on_its_own_timeout_when_the_upstream_never_answe assert list_fault_http_status(fault) == 504 +@pytest.mark.asyncio +async def test_list_tools_follows_tools_list_pagination_across_the_whole_catalog(): + """An upstream that pages tools/list (72 tools, 30 per page) must have every page read within the + one session, each request carrying the cursor the previous page returned. Reading only the first + page made 42 tools invisible to the proxy and every call to them fail as unknown.""" + tools = tuple( + MCPTool(name=f"tool_{i:02d}", inputSchema={"type": "object", "properties": {}}) for i in range(72) + ) + client = _ScriptedClient(timeout=30, tool_pages=(tools[:30], tools[30:60], tools[60:])) + + listed = await asyncio.wait_for(client.list_tools(raise_on_error=True), timeout=10) + + assert [tool.name for tool in listed] == [tool.name for tool in tools] + assert client._upstream.tools_list_cursors == [None, "1", "2"] + + @pytest.mark.asyncio async def test_upstream_json_rpc_error_408_is_not_reported_as_a_client_timeout(): """The SDK reports its own elapsed read timeout and relays an upstream JSON-RPC error through diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py new file mode 100644 index 00000000000..f588fdd1eee --- /dev/null +++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py @@ -0,0 +1,80 @@ +import logging + +import pytest +from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Tool as MCPTool + +import litellm.experimental_mcp_client.pagination as pagination_module +from litellm.experimental_mcp_client.pagination import collect_pages + + +def _tool(index: int) -> MCPTool: + return MCPTool(name=f"tool_{index:02d}", inputSchema={"type": "object", "properties": {}}) + + +class _PagedTools: + """A tools/list upstream serving ``total`` tools ``page_size`` at a time, cursors being offsets.""" + + def __init__(self, total: int, page_size: int): + self._tools = tuple(_tool(i) for i in range(total)) + self._page_size = page_size + self.cursors_seen: list[str | None] = [] + + async def fetch(self, params: PaginatedRequestParams | None) -> ListToolsResult: + cursor = params.cursor if params is not None else None + self.cursors_seen.append(cursor) + start = int(cursor) if cursor else 0 + end = start + self._page_size + return ListToolsResult( + tools=list(self._tools[start:end]), + nextCursor=str(end) if end < len(self._tools) else None, + ) + + +@pytest.mark.asyncio +async def test_collect_pages_follows_next_cursor_until_exhausted(): + upstream = _PagedTools(total=72, page_size=30) + + tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") + + assert [t.name for t in tools] == [f"tool_{i:02d}" for i in range(72)] + assert upstream.cursors_seen == [None, "30", "60"], "each page must be requested with the cursor the previous one returned" + + +@pytest.mark.asyncio +async def test_collect_pages_single_page_makes_one_request(): + upstream = _PagedTools(total=5, page_size=30) + + tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") + + assert len(tools) == 5 + assert upstream.cursors_seen == [None] + + +@pytest.mark.asyncio +async def test_collect_pages_stops_on_a_repeated_cursor_and_keeps_what_it_read(caplog): + calls: list[str | None] = [] + + async def fetch(params: PaginatedRequestParams | None) -> ListToolsResult: + calls.append(params.cursor if params else None) + return ListToolsResult(tools=[_tool(len(calls))], nextCursor="same") + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tools = await collect_pages(fetch, lambda page: page.tools, method="tools/list", server="s") + + assert calls == [None, "same"], "the cursor must be followed once and refused the second time it comes back" + assert len(tools) == 2 + assert any("repeated cursor" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_collect_pages_honors_the_page_cap(monkeypatch, caplog): + monkeypatch.setattr(pagination_module, "MCP_LIST_MAX_PAGES", 3) + upstream = _PagedTools(total=1000, page_size=10) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") + + assert len(upstream.cursors_seen) == 3 + assert len(tools) == 30 + assert any("LITELLM_MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records) From 26b48d58919e5021b9b251339bf5c720a3e3649e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:00:27 +0000 Subject: [PATCH 03/56] refactor(mcp): keep list pagination within type-discipline budget and ratchet LIT001 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 56 ++++++++++--------- .../mcp_server/rest_endpoints.py | 4 +- .../test_pagination.py | 4 +- type-discipline-budget.json | 2 +- 4 files changed, 35 insertions(+), 31 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 814074d35a4..46d8823e439 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -377,29 +377,31 @@ class MCPClient: return provided_env # Minimal allowlist of safe/standard environment variables - safe_keys: Final = { - "PATH", - "HOME", - "USER", - "LOGNAME", - "TMPDIR", - "TMP", - "TEMP", - "SHELL", - "LANG", - "LC_ALL", - # Node/Package manager caches - "NPM_CONFIG_CACHE", - "PNPM_HOME", - "XDG_CACHE_HOME", - "XDG_CONFIG_HOME", - "XDG_DATA_HOME", - # System info - "SYSTEMROOT", - "COMSPEC", - "PATHEXT", - "WINDIR", - } + safe_keys: Final = frozenset( + { + "PATH", + "HOME", + "USER", + "LOGNAME", + "TMPDIR", + "TMP", + "TEMP", + "SHELL", + "LANG", + "LC_ALL", + # Node/Package manager caches + "NPM_CONFIG_CACHE", + "PNPM_HOME", + "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + # System info + "SYSTEMROOT", + "COMSPEC", + "PATHEXT", + "WINDIR", + } + ) safe_env: Final = {} for key in safe_keys: @@ -615,7 +617,7 @@ class MCPClient: try: tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) tool_count: Final = len(tools) - tool_names: Final = [tool.name for tool in tools] + tool_names: Final = tuple(tool.name for tool in tools) verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) @@ -746,7 +748,7 @@ class MCPClient: try: prompts: Final = await self.run_with_session(_list_prompts_operation) prompt_count: Final = len(prompts) - prompt_names: Final = [prompt.name for prompt in prompts] + prompt_names: Final = tuple(prompt.name for prompt in prompts) verbose_logger.info( "MCP client listed %s prompts from %s: %s", prompt_count, self.server_url or "stdio", prompt_names ) @@ -823,7 +825,7 @@ class MCPClient: try: resources: Final = await self.run_with_session(_list_resources_operation) resource_count: Final = len(resources) - resource_names: Final = [resource.name for resource in resources] + resource_names: Final = tuple(resource.name for resource in resources) verbose_logger.info( "MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names ) @@ -859,7 +861,7 @@ class MCPClient: try: resource_templates: Final = await self.run_with_session(_list_resource_templates_operation) resource_template_count: Final = len(resource_templates) - resource_template_names: Final = [resource_template.name for resource_template in resource_templates] + resource_template_names: Final = tuple(resource_template.name for resource_template in resource_templates) verbose_logger.info( "MCP client listed %s resource templates from %s: %s", resource_template_count, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3ca6b6c5f90..beee7c1db78 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,6 @@ import asyncio import importlib -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -1402,7 +1402,7 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True) + list_tools_result: Final[Sequence[MCPTool]] = await client.list_tools(raise_on_error=True) model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py index f588fdd1eee..a76f410ac3c 100644 --- a/tests/test_litellm/experimental_mcp_client/test_pagination.py +++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py @@ -38,7 +38,9 @@ async def test_collect_pages_follows_next_cursor_until_exhausted(): tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") assert [t.name for t in tools] == [f"tool_{i:02d}" for i in range(72)] - assert upstream.cursors_seen == [None, "30", "60"], "each page must be requested with the cursor the previous one returned" + assert upstream.cursors_seen == [None, "30", "60"], ( + "each page must be requested with the cursor the previous one returned" + ) @pytest.mark.asyncio diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3d2e97d55a5..b0c3cc7f9fd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22367 + "limit": 22366 }, "LIT002": { "limit": 26777 From d32f8a07c88ed165e55ce944026504a1cd15a327 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:15:10 +0000 Subject: [PATCH 04/56] fix(mcp): make list page cap a plain constant and use a real ListToolsResult in the unit mock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 +- litellm/experimental_mcp_client/pagination.py | 2 +- tests/mcp_tests/test_mcp_client_unit.py | 6 ++---- .../test_litellm/experimental_mcp_client/test_pagination.py | 2 +- 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 11f35177636..07914934495 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -136,7 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) -MCP_LIST_MAX_PAGES: Final = int(os.getenv("LITELLM_MCP_LIST_MAX_PAGES", "100")) +MCP_LIST_MAX_PAGES: Final = 100 # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. diff --git a/litellm/experimental_mcp_client/pagination.py b/litellm/experimental_mcp_client/pagination.py index 8852715aba5..85f46268f92 100644 --- a/litellm/experimental_mcp_client/pagination.py +++ b/litellm/experimental_mcp_client/pagination.py @@ -42,7 +42,7 @@ async def collect_pages( return items if pages_read >= MCP_LIST_MAX_PAGES: verbose_logger.warning( - "MCP %s from %s still paginating after %s pages (LITELLM_MCP_LIST_MAX_PAGES); returning what was read", + "MCP %s from %s still paginating after %s pages (MCP_LIST_MAX_PAGES); returning what was read", method, server, pages_read, diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index aadaadd510e..ef4231fe1d9 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport -from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult +from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult, ListToolsResult def test_mcp_client_uses_configurable_default_timeout(): @@ -174,9 +174,7 @@ class TestMCPClientUnitTests: }, ) ] - mock_result = MagicMock() - mock_result.tools = mock_tools - mock_session_instance.list_tools.return_value = mock_result + mock_session_instance.list_tools.return_value = ListToolsResult(tools=mock_tools) client = MCPClient("http://example.com") result = await client.list_tools() diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py index a76f410ac3c..93952c176a8 100644 --- a/tests/test_litellm/experimental_mcp_client/test_pagination.py +++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py @@ -79,4 +79,4 @@ async def test_collect_pages_honors_the_page_cap(monkeypatch, caplog): assert len(upstream.cursors_seen) == 3 assert len(tools) == 30 - assert any("LITELLM_MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records) + assert any("MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records) From f2b6c0da81ca247a3b6e7e52c85c51d313310a78 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 16:41:12 -0700 Subject: [PATCH 05/56] feat(bedrock_mantle): serve /v1/messages for Claude models on Mantle's native Anthropic Messages API --- .../messages/handler.py | 3 +- .../llms/bedrock_mantle/messages/__init__.py | 0 .../bedrock_mantle/messages/transformation.py | 101 +++++ litellm/utils.py | 7 + ..._bedrock_mantle_messages_transformation.py | 346 ++++++++++++++++++ tests/test_litellm/test_utils.py | 22 ++ 6 files changed, 478 insertions(+), 1 deletion(-) create mode 100644 litellm/llms/bedrock_mantle/messages/__init__.py create mode 100644 litellm/llms/bedrock_mantle/messages/transformation.py create mode 100644 tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 87a4801f987..e1309ea4063 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -501,6 +501,7 @@ def anthropic_messages_handler( api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) + resolved_api_base: Final = dynamic_api_base if dynamic_api_base is not None else api_base # Store agentic loop params in logging object for agentic hooks # This provides original request context needed for follow-up calls @@ -662,7 +663,7 @@ def anthropic_messages_handler( litellm_params=litellm_params, logging_obj=litellm_logging_obj, api_key=api_key, - api_base=api_base, + api_base=resolved_api_base, stream=stream, kwargs=kwargs, ) diff --git a/litellm/llms/bedrock_mantle/messages/__init__.py b/litellm/llms/bedrock_mantle/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py new file mode 100644 index 00000000000..a4365cfa49b --- /dev/null +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -0,0 +1,101 @@ +from collections.abc import Mapping +from typing import Final + +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + DEFAULT_ANTHROPIC_API_VERSION, +) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH +from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig +from litellm.llms.bedrock_mantle.common_utils import ( + MANTLE_HOST_RE, + BedrockMantleAuthMixin, + resolve_mantle_region, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams + +_BASE_SUFFIXES_TO_STRIP: Final = ( + MANTLE_MESSAGES_PATH, + "/v1/messages", + "/messages", + "/anthropic/v1", + "/openai/v1", + "/v1", +) + + +def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str: + region: Final = resolve_mantle_region({**litellm_params, "api_base": api_base}) + configured: Final = ( + api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") or f"https://bedrock-mantle.{region}.api.aws" + ).rstrip("/") + stripped: Final = next( + (configured[: -len(suffix)] for suffix in _BASE_SUFFIXES_TO_STRIP if configured.endswith(suffix)), + configured, + ) + host: Final = f"https://bedrock-mantle.{region}.api.aws" if MANTLE_HOST_RE.match(stripped) else stripped + return f"{host}{MANTLE_MESSAGES_PATH}" + + +class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig): + def __init__(self, aws_signer: BaseAWSLLM | None = None) -> None: + AmazonMantleMessagesConfig.__init__(self) + self._aws_signer = aws_signer or self + + @property + def custom_llm_provider(self) -> str | None: + return "bedrock_mantle" + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + stream: bool | None = None, + ) -> str: + return build_mantle_native_messages_url(api_base=api_base, litellm_params=litellm_params) + + def validate_anthropic_messages_environment( + self, + headers: dict, + model: str, + messages: list[dict], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict, str | None]: + merged_headers, resolved_api_base = super().validate_anthropic_messages_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + if any(name.lower() == "anthropic-version" for name in merged_headers): + return merged_headers, resolved_api_base + return {**merged_headers, "anthropic-version": DEFAULT_ANTHROPIC_API_VERSION}, resolved_api_base + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict], + anthropic_messages_optional_request_params: dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> dict: + request: Final = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if "anthropic_version" in anthropic_messages_optional_request_params: + return request + return {key: value for key, value in request.items() if key != "anthropic_version"} diff --git a/litellm/utils.py b/litellm/utils.py index b724313641f..3439a21b560 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8681,6 +8681,13 @@ class ProviderConfigManager: from litellm.llms.bedrock.common_utils import BedrockModelInfo return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model) + elif litellm.LlmProviders.BEDROCK_MANTLE == provider: + if "claude" in model_lower: + from litellm.llms.bedrock_mantle.messages.transformation import ( + BedrockMantleAnthropicMessagesConfig, + ) + + return BedrockMantleAnthropicMessagesConfig() elif litellm.LlmProviders.VERTEX_AI == provider: if "claude" in model_lower: from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py new file mode 100644 index 00000000000..2961eee925c --- /dev/null +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -0,0 +1,346 @@ +""" +Unit tests for the bedrock_mantle native Anthropic Messages route. + +Mantle serves its Claude models only on `/anthropic/v1/messages` (the OpenAI +paths reject them), so `bedrock_mantle/anthropic.claude-*` requests on +/v1/messages must hit that endpoint directly instead of the chat-completions +bridge. These tests lock the dispatcher gate, the URL derivation from the +OpenAI-surface base that get_llm_provider pre-fills, the version header, the +Bearer/SigV4 auth chain, and the wire request through the public entrypoint. +""" + +import json +from unittest.mock import MagicMock + +import httpx +import pytest +import respx + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock_mantle.messages.transformation import ( + BedrockMantleAnthropicMessagesConfig, + build_mantle_native_messages_url, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager + +MESSAGES_PATH = "/anthropic/v1/messages" + + +@pytest.fixture(autouse=True) +def _httpx_transport_with_fresh_clients(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + + +@pytest.fixture(autouse=True) +def _no_ambient_mantle_env(monkeypatch): + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + + +def _anthropic_response() -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + ) + + +_SSE_EVENTS = ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_stream", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-sonnet-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + }, + ), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}), + ("message_stop", {"type": "message_stop"}), +) + + +def _sse_response() -> httpx.Response: + body = "".join(f"event: {event}\ndata: {json.dumps(payload)}\n\n" for event, payload in _SSE_EVENTS).encode() + return httpx.Response(status_code=200, content=body, headers={"content-type": "text/event-stream"}) + + +def _mantle_messages_route(region: str) -> respx.Route: + return respx.post(f"https://bedrock-mantle.{region}.api.aws{MESSAGES_PATH}") + + +def _sent_body(route: respx.Route) -> dict: + return json.loads(route.calls.last.request.content) + + +class TestDispatch: + def test_claude_models_get_the_native_messages_config(self): + config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="anthropic.claude-sonnet-5", provider=litellm.LlmProviders.BEDROCK_MANTLE + ) + assert isinstance(config, BedrockMantleAnthropicMessagesConfig) + assert config.custom_llm_provider == "bedrock_mantle" + + @pytest.mark.parametrize("model", ["openai.gpt-5.6-sol", "openai.gpt-oss-120b-1:0", "google.gemma-4-31b"]) + def test_non_claude_models_keep_the_bridge(self, model): + assert ( + ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=litellm.LlmProviders.BEDROCK_MANTLE + ) + is None + ) + + +class TestURL: + @pytest.mark.parametrize( + "api_base", + [ + "https://bedrock-mantle.us-east-1.api.aws/v1", + "https://bedrock-mantle.us-east-1.api.aws/openai/v1", + "https://bedrock-mantle.us-east-1.api.aws/openai/v1/", + "https://bedrock-mantle.us-east-1.api.aws", + "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages", + ], + ) + def test_prefilled_openai_base_becomes_the_messages_endpoint(self, api_base): + url = build_mantle_native_messages_url(api_base, {"aws_region_name": "us-east-1"}) + assert url == f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}" + + def test_aws_region_name_wins_over_the_prefilled_host_region(self): + url = build_mantle_native_messages_url( + "https://bedrock-mantle.us-east-1.api.aws/v1", {"aws_region_name": "us-east-2"} + ) + assert url == f"https://bedrock-mantle.us-east-2.api.aws{MESSAGES_PATH}" + + def test_host_region_is_used_when_no_region_param(self): + url = build_mantle_native_messages_url("https://bedrock-mantle.eu-west-1.api.aws/v1", {}) + assert url == f"https://bedrock-mantle.eu-west-1.api.aws{MESSAGES_PATH}" + + def test_custom_host_is_preserved(self): + url = build_mantle_native_messages_url("https://vpce-abc.bedrock-mantle.example.com/v1", {}) + assert url == f"https://vpce-abc.bedrock-mantle.example.com{MESSAGES_PATH}" + + def test_env_base_is_used_without_api_base(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://mantle-proxy.internal/openai/v1") + assert build_mantle_native_messages_url(None, {}) == f"https://mantle-proxy.internal{MESSAGES_PATH}" + + def test_default_host_comes_from_mantle_region_env(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "ap-northeast-1") + assert build_mantle_native_messages_url(None, {}) == f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}" + + def test_config_get_complete_url_reads_litellm_params(self): + config = BedrockMantleAnthropicMessagesConfig() + url = config.get_complete_url( + api_base="https://bedrock-mantle.us-east-1.api.aws/v1", + api_key=None, + model="anthropic.claude-sonnet-5", + optional_params={}, + litellm_params={"aws_region_name": "us-west-2"}, + ) + assert url == f"https://bedrock-mantle.us-west-2.api.aws{MESSAGES_PATH}" + + +class TestEnvironment: + def _validate(self, headers: dict, litellm_params: dict) -> dict: + config = BedrockMantleAnthropicMessagesConfig() + merged, _ = config.validate_anthropic_messages_environment( + headers=headers, + model="anthropic.claude-sonnet-5", + messages=[], + optional_params={}, + litellm_params=litellm_params, + ) + return merged + + def test_adds_the_anthropic_version_header(self): + assert self._validate({}, {})["anthropic-version"] == "2023-06-01" + + def test_keeps_a_caller_supplied_version_header(self): + merged = self._validate({"Anthropic-Version": "2024-01-01"}, {}) + assert merged["Anthropic-Version"] == "2024-01-01" + assert "anthropic-version" not in merged + + def test_project_id_becomes_the_workspace_header(self): + assert self._validate({}, {"aws_bedrock_project_id": "proj_123"})["anthropic-workspace"] == "proj_123" + + +class TestRequestBody: + def test_body_carries_model_and_stream_but_not_the_invoke_version(self): + config = BedrockMantleAnthropicMessagesConfig() + body = config.transform_anthropic_messages_request( + model="anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + anthropic_messages_optional_request_params={"max_tokens": 8, "stream": True}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["model"] == "anthropic.claude-sonnet-5" + assert body["stream"] is True + assert body["max_tokens"] == 8 + assert "anthropic_version" not in body + + def test_body_omits_stream_when_not_streaming(self): + config = BedrockMantleAnthropicMessagesConfig() + body = config.transform_anthropic_messages_request( + model="anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + anthropic_messages_optional_request_params={"max_tokens": 8}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert "stream" not in body + + +class TestAuth: + def test_bearer_from_api_key_skips_aws_credentials(self): + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=AssertionError("must not resolve AWS credentials")) + config = BedrockMantleAnthropicMessagesConfig(aws_signer=signer) + headers, signed = config.sign_request( + headers={"anthropic-version": "2023-06-01"}, + optional_params={}, + request_data={"model": "anthropic.claude-sonnet-5"}, + api_base=f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}", + api_key="arg-bearer", + ) + assert headers["Authorization"] == "Bearer arg-bearer" + assert headers["anthropic-version"] == "2023-06-01" + assert signed == b'{"model": "anthropic.claude-sonnet-5"}' + + def test_bearer_from_mantle_env_key(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + config = BedrockMantleAnthropicMessagesConfig() + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={}, + api_base=f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}", + api_key=None, + ) + assert headers["Authorization"] == "Bearer env-bearer" + + def test_sigv4_scope_is_pinned_to_the_url_host_region(self): + config = BedrockMantleAnthropicMessagesConfig() + headers, signed = config.sign_request( + headers={"anthropic-version": "2023-06-01"}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-1", + }, + request_data={"model": "anthropic.claude-sonnet-5"}, + api_base=f"https://bedrock-mantle.us-west-2.api.aws{MESSAGES_PATH}", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "/us-west-2/bedrock/aws4_request" in headers["Authorization"] + assert signed == b'{"model": "anthropic.claude-sonnet-5"}' + + +class TestWireRequest: + @pytest.mark.asyncio + @respx.mock + async def test_claude_request_hits_the_native_messages_endpoint(self): + route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response()) + + response = await litellm.anthropic_messages( + model="bedrock_mantle/anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + max_tokens=8, + api_key="test-bearer", + aws_region_name="us-east-1", + ) + + assert response["content"][0]["text"] == "pong" + assert route.call_count == 1 + sent = route.calls.last.request + assert sent.headers["authorization"] == "Bearer test-bearer" + assert sent.headers["anthropic-version"] == "2023-06-01" + assert "x-api-key" not in sent.headers + body = _sent_body(route) + assert body["model"] == "anthropic.claude-sonnet-5" + assert body["messages"] == [{"role": "user", "content": "ping"}] + assert "anthropic_version" not in body + assert "stream" not in body + + @pytest.mark.asyncio + @respx.mock + async def test_region_prefix_selects_the_host_and_is_not_sent_as_model(self): + route = _mantle_messages_route("us-east-2").mock(return_value=_anthropic_response()) + + await litellm.anthropic_messages( + model="bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "ping"}], + max_tokens=8, + api_key="test-bearer", + ) + + assert route.call_count == 1 + assert _sent_body(route)["model"] == "anthropic.claude-haiku-4-5" + + @pytest.mark.asyncio + @respx.mock + async def test_streaming_sends_stream_and_passes_the_sse_through(self): + route = _mantle_messages_route("us-east-1").mock(return_value=_sse_response()) + + response = await litellm.anthropic_messages( + model="bedrock_mantle/anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + max_tokens=8, + stream=True, + api_key="test-bearer", + aws_region_name="us-east-1", + ) + raw = b"".join([chunk async for chunk in response]) + + assert route.call_count == 1 + assert _sent_body(route)["stream"] is True + text = raw.decode() + assert "event: message_start" in text + assert '"text": "pong"' in text + assert "event: message_stop" in text + + @pytest.mark.asyncio + @respx.mock + async def test_sigv4_request_signs_against_the_messages_url(self): + route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response()) + + await litellm.anthropic_messages( + model="bedrock_mantle/anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + max_tokens=8, + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + aws_region_name="us-east-1", + ) + + assert route.call_count == 1 + authorization = route.calls.last.request.headers["authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "/us-east-1/bedrock/aws4_request" in authorization diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3336ad6d33a..a4c06122189 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -3640,6 +3640,28 @@ class TestGetOptionalParamsTencent: assert isinstance(config, TencentAnthropicMessagesConfig) assert config.custom_llm_provider == "tencent" + def test_bedrock_mantle_claude_messages_config_routing(self): + import litellm + from litellm.llms.bedrock_mantle.messages.transformation import ( + BedrockMantleAnthropicMessagesConfig, + ) + + config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="anthropic.claude-sonnet-5", + provider=litellm.LlmProviders.BEDROCK_MANTLE, + ) + assert isinstance(config, BedrockMantleAnthropicMessagesConfig) + assert config.custom_llm_provider == "bedrock_mantle" + + def test_bedrock_mantle_openai_models_keep_the_messages_bridge(self): + import litellm + + config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="openai.gpt-5.6-sol", + provider=litellm.LlmProviders.BEDROCK_MANTLE, + ) + assert config is None + class TestValidateEnvironmentTencent: """Tests that validate_environment resolves TENCENT_API_KEY for the tencent provider.""" From 368a8396400bdf5f986f8379d84ac44c43b808c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:08:41 -0700 Subject: [PATCH 06/56] fix(bedrock_mantle): send anthropic betas in the header Mantle reads on /v1/messages --- litellm/anthropic_beta_headers_config.json | 35 +++++ .../anthropic_claude3_transformation.py | 22 ++-- .../bedrock_mantle/messages/transformation.py | 36 +++-- ..._bedrock_mantle_messages_transformation.py | 123 +++++++++++++++++- 4 files changed, 196 insertions(+), 20 deletions(-) diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index eb31cc17a15..1331de4c266 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -131,6 +131,41 @@ "web-fetch-2025-09-10": null, "web-search-2025-03-05": null }, + "bedrock_mantle": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "advisor-tool-2026-03-01": null, + "bash_20241022": null, + "bash_20250124": null, + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-08-25": null, + "compact-2026-01-12": "compact-2026-01-12", + "computer-use-2025-01-24": "computer-use-2025-01-24", + "computer-use-2025-11-24": "computer-use-2025-11-24", + "context-1m-2025-08-07": "context-1m-2025-08-07", + "context-management-2025-06-27": "context-management-2025-06-27", + "effort-2025-11-24": "effort-2025-11-24", + "fast-mode-2026-02-01": null, + "files-api-2025-04-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", + "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", + "mcp-client-2025-04-04": null, + "mcp-client-2025-11-20": null, + "mcp-servers-2025-12-04": null, + "output-128k-2025-02-19": "output-128k-2025-02-19", + "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", + "prompt-caching-scope-2026-01-05": null, + "skills-2025-10-02": null, + "structured-output-2024-03-01": null, + "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", + "text_editor_20241022": null, + "text_editor_20250124": null, + "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", + "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "web-fetch-2025-09-10": null, + "web-search-2025-03-05": "web-search-2025-03-05" + }, "vertex_ai": { "advisor-tool-2026-03-01": null, "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index d2be1ad9156..4b52a3bafe6 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -1,4 +1,4 @@ -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast @@ -445,13 +445,16 @@ class AmazonAnthropicClaudeMessagesConfig( # Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the # ``context-management-2025-06-27`` beta. AWS docs: # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md - _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: dict[str, str] = { - "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, - "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, - } + _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType( + { + "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + } + ) - @staticmethod + @classmethod def _filter_context_management_for_bedrock_invoke( + cls, anthropic_messages_request: dict, beta_set: set, ) -> None: @@ -481,7 +484,7 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request.pop("context_management", None) return - supported: Final = AmazonAnthropicClaudeMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS + supported: Final = cls._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS retained_edits: Final = [e for e in edits if isinstance(e, dict) and e.get("type") in supported] if not retained_edits: anthropic_messages_request.pop("context_management", None) @@ -546,15 +549,16 @@ class AmazonAnthropicClaudeMessagesConfig( if "tool-search-tool-2025-10-19" in beta_set: beta_set.add("tool-examples-2025-10-29") + beta_provider: Final = self.custom_llm_provider or "bedrock" filtered_betas: Final = sorted( filter_and_transform_beta_headers( beta_headers=list(beta_set), - provider="bedrock", + provider=beta_provider, ) ) dropped_user_betas: Final = sorted( - b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider="bedrock") + b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider=beta_provider) ) if dropped_user_betas: verbose_logger.warning( diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py index a4365cfa49b..480c09a0476 100644 --- a/litellm/llms/bedrock_mantle/messages/transformation.py +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -1,6 +1,9 @@ from collections.abc import Mapping +from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( DEFAULT_ANTHROPIC_API_VERSION, ) @@ -13,6 +16,7 @@ from litellm.llms.bedrock_mantle.common_utils import ( resolve_mantle_region, ) from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.router import GenericLiteLLMParams _BASE_SUFFIXES_TO_STRIP: Final = ( @@ -23,6 +27,9 @@ _BASE_SUFFIXES_TO_STRIP: Final = ( "/openai/v1", "/v1", ) +_BODY_FIELDS_MANTLE_READS_FROM_HEADERS: Final = frozenset({"anthropic_version", "anthropic_beta"}) +_ANTHROPIC_BETAS: Final = TypeAdapter(tuple[str, ...]) +_MANTLE_REQUEST: Final = TypeAdapter(dict[str, object]) def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str: @@ -39,6 +46,13 @@ def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mappi class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig): + _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType( + { + **AmazonMantleMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS, + "clear_thinking_20251015": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + } + ) + def __init__(self, aws_signer: BaseAWSLLM | None = None) -> None: AmazonMantleMessagesConfig.__init__(self) self._aws_signer = aws_signer or self @@ -89,13 +103,17 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: - request: Final = super().transform_anthropic_messages_request( - model=model, - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params=litellm_params, - headers=headers, + request: Final = _MANTLE_REQUEST.validate_python( + super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ), ) - if "anthropic_version" in anthropic_messages_optional_request_params: - return request - return {key: value for key, value in request.items() if key != "anthropic_version"} + betas: Final = request.get("anthropic_beta") + if betas is not None: + header_betas: Final = ",".join(_ANTHROPIC_BETAS.validate_python(betas)) + headers["anthropic-beta"] = header_betas # rebind-ok: the handler signs and sends this same dict + return {key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS} diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py index 2961eee925c..3544262996c 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -79,7 +79,10 @@ _SSE_EVENTS = ( }, ), ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}), ("message_stop", {"type": "message_stop"}), @@ -152,7 +155,10 @@ class TestURL: def test_default_host_comes_from_mantle_region_env(self, monkeypatch): monkeypatch.setenv("BEDROCK_MANTLE_REGION", "ap-northeast-1") - assert build_mantle_native_messages_url(None, {}) == f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}" + assert ( + build_mantle_native_messages_url(None, {}) + == f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}" + ) def test_config_get_complete_url_reads_litellm_params(self): config = BedrockMantleAnthropicMessagesConfig() @@ -344,3 +350,116 @@ class TestWireRequest: authorization = route.calls.last.request.headers["authorization"] assert authorization.startswith("AWS4-HMAC-SHA256") assert "/us-east-1/bedrock/aws4_request" in authorization + + +def _sent_betas(route: respx.Route) -> list[str]: + return route.calls.last.request.headers["anthropic-beta"].split(",") + + +@pytest.mark.usefixtures("local_beta_headers_config") +class TestBetaHeadersOnTheWire: + async def _send(self, **request_params) -> respx.Route: + route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response()) + await litellm.anthropic_messages( + model="bedrock_mantle/anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": "ping"}], + max_tokens=8, + api_key="test-bearer", + aws_region_name="us-east-1", + **request_params, + ) + return route + + @pytest.mark.asyncio + @respx.mock + async def test_betas_mantle_accepts_reach_it_in_the_header(self): + route = await self._send( + extra_headers={ + "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14,context-management-2025-06-27" + } + ) + + assert _sent_betas(route) == [ + "claude-code-20250219", + "context-management-2025-06-27", + "interleaved-thinking-2025-05-14", + ] + + @pytest.mark.asyncio + @respx.mock + async def test_betas_mantle_rejects_are_dropped_before_the_request(self): + route = await self._send( + extra_headers={"anthropic-beta": "code-execution-2025-08-25,context-1m-2025-08-07,files-api-2025-04-14"} + ) + + assert _sent_betas(route) == ["context-1m-2025-08-07"] + + @pytest.mark.asyncio + @respx.mock + async def test_no_beta_header_is_sent_when_every_value_is_rejected(self): + route = await self._send(extra_headers={"anthropic-beta": "code-execution-2025-08-25"}) + + assert "anthropic-beta" not in route.calls.last.request.headers + + @pytest.mark.asyncio + @respx.mock + async def test_advanced_tool_use_is_renamed_to_the_beta_mantle_knows(self): + route = await self._send(extra_headers={"anthropic-beta": "advanced-tool-use-2025-11-20"}) + + assert "tool-search-tool-2025-10-19" in _sent_betas(route) + assert "advanced-tool-use-2025-11-20" not in _sent_betas(route) + + @pytest.mark.asyncio + @respx.mock + async def test_a_feature_beta_joins_the_callers_betas_in_the_header(self): + route = await self._send( + extra_headers={"anthropic-beta": "context-1m-2025-08-07"}, + context_management={"edits": [{"type": "clear_tool_uses_20250919"}]}, + ) + + assert _sent_betas(route) == ["context-1m-2025-08-07", "context-management-2025-06-27"] + assert _sent_body(route)["context_management"] == {"edits": [{"type": "clear_tool_uses_20250919"}]} + + @pytest.mark.asyncio + @respx.mock + async def test_betas_and_version_never_travel_in_the_body(self): + route = await self._send( + extra_headers={"anthropic-beta": "context-1m-2025-08-07"}, + context_management={"edits": [{"type": "clear_tool_uses_20250919"}]}, + anthropic_version="bedrock-2023-05-31", + ) + + body = _sent_body(route) + assert "anthropic_beta" not in body + assert "anthropic_version" not in body + assert route.calls.last.request.headers["anthropic-version"] == "2023-06-01" + + @pytest.mark.asyncio + @respx.mock + async def test_clear_thinking_edit_is_forwarded_with_thinking_on(self): + edits = [{"type": "clear_thinking_20251015", "keep": "all"}, {"type": "clear_tool_uses_20250919"}] + route = await self._send( + context_management={"edits": edits}, + thinking={"type": "adaptive"}, + ) + + body = _sent_body(route) + assert body["context_management"] == {"edits": edits} + assert body["thinking"] == {"type": "adaptive"} + assert "context-management-2025-06-27" in _sent_betas(route) + + @pytest.mark.asyncio + @respx.mock + async def test_tools_reach_mantle_unchanged(self): + tools = [ + { + "name": "get_weather", + "description": "Look up the weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ] + route = await self._send(tools=tools, tool_choice={"type": "auto"}) + + body = _sent_body(route) + assert body["tools"] == tools + assert body["tool_choice"] == {"type": "auto"} From 94b2fd827b4ea2678fa88ac0788e948f3b899348 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 19 Sep 2026 17:02:56 -0700 Subject: [PATCH 07/56] feat(ui): show prompt caching requests and net savings --- backend/routes/allowlist.py | 1 + .../prompt_caching_requests.py | 184 ++++++++++ litellm/proxy/proxy_server.py | 4 + litellm/proxy/spend_tracking/savings.py | 75 ++-- .../prompt_caching_requests.py | 35 ++ .../test_prompt_caching_requests.py | 321 ++++++++++++++++++ .../proxy/spend_tracking/test_savings.py | 37 ++ .../_components/CacheLeakageCard.tsx | 6 +- .../CostOptimizationView.activity.test.tsx | 1 + ...tCachingRequestsTable.integration.test.tsx | 248 ++++++++++++++ .../PromptCachingRequestsTable.tsx | 186 ++++++++++ .../_components/PromptCachingTab.test.tsx | 23 +- .../_components/PromptCachingTab.tsx | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 95 ++++++ 14 files changed, 1195 insertions(+), 28 deletions(-) create mode 100644 litellm/proxy/management_endpoints/prompt_caching_requests.py create mode 100644 litellm/types/management_endpoints/prompt_caching_requests.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 00c4e0070e6..c7f389c36a4 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/cache_settings", "/coordination_redis/", "/cost_tracking", + "/cost_optimization/", "/cost/", "/credentials", "/credential", diff --git a/litellm/proxy/management_endpoints/prompt_caching_requests.py b/litellm/proxy/management_endpoints/prompt_caching_requests.py new file mode 100644 index 00000000000..41255bd49b8 --- /dev/null +++ b/litellm/proxy/management_endpoints/prompt_caching_requests.py @@ -0,0 +1,184 @@ +from collections.abc import Callable, Mapping +from datetime import datetime, timezone +from types import MappingProxyType +from typing import TYPE_CHECKING, Annotated, Final + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel, Json, TypeAdapter + +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.spend_tracking.savings import ( + extract_cache_creation_tokens, + extract_cache_read_tokens, + marks_gateway_injection, + prompt_caching_savings_for_request, +) +from litellm.proxy.spend_tracking.spend_tracking_utils import ( + _query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below +) +from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY +from litellm.types.management_endpoints.prompt_caching_requests import ( + PromptCachingRequest, + PromptCachingRequestCursor, + PromptCachingRequestFilter, + PromptCachingRequestsResponse, +) + +if TYPE_CHECKING: + from litellm.router import Router + +router: Final = APIRouter() + + +def _numeric_token_sql(path: str) -> str: + value: Final = f"metadata #> '{{usage_object,{path}}}'" + return ( + f"CASE WHEN jsonb_typeof({value}) = 'number' THEN ({value} #>> '{{}}')::numeric " + f"WHEN {value} = 'true'::jsonb THEN 1 WHEN {value} = 'false'::jsonb THEN 0 END" + ) + + +def _cache_tokens_sql(*paths: str) -> str: + candidates: Final = ", ".join(f"NULLIF(({_numeric_token_sql(path)}), 0)" for path in paths) + return f"TRUNC(COALESCE({candidates}, 0))" + + +_CACHE_READ_SQL: Final = _cache_tokens_sql("cache_read_input_tokens", "prompt_tokens_details,cached_tokens") +_CACHE_CREATION_SQL: Final = _cache_tokens_sql( + "cache_creation_input_tokens", + "prompt_tokens_details,cache_write_tokens", + "prompt_tokens_details,cache_creation_tokens", +) +_GATEWAY_INJECTED_SQL: Final = ( + f"(jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' " + f"AND (metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = '' " + f"OR metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = model_id))" +) +_FILTER_SQL: Final = MappingProxyType( + { + "all": f"({_GATEWAY_INJECTED_SQL} OR {_CACHE_READ_SQL} > 0 OR {_CACHE_CREATION_SQL} > 0)", + "injected": _GATEWAY_INJECTED_SQL, + "hits": f"{_CACHE_READ_SQL} > 0", + } +) + + +def prompt_caching_requests_sql(filter: PromptCachingRequestFilter) -> str: + return f""" + SELECT request_id, "startTime" AS start_time, "endTime" AS end_time, + model, model_id, custom_llm_provider, spend, + CASE WHEN jsonb_typeof(metadata->'usage_object') = 'object' + THEN metadata->'usage_object' END AS usage_object, + CASE WHEN jsonb_typeof(metadata->'cost_breakdown') = 'object' + THEN metadata->'cost_breakdown' END AS cost_breakdown, + CASE WHEN jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' + THEN metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' END AS gateway_marker + FROM "LiteLLM_SpendLogs" + WHERE "startTime" >= ($1::text::timestamptz AT TIME ZONE 'UTC') + AND "startTime" <= ($2::text::timestamptz AT TIME ZONE 'UTC') + AND COALESCE(LOWER(cache_hit), 'false') != 'true' + AND {_FILTER_SQL[filter]} + AND ($4::text::timestamptz IS NULL OR + ("startTime", request_id) < (($4::text::timestamptz AT TIME ZONE 'UTC'), $5::text)) + ORDER BY "startTime" DESC, request_id DESC + LIMIT $3::integer + """ + + +class _PromptCachingRow(BaseModel): + request_id: str + start_time: datetime + end_time: datetime + model: str + model_id: str | None + custom_llm_provider: str | None + spend: float + usage_object: Json[Mapping[str, object]] | Mapping[str, object] | None + cost_breakdown: Json[Mapping[str, object]] | Mapping[str, object] | None + gateway_marker: str | None + + +_REQUEST_ROWS: Final = TypeAdapter(tuple[_PromptCachingRow, ...]) + + +def _request_result(row: _PromptCachingRow, llm_router: "Callable[[], Router | None]") -> PromptCachingRequest: + return PromptCachingRequest( + request_id=row.request_id, + start_time=row.start_time.replace(tzinfo=timezone.utc) if row.start_time.tzinfo is None else row.start_time, + model=row.model, + gateway_injected=marks_gateway_injection( + MappingProxyType({GATEWAY_INJECTED_CACHE_METADATA_KEY: row.gateway_marker}), row.model_id + ), + cache_read_tokens=extract_cache_read_tokens(row.usage_object), + cache_creation_tokens=extract_cache_creation_tokens(row.usage_object), + spend=row.spend, + net_savings=prompt_caching_savings_for_request( + model=row.model, + custom_llm_provider=row.custom_llm_provider, + usage_object=row.usage_object, + model_id=row.model_id, + llm_router=llm_router, + cost_breakdown=row.cost_breakdown, + billed_at=row.end_time, + ), + ) + + +@router.get( + "/cost_optimization/prompt_caching/requests", + tags=["Cost Optimization"], # mutable-ok: FastAPI's route API requires a list + response_model=PromptCachingRequestsResponse, +) +async def get_prompt_caching_requests( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_date: datetime, + end_date: datetime, + page_size: Annotated[int, Query(ge=1, le=100)] = 50, + filter: PromptCachingRequestFilter = "all", + cursor_start_time: datetime | None = None, + cursor_request_id: Annotated[str | None, Query(min_length=1)] = None, +) -> PromptCachingRequestsResponse: + from litellm.proxy.proxy_server import llm_router, prisma_client + + if not user_api_key_has_admin_view(user_api_key_dict): + raise HTTPException(status_code=403, detail="Only proxy admin roles can view prompt caching requests") + if (cursor_start_time is None) != (cursor_request_id is None): + raise HTTPException(status_code=400, detail="cursor_start_time and cursor_request_id must be provided together") + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + start: Final = start_date.replace(tzinfo=timezone.utc) if start_date.tzinfo is None else start_date + end: Final = end_date.replace(tzinfo=timezone.utc) if end_date.tzinfo is None else end_date + if end < start: + raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date") + cursor_time: Final = ( + cursor_start_time.replace(tzinfo=timezone.utc) + if cursor_start_time is not None and cursor_start_time.tzinfo is None + else cursor_start_time + ) + rows: Final = _REQUEST_ROWS.validate_python( + await _query_raw_rows( + prisma_client, + prompt_caching_requests_sql(filter), + start.isoformat(), + end.isoformat(), + page_size + 1, + cursor_time.isoformat() if cursor_time is not None else None, + cursor_request_id, + ) + or () + ) + + def current_router() -> "Router | None": + return llm_router + + requests: Final = tuple(_request_result(row, current_router) for row in rows[:page_size]) + has_more: Final = len(rows) > page_size + return PromptCachingRequestsResponse( + requests=requests, + page_size=page_size, + has_more=has_more, + next_cursor=PromptCachingRequestCursor(start_time=requests[-1].start_time, request_id=requests[-1].request_id) + if has_more + else None, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af25d418a63..f4a56e225cc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -587,6 +587,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) +from litellm.proxy.management_endpoints.prompt_caching_requests import ( + router as prompt_caching_requests_router, +) from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, ) @@ -19183,6 +19186,7 @@ app.include_router(workflow_management_router) app.include_router(memory_router) app.include_router(plugin_router) app.include_router(cost_tracking_settings_router) +app.include_router(prompt_caching_requests_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index b7a2ac62844..fbcf9c78d3e 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -578,6 +578,56 @@ def autorouter_savings_for_logging_payload( ) +def _request_savings_pricing( + model: str | None, + custom_llm_provider: str | None, + model_id: str | None, + llm_router: "Callable[[], Router | None] | None", +) -> tuple[str | None, ModelInfo | None]: + router_instance: Final = llm_router() if llm_router else None + identity: Final = _resolve_model(model, custom_llm_provider) + pricing: Final = _effective_model_info(router_instance, model_id, model or "") or ( + _model_info(identity) if identity else None + ) + return identity.provider if identity else custom_llm_provider, pricing + + +def _prompt_caching_savings( + pricing: ModelInfo | None, + provider: str | None, + usage_object: Mapping[str, object] | None, + cost_breakdown: Mapping[str, object] | None, + billed_at: datetime | str | None, +) -> float | None: + usage: Final = _usage_from_spend_log(usage_object) + if pricing is None or usage is None: + return None + basis: Final = _pricing_basis(cost_breakdown) + result: Final = calculate_prompt_caching_savings( + model_info=pricing, + usage=usage, + custom_llm_provider=provider, + service_tier=basis.service_tier, + data_residency=basis.data_residency, + vertex_location=basis.vertex_location, + billed_at=_coerce_billed_at(billed_at), + ) + return result if isfinite(result) else None + + +def prompt_caching_savings_for_request( + model: str | None, + custom_llm_provider: str | None, + usage_object: Mapping[str, object] | None, + model_id: str | None = None, + llm_router: "Callable[[], Router | None] | None" = None, + cost_breakdown: Mapping[str, object] | None = None, + billed_at: datetime | str | None = None, +) -> float | None: + request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router) + return _prompt_caching_savings(request_pricing[1], request_pricing[0], usage_object, cost_breakdown, billed_at) + + def compute_savings_spend( model: str | None, custom_llm_provider: str | None, @@ -639,29 +689,12 @@ def compute_savings_spend( # Deployment rates when the request came through one, public rates otherwise -- # `_effective_model_info` merges a deployment's configured prices over the built-in # map, so a negotiated price is not silently replaced by the list rate. - router_instance: Router | None = llm_router() if llm_router else None - identity: Final = _resolve_model(model, custom_llm_provider) - pricing: Final = _effective_model_info(router_instance, model_id, model or "") or ( - _model_info(identity) if identity else None - ) + request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router) + provider: Final = request_pricing[0] + pricing: Final = request_pricing[1] input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0 compression: Final = max(compression_saved_tokens, 0) * input_cost - usage: Final = _usage_from_spend_log(usage_object) - basis: Final = _pricing_basis(cost_breakdown) - billed_at_datetime: Final = _coerce_billed_at(billed_at) - prompt_caching: Final = ( - calculate_prompt_caching_savings( - model_info=pricing, - usage=usage, - custom_llm_provider=identity.provider if identity else custom_llm_provider, - service_tier=basis.service_tier, - data_residency=basis.data_residency, - vertex_location=basis.vertex_location, - billed_at=billed_at_datetime, - ) - if pricing is not None and usage is not None - else 0.0 - ) + prompt_caching: Final = _prompt_caching_savings(pricing, provider, usage_object, cost_breakdown, billed_at) or 0.0 gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0 # The figure the logging path recorded wins, before the usage gate on purpose: a row diff --git a/litellm/types/management_endpoints/prompt_caching_requests.py b/litellm/types/management_endpoints/prompt_caching_requests.py new file mode 100644 index 00000000000..e72183a113b --- /dev/null +++ b/litellm/types/management_endpoints/prompt_caching_requests.py @@ -0,0 +1,35 @@ +from datetime import datetime +from typing import Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict + +PromptCachingRequestFilter: TypeAlias = Literal["all", "injected", "hits"] + + +class PromptCachingRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + request_id: str + start_time: datetime + model: str + gateway_injected: bool + cache_read_tokens: int + cache_creation_tokens: int + spend: float + net_savings: float | None + + +class PromptCachingRequestCursor(BaseModel): + model_config = ConfigDict(frozen=True) + + start_time: datetime + request_id: str + + +class PromptCachingRequestsResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + requests: tuple[PromptCachingRequest, ...] + page_size: int + has_more: bool + next_cursor: PromptCachingRequestCursor | None diff --git a/tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py b/tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py new file mode 100644 index 00000000000..0995de6c39d --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py @@ -0,0 +1,321 @@ +import json +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Final + +import httpx +import psycopg +import pytest +import pytest_asyncio +from fastapi import FastAPI +from prisma import Prisma +from pydantic import TypeAdapter +from pytest_postgresql import factories + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.prompt_caching_requests import router +from litellm.proxy.spend_tracking.savings import ( + extract_cache_creation_tokens, + extract_cache_read_tokens, + marks_gateway_injection, +) +from litellm.types.management_endpoints.prompt_caching_requests import ( + PromptCachingRequestFilter, + PromptCachingRequestsResponse, +) + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + +_cache_postgresql_proc: Final = factories.postgresql_proc() # pyright: ignore[reportUnknownMemberType] # third-party fixture factory has incomplete callable types +_cache_postgresql: Final = factories.postgresql("_cache_postgresql_proc") +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_JSON_ROWS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_START: Final = "2026-09-01T00:00:00Z" +_END: Final = "2026-09-02T00:00:00Z" +_URL: Final = "/cost_optimization/prompt_caching/requests" +_MODEL: Final = "claude-sonnet-5" +_MARKER: Final = "litellm_gateway_injected_cache" +_DDL: Final = """ + CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT PRIMARY KEY, "startTime" TIMESTAMP, "endTime" TIMESTAMP, + model TEXT, model_id TEXT, custom_llm_provider TEXT, spend DOUBLE PRECISION, + metadata JSONB, cache_hit TEXT + ) +""" + + +@dataclass(frozen=True) +class _Case: + request_id: str + metadata: Mapping[str, object] + cache_hit: str | None = None + start_time: datetime = datetime(2026, 9, 1, 12, 0, 0, 123456) + + def matches(self, filter: PromptCachingRequestFilter) -> bool: + if self.cache_hit is not None and self.cache_hit.lower() == "true": + return False + if not datetime(2026, 9, 1) <= self.start_time <= datetime(2026, 9, 2): + return False + usage: Final = self.metadata.get("usage_object") + normalized: Final = _JSON_OBJECT.validate_python(usage) if isinstance(usage, Mapping) else None + injected: Final = marks_gateway_injection(self.metadata, "dep-a") + reads: Final = extract_cache_read_tokens(normalized) + writes: Final = extract_cache_creation_tokens(normalized) + match filter: + case "injected": + return injected + case "hits": + return reads > 0 + case "all": + return injected or reads > 0 or writes > 0 + + +_CASES: Final = ( + _Case("injected-empty", {_MARKER: ""}), + _Case("injected-deployment", {_MARKER: "dep-a"}), + _Case("wrong-deployment", {_MARKER: "dep-b"}), + _Case("legacy-read", {"usage_object": {"cache_read_input_tokens": 100}}), + _Case("nested-read", {"usage_object": {"prompt_tokens_details": {"cached_tokens": 100}}}), + _Case("write", {"usage_object": {"cache_creation_input_tokens": 100}}), + _Case("nested-write", {"usage_object": {"prompt_tokens_details": {"cache_write_tokens": 100}}}), + _Case("nested-creation", {"usage_object": {"prompt_tokens_details": {"cache_creation_tokens": 100}}}), + _Case( + "top-precedence", + {"usage_object": {"cache_read_input_tokens": -2, "prompt_tokens_details": {"cached_tokens": 100}}}, + ), + _Case( + "zero-fallback", + {"usage_object": {"cache_read_input_tokens": 0, "prompt_tokens_details": {"cached_tokens": 100}}}, + ), + _Case( + "fractional-precedence", + {"usage_object": {"cache_read_input_tokens": 0.5, "prompt_tokens_details": {"cached_tokens": 100}}}, + ), + _Case("malformed-number", {"usage_object": {"cache_read_input_tokens": "100"}}), + _Case("malformed-container", {"usage_object": [100]}), + _Case("boolean-number", {"usage_object": {"cache_read_input_tokens": True}}), + _Case("boolean-marker", {_MARKER: True}), + _Case("response-cache", {_MARKER: "", "usage_object": {"cache_read_input_tokens": 100}}, "True"), + _Case("outside-before", {_MARKER: ""}, start_time=datetime(2026, 8, 31, 23, 59, 59)), + _Case( + "outside-after", {"usage_object": {"cache_read_input_tokens": 100}}, start_time=datetime(2026, 9, 2, 0, 0, 1) + ), +) + + +@pytest_asyncio.fixture(loop_scope="function") +async def _cache_prisma( + _cache_postgresql: psycopg.Connection[tuple[object, ...]], +) -> AsyncIterator[Prisma]: + info: Final = _cache_postgresql.info + database: Final = Prisma(datasource={ + "url": f"postgresql://{info.user}@{info.host}:{info.port}/{info.dbname}?connection_limit=1", + }) + await database.connect() + try: + yield database + finally: + await database.disconnect() + + +def _seed(connection: psycopg.Connection[tuple[object, ...]], cases: tuple[_Case, ...] = _CASES) -> None: + with connection.cursor() as cursor: + cursor.execute(_DDL) + cursor.executemany( + """INSERT INTO "LiteLLM_SpendLogs" + VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s)""", + tuple( + ( + case.request_id, + case.start_time, + datetime(2026, 9, 1, 12, 0, 1), + _MODEL, + "dep-a", + "anthropic", + 0.01, + json.dumps(dict(case.metadata)), + case.cache_hit, + ) + for case in cases + ), + ) + connection.commit() + + +def _app(role: LitellmUserRoles | None) -> FastAPI: + application: Final = FastAPI() + application.include_router(router) + + def caller() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=role) + + application.dependency_overrides[user_api_key_auth] = caller + return application + + +@pytest.mark.asyncio +@pytest.mark.parametrize("filter", ["all", "injected", "hits"]) +@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) +async def test_request_filters_match_accounting_and_paginate_before_projection( + _cache_postgresql: psycopg.Connection[tuple[object, ...]], + _cache_prisma: Prisma, + monkeypatch: pytest.MonkeyPatch, + filter: PromptCachingRequestFilter, + role: LitellmUserRoles, +) -> None: + from litellm.proxy import proxy_server + + _seed(_cache_postgresql) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=_cache_prisma)) + monkeypatch.setattr(proxy_server, "llm_router", None) + expected: Final = tuple(sorted((case.request_id for case in _CASES if case.matches(filter)), reverse=True)) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(role)), base_url="http://test") as client: + first: Final = await client.get( + _URL, params={"start_date": _START, "end_date": _END, "filter": filter, "page_size": 2} + ) + assert first.status_code == 200 + first_page: Final = PromptCachingRequestsResponse.model_validate_json(first.content) + assert tuple(row.request_id for row in first_page.requests) == expected[:2] + assert first_page.has_more is (len(expected) > 2) + assert (first_page.next_cursor is not None) is first_page.has_more + if first_page.next_cursor is not None: + assert first_page.next_cursor.request_id == expected[1] + assert first_page.next_cursor.start_time == first_page.requests[-1].start_time + next_response: Final = await client.get( + _URL, params={ + "start_date": _START, "end_date": _END, "filter": filter, "page_size": 2, + "cursor_start_time": first_page.next_cursor.start_time.astimezone( + timezone(timedelta(hours=-7)) + ).isoformat(), + "cursor_request_id": first_page.next_cursor.request_id, + } + ) + assert next_response.status_code == 200 + next_page: Final = PromptCachingRequestsResponse.model_validate_json(next_response.content) + assert tuple(row.request_id for row in next_page.requests) == expected[2:4] + assert next_page.has_more is (len(expected) > 4) + assert (next_page.next_cursor is not None) is next_page.has_more + second: Final = await client.get( + _URL, params={"start_date": _START, "end_date": _END, "filter": filter, "page_size": 100} + ) + assert second.status_code == 200 + complete: Final = PromptCachingRequestsResponse.model_validate_json(second.content) + assert tuple(row.request_id for row in complete.requests) == expected + assert complete.has_more is False + assert complete.next_cursor is None + assert all(row.start_time.tzinfo == timezone.utc for row in complete.requests) + payload: Final = _JSON_OBJECT.validate_json(second.content) + assert set(payload) == {"requests", "page_size", "has_more", "next_cursor"} + serialized_rows: Final = _JSON_ROWS.validate_python(payload["requests"]) + assert set(serialized_rows[0]) == { + "request_id", + "start_time", + "model", + "gateway_injected", + "cache_read_tokens", + "cache_creation_tokens", + "spend", + "net_savings", + } + by_id: Final = {row.request_id: row for row in complete.requests} + if filter == "all": + assert by_id["injected-empty"].gateway_injected is True + assert by_id["injected-empty"].net_savings is None + assert by_id["legacy-read"].gateway_injected is False + assert by_id["legacy-read"].net_savings is not None and by_id["legacy-read"].net_savings > 0 + assert by_id["write"].net_savings is not None and by_id["write"].net_savings < 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role", [None, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +async def test_non_admin_is_denied_before_database_access( + role: LitellmUserRoles | None, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", None) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(role)), base_url="http://test") as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END}) + assert response.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("params", [ + {"filter": "savings"}, {"page_size": 0}, {"page_size": 101}, {"start_date": "invalid"}, + {"cursor_start_time": "invalid", "cursor_request_id": "request"}, + {"cursor_start_time": _START, "cursor_request_id": ""}, +]) +async def test_invalid_request_is_rejected(params: Mapping[str, str | int]) -> None: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params}) + assert response.status_code == 422 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("params", [{"cursor_start_time": _START}, {"cursor_request_id": "request"}]) +async def test_incomplete_cursor_is_rejected( + params: Mapping[str, str], monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", None) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params}) + assert response.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete_before_cursor", [False, True]) +async def test_cursor_keeps_remaining_requests_once_during_insertions_and_deletions( + _cache_postgresql: psycopg.Connection[tuple[object, ...]], + _cache_prisma: Prisma, + monkeypatch: pytest.MonkeyPatch, + delete_before_cursor: bool, +) -> None: + from litellm.proxy import proxy_server + + cases: Final = (*_CASES, _Case( + "older-cache-read", {"usage_object": {"cache_read_input_tokens": 100}}, start_time=datetime(2026, 9, 1, 11), + )) + _seed(_cache_postgresql, cases) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=_cache_prisma)) + monkeypatch.setattr(proxy_server, "llm_router", None) + expected: Final = (*sorted((case.request_id for case in _CASES if case.matches("all")), reverse=True), "older-cache-read") + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + first: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, "page_size": 2}) + assert first.status_code == 200 + first_page: Final = PromptCachingRequestsResponse.model_validate_json(first.content) + assert tuple(row.request_id for row in first_page.requests) == expected[:2] + assert first_page.next_cursor is not None + with _cache_postgresql.cursor() as cursor: + cursor.executemany( + """INSERT INTO "LiteLLM_SpendLogs" + SELECT %s, %s, "endTime", model, model_id, custom_llm_provider, spend, metadata, cache_hit + FROM "LiteLLM_SpendLogs" WHERE request_id = %s""", + ( + ("newer-request", datetime(2026, 9, 1, 13), expected[0]), + ("zz-higher-id", cases[0].start_time, expected[0]), + ), + ) + if delete_before_cursor: + cursor.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (expected[0],)) + _cache_postgresql.commit() + following: Final = await client.get(_URL, params={ + "start_date": _START, "end_date": _END, "page_size": 100, + "cursor_start_time": first_page.next_cursor.start_time.isoformat(), + "cursor_request_id": first_page.next_cursor.request_id, + }) + assert following.status_code == 200 + following_page: Final = PromptCachingRequestsResponse.model_validate_json(following.content) + assert tuple(row.request_id for row in following_page.requests) == expected[2:] + assert following_page.has_more is False + assert following_page.next_cursor is None diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index aae966022e3..004f07da431 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -11,6 +11,7 @@ from litellm.proxy.spend_tracking.savings import ( compute_autorouter_savings, compute_savings_spend, marks_gateway_injection, + prompt_caching_savings_for_request, ) from litellm.router import Router from litellm.types.utils import Usage @@ -18,6 +19,42 @@ from litellm.types.utils import Usage pytestmark = pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("model,usage", [ + (None, {"cache_read_input_tokens": 100}), + ("claude-sonnet-5", None), + ("claude-sonnet-5", {"prompt_tokens": "invalid"}), +]) +def test_prompt_cache_estimate_distinguishes_unknown_from_zero(model: str | None, usage: dict[str, object] | None) -> None: + assert prompt_caching_savings_for_request(model, "anthropic", usage) is None + assert compute_savings_spend(model, "anthropic", 0, False, usage_object=usage).prompt_caching == 0 + assert prompt_caching_savings_for_request("claude-sonnet-5", "anthropic", {"prompt_tokens": 100}) == 0 + + +def test_prompt_cache_estimate_uses_the_rollup_pricing_and_retains_write_premiums() -> None: + router: Final = Router(model_list=[{ + "model_name": "negotiated", + "litellm_params": { + "model": "anthropic/claude-sonnet-5", "input_cost_per_token": 1e-6, + "cache_creation_input_token_cost": 1.25e-6, "cache_read_input_token_cost": 1e-7, + }, + "model_info": {"id": "negotiated-cache-prices"}, + }]) + + def current_router() -> Router: + return router + + usage: Final = {"cache_read_input_tokens": 1000, "cache_creation_input_tokens": 20000} + estimate: Final = prompt_caching_savings_for_request( + "claude-sonnet-5", "anthropic", usage, model_id="negotiated-cache-prices", llm_router=current_router, + ) + rollup: Final = compute_savings_spend( + "claude-sonnet-5", "anthropic", 0, True, usage_object=usage, + model_id="negotiated-cache-prices", llm_router=current_router, + ) + assert estimate == pytest.approx(1000 * (1e-6 - 1e-7) - 20000 * (1.25e-6 - 1e-6)) + assert estimate == rollup.prompt_caching == rollup.gateway_injected_caching + + @pytest.mark.parametrize("modifier", [{"speed": "fast"}, {"inference_geo": "us"}]) @pytest.mark.parametrize("continuing", [False, True]) def test_baseline_preserves_anthropic_pricing_fields(modifier: dict[str, str], continuing: bool) -> None: diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index a0877b04648..f5b71a00061 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -3,7 +3,6 @@ import React, { useMemo, useState } from "react"; import { ArrowDown, ArrowUp, ArrowUpDown, Info } from "lucide-react"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -81,7 +80,7 @@ const SortableHead = ({ }; const CacheLeakageCard: React.FC = ({ activity }) => { - const { dateValue, onDateChange, results, loading, isFetchingMore, apiKeyTruncation } = activity; + const { results, loading, isFetchingMore, apiKeyTruncation } = activity; const [dimension, setDimension] = useState("key"); const [sort, setSort] = useState({ column: "potentialSavings", dir: "desc" }); const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]); @@ -111,9 +110,6 @@ const CacheLeakageCard: React.FC = ({ activity }) => { cached token, after cache-write premiums.

-
- -
setDimension(value === "model" ? "model" : "key")}> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 03250e3e53b..f8336f5ab56 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -42,6 +42,7 @@ vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => })); vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./PromptCachingRequestsTable", () => ({ default: () =>
})); import CostOptimizationView from "./CostOptimizationView"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx new file mode 100644 index 00000000000..833a46ce16f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx @@ -0,0 +1,248 @@ +import { Profiler } from "react"; +import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor, within } from "@/../tests/test-utils"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import type { components } from "@/lib/http/schema"; +import PromptCachingRequestsTable from "./PromptCachingRequestsTable"; +import type { DateRange } from "./useDailyActivityRange"; + +type CacheRequest = components["schemas"]["PromptCachingRequest"]; +type RequestsResponse = components["schemas"]["PromptCachingRequestsResponse"]; +const firstCursor = { start_time: "2026-09-01T11:59:59.123456Z", request_id: "first-boundary?&" }; +const secondCursor = { start_time: firstCursor.start_time, request_id: "second-boundary" }; +const fetchMock = vi.fn(); +const dates = { from: new Date(2026, 8, 1, 12), to: new Date(2026, 8, 2, 12) }; +const request = (overrides: Partial = {}): CacheRequest => ({ + request_id: "request-default", + start_time: "2026-09-01T12:00:00Z", + model: "cache-test-model", + gateway_injected: true, + cache_read_tokens: 0, + cache_creation_tokens: 1000, + spend: 0.0375, + net_savings: -0.0075, + ...overrides, +}); +const response = (requests: CacheRequest[], nextCursor: RequestsResponse["next_cursor"] = null) => { + const body: RequestsResponse = { requests, has_more: nextCursor !== null, next_cursor: nextCursor, page_size: 50 }; + return Response.json(body); +}; +const lastQuery = () => new URL(String(fetchMock.mock.calls.at(-1)?.[0]), "http://localhost").searchParams; + +describe("PromptCachingRequestsTable", () => { + beforeEach(() => { + fetchMock.mockReset(); + vi.stubGlobal("fetch", fetchMock); + }); + + afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); + vi.unstubAllEnvs(); + vi.useRealTimers(); + }); + + it("separates recorded injection from cache hits, retains write premiums and unknown savings, and links each request", async () => { + const clientHit = { + request_id: "client-hit", + gateway_injected: false, + cache_read_tokens: 10000, + cache_creation_tokens: 0, + net_savings: 0.27, + }; + fetchMock.mockResolvedValue( + response([ + request({ request_id: "injected/write?&", net_savings: -0.0075 }), + request(clientHit), + request({ request_id: "unknown-price", net_savings: null }), + request({ request_id: "no-benefit", net_savings: 0 }), + ]), + ); + renderWithProviders(); + + const table = await screen.findByRole("table", { name: "Prompt caching requests" }); + const write = within(table).getByRole("row", { name: /injected\/write/ }); + expect(within(write).getByText("Recorded")).toBeInTheDocument(); + expect(within(write).getByText("1,000")).toBeInTheDocument(); + expect(within(write).getByText("$0.0375")).toBeInTheDocument(); + expect(within(write).getByText("-$0.0075")).toBeInTheDocument(); + expect(within(write).getByText(new Date("2026-09-01T12:00:00Z").toLocaleString())).toBeInTheDocument(); + expect(within(write).getByText("cache-test-model")).toHaveAttribute("title", "cache-test-model"); + expect(within(write).getByRole("link")).toHaveAttribute("href", "/ui/logs?log_id=injected%2Fwrite%3F%26"); + + const hit = within(table).getByRole("row", { name: /client-hit/ }); + expect(within(hit).getByText("Not recorded")).toBeInTheDocument(); + expect(within(hit).getByText("10,000")).toBeInTheDocument(); + expect(within(hit).getByText("$0.2700")).toBeInTheDocument(); + expect(within(table).getByRole("row", { name: /unknown-price/ })).toHaveTextContent("Unavailable"); + expect(within(table).getByRole("row", { name: /no-benefit/ })).toHaveTextContent("$0.00"); + expect(screen.getByText(/after cache-write premiums/)).toBeInTheDocument(); + expect(lastQuery().get("start_date")).toBe("2026-09-01T00:00:00.000Z"); + expect(lastQuery().get("end_date")).toBe("2026-09-02T23:59:59.999Z"); + expect(fetchMock.mock.calls[0][1]?.headers).toEqual(expect.objectContaining({ Authorization: "Bearer token-a" })); + }); + + it("forwards complete server cursors, goes back to prior cursors, and clears them for each caching filter", async () => { + fetchMock.mockImplementation(async (input) => { + const query = new URL(String(input), "http://localhost").searchParams; + const pages = new Map([ + [null, 1], + [firstCursor.request_id, 2], + [secondCursor.request_id, 3], + ]); + const page = pages.get(query.get("cursor_request_id")); + const nextCursor = + new Map([ + [1, firstCursor], + [2, secondCursor], + ]).get(page ?? 0) ?? null; + return response([request({ request_id: `${query.get("filter")}-${page}` })], nextCursor); + }); + renderWithProviders(); + await screen.findByRole("link", { name: "all-1" }); + expect(screen.getByRole("button", { name: "Previous" })).toBeDisabled(); + expect(lastQuery().has("page")).toBe(false); + expect(lastQuery().has("cursor_request_id")).toBe(false); + + fireEvent.click(screen.getByRole("button", { name: "Next" })); + await screen.findByRole("link", { name: "all-2" }); + expect(screen.getByText("Page 2")).toBeInTheDocument(); + expect(lastQuery().get("cursor_start_time")).toBe(firstCursor.start_time); + expect(lastQuery().get("cursor_request_id")).toBe(firstCursor.request_id); + fireEvent.click(screen.getByRole("button", { name: "Next" })); + await screen.findByRole("link", { name: "all-3" }); + expect(screen.getByText("Page 3")).toBeInTheDocument(); + expect(lastQuery().get("cursor_start_time")).toBe(secondCursor.start_time); + expect(lastQuery().get("cursor_request_id")).toBe(secondCursor.request_id); + expect(screen.getByRole("button", { name: "Next" })).toBeDisabled(); + + await testQueryClient.invalidateQueries({ refetchType: "none" }); + fireEvent.click(screen.getByRole("button", { name: "Previous" })); + await screen.findByRole("link", { name: "all-2" }); + await waitFor(() => expect(lastQuery().get("cursor_request_id")).toBe(firstCursor.request_id)); + expect(lastQuery().get("cursor_start_time")).toBe(firstCursor.start_time); + expect(screen.getByText("Page 2")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Previous" })); + await screen.findByRole("link", { name: "all-1" }); + await waitFor(() => expect(lastQuery().has("cursor_request_id")).toBe(false)); + expect(lastQuery().has("cursor_start_time")).toBe(false); + fireEvent.click(screen.getByRole("button", { name: "Next" })); + await screen.findByRole("link", { name: "all-2" }); + + fireEvent.click(screen.getByRole("tab", { name: "LiteLLM injected" })); + await screen.findByRole("link", { name: "injected-1" }); + expect(screen.queryByRole("link", { name: "all-2" })).not.toBeInTheDocument(); + expect(lastQuery().get("filter")).toBe("injected"); + expect(lastQuery().has("cursor_request_id")).toBe(false); + expect(lastQuery().has("cursor_start_time")).toBe(false); + + fireEvent.click(screen.getByRole("button", { name: "Next" })); + await screen.findByRole("link", { name: "injected-2" }); + fireEvent.click(screen.getByRole("tab", { name: "Cache hits" })); + await screen.findByRole("link", { name: "hits-1" }); + expect(lastQuery().get("filter")).toBe("hits"); + expect(lastQuery().get("page_size")).toBe("50"); + expect(screen.getByText("Page 1")).toBeInTheDocument(); + }); + + it("includes the current UTC day for a range ending today, matching the activity totals", async () => { + vi.stubEnv("TZ", "America/Los_Angeles"); + vi.setSystemTime(new Date("2026-09-20T03:00:00Z")); + fetchMock.mockResolvedValue(response([])); + const today = { from: new Date(2026, 8, 19), to: new Date() }; + renderWithProviders(); + + await screen.findByText("No matching prompt caching requests in this range"); + expect(lastQuery().get("start_date")).toBe("2026-09-19T00:00:00.000Z"); + expect(lastQuery().get("end_date")).toBe("2026-09-20T23:59:59.999Z"); + }); + + it.each(["date", "authentication"])( + "hides every old-scope frame and resets pagination when %s changes", + async (change) => { + fetchMock.mockResolvedValueOnce(response([request({ request_id: "old-first" })], firstCursor)); + fetchMock.mockResolvedValueOnce(response([request({ request_id: "old-second" })])); + const committedOldRows: boolean[] = []; + const snapshot = () => { + committedOldRows.push(screen.queryByRole("link", { name: "old-second" }) !== null); + }; + const tree = (accessToken: string, dateValue: DateRange) => ( + + + + ); + const { rerender } = renderWithProviders(tree("token-a", dates)); + await screen.findByRole("link", { name: "old-first" }); + fireEvent.click(screen.getByRole("button", { name: "Next" })); + await screen.findByRole("link", { name: "old-second" }); + + const pending = Promise.withResolvers(); + fetchMock.mockReturnValueOnce(pending.promise); + committedOldRows.length = 0; + rerender( + tree( + change === "authentication" ? "token-b" : "token-a", + change === "date" ? { ...dates, to: new Date(2026, 8, 3) } : dates, + ), + ); + + expect(screen.getByRole("status")).toHaveTextContent("Loading requests"); + expect(committedOldRows.length).toBeGreaterThan(0); + expect(committedOldRows.every((visible) => !visible)).toBe(true); + expect(lastQuery().has("cursor_request_id")).toBe(false); + expect(lastQuery().has("cursor_start_time")).toBe(false); + if (change === "date") { + expect(lastQuery().get("end_date")).toBe("2026-09-03T23:59:59.999Z"); + } else { + expect(fetchMock.mock.calls.at(-1)?.[1]?.headers).toEqual( + expect.objectContaining({ Authorization: "Bearer token-b" }), + ); + } + + pending.resolve(response([request({ request_id: "new-first" })])); + await screen.findByRole("link", { name: "new-first" }); + expect(screen.getByText("Page 1")).toBeInTheDocument(); + expect(committedOldRows.every((visible) => !visible)).toBe(true); + }, + ); + + it("ignores a delayed response from the previous caching filter", async () => { + const stale = Promise.withResolvers(); + const current = Promise.withResolvers(); + fetchMock.mockReturnValueOnce(stale.promise).mockReturnValueOnce(current.promise); + renderWithProviders(); + fireEvent.click(screen.getByRole("tab", { name: "Cache hits" })); + expect(lastQuery().get("filter")).toBe("hits"); + + current.resolve(response([request({ request_id: "current-hit" })])); + await screen.findByRole("link", { name: "current-hit" }); + await act(async () => { + stale.resolve(response([request({ request_id: "stale-all" })], firstCursor)); + await stale.promise; + }); + + expect(screen.getByRole("link", { name: "current-hit" })).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "stale-all" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Next" })).toBeDisabled(); + }); + + it("offers retry after a failed read and shows the empty state after it succeeds", async () => { + fetchMock.mockRejectedValueOnce(new Error("offline")); + fetchMock.mockResolvedValueOnce(response([])); + renderWithProviders(); + + expect(await screen.findByRole("alert")).toHaveTextContent("Could not load prompt caching requests"); + fireEvent.click(screen.getByRole("button", { name: "Retry" })); + expect(await screen.findByText("No matching prompt caching requests in this range")).toBeInTheDocument(); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Next" })).toBeDisabled(); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it("does not request data for an incomplete date range", async () => { + renderWithProviders(); + expect(screen.getByText("Select a date range to view requests")).toBeInTheDocument(); + expect(screen.queryByRole("status")).not.toBeInTheDocument(); + await waitFor(() => expect(fetchMock).not.toHaveBeenCalled()); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx new file mode 100644 index 00000000000..29aa9252e7b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx @@ -0,0 +1,186 @@ +"use client"; + +import { useQuery, type UseQueryOptions } from "@tanstack/react-query"; +import Link from "next/link"; +import { useState } from "react"; + +import { apiClient } from "@/components/networking"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { LOG_ID_QUERY_PARAM } from "@/components/view_logs/logDetailRouting"; +import type { paths } from "@/lib/http/schema"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { uiHref } from "@/utils/uiHref"; +import { usd } from "./costOptimizationUtils"; +import { benchmarksWindow as activityWindow } from "./useAutoRouterBenchmarks"; +import type { DateRange } from "./useDailyActivityRange"; + +const REQUESTS_PATH = "/cost_optimization/prompt_caching/requests"; +type RequestsEndpoint = paths[typeof REQUESTS_PATH]["get"]; +type RequestsResponse = RequestsEndpoint["responses"][200]["content"]["application/json"]; +type RequestsQuery = NonNullable; +type RequestFilter = NonNullable; +type RequestCursor = RequestsResponse["next_cursor"]; + +interface PromptCachingRequestsTableProps { + accessToken: string; + dateValue: DateRange; +} + +export default function PromptCachingRequestsTable({ accessToken, dateValue }: PromptCachingRequestsTableProps) { + const [filter, setFilter] = useState("all"); + const window = activityWindow(dateValue, new Date()); + const startDate = window.start_date ? `${window.start_date}T00:00:00.000Z` : ""; + const endDate = window.end_date ? `${window.end_date}T23:59:59.999Z` : ""; + const scope = JSON.stringify([accessToken, startDate, endDate, filter]); + const [pagination, setPagination] = useState<{ scope: string; cursors: readonly RequestCursor[] }>({ + scope, + cursors: [null], + }); + const cursors = pagination.scope === scope ? pagination.cursors : [null]; + const cursor = cursors.at(-1); + const page = cursors.length; + + if (pagination.scope !== scope) { + setPagination({ scope, cursors: [null] }); + } + + const enabled = Boolean(accessToken && startDate && endDate); + const query: RequestsQuery = { + start_date: startDate, + end_date: endDate, + filter, + page_size: 50, + cursor_start_time: cursor?.start_time, + cursor_request_id: cursor?.request_id, + }; + const queryOptions: UseQueryOptions = { + queryKey: [REQUESTS_PATH, accessToken, query], + queryFn: ({ signal }) => apiClient.get(REQUESTS_PATH, { accessToken, query, signal }), + enabled, + retry: false, + }; + const requests = useQuery(queryOptions); + const nextCursor = requests.data?.next_cursor; + + const changeFilter = (value: unknown) => { + if (value === "all" || value === "injected" || value === "hits") { + setFilter(value); + } + }; + + return ( + + +
+ Prompt caching requests +

+ Requests with recorded LiteLLM injection or provider cache reads or writes. A cache hit alone does not + establish LiteLLM injection; older logs may not record it. +

+

+ Net savings are estimated from logged usage and current configured pricing, after cache-write premiums. + Negative values mean caching cost more; unavailable means the request could not be priced. +

+
+ + + All caching + LiteLLM injected + Cache hits + + +
+ + {!enabled &&

Select a date range to view requests

} + {enabled && requests.isPending && ( +

+ Loading requests... +

+ )} + {enabled && requests.isError && ( +
+

Could not load prompt caching requests

+ +
+ )} + {enabled && requests.isSuccess && ( + <> + {requests.data.requests.length === 0 ? ( +

+ No matching prompt caching requests in this range +

+ ) : ( + + + + Request + Model + LiteLLM injection + Cache reads + Cache writes + Actual cost + Net savings + + + + {requests.data.requests.map((request) => ( + + + + {request.request_id} + + + + + + {request.model} + + + {request.gateway_injected ? "Recorded" : "Not recorded"} + {formatNumberWithCommas(request.cache_read_tokens)} + + {formatNumberWithCommas(request.cache_creation_tokens)} + + {usd(request.spend)} + + {request.net_savings === null ? "Unavailable" : usd(request.net_savings)} + + + ))} + +
+ )} +
+ + Page {page} + +
+ + )} +
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx index 66db347e70f..35464c5852e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx @@ -1,4 +1,4 @@ -import { render, waitFor, screen } from "@testing-library/react"; +import { fireEvent, render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; const mockGetGeneralSettingsCall = vi.fn(); @@ -12,6 +12,21 @@ vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => })); const mockCacheLeakageCard = vi.fn(); +const mockRequestsTable = vi.fn(); +const nextDateRange = { from: new Date(2026, 8, 1), to: new Date(2026, 8, 2) }; + +vi.mock("./PromptCachingRequestsTable", () => ({ + default: (props: unknown) => { + mockRequestsTable(props); + return
; + }, +})); + +vi.mock("@/components/shared/advanced_date_picker", () => ({ + default: ({ onValueChange }: { onValueChange: (range: typeof nextDateRange) => void }) => ( + + ), +})); vi.mock("./CacheLeakageCard", () => ({ __esModule: true, @@ -24,7 +39,7 @@ vi.mock("./CacheLeakageCard", () => ({ import PromptCachingTab from "./PromptCachingTab"; describe("PromptCachingTab", () => { - it("renders the cache leakage table alongside the caching settings", async () => { + it("shares the selected dates between requests and cache leakage alongside caching settings", async () => { mockGetGeneralSettingsCall.mockResolvedValue([]); const activity = { @@ -42,6 +57,10 @@ describe("PromptCachingTab", () => { expect(screen.getByTestId("caching-settings")).toBeInTheDocument(); expect(screen.getByTestId("cache-leakage-card")).toBeInTheDocument(); + expect(screen.getByTestId("caching-requests")).toBeInTheDocument(); + expect(mockRequestsTable).toHaveBeenCalledWith({ accessToken: "test-token", dateValue: activity.dateValue }); + fireEvent.click(screen.getByRole("button", { name: "Change caching dates" })); + expect(activity.onDateChange).toHaveBeenCalledWith(nextDateRange); await waitFor(() => expect(mockCacheLeakageCard).toHaveBeenCalledWith(expect.objectContaining({ activity }))); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx index 59b38f272e0..4e43317998e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx @@ -3,12 +3,14 @@ import React, { useCallback, useEffect, useState } from "react"; import { getGeneralSettingsCall } from "@/components/networking"; +import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { toast } from "@/lib/toast"; import { PromptCachingPanel, generalSettingsItem, } from "@/app/(dashboard)/router-settings/_components/general_settings"; import CacheLeakageCard from "./CacheLeakageCard"; +import PromptCachingRequestsTable from "./PromptCachingRequestsTable"; import { DailyActivityRange } from "./useDailyActivityRange"; interface PromptCachingTabProps { @@ -48,6 +50,11 @@ const PromptCachingTab: React.FC = ({ accessToken, activi return (
+
+

Date range for requests and cache leakage

+ +
+
); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 81580c8bfb1..d916509c06f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -3534,6 +3534,23 @@ export interface paths { patch?: never; trace?: never; }; + "/cost_optimization/prompt_caching/requests": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Prompt Caching Requests */ + get: operations["get_prompt_caching_requests_cost_optimization_prompt_caching_requests_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/credentials": { parameters: { query?: never; @@ -35814,6 +35831,48 @@ export interface components { prompt_id: string; prompt_info?: components["schemas"]["PromptInfo"] | null; }; + /** PromptCachingRequest */ + PromptCachingRequest: { + /** Cache Creation Tokens */ + cache_creation_tokens: number; + /** Cache Read Tokens */ + cache_read_tokens: number; + /** Gateway Injected */ + gateway_injected: boolean; + /** Model */ + model: string; + /** Net Savings */ + net_savings: number | null; + /** Request Id */ + request_id: string; + /** Spend */ + spend: number; + /** + * Start Time + * Format: date-time + */ + start_time: string; + }; + /** PromptCachingRequestCursor */ + PromptCachingRequestCursor: { + /** Request Id */ + request_id: string; + /** + * Start Time + * Format: date-time + */ + start_time: string; + }; + /** PromptCachingRequestsResponse */ + PromptCachingRequestsResponse: { + /** Has More */ + has_more: boolean; + next_cursor: components["schemas"]["PromptCachingRequestCursor"] | null; + /** Page Size */ + page_size: number; + /** Requests */ + requests: components["schemas"]["PromptCachingRequest"][]; + }; /** PromptInfo */ PromptInfo: { /** @@ -47238,6 +47297,42 @@ export interface operations { }; }; }; + get_prompt_caching_requests_cost_optimization_prompt_caching_requests_get: { + parameters: { + query: { + start_date: string; + end_date: string; + page_size?: number; + filter?: "all" | "injected" | "hits"; + cursor_start_time?: string | null; + cursor_request_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["PromptCachingRequestsResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_credentials_credentials_get: { parameters: { query?: never; From 875f015e24219110dbad35691d80acd1c1a3c375 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:33:56 -0700 Subject: [PATCH 08/56] fix(token_counter): count replayed redacted_thinking blocks so prompt_caching keeps pinning A conversation that replays a redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in for a reasoning item that carries no summary) made _count_content_list raise, is_prompt_caching_valid_prompt swallowed that to False, and the prompt_caching pre-call check neither recorded nor pinned the serving deployment, so the conversation bounced across the group and paid a cache write on every deployment. The block now counts like a thinking block with no text: zero tokens for the encrypted payload. --- litellm/litellm_core_utils/token_counter.py | 11 ++-- .../litellm_core_utils/test_token_counter.py | 19 +++++++ .../test_prompt_caching_deployment_check.py | 52 +++++++++++++++++++ 3 files changed, 79 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6c1b7946394..bf37b1be2e4 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -46,6 +46,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionDocumentObject, ChatCompletionNamedToolChoiceParam, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ChatCompletionToolParam, OpenAIMessageContentListBlock, ) @@ -854,6 +856,8 @@ def _count_content_list( content_list: str | Iterable[ OpenAIMessageContentListBlock + | ChatCompletionThinkingBlock + | ChatCompletionRedactedThinkingBlock | AnthropicMessagesTextParam | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam @@ -898,9 +902,9 @@ def _count_content_list( use_default_image_token_count, default_token_count, ) - elif c["type"] == "thinking": + elif c["type"] in ("thinking", "redacted_thinking"): # Claude extended thinking content block - # Count the thinking text and skip signature (opaque signature blob) + # Count the thinking text and skip the opaque blobs (signature, redacted data) thinking_text = str(c.get("thinking", "")) if thinking_text: num_tokens += count_function(thinking_text) @@ -920,7 +924,8 @@ def _count_content_list( raise ValueError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " - f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)." + f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " + f"tool_reference)." ) return num_tokens except Exception as e: diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index ba3a6be609f..f19a8891609 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1257,6 +1257,25 @@ def test_token_counter_with_thinking_content(): ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}" + +def test_token_counter_with_redacted_thinking_content(): + """ + A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in + for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking + block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the + prompt_caching pre-call check stop pinning the deployment that held the cached prefix. + """ + model = "anthropic/claude-sonnet-4-5-20250929" + reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."} + redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30} + user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]} + follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]} + + without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up] + with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up] + + assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block) + def test_token_counter_with_tool_reference_block(): """ Regression test: a message containing an Anthropic tool-search diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 333e7b2ff31..267109c9164 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -197,6 +197,58 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" +@pytest.mark.asyncio +async def test_replayed_redacted_thinking_block_still_records_and_pins(): + """ + A model that returns no reasoning summary (gpt-5.x through the /v1/messages bridge, Anthropic with + redacted reasoning) hands the client a `redacted_thinking` block, and the client replays it on every + later turn. The token count behind `is_prompt_caching_valid_prompt` raised on that block, the helper + swallowed it to False, and the check neither recorded the serving deployment nor pinned it, so the + conversation bounced across the group and paid a cache write on each deployment. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + model = "openai/gpt-5.6-sol" + deployments = _deployments(model, model, model) + messages = cast( + List[AllMessageValues], + [ + *_messages(word_count=3000), + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:" + "Z" * 400}, + {"type": "text", "text": "Draw from the box labeled Mixed."}, + ], + }, + {"role": "user", "content": "Restate that in one sentence."}, + ], + ) + + assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True + + await check.async_log_success_event( + kwargs={ + "standard_logging_object": { + "call_type": "anthropic_messages", + "model": model, + "messages": messages, + "model_id": "dep-2", + } + }, + response_obj=None, + start_time=None, + end_time=None, + ) + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + ) + + assert filtered == [deployments[1]] + + def _auto_caching_messages() -> List[AllMessageValues]: """A prompt over the model minimum that carries no client cache_control.""" return cast( From 3772993032e93d283c9c0b0cf5a80909feae52f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:42:05 -0700 Subject: [PATCH 09/56] fix(anthropic_messages): only Mantle consumes get_llm_provider's api_base The /v1/messages handler passed the api_base get_llm_provider resolved to every provider's native messages config, which shadowed DEEPSEEK_ANTHROPIC_API_BASE and TENCENT_ANTHROPIC_API_BASE with the chat default and changed the azure_ai precedence. Messages configs now opt in through uses_get_llm_provider_api_base(), true only for Bedrock Mantle, whose region-prefixed model must resolve to a region host before the prefix is stripped. Also registers BedrockMantleAnthropicMessagesConfig in the lazy import registry. --- litellm/__init__.py | 3 ++ litellm/_lazy_imports_registry.py | 5 +++ .../messages/handler.py | 6 ++- .../anthropic_messages/transformation.py | 3 ++ .../bedrock_mantle/messages/transformation.py | 3 ++ ...erimental_pass_through_messages_handler.py | 42 +++++++++++++++++++ 6 files changed, 61 insertions(+), 1 deletion(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index e17ab613dac..d2bbc107205 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1684,6 +1684,9 @@ if TYPE_CHECKING: from .llms.bedrock.messages.mantle_transformation import ( AmazonMantleMessagesConfig as AmazonMantleMessagesConfig, ) + from .llms.bedrock_mantle.messages.transformation import ( + BedrockMantleAnthropicMessagesConfig as BedrockMantleAnthropicMessagesConfig, + ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig from .llms.together_ai.chat.transformation import ( TogetherAIChatConfig as TogetherAIChatConfig, diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 9cfcb9e41f7..bca04a17250 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = ( "BedrockClaudePlatformMessagesConfig", "AmazonAnthropicClaudeMessagesConfig", "AmazonMantleMessagesConfig", + "BedrockMantleAnthropicMessagesConfig", "TogetherAIConfig", "TogetherAIChatConfig", "NLPCloudConfig", @@ -746,6 +747,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.bedrock.messages.mantle_transformation", "AmazonMantleMessagesConfig", ), + "BedrockMantleAnthropicMessagesConfig": ( + ".llms.bedrock_mantle.messages.transformation", + "BedrockMantleAnthropicMessagesConfig", + ), "TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"), "TogetherAIChatConfig": ( ".llms.together_ai.chat.transformation", diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index e1309ea4063..d87cb0a64f5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -501,7 +501,6 @@ def anthropic_messages_handler( api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) - resolved_api_base: Final = dynamic_api_base if dynamic_api_base is not None else api_base # Store agentic loop params in logging object for agentic hooks # This provides original request context needed for follow-up calls @@ -652,6 +651,11 @@ def anthropic_messages_handler( "display": "summarized", } + resolved_api_base: Final = ( + dynamic_api_base + if dynamic_api_base is not None and anthropic_messages_provider_config.uses_get_llm_provider_api_base() + else api_base + ) return base_llm_http_handler.anthropic_messages_handler( model=model, messages=strip_provider_specific_fields_from_anthropic_messages(messages), diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 8e7c22930fa..101a5e6c58c 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -128,6 +128,9 @@ class BaseAnthropicMessagesConfig(ABC): """ return True + def uses_get_llm_provider_api_base(self) -> bool: + return False + def get_async_streaming_response_iterator( self, model: str, diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py index 480c09a0476..480fe82ef4c 100644 --- a/litellm/llms/bedrock_mantle/messages/transformation.py +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -61,6 +61,9 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM def custom_llm_provider(self) -> str | None: return "bedrock_mantle" + def uses_get_llm_provider_api_base(self) -> bool: + return True + def get_complete_url( self, api_base: str | None, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 997a97c6fd3..9fa3ef153be 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1438,3 +1438,45 @@ async def test_anthropic_messages_leaves_non_provider_failures_unmapped(): ) assert "Traceback" not in str(excinfo.value) + + +def _recording_client(seen_urls: list[str]) -> AsyncHTTPHandler: + def record_and_answer(request: httpx.Request) -> httpx.Response: + seen_urls.append(str(request.url)) + return httpx.Response( + 200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "deepseek-chat", + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + ) + + upstream = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(record_and_answer)) + return upstream + + +@pytest.mark.asyncio +async def test_provider_messages_api_base_env_is_not_shadowed_by_the_chat_default(monkeypatch): + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + monkeypatch.delenv("DEEPSEEK_API_BASE", raising=False) + monkeypatch.setenv("DEEPSEEK_ANTHROPIC_API_BASE", "https://deepseek.internal.example/anthropic") + seen_urls: list[str] = [] + + await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "ping"}], + model="deepseek/deepseek-chat", + api_key="sk-test", + client=_recording_client(seen_urls), + ) + + assert seen_urls == ["https://deepseek.internal.example/anthropic/v1/messages"] + From 517fff5bb7bbbd397ad1942cba5a3a1b35e0640a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:52:22 -0700 Subject: [PATCH 10/56] fix(router): keep prompt caching affinity when the breakpoint moves The prompt_caching pre-call check keyed a deployment pin on a hash of the whole cacheable prefix, cache_control markers included. Agent clients such as Claude Code move the marker to the newest user turn on every request, so the key changed every turn, the pin never matched, and a multi-turn session drifted across deployments and lost its provider cache. Hash the prefix per content block with the markers stripped, chained so every block position has a key, and write the pin at the breakpoint block. Lookup walks back over the last PROMPT_CACHE_LOOKBACK_POSITIONS positions (a run of tool_use or tool_result blocks counting as one), the same window the provider probes for a cached prefix, in one batch cache read. Both sides hash the prefix after base64 truncation so a request carrying raw image bytes derives the keys the success event stored. --- litellm/constants.py | 3 + litellm/router_utils/prompt_caching_cache.py | 250 +++++++++++----- .../test_prompt_caching_deployment_check.py | 273 +++++++++++++++++- 3 files changed, 450 insertions(+), 76 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index bbeb4846e27..e4576ad4d5c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -399,6 +399,9 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = ( if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT ) +# Anthropic checks at most 20 block positions behind a breakpoint for a cached prefix, a run of tool_use +# or tool_result blocks counting as one position, so deployment affinity probes the same window +PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20 DEFAULT_TRIM_RATIO: Final = float( os.getenv("DEFAULT_TRIM_RATIO", 0.75) ) # default ratio of tokens to trim from the end of a prompt diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 39708e168f5..0b784e1fa91 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -4,12 +4,21 @@ Wrapper around router cache. Meant to store model id when prompt caching support import hashlib import json +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import JsonValue, TypeAdapter +from pydantic_core import to_jsonable_python from typing_extensions import TypedDict from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS +from litellm.litellm_core_utils.logging_utils import ( + truncate_base64_in_messages, + truncate_base64_in_messages_async, +) from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam if TYPE_CHECKING: @@ -28,10 +37,100 @@ class PromptCachingCacheValue(TypedDict): model_id: str +PROMPT_CACHE_PIN_TTL_SECONDS: Final = 300 +_TOOL_RUN_BLOCK_TYPES: Final = frozenset({"tool_use", "tool_result"}) +_PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) +_TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...]) +_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None) + + +@dataclass(frozen=True, slots=True) +class PrefixPosition: + cache_key: str + position: int + + +def _sorted_pairs(pairs: Iterable[tuple[str, JsonValue]]) -> tuple[tuple[str, JsonValue], ...]: + return tuple(sorted(pairs, key=lambda pair: pair[0])) + + +def _canonical_bytes(value: object) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + + +def _block_unit( + envelope: tuple[tuple[str, JsonValue], ...], message_run_type: str | None, block: JsonValue +) -> tuple[bytes, str | None]: + if not isinstance(block, dict): + return _canonical_bytes((envelope, block)), message_run_type + block_type: Final = block.get("type") + block_run_type: Final = block_type if isinstance(block_type, str) and block_type in _TOOL_RUN_BLOCK_TYPES else None + stripped: Final = _sorted_pairs(item for item in block.items() if item[0] != "cache_control") + return _canonical_bytes((envelope, stripped)), message_run_type or block_run_type + + +def _message_units(message: Mapping[str, JsonValue]) -> tuple[tuple[bytes, str | None], ...]: + envelope: Final = _sorted_pairs(item for item in message.items() if item[0] not in ("content", "cache_control")) + message_run_type: Final = "tool_result" if message.get("role") == "tool" else None + content: Final = message.get("content") + if isinstance(content, list) and content: + return tuple(_block_unit(envelope, message_run_type, block) for block in content) + if isinstance(content, str) and content: + return ((_canonical_bytes((envelope, (("text", content), ("type", "text")))), message_run_type),) + return ((_canonical_bytes((envelope, None)), message_run_type),) + + +def _chain_digest(digest: bytes, unit: bytes) -> bytes: + return hashlib.sha256(digest + unit).digest() + + +def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes: + if tools is None: + return hashlib.sha256(b"").digest() + return hashlib.sha256( + _canonical_bytes(_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True))) + ).digest() + + +def _positions_of( + prefix: tuple[Mapping[str, JsonValue], ...], tools: Sequence[ChatCompletionToolParam] | None +) -> tuple[PrefixPosition, ...]: + units: Final = tuple(unit for message in prefix for unit in _message_units(message)) + digests: Final = tuple(accumulate((unit_bytes for unit_bytes, _ in units), _chain_digest, initial=_seed(tools)))[1:] + run_types: Final = tuple(run_type for _, run_type in units) + positions: Final = accumulate( + 0 if run_type is not None and run_type == previous else 1 + for run_type, previous in zip(run_types, (None, *run_types[:-1])) + ) + return tuple( + PrefixPosition(cache_key=f"deployment:{digest.hex()}:prompt_caching", position=position) + for digest, position in zip(digests, positions) + ) + + +def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]: + if not positions: + return () + oldest_probed_position: Final = positions[-1].position - PROMPT_CACHE_LOOKBACK_POSITIONS + return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position) + + +def _pinned_value(value: JsonValue) -> PromptCachingCacheValue | None: + if not isinstance(value, dict): + return None + model_id: Final = value.get("model_id") + return PromptCachingCacheValue(model_id=model_id) if isinstance(model_id, str) else None + + +def _first_pin(values: tuple[JsonValue, ...] | None) -> PromptCachingCacheValue | None: + if values is None: + return None + return next((pin for pin in map(_pinned_value, values) if pin is not None), None) + + class PromptCachingCache: def __init__(self, cache: DualCache): self.cache = cache - self.in_memory_cache = InMemoryCache() @staticmethod def serialize_object(obj: Any) -> object: @@ -140,114 +239,123 @@ class PromptCachingCache: return cacheable_prefix @staticmethod - def get_prompt_caching_cache_key( + def prefix_positions( messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, - ) -> str | None: - if messages is None and tools is None: - return None + tools: Sequence[ChatCompletionToolParam] | None, + ) -> tuple[PrefixPosition, ...]: + """ + One cache key per content block of the cacheable prefix, oldest block first. - # Extract cacheable prefix from messages (only include up to last cache_control block) - cacheable_messages = None - if messages is not None: - cacheable_messages = PromptCachingCache.extract_cacheable_prefix(messages) - # If no cacheable prefix found, return None (can't cache) - if not cacheable_messages: - return None + Each key hashes the prefix content up to and including that block, with cache_control markers + left out, so the key of a block is the same whichever turn's breakpoint the prefix ends at. + String content hashes like a single text block, which is how the provider treats it and how + Claude Code re-sends a previously marked message. `position` counts a run of consecutive + tool_use (or tool_result) blocks as one, matching the provider's lookback window. - # Use serialize_object for consistent and stable serialization - data_to_hash: Final = {} - if cacheable_messages is not None: - serialized_messages: Final = PromptCachingCache.serialize_object(cacheable_messages) - data_to_hash["messages"] = serialized_messages - if tools is not None: - serialized_tools: Final = PromptCachingCache.serialize_object(tools) - data_to_hash["tools"] = serialized_tools - - # Combine serialized data into a single string - data_to_hash_str: Final = json.dumps( - data_to_hash, - sort_keys=True, - separators=(",", ":"), + The prefix is hashed in the shape the success event sees it, with long base64 data URIs + already replaced by their size placeholder, so a request carrying the raw image bytes + derives the same keys the write side stored. + """ + if not messages: + return () + return _positions_of( + _PREFIX_ADAPTER.validate_python( + to_jsonable_python( + truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)), + serialize_unknown=True, + ) + ), + tools, ) - # Create a hash of the serialized data for a stable cache key - hashed_data: Final = hashlib.sha256(data_to_hash_str.encode()).hexdigest() - return f"deployment:{hashed_data}:prompt_caching" + @staticmethod + async def async_prefix_positions( + messages: list[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + ) -> tuple[PrefixPosition, ...]: + if not messages: + return () + return _positions_of( + _PREFIX_ADAPTER.validate_python( + to_jsonable_python( + await truncate_base64_in_messages_async(PromptCachingCache.extract_cacheable_prefix(messages)), + serialize_unknown=True, + ) + ), + tools, + ) + + @staticmethod + def get_prompt_caching_cache_key( + messages: list[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + ) -> str | None: + positions: Final = PromptCachingCache.prefix_positions(messages, tools) + return positions[-1].cache_key if positions else None def add_model_id( self, model_id: str, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - if messages is None and tools is None: - return - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - # If no cacheable prefix found, don't cache (can't generate cache key) if cache_key is None: return - self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300) - return + self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=PROMPT_CACHE_PIN_TTL_SECONDS) async def async_add_model_id( self, model_id: str, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - if messages is None and tools is None: - return - - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - # If no cacheable prefix found, don't cache (can't generate cache key) - if cache_key is None: + positions: Final = await PromptCachingCache.async_prefix_positions(messages, tools) + if not positions: return await self.cache.async_set_cache( - cache_key, + positions[-1].cache_key, PromptCachingCacheValue(model_id=model_id), - ttl=300, # store for 5 minutes + ttl=PROMPT_CACHE_PIN_TTL_SECONDS, ) - return async def async_get_model_id( self, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> PromptCachingCacheValue | None: """ - Get model ID from cache using the cacheable prefix. - - The cache key is based on the cacheable prefix (everything up to and including - the last cache_control block), so requests with the same cacheable prefix but - different user messages will have the same cache key. + Find the deployment that last served this prefix, walking back from the breakpoint the + same way the provider cache does, so a breakpoint that moved forward since the last + turn still lands on the deployment whose cache holds the earlier prefix. """ - if messages is None and tools is None: + cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools)) + if not cache_keys: return None - # Generate cache key using cacheable prefix - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - if cache_key is None: - return None - - # Perform cache lookup - cache_result: Final = await self.cache.async_get_cache(key=cache_key) - return cache_result + return _first_pin( + _PINS_ADAPTER.validate_python( + await self.cache.async_batch_get_cache( + keys=list(cache_keys), # mutable-ok: DualCache.async_batch_get_cache only takes a list + ) + ) + ) def get_model_id( self, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> PromptCachingCacheValue | None: - if messages is None and tools is None: + cache_keys: Final = _lookback_keys(PromptCachingCache.prefix_positions(messages, tools)) + if not cache_keys: return None - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - # If no cacheable prefix found, return None (can't cache) - if cache_key is None: - return None - - return self.cache.get_cache(cache_key) + return _first_pin( + _PINS_ADAPTER.validate_python( + self.cache.batch_get_cache( + keys=list(cache_keys), # mutable-ok: DualCache.batch_get_cache only takes a list + ) + ) + ) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 333e7b2ff31..d0a9223dfa7 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1,5 +1,6 @@ import asyncio import copy +import functools from typing import List, cast import pytest @@ -7,7 +8,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache -from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT +from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( @@ -30,7 +31,6 @@ def _local_model_cost_map_autouse(local_model_cost_map): yield - def _deployments(*models: str) -> List[dict]: return [ { @@ -84,7 +84,9 @@ def test_write_gate_is_what_prevents_a_pin_below_the_model_minimum(): """ messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True + ) assert 1024 < token_count < 4096 assert is_prompt_caching_valid_prompt(model="anthropic/claude-opus-4-5", messages=messages) is False @@ -110,7 +112,9 @@ async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minim deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT < token_count < OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -136,7 +140,9 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=5000) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert token_count > OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -539,3 +545,260 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): "model_id": "dep-1" } assert_loop_stayed_free(took, lags) + + +LONG_PROMPT = "word " * 3000 +ONE_PIXEL_PNG = ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +def _turn(*messages: dict) -> List[AllMessageValues]: + return cast(List[AllMessageValues], list(messages)) + + +def _text(text: str) -> dict: + return {"type": "text", "text": text} + + +def _marked(text: str) -> dict: + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + + +@pytest.mark.asyncio +async def test_pin_survives_the_breakpoint_moving_to_the_next_turn(): + """ + The regression. Claude Code marks only the newest user message each turn, so the last breakpoint + moves forward every turn. The key hashed the prefix up to that moving breakpoint, markers + included, so no turn after the first ever found the pin the previous turn wrote, and a + multi-deployment group re-rolled the deployment mid-session, paying a cache write on a + deployment whose provider cache held nothing of the conversation. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + turn_one = _turn({"role": "user", "content": [_marked(LONG_PROMPT)]}) + turn_two = _turn( + {"role": "user", "content": [_text(LONG_PROMPT)]}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": [_marked("next")]}, + ) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=turn_one, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[1]] + + +@pytest.mark.asyncio +async def test_pin_survives_the_marked_message_coming_back_as_string_content(): + """ + Claude Code sends the message that carries a breakpoint as a one-block content list and re-sends + it next turn as plain string content once the marker has moved on. The provider caches both + shapes identically, so the key has to as well, or the walk-back never lands on the turn-one write. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + turn_one = _turn( + {"role": "system", "content": [_marked(LONG_PROMPT)]}, + {"role": "user", "content": [_marked("hello")]}, + ) + turn_two = _turn( + {"role": "system", "content": LONG_PROMPT}, + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": [_marked("again")]}, + ) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=turn_one, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[0]] + + +@pytest.mark.asyncio +async def test_lookback_stops_where_the_provider_cache_stops(): + """ + Anthropic finds a cached prefix at most PROMPT_CACHE_LOOKBACK_POSITIONS block positions behind a + breakpoint, the breakpoint block included. Probing further would pin to a deployment whose cache + the provider will not consult, and probing less would drop pins the provider still honors. + """ + prompt_cache = PromptCachingCache(cache=DualCache()) + await prompt_cache.async_add_model_id( + model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("block 0")]}), tools=None + ) + + def turn_with_blocks_after(count: int) -> List[AllMessageValues]: + later = [_text(f"block {index}") for index in range(1, count)] + [_marked(f"block {count}")] + return _turn({"role": "user", "content": [_text("block 0"), *later]}) + + inside_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS - 1) + past_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS) + + assert await prompt_cache.async_get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"} + assert prompt_cache.get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"} + assert await prompt_cache.async_get_model_id(messages=past_window, tools=None) is None + assert prompt_cache.get_model_id(messages=past_window, tools=None) is None + + +@pytest.mark.asyncio +async def test_a_run_of_tool_blocks_counts_as_one_lookback_position(): + """ + The provider counts consecutive tool_use blocks as one lookback position, and consecutive + tool_result blocks as one, in both the Anthropic and the OpenAI message shapes. An agent turn that + fans out into many tool calls would otherwise push the previous breakpoint out of the window + after a single turn, which is exactly when the conversation is longest and the cache matters most. + """ + prompt_cache = PromptCachingCache(cache=DualCache()) + await prompt_cache.async_add_model_id( + model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("task")]}), tools=None + ) + fan_out = PROMPT_CACHE_LOOKBACK_POSITIONS + 5 + + def anthropic_shaped(tool_use_type: str, tool_result_type: str) -> List[AllMessageValues]: + return _turn( + {"role": "user", "content": [_text("task")]}, + { + "role": "assistant", + "content": [ + {"type": tool_use_type, "id": f"call-{index}", "name": "read", "input": {"index": index}} + for index in range(fan_out) + ], + }, + { + "role": "user", + "content": [ + *( + {"type": tool_result_type, "tool_use_id": f"call-{index}", "content": "ok"} + for index in range(fan_out) + ), + _marked("continue"), + ], + }, + ) + + openai_shaped = _turn( + {"role": "user", "content": [_text("task")]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": f"call-{index}", "type": "function", "function": {"name": "read", "arguments": "{}"}} + for index in range(fan_out) + ], + }, + *({"role": "tool", "tool_call_id": f"call-{index}", "content": "ok"} for index in range(fan_out)), + {"role": "user", "content": [_marked("continue")]}, + ) + + assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("tool_use", "tool_result"), tools=None) == { + "model_id": "dep-1" + } + assert await prompt_cache.async_get_model_id(messages=openai_shaped, tools=None) == {"model_id": "dep-1"} + assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("text", "text"), tools=None) is None + + +@pytest.mark.asyncio +async def test_an_edited_earlier_block_does_not_inherit_the_pin(): + """Walking back must still bind every block's content, or an edited conversation pins to a stale cache.""" + prompt_cache = PromptCachingCache(cache=DualCache()) + await prompt_cache.async_add_model_id( + model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("original")]}), tools=None + ) + edited = _turn( + {"role": "user", "content": [_text("edited")]}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": [_marked("next")]}, + ) + + assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None + + +class _BrokenBatchReadCache(DualCache): + async def async_batch_get_cache(self, keys, parent_otel_span=None, local_only=False, **kwargs): + return None + + +@pytest.mark.asyncio +async def test_a_failed_batch_read_pins_nothing(): + """DualCache answers None rather than a list when the batch read raises, and routing must fall through.""" + prompt_cache = PromptCachingCache(cache=_BrokenBatchReadCache()) + + assert ( + await prompt_cache.async_get_model_id(messages=_turn({"role": "user", "content": [_marked("x")]}), tools=None) + is None + ) + + +@pytest.mark.asyncio +async def test_pin_matches_when_the_success_event_truncated_an_image_payload(monkeypatch, local_model_cost_map): + """ + The success event only ever sees the standard logging payload, whose long base64 data URIs are + replaced by size placeholders, while routing sees the raw request. Hashing the raw bytes on the + read side would key every image-carrying session past its own pin. + """ + capture = _SentMessagesCapture() + monkeypatch.setattr(litellm, "callbacks", [capture]) + image = {"type": "image_url", "image_url": {"url": ONE_PIXEL_PNG}} + turn_one = _turn({"role": "user", "content": [image, _marked(LONG_PROMPT)]}) + + await litellm.acompletion( + model=AUTO_CACHING_MODEL, messages=copy.deepcopy(turn_one), mock_response="ok", api_key="sk-fake" + ) + logged = await _eventually(lambda: capture.messages) + assert logged is not None + assert logged != turn_one + + cache = DualCache() + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=logged, tools=None) + turn_two = _turn( + {"role": "user", "content": [image, _text(LONG_PROMPT)]}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": [_marked("next")]}, + ) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + + filtered = await PromptCachingDeploymentCheck(cache=cache).async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[1]] + + +@pytest.mark.asyncio +async def test_claude_code_style_session_stays_on_one_deployment_across_turns(local_model_cost_map): + """ + End to end over the router with a client that marks only the newest user message each turn, the + way Claude Code does. Every turn has to land on the deployment that served the first one. + """ + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2", "dep-3") + ], + optional_pre_call_checks=["prompt_caching"], + ) + user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 6))] + history: List[AllMessageValues] = [] + served: List[str] = [] + for text in user_turns: + request = cast(List[AllMessageValues], [*history, {"role": "user", "content": [_marked(text)]}]) + response = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=request, mock_response="ok") + served.append(response._hidden_params["model_id"]) + pin_key = PromptCachingCache.get_prompt_caching_cache_key(request, None) + assert await _eventually(functools.partial(router.cache.get_cache, key=pin_key)) is not None + history = [*history, {"role": "user", "content": [_text(text)]}, {"role": "assistant", "content": "ok"}] + + assert served == [served[0]] * len(user_turns) From c13dcb0abfe3de7b6722e18d7acf0f59eaa39fc8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:11:14 -0700 Subject: [PATCH 11/56] fix(proxy): forward a client's anthropic-beta and anthropic-version headers to bedrock_mantle --- litellm/proxy/litellm_pre_call_utils.py | 7 +++++- ..._bedrock_mantle_messages_transformation.py | 19 +++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 24 ++++++++++++++++++- 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9a973755894..e415a78f412 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3418,7 +3418,12 @@ async def add_guardrails_from_policy_engine( _ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( - (LlmProviders.ANTHROPIC.value, LlmProviders.BEDROCK.value, LlmProviders.VERTEX_AI.value) + ( + LlmProviders.ANTHROPIC.value, + LlmProviders.BEDROCK.value, + LlmProviders.BEDROCK_MANTLE.value, + LlmProviders.VERTEX_AI.value, + ) ) _ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.ANTHROPIC.value diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py index 3544262996c..6bacf8f3d94 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -385,6 +385,25 @@ class TestBetaHeadersOnTheWire: "interleaved-thinking-2025-05-14", ] + @pytest.mark.asyncio + @respx.mock + async def test_betas_a_proxy_client_sends_reach_mantle_filtered(self): + from litellm.proxy.litellm_pre_call_utils import add_provider_specific_headers_to_request + + proxy_request_data: dict = {} + add_provider_specific_headers_to_request( + data=proxy_request_data, + headers={ + "anthropic-beta": "claude-code-20250219,fast-mode-2026-02-01,interleaved-thinking-2025-05-14", + "anthropic-version": "2023-06-01", + "user-agent": "claude-cli/2.1.239", + }, + ) + + route = await self._send(**proxy_request_data) + + assert _sent_betas(route) == ["claude-code-20250219", "interleaved-thinking-2025-05-14"] + @pytest.mark.asyncio @respx.mock async def test_betas_mantle_rejects_are_dropped_before_the_request(self): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 88d38d74f49..9257a2dd23d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7249,7 +7249,7 @@ CROSS_ACCOUNT_AUTHORIZATION = "Bearer deliberately-configured-pass-through-token SIGV4_PREFIX = "AWS4-HMAC-SHA256" AUTHORIZATION_HEADER_CASINGS = ["authorization", "Authorization", "AUTHORIZATION"] -LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "vertex_ai"] +LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "bedrock_mantle", "vertex_ai"] BEDROCK_ENDPOINT = ( "https://bedrock-runtime.us-west-2.amazonaws.com/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke" @@ -7342,6 +7342,28 @@ def test_oauth_credential_entry_is_scoped_to_anthropic_alone(): assert [entry["custom_llm_provider"] for entry in credential_entries] == ["anthropic"] +@pytest.mark.parametrize("custom_llm_provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"]) +def test_client_anthropic_api_headers_reach_every_anthropic_messages_provider(custom_llm_provider): + client_headers = { + "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14", + "anthropic-version": "2023-06-01", + "user-agent": "claude-cli/2.1.239", + } + + forwarded = _headers_forwarded_to(client_headers, custom_llm_provider) + + assert forwarded == { + "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14", + "anthropic-version": "2023-06-01", + } + + +def test_client_anthropic_api_headers_stay_off_openai_compatible_providers(): + forwarded = _headers_forwarded_to({"anthropic-beta": "claude-code-20250219"}, "openai") + + assert forwarded == {} + + def test_no_provider_specific_header_when_client_sends_nothing_anthropic(): data: dict = {} add_provider_specific_headers_to_request( From 0f0c0fe499fc12856273f6094e622a8f9dc72311 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:13:59 -0700 Subject: [PATCH 12/56] fix: drop a blank anthropic-beta header before it reaches the provider --- litellm/anthropic_beta_headers_manager.py | 2 +- .../test_anthropic_beta_headers_filtering.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index abce47c191e..7e7099a53b0 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -334,7 +334,7 @@ def update_headers_with_filtered_beta( Updated headers dict """ existing_beta: Final = headers.get("anthropic-beta") - if not existing_beta: + if existing_beta is None: return headers # Parse existing beta headers diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index 3c967283abf..d404edb1281 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -18,6 +18,7 @@ import pytest import litellm from litellm.anthropic_beta_headers_manager import ( filter_and_transform_beta_headers, + update_headers_with_filtered_beta, update_request_with_filtered_beta, ) @@ -511,3 +512,20 @@ class TestAnthropicBetaHeadersFiltering: assert ( "unknown-header-123" not in filtered ), f"Unknown header should not be in result for {provider}" + + @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"]) + def test_blank_anthropic_beta_header_is_removed(self, provider): + headers = {"anthropic-beta": "", "anthropic-version": "2023-06-01"} + + assert update_headers_with_filtered_beta(headers, provider) == {"anthropic-version": "2023-06-01"} + + @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"]) + def test_whitespace_only_anthropic_beta_header_is_removed(self, provider): + headers = {"anthropic-beta": " , ", "anthropic-version": "2023-06-01"} + + assert update_headers_with_filtered_beta(headers, provider) == {"anthropic-version": "2023-06-01"} + + def test_absent_anthropic_beta_header_is_left_alone(self): + headers = {"anthropic-version": "2023-06-01"} + + assert update_headers_with_filtered_beta(headers, "bedrock_mantle") == {"anthropic-version": "2023-06-01"} From f24208f9ca8c0c5842e92eba09d6bc9b35b8a66f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:17:55 -0700 Subject: [PATCH 13/56] fix(bedrock_mantle): price region-prefixed Claude responses from the bare Bedrock row --- litellm/utils.py | 17 ++++++++++++--- tests/test_litellm/test_cost_calculator.py | 25 ++++++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 3439a21b560..f3b9fcfd1ed 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5624,6 +5624,12 @@ def _get_model_info_from_generalization( return None +def _strip_mantle_region_prefix(model: str) -> str: + from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix + + return split_mantle_region_prefix(model)[1] + + def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> PotentialModelNamesAndCustomLLMProvider: if custom_llm_provider is None: # Get custom_llm_provider @@ -5656,17 +5662,22 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P split_model = strip_bedrock_routing_prefix(split_model) + region_free_split_model: Final = ( + _strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model + ) provider_model_info: Final = ( - ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider)) + ProviderConfigManager.get_provider_model_info( + model=region_free_split_model, provider=LlmProviders(custom_llm_provider) + ) if custom_llm_provider in LlmProvidersSet else None ) provider_cost_key: Final = ( - provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None + provider_model_info.get_model_cost_key(region_free_split_model) if provider_model_info is not None else None ) return PotentialModelNamesAndCustomLLMProvider( - split_model=split_model, + split_model=region_free_split_model, combined_model_name=combined_model_name, stripped_model_name=stripped_model_name, combined_stripped_model_name=combined_stripped_model_name, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index aef17f3d5d0..fe52b9993f2 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3522,6 +3522,31 @@ def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_mo ) +def test_completion_cost_mantle_native_messages_prices_claude_from_the_bedrock_row(_local_model_cost_map): + """Mantle's native Messages API answers with Anthropic's canonical model name and the proxy + resolves a Mantle region for every call, so the first cost candidate is + bedrock_mantle//claude-sonnet-5. That name has no row of its own and must fall through to + the deployment's bare Bedrock row instead of stopping on an unpriced capability rule at $0.""" + + response = litellm.ModelResponse( + id="msg_x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="claude-sonnet-5", + usage={"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110}, + ) + row = litellm.model_cost["anthropic.claude-sonnet-5"] + expected = 100 * row["input_cost_per_token"] + 10 * row["output_cost_per_token"] + assert expected > 0 + + for region_name in ("us-east-1", None): + assert litellm.completion_cost( + completion_response=response, + model="bedrock_mantle/anthropic.claude-sonnet-5", + custom_llm_provider="bedrock_mantle", + region_name=region_name, + ) == pytest.approx(expected) + + def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map): """An explicit base_model keeps pricing on that model's own key even when the request carries a region with different regional rates, so the private provider model never widens region pricing.""" From 3ffe6272c96c08f54f972ef43a2541d73222f2ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:22:28 -0700 Subject: [PATCH 14/56] fix(router): hash the prompt caching affinity prefix off the event loop Offload the per-block hashing through offload_token_count on both the pre-call read and the success-event write, hash raw bytes as base64 instead of raising, drop the unused serialize_object helper, and bind the chained digest, the message envelope, and the bytes path in the regression tests --- litellm/constants.py | 2 - litellm/router_utils/prompt_caching_cache.py | 38 +++------------ .../test_router_prompt_caching.py | 48 ------------------- .../test_prompt_caching_deployment_check.py | 40 ++++++++++++++-- 4 files changed, 43 insertions(+), 85 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e4576ad4d5c..215f25bccd1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -399,8 +399,6 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = ( if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT ) -# Anthropic checks at most 20 block positions behind a breakpoint for a cached prefix, a run of tool_use -# or tool_result blocks counting as one position, so deployment affinity probes the same window PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20 DEFAULT_TRIM_RATIO: Final = float( os.getenv("DEFAULT_TRIM_RATIO", 0.75) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 0b784e1fa91..78fc5e3fe6d 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -15,10 +15,8 @@ from typing_extensions import TypedDict from litellm.caching.caching import DualCache from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS -from litellm.litellm_core_utils.logging_utils import ( - truncate_base64_in_messages, - truncate_base64_in_messages_async, -) +from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam if TYPE_CHECKING: @@ -88,7 +86,9 @@ def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes: if tools is None: return hashlib.sha256(b"").digest() return hashlib.sha256( - _canonical_bytes(_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True))) + _canonical_bytes( + _TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64")) + ) ).digest() @@ -132,23 +132,6 @@ class PromptCachingCache: def __init__(self, cache: DualCache): self.cache = cache - @staticmethod - def serialize_object(obj: Any) -> object: - """Helper function to serialize Pydantic objects, dictionaries, or fallback to string.""" - if hasattr(obj, "dict"): - # If the object is a Pydantic model, use its `dict()` method - return obj.dict() - elif isinstance(obj, dict): - # If the object is a dictionary, serialize it with sorted keys - return json.dumps(obj, sort_keys=True, separators=(",", ":")) # Standardize serialization - - elif isinstance(obj, list): - # Serialize lists by ensuring each element is handled properly - return [PromptCachingCache.serialize_object(item) for item in obj] - elif isinstance(obj, (int, float, bool)): - return obj # Keep primitive types as-is - return str(obj) - @staticmethod def extract_cacheable_prefix( messages: list[AllMessageValues], @@ -263,6 +246,7 @@ class PromptCachingCache: to_jsonable_python( truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)), serialize_unknown=True, + bytes_mode="base64", ) ), tools, @@ -275,15 +259,7 @@ class PromptCachingCache: ) -> tuple[PrefixPosition, ...]: if not messages: return () - return _positions_of( - _PREFIX_ADAPTER.validate_python( - to_jsonable_python( - await truncate_base64_in_messages_async(PromptCachingCache.extract_cacheable_prefix(messages)), - serialize_unknown=True, - ) - ), - tools, - ) + return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools) @staticmethod def get_prompt_caching_cache_key( diff --git a/tests/router_unit_tests/test_router_prompt_caching.py b/tests/router_unit_tests/test_router_prompt_caching.py index 5c36c30e818..879264ca502 100644 --- a/tests/router_unit_tests/test_router_prompt_caching.py +++ b/tests/router_unit_tests/test_router_prompt_caching.py @@ -11,57 +11,9 @@ from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload import unittest -from pydantic import BaseModel from litellm.router_utils.prompt_caching_cache import PromptCachingCache -class ExampleModel(BaseModel): - field1: str - field2: int - - -def test_serialize_pydantic_object(): - model = ExampleModel(field1="value", field2=42) - serialized = PromptCachingCache.serialize_object(model) - assert serialized == {"field1": "value", "field2": 42} - - -def test_serialize_dict(): - obj = {"b": 2, "a": 1} - serialized = PromptCachingCache.serialize_object(obj) - assert serialized == '{"a":1,"b":2}' # JSON string with sorted keys - - -def test_serialize_nested_dict(): - obj = {"z": {"b": 2, "a": 1}, "x": [1, 2, {"c": 3}]} - serialized = PromptCachingCache.serialize_object(obj) - expected = '{"x":[1,2,{"c":3}],"z":{"a":1,"b":2}}' # JSON string with sorted keys - assert serialized == expected - - -def test_serialize_list(): - obj = ["item1", {"a": 1, "b": 2}, 42] - serialized = PromptCachingCache.serialize_object(obj) - expected = ["item1", '{"a":1,"b":2}', 42] - assert serialized == expected - - -def test_serialize_fallback(): - obj = 12345 # Simple non-serializable object - serialized = PromptCachingCache.serialize_object(obj) - assert serialized == 12345 - - -def test_serialize_non_serializable(): - class CustomClass: - def __str__(self): - return "custom_object" - - obj = CustomClass() - serialized = PromptCachingCache.serialize_object(obj) - assert serialized == "custom_object" # Fallback to string conversion - - @pytest.mark.asyncio async def test_router_prompt_caching_same_cacheable_prefix_routes_to_same_deployment(): """ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index d0a9223dfa7..ad92f442a6e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -708,7 +708,10 @@ async def test_a_run_of_tool_blocks_counts_as_one_lookback_position(): @pytest.mark.asyncio async def test_an_edited_earlier_block_does_not_inherit_the_pin(): - """Walking back must still bind every block's content, or an edited conversation pins to a stale cache.""" + """ + Every key must bind the whole prefix before its block, not the block alone, or a conversation + that repeats a pinned block after an edit walks back onto a cache the provider no longer holds. + """ prompt_cache = PromptCachingCache(cache=DualCache()) await prompt_cache.async_add_model_id( model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("original")]}), tools=None @@ -716,12 +719,41 @@ async def test_an_edited_earlier_block_does_not_inherit_the_pin(): edited = _turn( {"role": "user", "content": [_text("edited")]}, {"role": "assistant", "content": "ok"}, - {"role": "user", "content": [_marked("next")]}, + {"role": "user", "content": [_marked("original")]}, ) assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None +@pytest.mark.asyncio +async def test_swapped_roles_do_not_inherit_the_pin(): + """The message envelope is part of what the provider caches, so the same blocks under other roles key apart.""" + prompt_cache = PromptCachingCache(cache=DualCache()) + pinned = _turn( + {"role": "user", "content": [_text("question")]}, + {"role": "assistant", "content": [_marked("answer")]}, + ) + swapped = _turn( + {"role": "assistant", "content": [_text("question")]}, + {"role": "user", "content": [_marked("answer")]}, + ) + await prompt_cache.async_add_model_id(model_id="dep-1", messages=pinned, tools=None) + + assert await prompt_cache.async_get_model_id(messages=pinned, tools=None) == {"model_id": "dep-1"} + assert await prompt_cache.async_get_model_id(messages=swapped, tools=None) is None + + +@pytest.mark.asyncio +async def test_raw_bytes_in_a_block_hash_instead_of_failing_the_request(): + """A block carrying raw bytes must key like any other block rather than raising out of the router filter.""" + prompt_cache = PromptCachingCache(cache=DualCache()) + binary_block = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": b"\xff\xfe"}} + turn = _turn({"role": "user", "content": [binary_block, _marked("describe")]}) + await prompt_cache.async_add_model_id(model_id="dep-1", messages=turn, tools=None) + + assert await prompt_cache.async_get_model_id(messages=turn, tools=None) == {"model_id": "dep-1"} + + class _BrokenBatchReadCache(DualCache): async def async_batch_get_cache(self, keys, parent_otel_span=None, local_only=False, **kwargs): return None @@ -786,11 +818,11 @@ async def test_claude_code_style_session_stays_on_one_deployment_across_turns(lo "litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"}, "model_info": {"id": model_id}, } - for model_id in ("dep-1", "dep-2", "dep-3") + for model_id in (f"dep-{number}" for number in range(1, 7)) ], optional_pre_call_checks=["prompt_caching"], ) - user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 6))] + user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 9))] history: List[AllMessageValues] = [] served: List[str] = [] for text in user_turns: From 0c68c58eb1d63f0d857bba7ebdd8c4c5dbea992a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:39:53 -0700 Subject: [PATCH 15/56] test(proxy): expect bedrock_mantle in the anthropic header provider list --- tests/proxy_unit_tests/test_proxy_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 160753e3442..c62aab11930 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2004,7 +2004,7 @@ def test_provider_specific_header(): ) # Verify multi-provider support: anthropic headers work across multiple providers assert data["provider_specific_header"] == { - "custom_llm_provider": "anthropic,bedrock,vertex_ai", + "custom_llm_provider": "anthropic,bedrock,bedrock_mantle,vertex_ai", "extra_headers": { "anthropic-beta": "prompt-caching-2024-07-31", }, @@ -2076,7 +2076,7 @@ def test_provider_specific_header_multi_provider(): assert "provider_specific_header" in data assert ( data["provider_specific_header"]["custom_llm_provider"] - == "anthropic,bedrock,vertex_ai" + == "anthropic,bedrock,bedrock_mantle,vertex_ai" ) assert data["provider_specific_header"]["extra_headers"] == { "anthropic-beta": "context-1m-2025-08-07", From 7fc114c24f393d9329c5b6cd919cee459a96fba7 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:24:47 +0000 Subject: [PATCH 16/56] feat(policy_engine): add default fallback policy attachments A policy attachment with default: true applies only when no non-default attachment matches the request, so an opt-in guardrail policy replaces the fallback one instead of running alongside it. Supported in config.yaml, /policies/attachments, the Admin UI Attachments tab and the resolver (matched_via is prefixed with default:). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 1 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 18 +++ .../policy_engine/attachment_registry.py | 24 ++- .../proxy/policy_engine/policy_endpoints.py | 1 + litellm/proxy/schema.prisma | 1 + .../types/proxy/policy_engine/policy_types.py | 4 + .../proxy/policy_engine/resolver_types.py | 8 + schema.prisma | 1 + .../policy_engine/test_attachment_registry.py | 152 ++++++++++++------ .../_components/AttachmentTable.test.tsx | 13 ++ .../_components/AttachmentTableColumns.tsx | 14 ++ .../_components/add_attachment_form.test.tsx | 15 ++ .../_components/add_attachment_form.tsx | 18 +++ .../_components/build_attachment_data.test.ts | 10 ++ .../_components/build_attachment_data.ts | 2 + .../src/components/policies/types.ts | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 ++ 18 files changed, 241 insertions(+), 56 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql new file mode 100644 index 00000000000..a6c45448d03 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d2032cec0d0..2d7e557a9d1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable { models String[] @default([]) // Model names or patterns tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) priority Int? // Explicit execution order + is_default Boolean @default(false) // Applied only when no non-default attachment matches created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 06e157498aa..6f9a2d8c96d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -34982,6 +34982,12 @@ "PolicyAttachmentCreateRequest": { "description": "Request body for creating a policy attachment.", "properties": { + "default": { + "default": false, + "description": "Apply this attachment only when no non-default attachment matches the request.", + "title": "Default", + "type": "boolean" + }, "keys": { "anyOf": [ { @@ -35113,6 +35119,12 @@ "description": "Who created the attachment.", "title": "Created By" }, + "default": { + "default": false, + "description": "Apply this attachment only when no non-default attachment matches the request.", + "title": "Default", + "type": "boolean" + }, "definition_location": { "default": "db", "description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).", @@ -37141,6 +37153,12 @@ "PolicyAttachmentCreateRequest": { "description": "Request body for creating a policy attachment.", "properties": { + "default": { + "default": false, + "description": "Apply this attachment only when no non-default attachment matches the request.", + "title": "Default", + "type": "boolean" + }, "keys": { "anyOf": [ { diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 3735c335bd4..04009151487 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -119,6 +119,7 @@ class AttachmentRegistry: models=attachment_data.get("models"), tags=attachment_data.get("tags"), priority=attachment_data.get("priority"), + default=attachment_data.get("default", False), ) def get_attached_policies(self, context: PolicyMatchContext) -> list[str]: @@ -142,12 +143,14 @@ class AttachmentRegistry: """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher + in_scope: Final = tuple( + attachment + for attachment in self._attachments + if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) + ) + non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default) matching_attachments: Final = sorted( - ( - attachment - for attachment in self._attachments - if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) - ), + non_default or tuple(attachment for attachment in in_scope if attachment.default), key=_attachment_sort_key, ) broadest_attachment_by_policy: Final = MappingProxyType( @@ -169,6 +172,11 @@ class AttachmentRegistry: @staticmethod def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: """Describe why an attachment matched the context.""" + reason: Final = AttachmentRegistry._describe_scope_match(attachment, context) + return f"default:{reason}" if attachment.default else reason + + @staticmethod + def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher if attachment.is_global(): @@ -324,6 +332,7 @@ class AttachmentRegistry: "models": attachment_request.models or [], "tags": attachment_request.tags or [], "priority": attachment_request.priority, + "is_default": attachment_request.default, "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), "created_by": created_by, @@ -340,6 +349,7 @@ class AttachmentRegistry: models=attachment_request.models, tags=attachment_request.tags, priority=attachment_request.priority, + default=attachment_request.default, ) self.add_attachment(attachment) @@ -352,6 +362,7 @@ class AttachmentRegistry: models=created_attachment.models or [], tags=created_attachment.tags or [], priority=created_attachment.priority, + default=created_attachment.is_default, created_at=created_attachment.created_at, updated_at=created_attachment.updated_at, created_by=created_attachment.created_by, @@ -429,6 +440,7 @@ class AttachmentRegistry: models=attachment.models or [], tags=attachment.tags or [], priority=attachment.priority, + default=attachment.is_default, created_at=attachment.created_at, updated_at=attachment.updated_at, created_by=attachment.created_by, @@ -468,6 +480,7 @@ class AttachmentRegistry: models=a.models or [], tags=a.tags or [], priority=a.priority, + default=a.is_default, created_at=a.created_at, updated_at=a.updated_at, created_by=a.created_by, @@ -502,6 +515,7 @@ class AttachmentRegistry: models=(attachment_response.models if attachment_response.models else None), tags=attachment_response.tags if attachment_response.tags else None, priority=attachment_response.priority, + default=attachment_response.default, ) for attachment_response in attachments ] diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 1e30238c8b4..f4b38bea14e 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -61,6 +61,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment) models=attachment.models or [], tags=attachment.tags or [], priority=attachment.priority, + default=attachment.default, definition_location="config", ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d2032cec0d0..2d7e557a9d1 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable { models String[] @default([]) // Model names or patterns tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) priority Int? // Explicit execution order + is_default Boolean @default(false) // Applied only when no non-default attachment matches created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/types/proxy/policy_engine/policy_types.py b/litellm/types/proxy/policy_engine/policy_types.py index 66e5fbb4b49..73eeffa3585 100644 --- a/litellm/types/proxy/policy_engine/policy_types.py +++ b/litellm/types/proxy/policy_engine/policy_types.py @@ -294,6 +294,10 @@ class PolicyAttachment(BaseModel): le=2147483647, description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.", ) + default: bool = Field( + default=False, + description="Apply this attachment only when no non-default attachment matches the request.", + ) model_config = ConfigDict(extra="forbid") diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index e6f501ed4b5..ebdedb98b12 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -311,6 +311,10 @@ class PolicyAttachmentCreateRequest(BaseModel): le=2147483647, description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.", ) + default: bool = Field( + default=False, + description="Apply this attachment only when no non-default attachment matches the request.", + ) class PolicyAttachmentDBResponse(BaseModel): @@ -327,6 +331,10 @@ class PolicyAttachmentDBResponse(BaseModel): default=None, description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.", ) + default: bool = Field( + default=False, + description="Apply this attachment only when no non-default attachment matches the request.", + ) created_at: datetime | None = Field(default=None, description="When the attachment was created.") updated_at: datetime | None = Field(default=None, description="When the attachment was last updated.") created_by: str | None = Field(default=None, description="Who created the attachment.") diff --git a/schema.prisma b/schema.prisma index d2032cec0d0..2d7e557a9d1 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable { models String[] @default([]) // Model names or patterns tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) priority Int? // Explicit execution order + is_default Boolean @default(false) // Applied only when no non-default attachment matches created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 089bec59583..b419f3db060 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -30,9 +30,7 @@ class TestGetAttachedPolicies: ) # Should match any context - context = PolicyMatchContext( - team_alias="any-team", key_alias="any-key", model="any-model" - ) + context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="any-model") attached = registry.get_attached_policies(context) assert "global-baseline" in attached @@ -46,15 +44,11 @@ class TestGetAttachedPolicies: ) # Match - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") assert "healthcare-policy" in registry.get_attached_policies(context) # No match - different team - context_other = PolicyMatchContext( - team_alias="finance-team", key_alias="key", model="gpt-4" - ) + context_other = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4") assert "healthcare-policy" not in registry.get_attached_policies(context_other) def test_key_wildcard_pattern_attachment(self): @@ -67,15 +61,11 @@ class TestGetAttachedPolicies: ) # Match - key starts with dev-key- - context = PolicyMatchContext( - team_alias="team", key_alias="dev-key-123", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="team", key_alias="dev-key-123", model="gpt-4") assert "dev-policy" in registry.get_attached_policies(context) # No match - different prefix - context_prod = PolicyMatchContext( - team_alias="team", key_alias="prod-key-123", model="gpt-4" - ) + context_prod = PolicyMatchContext(team_alias="team", key_alias="prod-key-123", model="gpt-4") assert "dev-policy" not in registry.get_attached_policies(context_prod) def test_model_specific_attachment(self): @@ -92,9 +82,7 @@ class TestGetAttachedPolicies: assert "gpt4-policy" in registry.get_attached_policies(context) # No match - context_other = PolicyMatchContext( - team_alias="team", key_alias="key", model="gpt-3.5" - ) + context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-3.5") assert "gpt4-policy" not in registry.get_attached_policies(context_other) def test_model_wildcard_pattern(self): @@ -107,15 +95,11 @@ class TestGetAttachedPolicies: ) # Match - context = PolicyMatchContext( - team_alias="team", key_alias="key", model="bedrock/claude-3" - ) + context = PolicyMatchContext(team_alias="team", key_alias="key", model="bedrock/claude-3") assert "bedrock-policy" in registry.get_attached_policies(context) # No match - context_other = PolicyMatchContext( - team_alias="team", key_alias="key", model="openai/gpt-4" - ) + context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="openai/gpt-4") assert "bedrock-policy" not in registry.get_attached_policies(context_other) def test_multiple_attachments_match_same_context(self): @@ -129,9 +113,7 @@ class TestGetAttachedPolicies: ] ) - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") attached = registry.get_attached_policies(context) # All three should match @@ -277,9 +259,7 @@ class TestGetAttachedPolicies: ] ) - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") attached = registry.get_attached_policies(context) # Should only appear once @@ -288,9 +268,7 @@ class TestGetAttachedPolicies: def test_many_distinct_policies_resolve_in_linear_time(self): policy_count = 20_000 registry = AttachmentRegistry() - registry.load_attachments( - [{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)] - ) + registry.load_attachments([{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)]) context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") started = time.perf_counter() @@ -318,9 +296,7 @@ class TestGetAttachedPolicies: ] ) - context = PolicyMatchContext( - team_alias="finance-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4") attached = registry.get_attached_policies(context) assert attached == [] @@ -338,23 +314,15 @@ class TestGetAttachedPolicies: ) # Match - both team and model match - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") assert "strict-policy" in registry.get_attached_policies(context) # No match - team matches but model doesn't - context_wrong_model = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-3.5" - ) - assert "strict-policy" not in registry.get_attached_policies( - context_wrong_model - ) + context_wrong_model = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-3.5") + assert "strict-policy" not in registry.get_attached_policies(context_wrong_model) # No match - model matches but team doesn't - context_wrong_team = PolicyMatchContext( - team_alias="finance-team", key_alias="key", model="gpt-4" - ) + context_wrong_team = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4") assert "strict-policy" not in registry.get_attached_policies(context_wrong_team) @@ -527,6 +495,79 @@ class TestMatchAttribution: assert "catch-all" in attached +class TestDefaultAttachments: + """`default: true` attachments apply only when no non-default attachment matches.""" + + @staticmethod + def _registry() -> AttachmentRegistry: + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "guardrail-y", "scope": "*", "default": True}, + {"policy": "guardrail-x", "tags": ["opt-in"]}, + ] + ) + return registry + + def test_opted_in_request_gets_only_the_opt_in_policy(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + + assert self._registry().get_attached_policies(context) == ["guardrail-x"] + + def test_request_without_opt_in_falls_back_to_default_policy(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2") + + assert self._registry().get_attached_policies(context) == ["guardrail-y"] + + def test_default_attachment_still_honors_its_own_scope(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "team-default", "teams": ["team-a"], "default": True}]) + + assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-a", key_alias="k", model="m")) == [ + "team-default" + ] + assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-b", key_alias="k", model="m")) == [] + + def test_all_matching_defaults_apply_when_nothing_else_matches(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "default-a", "scope": "*", "default": True}, + {"policy": "default-b", "teams": ["team-a"], "default": True}, + {"policy": "opt-in", "tags": ["opt-in"]}, + ] + ) + context = PolicyMatchContext(team_alias="team-a", key_alias="k", model="m") + + assert registry.get_attached_policies(context) == ["default-a", "default-b"] + + def test_non_default_attachments_remain_additive(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "baseline", "scope": "*"}, + {"policy": "opt-in", "tags": ["opt-in"]}, + {"policy": "fallback", "scope": "*", "default": True}, + ] + ) + context = PolicyMatchContext(team_alias="t", key_alias="k", model="m", tags=["opt-in"]) + + assert registry.get_attached_policies(context) == ["baseline", "opt-in"] + + def test_default_match_reason_is_labelled(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="m") + + results = self._registry().get_attached_policies_with_reasons(context) + + assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] + + def test_default_defaults_to_false_when_omitted(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "p"}]) + + assert registry.get_all_attachments()[0].default is False + + class TestAttachmentRegistrySingleton: """Test global singleton behavior.""" @@ -557,6 +598,7 @@ def _make_db_attachment_row( scope: str | None = None, teams: list[str] | None = None, priority: int | None = None, + is_default: bool = False, ) -> MagicMock: row = MagicMock() row.attachment_id = attachment_id @@ -567,6 +609,7 @@ def _make_db_attachment_row( row.models = [] row.tags = [] row.priority = priority + row.is_default = is_default row.created_at = datetime.now(timezone.utc) row.updated_at = datetime.now(timezone.utc) row.created_by = None @@ -576,9 +619,7 @@ def _make_db_attachment_row( def _prisma_with_attachment_rows(rows: list[MagicMock]) -> MagicMock: prisma = MagicMock() - prisma.configure_mock( - **{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)} - ) + prisma.configure_mock(**{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)}) return prisma @@ -629,6 +670,15 @@ class TestConfigAttachmentsPreservedAcrossDbSync: assert registry.get_all_attachments()[0].priority == 7 + @pytest.mark.asyncio + async def test_sync_round_trips_db_attachment_default_flag(self): + registry = AttachmentRegistry() + db_row = _make_db_attachment_row(is_default=True) + + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row])) + + assert registry.get_all_attachments()[0].default is True + @pytest.mark.asyncio async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self): registry = AttachmentRegistry() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx index 43ad6a7cc9e..be83f73bb2e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx @@ -65,6 +65,19 @@ describe("AttachmentTable", () => { ); }); + it("should show a Default badge only for default attachments", () => { + const attachments = [ + makeAttachment({ attachment_id: "att-def00001", policy_name: "fallback", default: true }), + makeAttachment({ attachment_id: "att-def00002", policy_name: "regular" }), + ]; + renderWithProviders(); + const rows = screen.getAllByRole("row").slice(1); + const fallbackRow = rows.find((row) => within(row).queryByText("fallback")); + const regularRow = rows.find((row) => within(row).queryByText("regular")); + expect(within(fallbackRow!).getByText("Default")).toBeInTheDocument(); + expect(within(regularRow!).queryByText("Default")).not.toBeInTheDocument(); + }); + it("should show skeleton rows when isLoading is true", () => { renderWithProviders(); expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx index 9a190401d08..3265b9db834 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx @@ -181,6 +181,20 @@ export const getAttachmentTableColumns = ({ {row.original.priority} ), }, + { + id: "default", + accessorFn: (row) => (row.default ? 1 : 0), + meta: { title: "Default" }, + header: ({ column }) => , + size: 100, + enableSorting: true, + cell: ({ row }) => + row.original.default ? ( + + ) : ( + - + ), + }, { id: "created_at", accessorFn: (row) => row.created_at ?? "", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx index dfc023d428e..14af4a2b8f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx @@ -237,6 +237,21 @@ describe("AddAttachmentForm", () => { expect(createAttachment).toHaveBeenCalledWith("test-token", { policy_name: "policy-alpha", scope: "*" }); }); + it("sends default: true when the Default switch is turned on", async () => { + const user = userEvent.setup(); + const createAttachment = vi.fn().mockResolvedValue({}); + renderWithProviders(); + await selectPolicy(user, "policy-alpha"); + await user.click(screen.getByRole("switch", { name: /default/i })); + await submit(user); + await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1)); + expect(createAttachment).toHaveBeenCalledWith("test-token", { + policy_name: "policy-alpha", + scope: "*", + default: true, + }); + }); + it.each([ ["2147483648", /at most 2147483647/i], ["-2147483649", /at least -2147483648/i], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx index 02463a89139..74c4978392f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx @@ -11,6 +11,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Separator } from "@/components/ui/separator"; +import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; @@ -38,6 +39,7 @@ interface AttachmentFormValues { models: string[]; tags: string[]; priority: number | null; + default: boolean; } const EMPTY_VALUES: AttachmentFormValues = { @@ -47,6 +49,7 @@ const EMPTY_VALUES: AttachmentFormValues = { models: [], tags: [], priority: null, + default: false, }; const INT32_MIN = -2147483648; @@ -64,6 +67,7 @@ const attachmentShape = { .min(INT32_MIN, `Priority must be at least ${INT32_MIN}`) .max(INT32_MAX, `Priority must be at most ${INT32_MAX}`) .nullable(), + default: z.boolean(), }; const buildAttachmentSchema = (scopeType: ScopeType, teamsLoaded: boolean, availableTeams: string[]) => @@ -453,6 +457,20 @@ const AddAttachmentForm: React.FC = ({ /> )} + + + {({ value, onChange, ref, ...field }) => ( + + )} + {impactResult && } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts index 930e755f242..80617a40f13 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts @@ -80,6 +80,16 @@ describe("buildAttachmentData", () => { }); }); + describe("default", () => { + it.each(["global", "specific"] as const)("should send default: true for a %s scope", (scopeType) => { + expect(buildAttachmentData({ policy_name: "p", default: true }, scopeType).default).toBe(true); + }); + + it.each([undefined, false])("should omit default when it is %s", (value) => { + expect(buildAttachmentData({ policy_name: "p", default: value }, "specific")).not.toHaveProperty("default"); + }); + }); + describe("priority", () => { it.each(["global", "specific"] as const)("should include priority for a %s scope", (scopeType) => { expect(buildAttachmentData({ policy_name: "p", priority: 0 }, scopeType).priority).toBe(0); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts index 8b21142df74..8b50cd7bdc1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts @@ -7,6 +7,7 @@ export interface AttachmentFormInput { models?: string[]; tags?: string[]; priority?: number | null; + default?: boolean; } export function buildAttachmentData( @@ -25,5 +26,6 @@ export function buildAttachmentData( if (formValues.tags && formValues.tags.length > 0) data.tags = formValues.tags; } if (typeof formValues.priority === "number") data.priority = formValues.priority; + if (formValues.default === true) data.default = true; return data; } diff --git a/ui/litellm-dashboard/src/components/policies/types.ts b/ui/litellm-dashboard/src/components/policies/types.ts index 9f3ef02ba5d..430864f93df 100644 --- a/ui/litellm-dashboard/src/components/policies/types.ts +++ b/ui/litellm-dashboard/src/components/policies/types.ts @@ -45,6 +45,7 @@ export interface PolicyAttachment { models: string[]; tags: string[]; priority?: number | null; + default?: boolean; created_at?: string; updated_at?: string; created_by?: string; @@ -80,6 +81,7 @@ export interface PolicyAttachmentCreateRequest { models?: string[]; tags?: string[]; priority?: number; + default?: boolean; } export interface PolicyListResponse { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e33764c3d1a..ccb5fda2149 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35164,6 +35164,12 @@ export interface components { * @description Request body for creating a policy attachment. */ PolicyAttachmentCreateRequest: { + /** + * Default + * @description Apply this attachment only when no non-default attachment matches the request. + * @default false + */ + default: boolean; /** * Keys * @description Key aliases or patterns this attachment applies to. @@ -35220,6 +35226,12 @@ export interface components { * @description Who created the attachment. */ created_by?: string | null; + /** + * Default + * @description Apply this attachment only when no non-default attachment matches the request. + * @default false + */ + default: boolean; /** * Definition Location * @description Where this attachment is defined: 'db' (database) or 'config' (config.yaml). From ef55eb6bdf0ccb69533d9541266c5108e0176609 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 09:07:08 +0000 Subject: [PATCH 17/56] fix(policy_engine): ignore inapplicable non-default attachments when selecting defaults A non-default attachment whose policy is missing or whose condition does not match the request no longer suppresses default attachments. The impact preview marks default counts as an upper bound Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 4 ++- .../policy_engine/attachment_registry.py | 18 ++++++++-- litellm/proxy/policy_engine/policy_matcher.py | 15 ++++++++ .../policy_engine/policy_resolve_endpoints.py | 4 ++- .../proxy/policy_engine/response_retrieval.py | 4 ++- .../policy_engine/test_attachment_registry.py | 35 ++++++++++++++++++- .../_components/add_attachment_form.tsx | 2 +- .../_components/impact_preview_alert.test.tsx | 11 ++++++ .../_components/impact_preview_alert.tsx | 11 ++++-- 9 files changed, 94 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9a973755894..8aa861e574f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3216,7 +3216,9 @@ def _match_and_track_policies( attachment_registry: Final = ( attachment_registry_override if attachment_registry_override is not None else get_attachment_registry() ) - matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context) + matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons( + context, PolicyMatcher.policy_applies(context, policies_override) + ) matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons] policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 04009151487..d81471b3c1a 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -5,6 +5,7 @@ Attachments define WHERE policies apply, separate from the policy definitions. This allows the same policy to be attached to multiple scopes. """ +from collections.abc import Callable from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict @@ -122,24 +123,34 @@ class AttachmentRegistry: default=attachment_data.get("default", False), ) - def get_attached_policies(self, context: PolicyMatchContext) -> list[str]: + def get_attached_policies( + self, + context: PolicyMatchContext, + policy_applies: Callable[[str], bool] | None = None, + ) -> list[str]: """ Get list of policy names attached to the given context. Args: context: The request context to match against + policy_applies: Optional predicate; attachments whose policy does not apply are ignored Returns: List of policy names that are attached to matching scopes """ - return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)] + return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)] - def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]: + def get_attached_policies_with_reasons( + self, + context: PolicyMatchContext, + policy_applies: Callable[[str], bool] | None = None, + ) -> list[PolicyAttachmentMatch]: """ Get list of policy names and match reasons for the given context. Returns a list of dicts with 'policy_name' and 'matched_via' keys. The 'matched_via' describes which dimension caused the match. + Attachments whose policy fails `policy_applies` are dropped before defaults are considered. """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher @@ -147,6 +158,7 @@ class AttachmentRegistry: attachment for attachment in self._attachments if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) + and (policy_applies is None or policy_applies(attachment.policy)) ) non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default) matching_attachments: Final = sorted( diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index 001e4115374..2f7fcd23b75 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -7,6 +7,7 @@ apply to a given request based on team alias, key alias, and model. Policies are matched via policy_attachments which define WHERE each policy applies. """ +from collections.abc import Callable from typing import Final from litellm._logging import verbose_proxy_logger @@ -130,6 +131,20 @@ class PolicyMatcher: """ return PolicyMatcher.get_matching_policies(context=context) + @staticmethod + def policy_applies( + context: PolicyMatchContext, + policies: dict[str, Policy] | None = None, + ) -> Callable[[str], bool]: + """Predicate telling whether a policy exists and its condition matches the context.""" + return lambda policy_name: bool( + PolicyMatcher.get_policies_with_matching_conditions( + policy_names=[policy_name], # mutable-ok: the matcher takes a list + context=context, + policies=policies, + ) + ) + @staticmethod def get_policies_with_matching_conditions( policy_names: list[str], diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index a8a9856b833..898e42635c5 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -265,7 +265,9 @@ async def resolve_policies_for_context( ) # Get matching policies with reasons - match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(context=context) + match_results: Final = get_attachment_registry().get_attached_policies_with_reasons( + context=context, policy_applies=PolicyMatcher.policy_applies(context) + ) if not match_results: return PolicyResolveResponse( diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index d284c44397e..0f373b08056 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -84,7 +84,9 @@ def _retrieval_context( def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: - matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context) + matches: Final = get_attachment_registry().get_attached_policies_with_reasons( + context, PolicyMatcher.policy_applies(context) + ) if not matches: return (), MappingProxyType({}) applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index b419f3db060..faa8d67fe3a 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -14,7 +14,8 @@ from litellm.proxy.policy_engine.attachment_registry import ( AttachmentRegistry, get_attachment_registry, ) -from litellm.types.proxy.policy_engine import PolicyMatchContext +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails, PolicyMatchContext class TestGetAttachedPolicies: @@ -561,6 +562,38 @@ class TestDefaultAttachments: assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] + def test_inapplicable_opt_in_policy_does_not_suppress_default(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + policies = { + "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])), + "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="claude.*")), + } + + results = self._registry().get_attached_policies_with_reasons( + context, PolicyMatcher.policy_applies(context, policies) + ) + + assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] + + def test_attachment_to_missing_policy_does_not_suppress_default(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + policies = {"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"]))} + + assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [ + "guardrail-y" + ] + + def test_applicable_opt_in_policy_still_wins_with_predicate(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + policies = { + "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])), + "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="gpt.*")), + } + + assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [ + "guardrail-x" + ] + def test_default_defaults_to_false_when_omitted(self): registry = AttachmentRegistry() registry.load_attachments([{"policy": "p"}]) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx index 74c4978392f..5cd240a0838 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx @@ -473,7 +473,7 @@ const AddAttachmentForm: React.FC = ({ - {impactResult && } + {impactResult && }
); -const ImpactPreviewAlert: React.FC = ({ impactResult }) => { +const ImpactPreviewAlert: React.FC = ({ impactResult, isDefault = false }) => { const isGlobal = impactResult.affected_keys_count === -1; + const qualifier = isDefault ? "up to " : ""; return ( @@ -47,7 +49,7 @@ const ImpactPreviewAlert: React.FC = ({ impactResult }) ) : (
- This attachment would affect{" "} + This attachment would affect {qualifier} {impactResult.affected_keys_count} key{impactResult.affected_keys_count !== 1 ? "s" : ""} {" "} @@ -57,6 +59,11 @@ const ImpactPreviewAlert: React.FC = ({ impactResult }) . + {isDefault && ( +
+ Default attachments only apply to requests no non-default attachment matches, so fewer may be affected. +
+ )} {impactResult.sample_keys.length > 0 && ( Date: Sun, 20 Sep 2026 09:09:27 +0000 Subject: [PATCH 18/56] refactor(policy_engine): accept any sequence of policy names in condition matching Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/policy_engine/policy_matcher.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index 2f7fcd23b75..2b54b5dbe41 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -7,7 +7,7 @@ apply to a given request based on team alias, key alias, and model. Policies are matched via policy_attachments which define WHERE each policy applies. """ -from collections.abc import Callable +from collections.abc import Callable, Sequence from typing import Final from litellm._logging import verbose_proxy_logger @@ -139,7 +139,7 @@ class PolicyMatcher: """Predicate telling whether a policy exists and its condition matches the context.""" return lambda policy_name: bool( PolicyMatcher.get_policies_with_matching_conditions( - policy_names=[policy_name], # mutable-ok: the matcher takes a list + policy_names=(policy_name,), context=context, policies=policies, ) @@ -147,7 +147,7 @@ class PolicyMatcher: @staticmethod def get_policies_with_matching_conditions( - policy_names: list[str], + policy_names: Sequence[str], context: PolicyMatchContext, policies: dict[str, Policy] | None = None, ) -> list[str]: From 95cf7066d16186e94a8f27828ac38db55ef45cf7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 20:38:13 -0700 Subject: [PATCH 19/56] refactor(cache): use static dispatch and typed backend codecs --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/cache-memory/src/cache.rs | 23 +-- .../crates/cache-memory/tests/cache.rs | 47 ++++- litellm-rust/crates/cache-redis/Cargo.toml | 2 +- litellm-rust/crates/cache-redis/src/cache.rs | 171 ++++++++---------- .../crates/cache-redis/tests/cache.rs | 152 +++++++++++++++- litellm-rust/crates/cache/Cargo.toml | 1 + litellm-rust/crates/cache/src/base_cache.rs | 53 +++--- litellm-rust/crates/cache/src/caching.rs | 16 +- litellm-rust/crates/cache/src/codec.rs | 42 +++++ litellm-rust/crates/cache/src/lib.rs | 6 +- litellm-rust/crates/cache/tests/caching.rs | 70 ++++++- litellm-rust/crates/cache/tests/codec.rs | 53 ++++++ 13 files changed, 480 insertions(+), 157 deletions(-) create mode 100644 litellm-rust/crates/cache/src/codec.rs create mode 100644 litellm-rust/crates/cache/tests/codec.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8c35a0be0b4..ccfaee44f50 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2462,6 +2462,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "thiserror 2.0.19", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 1908ff44a81..974cdbe9760 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -4,8 +4,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs, - Error, + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, }; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; @@ -214,8 +213,8 @@ impl InMemoryCache { } } -impl BaseCache for InMemoryCache { - type Value = CacheEntry; +impl BaseCache for InMemoryCache { + type Value = V; fn default_ttl(&self) -> Duration { self.default_ttl @@ -238,17 +237,15 @@ impl BaseCache for InMemoryCache { self.flush_cache() } - fn disconnect(&self) -> CacheFuture<'_, ()> { - Box::pin(async { Ok(()) }) + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) } - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { - Box::pin(async { - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "In-memory cache connection test successful".into(), - error: None, - }) + async fn test_connection(&self) -> Result { + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "In-memory cache connection test successful".into(), + error: None, }) } } diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index aaf82641db7..ffac9d8ae64 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -2,7 +2,10 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; -use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error}; +use litellm_cache::{ + BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache, + set_cache, +}; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -156,3 +159,45 @@ async fn connection_test_matches_python_result_contract() { }) ); } + +#[tokio::test] +async fn generic_consumers_share_typed_values_and_honor_expiration() { + let clock = clock(); + let cache: CacheBackend> = Arc::new(cache(clock.clone(), 4)); + let reader = Arc::clone(&cache); + let kwargs = CacheKwargs { + ttl: Some(Duration::from_secs(5)), + ..Default::default() + }; + set_cache(cache.as_ref(), "sync", "first".into(), kwargs.clone()).unwrap(); + assert_eq!( + get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), + Some("first".into()) + ); + cache + .batch_cache_write("async", "second".into(), kwargs.clone()) + .await + .unwrap(); + cache + .async_set_cache_pipeline(vec![("batch".into(), "third".into())], kwargs.clone()) + .await + .unwrap(); + drop(cache); + for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] { + assert_eq!( + reader.async_get_cache(key, &kwargs).await.unwrap(), + Some(value.into()) + ); + } + reader.async_delete_cache("async").await.unwrap(); + assert_eq!( + reader.async_get_cache("async", &kwargs).await.unwrap(), + None + ); + clock.store(106, Ordering::SeqCst); + assert_eq!(get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), None); + assert_eq!( + reader.async_get_cache("batch", &kwargs).await.unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 933b0feaae4..a60813b6260 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -8,8 +8,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true redis = "1.7.0" -serde_json.workspace = true tokio.workspace = true [dev-dependencies] redis-test = "1.0.4" +serde_json.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 69dee6c6363..8d92fdc8f75 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -2,35 +2,37 @@ use std::sync::{Arc, Mutex, MutexGuard}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs, - Error, + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error, }; use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); const KEY_PREFIX: &str = "litellm-cache:"; -pub struct RedisCache { +pub struct RedisCache { connection: Arc>, default_ttl: Duration, + codec: S, } -impl RedisCache { - pub fn new(url: &str, default_ttl: Option) -> Result { +impl RedisCache { + pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; let connection = client.get_connection().map_err(|_| Error::Unavailable)?; - Ok(Self::with_connection(connection, default_ttl)) + Ok(Self::with_connection(connection, default_ttl, codec)) } } -impl RedisCache +impl RedisCache where + S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { - fn with_connection(connection: C, default_ttl: Option) -> Self { + pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { connection: Arc::new(Mutex::new(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + codec, } } @@ -47,48 +49,39 @@ where PATTERN } - fn encode(value: &CacheEntry) -> Result, Error> { - serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) - } - - fn decode(value: Vec) -> Result { - serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry) - } - fn ttl_seconds(ttl: Duration) -> u64 { ttl.as_secs() .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - fn run_blocking(connection: Arc>, operation: F) -> CacheFuture<'static, T> + async fn run_blocking(connection: Arc>, operation: F) -> Result where T: Send + 'static, F: FnOnce(&mut C) -> Result + Send + 'static, { - Box::pin(async move { - tokio::task::spawn_blocking(move || { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; - operation(&mut connection) - }) - .await - .map_err(|_| Error::Unavailable)? + tokio::task::spawn_blocking(move || { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut connection) }) + .await + .map_err(|_| Error::Unavailable)? } } -impl BaseCache for RedisCache +impl BaseCache for RedisCache where + S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { - type Value = CacheEntry; + type Value = S::Value; fn default_ttl(&self) -> Duration { self.default_ttl } fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { - let payload = Self::encode(&value)?; + let payload = self.codec.encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); self.connection()? .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl) @@ -96,11 +89,11 @@ where } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - self.connection()? + let bytes = self + .connection()? .get::<_, Option>>(Self::namespaced_key(key)) - .map_err(|_| Error::Unavailable)? - .map(Self::decode) - .transpose() + .map_err(|_| Error::Unavailable)?; + bytes.map(|bytes| self.codec.decode(&bytes)).transpose() } fn delete_cache(&self, key: &str) -> Result<(), Error> { @@ -125,86 +118,87 @@ where .map_err(|_| Error::Unavailable) } - fn async_set_cache<'a>( - &'a self, - key: &'a str, + async fn async_set_cache( + &self, + key: &str, value: Self::Value, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { - let payload = Self::encode(&value); + ) -> Result<(), Error> { + let payload = self.codec.encode(&value)?; let key = Self::namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection - .set_ex::<_, _, ()>(key, payload?, ttl) + .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) }) + .await } - fn async_get_cache<'a>( - &'a self, - key: &'a str, - _: &'a CacheKwargs, - ) -> CacheFuture<'a, Option> { + async fn async_get_cache( + &self, + key: &str, + _: &CacheKwargs, + ) -> Result, Error> { let key = Self::namespaced_key(key); - Box::pin(async move { - Self::run_blocking(Arc::clone(&self.connection), move |connection| { - connection - .get::<_, Option>>(key) - .map_err(|_| Error::Unavailable) - }) - .await? - .map(Self::decode) - .transpose() + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection + .get::<_, Option>>(key) + .map_err(|_| Error::Unavailable) }) + .await? + .map(|bytes| self.codec.decode(&bytes)) + .transpose() } - fn async_set_cache_pipeline<'a>( - &'a self, + async fn async_set_cache_pipeline( + &self, cache_list: Vec<(String, Self::Value)>, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { + ) -> Result<(), Error> { let entries = cache_list .into_iter() .map(|(key, value)| { - Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload)) + self.codec + .encode(&value) + .map(|payload| (Self::namespaced_key(&key), payload)) }) - .collect::, _>>(); + .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); Self::run_blocking(Arc::clone(&self.connection), move |connection| { - for (key, payload) in entries? { + for (key, payload) in entries { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable)?; } Ok(()) }) + .await } - fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> { + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = Self::namespaced_key(key); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) + .await } - fn disconnect(&self) -> CacheFuture<'_, ()> { - Box::pin(async { Ok(()) }) + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) } - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { - Box::pin(async move { - Self::run_blocking(Arc::clone(&self.connection), |connection| { - redis::cmd("PING") - .query::(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "Redis cache connection test successful".into(), - error: None, - }) + async fn test_connection(&self) -> Result { + Self::run_blocking(Arc::clone(&self.connection), |connection| { + redis::cmd("PING") + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, }) } } @@ -212,7 +206,7 @@ where #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::{BaseCache, CacheEntry, CacheKwargs}; + use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; @@ -224,33 +218,18 @@ mod tests { } } - #[test] - fn cache_entries_round_trip_through_json() { - let entry = entry(); - let encoded = RedisCache::::encode(&entry).unwrap(); - assert_eq!( - RedisCache::::decode(encoded).unwrap(), - entry - ); - } - - #[test] - fn invalid_json_is_rejected() { - assert!(RedisCache::::decode(b"not json".to_vec()).is_err()); - } - #[test] fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { assert_eq!( - RedisCache::::ttl_seconds(Duration::ZERO), + RedisCache::>::ttl_seconds(Duration::ZERO), 1 ); assert_eq!( - RedisCache::::ttl_seconds(Duration::from_millis(1500)), + RedisCache::>::ttl_seconds(Duration::from_millis(1500)), 2 ); assert_eq!( - RedisCache::::ttl_seconds(Duration::from_secs(15)), + RedisCache::>::ttl_seconds(Duration::from_secs(15)), 15 ); } @@ -258,7 +237,7 @@ mod tests { #[test] fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { let value = entry(); - let payload = RedisCache::::encode(&value).unwrap(); + let payload = JsonCodec::::new().encode(&value).unwrap(); let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") @@ -271,7 +250,7 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -296,7 +275,7 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); cache.flush_cache().unwrap(); } @@ -305,7 +284,7 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); assert_eq!( cache.test_connection().await.unwrap().status, diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 76f73145da8..fe15fcd975b 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,6 +1,156 @@ +use std::time::Duration; + +use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, Error, JsonCodec, get_cache, set_cache}; use litellm_cache_redis::RedisCache; +use redis_test::{MockCmd, MockRedisConnection}; + +struct TaggedByteCodec(u8); + +impl CacheCodec for TaggedByteCodec { + type Value = u8; + + fn encode(&self, value: &u8) -> Result, Error> { + if *value > 127 { + return Err(Error::InvalidEntry); + } + Ok(vec![self.0, *value]) + } + + fn decode(&self, bytes: &[u8]) -> Result { + match bytes { + [tag, value] if *tag == self.0 => Ok(*value), + _ => Err(Error::InvalidEntry), + } + } +} #[test] fn constructor_rejects_invalid_urls() { - assert!(RedisCache::new("not a redis url", None).is_err()); + assert!(RedisCache::new("not a redis url", None, JsonCodec::::new()).is_err()); +} + +#[test] +fn generic_helpers_use_the_injected_codec_and_ttl() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:counter") + .arg(2) + .arg([42u8, 7].as_slice()), + Ok("OK"), + ), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:counter"), + Ok(vec![42u8, 7]), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + let kwargs = CacheKwargs { + ttl: Some(Duration::from_millis(1500)), + ..Default::default() + }; + set_cache(&cache, "counter", 7, kwargs.clone()).unwrap(); + assert_eq!(get_cache(&cache, "counter", &kwargs).unwrap(), Some(7)); +} + +#[tokio::test] +async fn async_operations_preserve_codec_ttl_and_missing_values() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:counter") + .arg(9) + .arg([42u8, 7].as_slice()), + Ok("OK"), + ), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:counter"), + Ok(vec![42u8, 7]), + ), + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:batch") + .arg(2) + .arg([42u8, 8].as_slice()), + Ok("OK"), + ), + MockCmd::new(redis::cmd("DEL").arg("litellm-cache:counter"), Ok(1u32)), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:counter"), + Ok(redis::Value::Nil), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection( + connection, + Some(Duration::from_secs(9)), + TaggedByteCodec(42), + ); + let kwargs = CacheKwargs::default(); + cache + .batch_cache_write("counter", 7, kwargs.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("counter", &kwargs).await.unwrap(), + Some(7) + ); + cache + .async_set_cache_pipeline( + vec![("batch".into(), 8)], + CacheKwargs { + ttl: Some(Duration::from_millis(1500)), + ..Default::default() + }, + ) + .await + .unwrap(); + cache.async_delete_cache("counter").await.unwrap(); + assert_eq!( + cache.async_get_cache("counter", &kwargs).await.unwrap(), + None + ); +} + +#[tokio::test] +async fn codec_errors_propagate_without_writing_partial_batches() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:invalid"), + Ok(vec![99u8, 7]), + ), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:invalid"), + Ok(vec![99u8, 7]), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + let kwargs = CacheKwargs::default(); + assert_eq!( + cache.set_cache("invalid", 255, kwargs.clone()), + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_set_cache("invalid", 255, kwargs.clone()).await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache + .async_set_cache_pipeline( + vec![("valid".into(), 7), ("invalid".into(), 255)], + kwargs.clone(), + ) + .await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.get_cache("invalid", &kwargs), + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_get_cache("invalid", &kwargs).await, + Err(Error::InvalidEntry) + ); } diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index a14c4294aa0..350db4b1adb 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -13,3 +13,4 @@ thiserror.workspace = true [dev-dependencies] rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 2ba8ff92ebd..1891e417cb0 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -1,5 +1,4 @@ use std::future::Future; -use std::pin::Pin; use std::time::Duration; use serde::{Deserialize, Serialize}; @@ -7,8 +6,6 @@ use serde_json::{Map, Value}; use crate::Error; -pub type CacheFuture<'a, T> = Pin> + Send + 'a>>; - #[derive(Clone, Debug, Default, PartialEq)] pub struct CacheKwargs { pub ttl: Option, @@ -45,54 +42,54 @@ pub trait BaseCache: Send + Sync { fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>; - fn async_set_cache<'a>( - &'a self, - key: &'a str, + fn async_set_cache( + &self, + key: &str, value: Self::Value, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { - Box::pin(async move { self.set_cache(key, value, kwargs) }) + ) -> impl Future> + Send { + async move { self.set_cache(key, value, kwargs) } } - fn async_get_cache<'a>( - &'a self, - key: &'a str, - kwargs: &'a CacheKwargs, - ) -> CacheFuture<'a, Option> { - Box::pin(async move { self.get_cache(key, kwargs) }) + fn async_get_cache( + &self, + key: &str, + kwargs: &CacheKwargs, + ) -> impl Future, Error>> + Send { + async move { self.get_cache(key, kwargs) } } - fn async_set_cache_pipeline<'a>( - &'a self, + fn async_set_cache_pipeline( + &self, cache_list: Vec<(String, Self::Value)>, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { - Box::pin(async move { + ) -> impl Future> + Send { + async move { for (key, value) in cache_list { - self.set_cache(&key, value, kwargs.clone())?; + self.async_set_cache(&key, value, kwargs.clone()).await?; } Ok(()) - }) + } } - fn batch_cache_write<'a>( - &'a self, - key: &'a str, + fn batch_cache_write( + &self, + key: &str, value: Self::Value, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { + ) -> impl Future> + Send { self.async_set_cache(key, value, kwargs) } fn delete_cache(&self, key: &str) -> Result<(), Error>; - fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> { - Box::pin(async move { self.delete_cache(key) }) + fn async_delete_cache(&self, key: &str) -> impl Future> + Send { + async move { self.delete_cache(key) } } fn flush_cache(&self) -> Result<(), Error>; - fn disconnect(&self) -> CacheFuture<'_, ()>; + fn disconnect(&self) -> impl Future> + Send; - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>; + fn test_connection(&self) -> impl Future> + Send; } diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 1aab6ee8e91..21ebcce29bb 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -146,21 +146,21 @@ impl CacheEntry { } } -pub fn get_cache( - cache: &dyn BaseCache, +pub fn get_cache( + cache: &B, key: &str, kwargs: &CacheKwargs, -) -> Result, Error> { +) -> Result, Error> { cache.get_cache(key, kwargs) } -pub fn set_cache( - cache: &dyn BaseCache, +pub fn set_cache( + cache: &B, key: &str, - entry: CacheEntry, + value: B::Value, kwargs: CacheKwargs, ) -> Result<(), Error> { - cache.set_cache(key, entry, kwargs) + cache.set_cache(key, value, kwargs) } -pub type CacheBackend = Arc>; +pub type CacheBackend = Arc; diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs new file mode 100644 index 00000000000..09bee6032f6 --- /dev/null +++ b/litellm-rust/crates/cache/src/codec.rs @@ -0,0 +1,42 @@ +use std::marker::PhantomData; + +use serde::{Serialize, de::DeserializeOwned}; + +use crate::Error; + +pub trait CacheCodec: Send + Sync { + type Value: Clone + Send + Sync + 'static; + + fn encode(&self, value: &Self::Value) -> Result, Error>; + + fn decode(&self, bytes: &[u8]) -> Result; +} + +pub struct JsonCodec(PhantomData V>); + +impl Default for JsonCodec { + fn default() -> Self { + Self::new() + } +} + +impl JsonCodec { + pub const fn new() -> Self { + Self(PhantomData) + } +} + +impl CacheCodec for JsonCodec +where + V: Clone + Send + Sync + Serialize + DeserializeOwned + 'static, +{ + type Value = V; + + fn encode(&self, value: &Self::Value) -> Result, Error> { + serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + } + + fn decode(&self, bytes: &[u8]) -> Result { + serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry) + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index d0fe3de15cd..a1d9d1402bb 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -1,12 +1,12 @@ mod base_cache; mod caching; +mod codec; mod error; -pub use base_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs, -}; +pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; pub use caching::{ Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache, }; +pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 1192fc9a2b0..5c250c6b3c9 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,12 +1,13 @@ use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext, - CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, + BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, + CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, }; use sha2::{Digest, Sha256}; -use std::time::Duration; +use std::{sync::Mutex, time::Duration}; struct TestCache { default_ttl: Duration, + writes: Mutex>, } impl BaseCache for TestCache { @@ -17,6 +18,22 @@ impl BaseCache for TestCache { } fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> { + Err(Error::Unavailable) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + if key == "unavailable" { + return Err(Error::Unavailable); + } + self.writes + .lock() + .unwrap() + .push((key.into(), value, kwargs)); Ok(()) } @@ -32,11 +49,11 @@ impl BaseCache for TestCache { Ok(()) } - fn disconnect(&self) -> CacheFuture<'_, ()> { - Box::pin(async { Ok(()) }) + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) } - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { + async fn test_connection(&self) -> Result { unreachable!() } } @@ -45,6 +62,7 @@ impl BaseCache for TestCache { fn ttl_uses_default_and_allows_per_call_override() { let cache = TestCache { default_ttl: Duration::from_secs(60), + writes: Mutex::default(), }; assert_eq!( cache.get_ttl(&CacheKwargs::default()), @@ -59,6 +77,46 @@ fn ttl_uses_default_and_allows_per_call_override() { ); } +#[tokio::test] +async fn default_batch_operations_use_async_writes_and_stop_on_failure() { + let cache = TestCache { + default_ttl: Duration::from_secs(60), + writes: Mutex::default(), + }; + let entry = CacheEntry { + timestamp: 123.0, + response: serde_json::json!("cached"), + }; + let kwargs = CacheKwargs { + ttl: Some(Duration::from_secs(5)), + ..Default::default() + }; + cache + .batch_cache_write("single", entry.clone(), kwargs.clone()) + .await + .unwrap(); + assert_eq!( + cache + .async_set_cache_pipeline( + vec![ + ("first".into(), entry.clone()), + ("unavailable".into(), entry.clone()), + ("skipped".into(), entry.clone()), + ], + kwargs.clone(), + ) + .await, + Err(Error::Unavailable) + ); + assert_eq!( + *cache.writes.lock().unwrap(), + vec![ + ("single".into(), entry.clone(), kwargs.clone()), + ("first".into(), entry, kwargs), + ] + ); +} + #[test] fn keys_match_python_order_groups_files_presets_and_namespaces() { let mut input = CacheKeyInput { diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs new file mode 100644 index 00000000000..dad5398a879 --- /dev/null +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -0,0 +1,53 @@ +use std::collections::BTreeMap; + +use litellm_cache::{CacheCodec, CacheEntry, Error, JsonCodec}; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +struct RoutingState { + deployment: String, + cooldown_seconds: u64, +} + +#[test] +fn json_codec_round_trips_typed_domain_values() { + let codec = JsonCodec::::new(); + let value = RoutingState { + deployment: "deployment-a".into(), + cooldown_seconds: 30, + }; + let bytes = codec.encode(&value).unwrap(); + assert_eq!(codec.decode(&bytes).unwrap(), value); + assert_eq!( + serde_json::from_slice::(&bytes).unwrap(), + json!({"deployment": "deployment-a", "cooldown_seconds": 30}) + ); +} + +#[test] +fn response_entries_preserve_the_existing_json_representation() { + let codec = JsonCodec::::new(); + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + let bytes = codec.encode(&entry).unwrap(); + assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); + assert_eq!(codec.decode(&bytes).unwrap(), entry); +} + +#[test] +fn json_codec_rejects_malformed_and_wrongly_typed_entries() { + let codec = JsonCodec::::new(); + for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] { + assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry); + } +} + +#[test] +fn json_codec_propagates_encoding_errors() { + let codec = JsonCodec::>::new(); + let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]); + assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry); +} From 081c93908f1e2750c37beb0aa4e87660aa983d4e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 20:55:45 -0700 Subject: [PATCH 20/56] feat(cache): add native response cache and Python binding foundations --- litellm-rust/Cargo.lock | 100 ++++- litellm-rust/Cargo.toml | 2 + litellm-rust/crates/cache-redis/src/cache.rs | 83 +++-- .../crates/cache-redis/tests/cache.rs | 78 ++-- litellm-rust/crates/cache-response/Cargo.toml | 19 + .../crates/cache-response/src/codec.rs | 100 +++++ litellm-rust/crates/cache-response/src/lib.rs | 7 + .../crates/cache-response/src/native.rs | 97 +++++ .../crates/cache-response/src/response.rs | 124 +++++++ .../crates/cache-response/tests/response.rs | 290 +++++++++++++++ litellm-rust/crates/cache/src/caching.rs | 1 + litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 3 + .../crates/python-bridge/src/cache/facade.rs | 210 +++++++++++ .../crates/python-bridge/src/cache/mod.rs | 350 ++++++++++++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 6 + litellm/rust_bridge/_native.pyi | 49 ++- tests/test_litellm_rust/test_cache.py | 231 ++++++++++++ 18 files changed, 1702 insertions(+), 50 deletions(-) create mode 100644 litellm-rust/crates/cache-response/Cargo.toml create mode 100644 litellm-rust/crates/cache-response/src/codec.rs create mode 100644 litellm-rust/crates/cache-response/src/lib.rs create mode 100644 litellm-rust/crates/cache-response/src/native.rs create mode 100644 litellm-rust/crates/cache-response/src/response.rs create mode 100644 litellm-rust/crates/cache-response/tests/response.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/facade.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/mod.rs create mode 100644 tests/test_litellm_rust/test_cache.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ccfaee44f50..aa62e4f3770 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2486,6 +2486,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-response" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-memory", + "litellm-cache-redis", + "py_literal", + "redis", + "redis-test", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" @@ -2649,6 +2664,8 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-gcp", + "litellm-cache", + "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", @@ -2660,6 +2677,7 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "rstest", + "serde", "serde_json", "tokio", "tokio-tungstenite", @@ -2948,6 +2966,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.5.1" @@ -2958,6 +2986,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -3131,6 +3168,48 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pest" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d45aeb61b4bf818e12d4205f2466f8c4748f85f4fce0146d1c03d69d753f0ad" +dependencies = [ + "memchr", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89cc5a242e25ed4e7704d0be240f2cfbe20a8c27e7e252d94835be93d92dc39f" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7abf21475cc3820fe4b2ca2dc2142902f67a02189f3b5b3a229f4febc01a43e5" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pest_meta" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adba4db388f687393c18c51348d44a41d870ca9df71a2c98172ea3035dc6936e" +dependencies = [ + "pest", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -3305,6 +3384,19 @@ dependencies = [ "prost", ] +[[package]] +name = "py_literal" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1" +dependencies = [ + "num-bigint 0.4.8", + "num-complex", + "num-traits", + "pest", + "pest_derive", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -3604,7 +3696,7 @@ dependencies = [ "arcstr", "combine", "itoa", - "num-bigint", + "num-bigint 0.5.1", "percent-encoding", "ryu", "sha1_smol", @@ -4955,6 +5047,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + [[package]] name = "unarray" version = "0.1.4" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 2f6f5feb4ad..570d0dd3568 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -28,6 +28,8 @@ litellm-types = { path = "crates/types" } 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-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 8d92fdc8f75..0faca6cdaaf 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -7,12 +7,12 @@ use litellm_cache::{ use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); -const KEY_PREFIX: &str = "litellm-cache:"; pub struct RedisCache { connection: Arc>, default_ttl: Duration, codec: S, + namespace: Option, } impl RedisCache { @@ -33,6 +33,7 @@ where connection: Arc::new(Mutex::new(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, + namespace: None, } } @@ -40,13 +41,44 @@ where self.connection.lock().map_err(|_| Error::Unavailable) } - fn namespaced_key(key: &str) -> String { - format!("{KEY_PREFIX}{key}") + pub fn with_namespace(self, namespace: Option) -> Self { + Self { + namespace: namespace.filter(|value| !value.is_empty()), + ..self + } } - fn namespaced_pattern() -> &'static str { - const PATTERN: &str = "litellm-cache:*"; - PATTERN + fn namespaced_key(&self, key: &str) -> String { + match &self.namespace { + Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { + format!("{namespace}:{key}") + } + _ => key.into(), + } + } + + fn namespaced_pattern(&self) -> Result { + let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?; + let escaped: String = namespace + .chars() + .flat_map(|ch| { + if matches!(ch, '*' | '?' | '[' | ']' | '\\') { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect(); + Ok(format!("{escaped}:*")) + } + + fn decode_response(&self, value: redis::Value) -> Result, Error> { + match value { + redis::Value::Nil => Ok(None), + redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some), + redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some), + _ => Err(Error::InvalidEntry), + } } fn ttl_seconds(ttl: Duration) -> u64 { @@ -84,28 +116,29 @@ where let payload = self.codec.encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); self.connection()? - .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl) + .set_ex::<_, _, ()>(self.namespaced_key(key), payload, ttl) .map_err(|_| Error::Unavailable) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - let bytes = self + let value = self .connection()? - .get::<_, Option>>(Self::namespaced_key(key)) + .get::<_, redis::Value>(self.namespaced_key(key)) .map_err(|_| Error::Unavailable)?; - bytes.map(|bytes| self.codec.decode(&bytes)).transpose() + self.decode_response(value) } fn delete_cache(&self, key: &str) -> Result<(), Error> { self.connection()? - .del::<_, ()>(Self::namespaced_key(key)) + .del::<_, ()>(self.namespaced_key(key)) .map_err(|_| Error::Unavailable) } fn flush_cache(&self) -> Result<(), Error> { + let pattern = self.namespaced_pattern()?; let mut connection = self.connection()?; let keys = connection - .scan_match(Self::namespaced_pattern()) + .scan_match(pattern) .map_err(|_| Error::Unavailable)? .collect::>>() .map_err(|_| Error::Unavailable)?; @@ -125,7 +158,7 @@ where kwargs: CacheKwargs, ) -> Result<(), Error> { let payload = self.codec.encode(&value)?; - let key = Self::namespaced_key(key); + let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection @@ -140,15 +173,14 @@ where key: &str, _: &CacheKwargs, ) -> Result, Error> { - let key = Self::namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + let key = self.namespaced_key(key); + let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection - .get::<_, Option>>(key) + .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) }) - .await? - .map(|bytes| self.codec.decode(&bytes)) - .transpose() + .await?; + self.decode_response(value) } async fn async_set_cache_pipeline( @@ -161,7 +193,7 @@ where .map(|(key, value)| { self.codec .encode(&value) - .map(|payload| (Self::namespaced_key(&key), payload)) + .map(|payload| (self.namespaced_key(&key), payload)) }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); @@ -177,7 +209,7 @@ where } async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - let key = Self::namespaced_key(key); + let key = self.namespaced_key(key); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) @@ -250,7 +282,8 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -275,7 +308,8 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache.flush_cache().unwrap(); } @@ -284,7 +318,8 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); assert_eq!( cache.test_connection().await.unwrap().status, diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index fe15fcd975b..d5bba19a8bd 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -34,15 +34,12 @@ fn generic_helpers_use_the_injected_codec_and_ttl() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") - .arg("litellm-cache:counter") + .arg("counter") .arg(2) .arg([42u8, 7].as_slice()), Ok("OK"), ), - MockCmd::new( - redis::cmd("GET").arg("litellm-cache:counter"), - Ok(vec![42u8, 7]), - ), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); @@ -59,27 +56,21 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") - .arg("litellm-cache:counter") + .arg("counter") .arg(9) .arg([42u8, 7].as_slice()), Ok("OK"), ), - MockCmd::new( - redis::cmd("GET").arg("litellm-cache:counter"), - Ok(vec![42u8, 7]), - ), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), MockCmd::new( redis::cmd("SETEX") - .arg("litellm-cache:batch") + .arg("batch") .arg(2) .arg([42u8, 8].as_slice()), Ok("OK"), ), - MockCmd::new(redis::cmd("DEL").arg("litellm-cache:counter"), Ok(1u32)), - MockCmd::new( - redis::cmd("GET").arg("litellm-cache:counter"), - Ok(redis::Value::Nil), - ), + MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)), ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection( @@ -116,14 +107,8 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() { #[tokio::test] async fn codec_errors_propagate_without_writing_partial_batches() { let connection = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("GET").arg("litellm-cache:invalid"), - Ok(vec![99u8, 7]), - ), - MockCmd::new( - redis::cmd("GET").arg("litellm-cache:invalid"), - Ok(vec![99u8, 7]), - ), + MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), + MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); @@ -154,3 +139,48 @@ async fn codec_errors_propagate_without_writing_partial_batches() { Err(Error::InvalidEntry) ); } + +#[test] +fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() { + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)), + MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + assert_eq!( + cache.get_cache("key", &CacheKwargs::default()).unwrap(), + None + ); + assert_eq!( + cache + .get_cache("team:key", &CacheKwargs::default()) + .unwrap(), + None + ); +} + +#[test] +fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { + let unscoped = RedisCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + None, + JsonCodec::::new(), + ); + assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush)); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team\\*:*"), + Ok(redis_test::redis_value!(["0", ["team*:key"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let scoped = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team*".into())); + scoped.flush_cache().unwrap(); +} diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml new file mode 100644 index 00000000000..a0c4a1f74ef --- /dev/null +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-cache-response" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true +py_literal = "0.4.0" +redis = "1.7.0" +serde.workspace = true +serde_json.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +tokio.workspace = true diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs new file mode 100644 index 00000000000..137cf61267a --- /dev/null +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -0,0 +1,100 @@ +use litellm_cache::{CacheCodec, CacheEntry, Error}; +use serde_json::Value; + +pub struct ResponseCacheCodec; + +impl CacheCodec for ResponseCacheCodec { + type Value = CacheEntry; + + fn encode(&self, value: &CacheEntry) -> Result, Error> { + if !value.timestamp.is_finite() { + return Err(Error::InvalidEntry); + } + serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + } + + fn decode(&self, bytes: &[u8]) -> Result { + let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; + let entry: CacheEntry = + serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?; + if !entry.timestamp.is_finite() { + return Err(Error::InvalidEntry); + } + Ok(entry) + } +} + +pub(crate) fn decode_value(text: &str) -> Result { + if let Ok(value) = serde_json::from_str(text) { + return Ok(value); + } + check_literal_depth(text)?; + let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?; + literal_value(literal, 0) +} + +fn literal_value(value: py_literal::Value, depth: usize) -> Result { + use py_literal::Value as Literal; + if depth > 128 { + return Err(Error::InvalidEntry); + } + match value { + Literal::String(text) => Ok(Value::String(text)), + Literal::Boolean(value) => Ok(Value::Bool(value)), + Literal::None => Ok(Value::Null), + Literal::Integer(value) => { + serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry) + } + Literal::Float(value) => serde_json::Number::from_f64(value) + .map(Value::Number) + .ok_or(Error::InvalidEntry), + Literal::List(values) | Literal::Tuple(values) => values + .into_iter() + .map(|value| literal_value(value, depth + 1)) + .collect::, _>>() + .map(Value::Array), + Literal::Dict(entries) => entries + .into_iter() + .map(|(key, value)| { + let Literal::String(key) = key else { + return Err(Error::InvalidEntry); + }; + Ok((key, literal_value(value, depth + 1)?)) + }) + .collect::, _>>() + .map(Value::Object), + _ => Err(Error::InvalidEntry), + } +} + +fn check_literal_depth(text: &str) -> Result<(), Error> { + let mut quote = None; + let mut escaped = false; + let mut depth = 0usize; + for ch in text.chars() { + if escaped { + escaped = false; + continue; + } + if let Some(delimiter) = quote { + if ch == '\\' { + escaped = true; + } else if ch == delimiter { + quote = None; + } + continue; + } + match ch { + '\'' | '"' => quote = Some(ch), + '[' | '{' | '(' => { + depth += 1; + if depth > 128 { + return Err(Error::InvalidEntry); + } + } + ']' | '}' | ')' => depth = depth.saturating_sub(1), + _ => {} + } + } + Ok(()) +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs new file mode 100644 index 00000000000..454a82e76de --- /dev/null +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -0,0 +1,7 @@ +mod codec; +mod native; +mod response; + +pub use codec::ResponseCacheCodec; +pub use native::NativeResponseCache; +pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/native.rs b/litellm-rust/crates/cache-response/src/native.rs new file mode 100644 index 00000000000..c25c61cee82 --- /dev/null +++ b/litellm-rust/crates/cache-response/src/native.rs @@ -0,0 +1,97 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{CacheEntry, Error}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_redis::RedisCache; +use serde_json::Value; + +use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; + +pub enum NativeResponseCache +where + C: redis::ConnectionLike + Send + 'static, +{ + Memory(Arc>>), + Redis(Arc>>), +} + +impl Clone for NativeResponseCache { + fn clone(&self) -> Self { + match self { + Self::Memory(cache) => Self::Memory(Arc::clone(cache)), + Self::Redis(cache) => Self::Redis(Arc::clone(cache)), + } + } +} + +impl NativeResponseCache { + pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { + Self::Memory(Arc::new(ResponseCache::new(Arc::new( + InMemoryCache::response_cache(capacity, ttl, max_entry_bytes), + )))) + } + + pub fn redis( + url: &str, + ttl: Option, + namespace: Option, + ) -> Result { + let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace); + Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend))))) + } +} + +impl NativeResponseCache { + pub fn kind(&self) -> &'static str { + match self { + Self::Memory(_) => "memory", + Self::Redis(_) => "redis", + } + } + + pub fn lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + match self { + Self::Memory(cache) => cache.lookup(request, now), + Self::Redis(cache) => cache.lookup(request, now), + } + } + + pub fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.store(request, response, now), + Self::Redis(cache) => cache.store(request, response, now), + } + } + + pub async fn async_lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + match self { + Self::Memory(cache) => cache.async_lookup(request, now).await, + Self::Redis(cache) => cache.async_lookup(request, now).await, + } + } + + pub async fn async_store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Redis(cache) => cache.async_store(request, response, now).await, + } + } +} diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs new file mode 100644 index 00000000000..3e9807b6d1b --- /dev/null +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -0,0 +1,124 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key, +}; +use serde_json::Value; + +#[derive(Clone)] +pub struct ResponseCacheRequest { + pub key: CacheKeyInput, + pub controls: CacheControls, + pub kwargs: CacheKwargs, + pub max_age: Option, +} + +impl ResponseCacheRequest { + pub fn new(key: CacheKeyInput) -> Self { + Self { + key, + controls: CacheControls { + configured: true, + supported_call_type: true, + native_backend: true, + default_on: true, + ..Default::default() + }, + kwargs: CacheKwargs::default(), + max_age: None, + } + } +} + +pub struct ResponseCache> { + backend: Arc, +} + +impl> ResponseCache { + pub fn new(backend: Arc) -> Self { + Self { backend } + } + + pub fn lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + if !request.controls.reads() { + return Ok(None); + } + let entry = self + .backend + .get_cache(&cache_key(&request.key), &request.kwargs)?; + Self::fresh_response(entry, now, request.max_age) + } + + pub async fn async_lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + if !request.controls.reads() { + return Ok(None); + } + let entry = self + .backend + .async_get_cache(&cache_key(&request.key), &request.kwargs) + .await?; + Self::fresh_response(entry, now, request.max_age) + } + + pub fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + if !request.controls.writes() { + return Ok(()); + } + self.backend.set_cache( + &cache_key(&request.key), + CacheEntry { + timestamp: now.as_secs_f64(), + response, + }, + request.kwargs.clone(), + ) + } + + pub async fn async_store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + if !request.controls.writes() { + return Ok(()); + } + self.backend + .async_set_cache( + &cache_key(&request.key), + CacheEntry { + timestamp: now.as_secs_f64(), + response, + }, + request.kwargs.clone(), + ) + .await + } + + fn fresh_response( + entry: Option, + now: Duration, + max_age: Option, + ) -> Result, Error> { + entry + .filter(|entry| entry.fresh(now, max_age)) + .map(|entry| match entry.response { + Value::String(text) => crate::codec::decode_value(&text), + value => Ok(value), + }) + .transpose() + } +} diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs new file mode 100644 index 00000000000..a4beda9fa4c --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -0,0 +1,290 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_redis::RedisCache; +use litellm_cache_response::{ + NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::json; + +fn request() -> ResponseCacheRequest { + ResponseCacheRequest::new(CacheKeyInput { + preset: Some("tenant:key".into()), + ..Default::default() + }) +} + +#[tokio::test] +async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { + let clock = Arc::new(AtomicU64::new(100)); + let backend = Arc::new(InMemoryCache::with_clock( + Some(8), + Some(Duration::from_secs(600)), + { + let clock = clock.clone(); + move || Duration::from_secs(clock.load(Ordering::SeqCst)) + }, + )); + let cache = ResponseCache::new(backend.clone()); + let mut request = request(); + request.kwargs.ttl = Some(Duration::from_secs(10)); + request.max_age = Some(Duration::from_secs(5)); + cache + .store( + &request, + json!({"choices": [1], "usage": {"total_tokens": 7}}), + Duration::from_secs(100), + ) + .unwrap(); + assert_eq!( + backend.expires_at("tenant:key").unwrap(), + Some(Duration::from_secs(110)) + ); + assert!( + cache + .async_lookup(&request, Duration::from_secs(105)) + .await + .unwrap() + .is_some() + ); + assert_eq!( + cache.lookup(&request, Duration::from_secs(106)).unwrap(), + None + ); + request.max_age = None; + assert_eq!( + cache + .lookup(&request, Duration::from_secs(106)) + .unwrap() + .unwrap()["usage"]["total_tokens"], + 7 + ); + clock.store(111, Ordering::SeqCst); + assert_eq!( + cache + .async_lookup(&request, Duration::from_secs(111)) + .await + .unwrap(), + None + ); + cache + .async_store(&request, json!({"choices": [2]}), Duration::from_secs(111)) + .await + .unwrap(); + assert_eq!( + cache.lookup(&request, Duration::from_secs(111)).unwrap(), + Some(json!({"choices": [2]})) + ); +} + +#[tokio::test] +async fn directives_skip_io_and_keep_reads_and_writes_independent() { + let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let mut request = request(); + let now = Duration::from_secs(100); + request.controls.no_store = true; + cache + .async_store(&request, json!({"v": 1}), now) + .await + .unwrap(); + assert_eq!(cache.lookup(&request, now).unwrap(), None); + request.controls.no_store = false; + request.controls.no_cache = true; + cache.store(&request, json!({"v": 2}), now).unwrap(); + assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None); + request.controls.no_cache = false; + assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2}))); + request.controls.default_on = false; + cache.store(&request, json!({"v": 3}), now).unwrap(); + assert_eq!(cache.lookup(&request, now).unwrap(), None); + request.controls.use_cache = true; + assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2}))); + request.controls.supported_call_type = false; + assert_eq!(cache.lookup(&request, now).unwrap(), None); +} + +#[tokio::test] +async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("GET").arg("tenant:key"), + Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()), + ), + MockCmd::new( + redis::cmd("GET").arg("tenant:key"), + Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()), + ), + MockCmd::new( + redis::cmd("SETEX") + .arg("tenant:key") + .arg(600) + .arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()), + Ok("OK"), + ), + ]) + .assert_all_commands_consumed(); + let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec) + .with_namespace(Some("tenant".into())); + let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))); + let request = request(); + let expected = json!({"ok": true, "text": "cached"}); + assert_eq!( + cache.lookup(&request, Duration::from_secs(101)).unwrap(), + Some(expected.clone()) + ); + assert_eq!( + cache + .async_lookup(&request, Duration::from_secs(101)) + .await + .unwrap(), + Some(expected.clone()) + ); + cache + .async_store(&request, expected, Duration::from_secs(100)) + .await + .unwrap(); +} + +#[tokio::test] +async fn captured_enum_keeps_the_selected_backend_for_background_writes() { + let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let captured = original.clone(); + let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let request = request(); + let writer = tokio::spawn({ + let request = request.clone(); + async move { + captured + .async_store( + &request, + json!({"selected": "original"}), + Duration::from_secs(100), + ) + .await + } + }); + writer.await.unwrap().unwrap(); + assert_eq!( + original.lookup(&request, Duration::from_secs(100)).unwrap(), + Some(json!({"selected":"original"})) + ); + assert_eq!( + replacement + .lookup(&request, Duration::from_secs(100)) + .unwrap(), + None + ); +} + +#[test] +fn generated_keys_preserve_namespace_and_explicit_keys() { + let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let key = CacheKeyInput { + fields: vec![CacheKeyField { + name: "model".into(), + value: Some("a".into()), + api_parameter: true, + internal_parameter: false, + }], + namespace: Some("tenant".into()), + ..Default::default() + }; + let generated = ResponseCacheRequest::new(key.clone()); + let explicit = ResponseCacheRequest::new(CacheKeyInput { + preset: Some(litellm_cache::cache_key(&key)), + ..Default::default() + }); + cache + .store(&generated, json!({"value": 7}), Duration::from_secs(100)) + .unwrap(); + assert_eq!( + cache.lookup(&explicit, Duration::from_secs(100)).unwrap(), + Some(json!({"value":7})) + ); +} + +#[test] +fn response_codec_accepts_python_literals_without_executing_code() { + let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#; + let entry = ResponseCacheCodec.decode(bytes).unwrap(); + assert_eq!( + entry.response, + json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]}) + ); + for bytes in [ + b"__import__('os').system('false')".as_slice(), + b"{'timestamp': 'invalid', 'response': {}}", + b"{'timestamp': 1e9999, 'response': {}}", + ] { + assert_eq!( + ResponseCacheCodec.decode(bytes).unwrap_err(), + Error::InvalidEntry + ); + } + let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000)); + assert_eq!( + ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(), + Error::InvalidEntry + ); + assert_eq!( + ResponseCacheCodec + .encode(&CacheEntry { + timestamp: f64::NAN, + response: json!({}) + }) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[tokio::test] +async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() { + let connection = MockRedisConnection::new([MockCmd::new( + redis::cmd("GET").arg("tenant:key"), + Ok(b"invalid".to_vec()), + )]) + .assert_all_commands_consumed(); + let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec); + let cache = ResponseCache::new(Arc::new(backend)); + let mut request = request(); + request.controls.no_cache = true; + assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); + request.controls.no_cache = false; + assert_eq!( + cache + .async_lookup(&request, Duration::ZERO) + .await + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn malformed_memory_entries_are_rejected_by_the_response_consumer() { + let backend = Arc::new(InMemoryCache::default()); + BaseCache::set_cache( + backend.as_ref(), + "tenant:key", + CacheEntry { + timestamp: 100.0, + response: json!("not a serialized response"), + }, + Default::default(), + ) + .unwrap(); + let cache = ResponseCache::new(backend); + assert_eq!( + cache + .lookup(&request(), Duration::from_secs(100)) + .unwrap_err(), + Error::InvalidEntry + ); +} diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 21ebcce29bb..1d1df966ba2 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -27,6 +27,7 @@ pub struct CacheKeyField { } #[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(default)] pub struct CacheKeyInput { pub fields: Vec, pub preset: Option, diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index d447c80f62d..ff3ff6572d4 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -4,4 +4,6 @@ pub enum Error { Unavailable, #[error("invalid cache entry")] InvalidEntry, + #[error("flushing Redis requires an explicit namespace")] + UnscopedFlush, } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a76b069935f..308dfd2dd7a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,9 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true +litellm-cache.workspace = true +litellm-cache-response.workspace = true +serde.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy-python.workspace = true litellm-core.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs new file mode 100644 index 00000000000..32360e72469 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -0,0 +1,210 @@ +use litellm_cache_response::NativeResponseCache; +use litellm_host_python::from_py; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::PyTypeError, + prelude::*, + types::{PyDict, PyTuple, PyType}, +}; +use serde_json::Value; + +use super::NativeCacheHandle; + +struct ClassGuard { + class: Py, + attributes: Vec<(String, Py)>, +} + +struct ObjectGuard { + reference: Py, + classes: Vec, + config_names: &'static [&'static str], + config: Vec, +} + +pub(super) struct FacadeGuard { + outer: ObjectGuard, + backend: ObjectGuard, +} + +impl ObjectGuard { + fn capture( + py: Python<'_>, + object: &Bound<'_, PyAny>, + config_names: &'static [&'static str], + ) -> PyResult { + let classes = object + .get_type() + .getattr("__mro__")? + .cast_into::()? + .iter() + .map(|class| { + let class = class.cast_into::()?; + let attributes = class + .getattr("__dict__")? + .call_method0("items")? + .try_iter()? + .map(|item| item?.extract::<(String, Py)>()) + .collect::>>()?; + Ok(ClassGuard { + class: class.unbind(), + attributes, + }) + }) + .collect::>>()?; + let guard = Self { + reference: py + .import("weakref")? + .getattr("ref")? + .call1((object,))? + .unbind(), + classes, + config_names, + config: Self::config(object, config_names)?, + }; + if !guard.matches(py, object)? { + return Err(PyTypeError::new_err( + "native facade registration requires unmodified built-in methods", + )); + } + Ok(guard) + } + + fn config(object: &Bound<'_, PyAny>, names: &[&str]) -> PyResult> { + names + .iter() + .map(|name| match object.getattr(*name) { + Ok(value) => from_py(&value), + Err(error) + if error.is_instance_of::(object.py()) => + { + Ok(Value::Null) + } + Err(error) => Err(error), + }) + .collect() + } + + fn matches(&self, py: Python<'_>, object: &Bound<'_, PyAny>) -> PyResult { + if !self.reference.bind(py).call0()?.is(object) { + return Ok(false); + } + let mro = object + .get_type() + .getattr("__mro__")? + .cast_into::()?; + if mro.len() != self.classes.len() { + return Ok(false); + } + let instance = object.getattr("__dict__")?.cast_into::()?; + for (class, expected) in mro.iter().zip(&self.classes) { + if !class.is(expected.class.bind(py)) { + return Ok(false); + } + let attributes = class.getattr("__dict__")?; + if attributes.len()? != expected.attributes.len() { + return Ok(false); + } + for (name, value) in &expected.attributes { + if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + return Ok(false); + } + } + } + Ok(Self::config(object, self.config_names)? == self.config) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference)?; + for class in &self.classes { + visit.call(&class.class)?; + for (_, value) in &class.attributes { + visit.call(value)?; + } + } + Ok(()) + } +} + +impl FacadeGuard { + pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult { + let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; + if !facade.get_type().is(&cache_type) { + return Err(PyTypeError::new_err( + "only exact built-in Cache facades can be registered", + )); + } + let (module, name, cache_kind) = match kind { + "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), + "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + _ => unreachable!(), + }; + let backend = facade.getattr("cache")?; + if facade.getattr("type")?.extract::()? != cache_kind + || !backend.get_type().is(&py.import(module)?.getattr(name)?) + { + return Err(PyTypeError::new_err( + "facade and native backend types must match", + )); + } + Ok(Self { + outer: ObjectGuard::capture( + py, + facade, + &[ + "type", + "mode", + "ttl", + "namespace", + "supported_call_types", + "redis_flush_size", + ], + )?, + backend: ObjectGuard::capture( + py, + &backend, + &[ + "namespace", + "default_ttl", + "max_size_in_memory", + "max_size_per_item", + ], + )?, + }) + } + + fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { + Ok(self.outer.matches(py, facade)? + && self.backend.matches(py, &facade.getattr("cache")?)?) + } + + pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.outer.traverse(&visit)?; + self.backend.traverse(&visit) + } +} + +pub(super) fn resolve( + py: Python<'_>, + facade: &Bound<'_, PyAny>, +) -> PyResult> { + let Ok(dict) = facade + .getattr("__dict__") + .and_then(|dict| dict.cast_into::().map_err(Into::into)) + else { + return Ok(None); + }; + let Some(handle) = dict.get_item("_native_cache_handle")? else { + return Ok(None); + }; + let Ok(handle) = handle.extract::>() else { + return Ok(None); + }; + let Some(guard) = &handle.guard else { + return Ok(None); + }; + if !guard.matches(py, facade).unwrap_or(false) { + return Ok(None); + } + handle.service().map(Some) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs new file mode 100644 index 00000000000..f00cceeb86e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -0,0 +1,350 @@ +mod facade; + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache::{CacheControls, CacheKeyInput, Error}; +use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest}; +use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + prelude::*, + types::PyDict, +}; +use serde::Deserialize; +use serde_json::Value; + +use facade::FacadeGuard; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestInput { + key: CacheKeyInput, + controls: Option, + ttl_seconds: Option, + max_age_seconds: Option, +} + +fn request(value: &Bound<'_, PyAny>) -> PyResult { + let input: RequestInput = from_py(value)?; + let mut request = ResponseCacheRequest::new(input.key); + if let Some(controls) = input.controls { + request.controls = controls; + } + request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; + request.max_age = input.max_age_seconds.map(duration).transpose()?; + Ok(request) +} + +fn duration(seconds: f64) -> PyResult { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) +} + +fn now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() +} + +fn cache_error(error: Error) -> PyErr { + match error { + Error::InvalidEntry => PyValueError::new_err(error.to_string()), + _ => PyRuntimeError::new_err(error.to_string()), + } +} + +#[pyclass(frozen)] +pub(crate) struct NativeCacheHandle { + service: NativeResponseCache, + guard: Option, + pid: u32, +} + +impl NativeCacheHandle { + fn service(&self) -> PyResult { + if self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache handles must be recreated after fork", + )); + } + Ok(self.service.clone()) + } +} + +#[pymethods] +impl NativeCacheHandle { + #[staticmethod] + #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] + fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { + Ok(Self { + service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), + guard: None, + pid: std::process::id(), + }) + } + + #[staticmethod] + #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + fn redis( + py: Python<'_>, + url: String, + ttl_seconds: Option, + namespace: Option, + ) -> PyResult { + let ttl = ttl_seconds.map(duration).transpose()?; + let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + + #[getter] + fn backend(&self) -> &'static str { + self.service.kind() + } + + fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { + let service = self.service()?; + let guard = FacadeGuard::capture(py, facade, self.backend())?; + let handle = Py::new( + py, + Self { + service, + guard: Some(guard), + pid: self.pid, + }, + )?; + facade.setattr("_native_cache_handle", handle) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(guard) = &self.guard { + guard.traverse(visit)?; + } + Ok(()) + } +} + +enum CacheBinding { + Disabled, + Native(NativeResponseCache), + PythonCallback(Py), +} + +#[pyclass(frozen, name = "CacheBinding")] +pub(crate) struct ResolvedCache { + binding: CacheBinding, + pid: u32, +} + +impl ResolvedCache { + fn check_process(&self) -> PyResult<()> { + if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache bindings must be resolved again after fork", + )); + } + Ok(()) + } + + pub(crate) fn lookup_step( + &self, + py: Python<'_>, + input: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.check_process()?; + let awaitable = match &self.binding { + CacheBinding::Disabled => ready_none(py)?, + CacheBinding::Native(service) => { + let request = request(input)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup(&request, now()).await }, + cache_error, + )? + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_get_cache", + (), + Some(callback_kwargs(kwargs)?), + )?, + }; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +#[pymethods] +impl ResolvedCache { + #[getter] + fn kind(&self) -> &'static str { + match self.binding { + CacheBinding::Disabled => "disabled", + CacheBinding::Native(_) => "native", + CacheBinding::PythonCallback(_) => "python_callback", + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn lookup( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(py.None()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup(&request, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(Bound::unbind), + } + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn store( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + response: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + release_gil(py, move || service.store(&request, response, now())) + .map_err(cache_error) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "add_cache", + (response,), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(|_| ()), + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn async_lookup<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? + else { + unreachable!() + }; + Ok(awaitable.into_bound(py)) + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn async_store<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + response: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + run_async( + py, + async move { service.async_store(&request, response, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_add_cache", + (response,), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let CacheBinding::PythonCallback(object) = &self.binding { + visit.call(object)?; + } + Ok(()) + } +} + +fn callback_kwargs<'a, 'py>( + kwargs: Option<&'a Bound<'py, PyDict>>, +) -> PyResult<&'a Bound<'py, PyDict>> { + kwargs.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") + }) +} + +fn ready_none(py: Python<'_>) -> PyResult> { + let future = py + .import("asyncio")? + .call_method0("get_running_loop")? + .call_method0("create_future")?; + future.call_method1("set_result", (py.None(),))?; + Ok(future) +} + +#[pyclass(frozen)] +pub(crate) struct CacheResolver { + namespace: Py, +} + +#[pymethods] +impl CacheResolver { + #[new] + fn new(namespace: Py) -> Self { + Self { namespace } + } + + pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { + let object = self.namespace.bind(py).getattr("cache")?; + let binding = if object.is_none() { + CacheBinding::Disabled + } else if let Ok(handle) = object.extract::>() { + CacheBinding::Native(handle.service()?) + } else if let Some(service) = facade::resolve(py, &object)? { + CacheBinding::Native(service) + } else { + CacheBinding::PythonCallback(object.unbind()) + }; + Ok(ResolvedCache { + binding, + pid: std::process::id(), + }) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.namespace) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 46f98736aa1..621c111a35b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,3 +1,4 @@ +mod cache; mod credentials; mod diagnostics; mod errors; @@ -9,6 +10,8 @@ mod token_counter; #[pymodule(gil_used = true)] mod _native { + #[pymodule_export] + use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache}; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -65,6 +68,9 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "TokenCounter", + "CacheResolver", + "NativeCacheHandle", + "CacheBinding", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..4fd2f0829a3 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,5 +1,5 @@ from asyncio import Future -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -93,6 +93,50 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... +@final +class NativeCacheHandle: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @staticmethod + def memory( + *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576 + ) -> NativeCacheHandle: ... + @staticmethod + def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ... + @property + def backend(self) -> str: ... + def bind_facade(self, facade: object) -> None: ... + +@final +class CacheResolver: + def __new__(cls, namespace: object) -> CacheResolver: ... + def resolve(self) -> CacheBinding: ... + +@final +class CacheBinding: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @property + def kind(self) -> str: ... + def lookup( + self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None + ) -> object: ... + def store( + self, + request: Mapping[str, object] | None, + response: object, + *, + callback_kwargs: dict[str, object] | None = None, + ) -> None: ... + def async_lookup( + self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None + ) -> Awaitable[object]: ... + def async_store( + self, + request: Mapping[str, object] | None, + response: object, + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @@ -109,7 +153,10 @@ def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... __all__ = [ + "CacheBinding", + "CacheResolver", "ForkedAfterNativeRuntimeStarted", + "NativeCacheHandle", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py new file mode 100644 index 00000000000..d1cea860cb0 --- /dev/null +++ b/tests/test_litellm_rust/test_cache.py @@ -0,0 +1,231 @@ +import asyncio +import contextvars +import gc +import json +import threading +import time +import weakref +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final, Protocol, cast + +import fakeredis +import pytest +import redis + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +class CacheLookup(Protocol): + def get_cache(self, **kwargs: object) -> object: ... + + +def request(key: str = "key") -> dict[str, object]: + return {"key": {"preset": key}} + + +@pytest.fixture +def redis_url() -> Generator[str]: + server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + try: + yield f"redis://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + +def test_existing_constructor_and_global_are_unchanged() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + assert type(facade.cache) is InMemoryCache + assert "_native_cache_handle" not in vars(facade) + with rebound(litellm, "cache", facade): + resolver: Final = _native.CacheResolver(litellm) + assert resolver.resolve().kind == "python_callback" + resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) + assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} + + +async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: + namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory()) + resolver: Final = _native.CacheResolver(namespace) + selected: Final = resolver.resolve() + assert selected.kind == "native" + selected.store(request(), {"answer": 1}) + assert await selected.async_lookup(request()) == {"answer": 1} + with rebound(namespace, "cache", _native.NativeCacheHandle.memory()): + replacement: Final = resolver.resolve() + await selected.async_store(request(), {"answer": 2}) + assert replacement.lookup(request()) is None + assert selected.lookup(request()) == {"answer": 2} + with rebound(namespace, "cache", None): + disabled: Final = resolver.resolve() + assert disabled.kind == "disabled" + assert disabled.lookup(None) is None + await disabled.async_store(None, object()) + assert await disabled.async_lookup(None) is None + assert selected.lookup(request()) == {"answer": 2} + + +async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None: + context: Final = contextvars.ContextVar("cache_context", default="caller") + caller: Final = asyncio.current_task() + sentinel: Final = object() + failure: Final = RuntimeError("callback failed") + + class CustomCache: + async def async_get_cache(self, *, marker: object) -> object: + assert marker is sentinel + assert asyncio.current_task() is caller + context.set("callback") + return marker + + async def async_add_cache(self, response: object, *, marker: object) -> None: + assert response is sentinel + assert marker is sentinel + raise failure + + namespace: Final = SimpleNamespace(cache=CustomCache()) + binding: Final = _native.CacheResolver(namespace).resolve() + assert binding.kind == "python_callback" + assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel + assert context.get() == "callback" + with pytest.raises(RuntimeError) as caught: + await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel}) + assert caught.value is failure + + +async def test_callback_cancellation_stays_in_the_callers_task() -> None: + entered: Final = asyncio.Event() + finished: Final = asyncio.Event() + + class CustomCache: + async def async_get_cache(self) -> None: + entered.set() + try: + await asyncio.Future() + finally: + finished.set() + + binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() + + async def lookup() -> object: + return await binding.async_lookup(None, callback_kwargs={}) + + task: Final = asyncio.create_task(lookup()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + + +def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle: Final = _native.NativeCacheHandle.memory() + handle.bind_facade(facade) + resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + native.store(request(), {"source": "native"}) + assert native.lookup(request()) == {"source": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="key") is None + sentinel: Final = object() + + def outer_override(**_kwargs: object) -> object: + return sentinel + + def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]: + return {"source": "override"} + + with rebound(facade, "get_cache", outer_override): + fallback: Final = resolver.resolve() + assert fallback.kind == "python_callback" + assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache") + assert resolver.resolve().kind == "native" + with rebound(facade.cache, "get_cache", backend_override): + backend_fallback: Final = resolver.resolve() + assert backend_fallback.kind == "python_callback" + assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"} + + +def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None: + class CustomCache(Cache): + pass + + handle: Final = _native.NativeCacheHandle.memory() + with pytest.raises(TypeError): + handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle.bind_facade(facade) + resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + with rebound(facade, "cache", InMemoryCache()): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + def custom_key(**_kwargs: object) -> str: + return "custom" + + with rebound(facade, "get_cache_key", custom_key): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache_key") + assert resolver.resolve().kind == "native" + + +def test_resolver_and_callback_cycles_can_be_collected() -> None: + class CustomCache: + pass + + def cyclic_reference() -> weakref.ReferenceType[CustomCache]: + callback: Final = CustomCache() + namespace: Final = SimpleNamespace(cache=callback) + binding: Final = _native.CacheResolver(namespace).resolve() + setattr(callback, "binding", binding) + return weakref.ref(callback) + + reference: Final = cyclic_reference() + gc.collect() + assert reference() is None + + +async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team")) + binding: Final = _native.CacheResolver(namespace).resolve() + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} + client.set("team:sync", str(envelope)) + client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("team:async")) == response + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = client.get("team:native") + assert isinstance(stored, bytes) + assert json.loads(stored)["response"] == response + assert 0 < client.ttl("team:native") <= 12 + assert client.get("litellm-cache:team:native") is None + assert client.get("team:team:async") is None + client.close() + + +def test_invalid_duration_and_request_shape_fail_before_storage() -> None: + binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + for seconds in (-1.0, float("nan"), float("inf")): + with pytest.raises(ValueError): + binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) + assert binding.lookup(request()) is None + with pytest.raises(ValueError): + _native.NativeCacheHandle.memory(ttl_seconds=-1) From 0c3a0a208948e97d0da05ec6c7e28672205c2f00 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:11:48 -0700 Subject: [PATCH 21/56] refactor(cache): separate response policy and host selection --- litellm-rust/Cargo.lock | 4 +- litellm-rust/crates/cache-memory/Cargo.toml | 2 +- litellm-rust/crates/cache-memory/src/cache.rs | 47 +----- .../crates/cache-memory/tests/cache.rs | 80 ++++------ litellm-rust/crates/cache-redis/src/cache.rs | 34 +++-- litellm-rust/crates/cache-response/Cargo.toml | 7 +- litellm-rust/crates/cache-response/README.md | 55 +++++++ .../crates/cache-response/src/caching.rs | 143 ++++++++++++++++++ .../crates/cache-response/src/codec.rs | 4 +- litellm-rust/crates/cache-response/src/lib.rs | 7 +- .../crates/cache-response/src/response.rs | 6 +- .../crates/cache-response/tests/caching.rs | 83 ++++++++++ .../crates/cache-response/tests/response.rs | 40 +++-- litellm-rust/crates/cache/Cargo.toml | 1 - litellm-rust/crates/cache/src/caching.rs | 143 ------------------ litellm-rust/crates/cache/src/lib.rs | 5 +- litellm-rust/crates/cache/tests/caching.rs | 94 +----------- litellm-rust/crates/cache/tests/codec.rs | 14 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/facade.rs | 3 +- .../crates/python-bridge/src/cache/mod.rs | 6 +- .../src => python-bridge/src/cache}/native.rs | 33 ++-- tests/test_litellm_rust/test_cache.py | 20 ++- 23 files changed, 426 insertions(+), 407 deletions(-) create mode 100644 litellm-rust/crates/cache-response/README.md create mode 100644 litellm-rust/crates/cache-response/src/caching.rs create mode 100644 litellm-rust/crates/cache-response/tests/caching.rs rename litellm-rust/crates/{cache-response/src => python-bridge/src/cache}/native.rs (75%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index aa62e4f3770..8b299f5455b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2460,7 +2460,6 @@ dependencies = [ "rstest", "serde", "serde_json", - "sha2 0.10.9", "thiserror 2.0.19", "tokio", ] @@ -2498,6 +2497,7 @@ dependencies = [ "redis-test", "serde", "serde_json", + "sha2 0.10.9", "tokio", ] @@ -2665,6 +2665,8 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-memory", + "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", diff --git a/litellm-rust/crates/cache-memory/Cargo.toml b/litellm-rust/crates/cache-memory/Cargo.toml index d4487573a9a..86ab01564c8 100644 --- a/litellm-rust/crates/cache-memory/Cargo.toml +++ b/litellm-rust/crates/cache-memory/Cargo.toml @@ -7,8 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -serde_json.workspace = true [dev-dependencies] +serde_json.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 974cdbe9760..43186faf3f7 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -3,15 +3,12 @@ use std::collections::{BinaryHeap, HashMap}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, -}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); type ValueMeasure = Arc Result + Send + Sync>; -type ValueValidator = Arc Result<(), Error> + Send + Sync>; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CacheWrite { @@ -32,7 +29,6 @@ pub struct InMemoryCache { default_ttl: Duration, max_entry_bytes: Option, measure_value: Option>, - validate_value: Option>, now: Arc Duration + Send + Sync>, } @@ -76,7 +72,6 @@ impl InMemoryCache { default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), max_entry_bytes, measure_value, - validate_value: None, now: Arc::new(now), } } @@ -90,9 +85,6 @@ impl InMemoryCache { if self.max_size_in_memory == 0 { return Ok(CacheWrite::Disabled); } - if let Some(validate) = &self.validate_value { - validate(&value)?; - } if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value) && measure(&value)? > limit { @@ -176,43 +168,6 @@ impl InMemoryCache { } } -impl InMemoryCache { - pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { - Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - }) - } - - pub fn response_cache_with_clock( - capacity: usize, - ttl: Duration, - max_entry_bytes: usize, - now: impl Fn() -> Duration + Send + Sync + 'static, - ) -> Self { - let mut cache = Self::with_clock_and_size_measurement( - Some(capacity), - Some(ttl), - Some(max_entry_bytes), - Some(Arc::new(|entry: &CacheEntry| { - serde_json::to_vec(entry) - .map(|bytes| bytes.len()) - .map_err(|_| Error::InvalidEntry) - })), - now, - ); - cache.validate_value = Some(Arc::new(|entry: &CacheEntry| { - entry - .timestamp - .is_finite() - .then_some(()) - .ok_or(Error::InvalidEntry) - })); - cache - } -} - impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index ffac9d8ae64..370145bebef 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -3,8 +3,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache, - set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -87,66 +86,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc } #[test] -fn disabled_size_limited_and_synchronized_response_writes_are_observable() { - let disabled = InMemoryCache::::response_cache(0, Duration::from_secs(60), 80); +fn disabled_size_limited_and_validated_writes_are_observable() { + let cache = |capacity| { + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(Duration::from_secs(60)), + Some(4), + Some(Arc::new(|value: &String| { + if value.is_empty() { + return Err(Error::InvalidEntry); + } + Ok(value.len()) + })), + || Duration::from_secs(100), + ) + }; + let disabled = cache(0); assert_eq!( - disabled - .set_cache( - "a", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("x") - }, - None - ) - .unwrap(), + disabled.set_cache("a", "x".into(), None).unwrap(), CacheWrite::Disabled ); - let cache = InMemoryCache::::response_cache(2, Duration::from_secs(60), 80); + let cache = cache(2); assert_eq!( - cache - .set_cache( - "large", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("x".repeat(100)) - }, - None - ) - .unwrap(), + cache.set_cache("large", "oversized".into(), None).unwrap(), CacheWrite::TooLarge ); - cache - .set_cache( - "small", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("ok"), - }, - None, - ) - .unwrap(); - assert!(cache.get_cache("small").unwrap().is_some()); + assert_eq!(cache.get_cache("large").unwrap(), None); assert_eq!( - cache - .set_cache( - "invalid", - CacheEntry { - timestamp: f64::NAN, - response: serde_json::json!("bad"), - }, - None, - ) - .unwrap_err(), - Error::InvalidEntry + cache.set_cache("small", "ok".into(), None).unwrap(), + CacheWrite::Stored ); + assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into())); + assert_eq!( + cache.set_cache("invalid", String::new(), None), + Err(Error::InvalidEntry) + ); + assert_eq!(cache.get_cache("invalid").unwrap(), None); cache.delete_cache("small").unwrap(); - cache.flush_cache().unwrap(); + assert_eq!(cache.get_cache("small").unwrap(), None); } #[tokio::test] async fn connection_test_matches_python_result_contract() { - let cache = InMemoryCache::::default(); + let cache = InMemoryCache::::default(); let result = BaseCache::test_connection(&cache).await.unwrap(); assert_eq!(result.status, CacheConnectionStatus::Success); assert_eq!(result.message, "In-memory cache connection test successful"); diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 0faca6cdaaf..d4ca0cf0522 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -238,30 +238,27 @@ where #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec}; + use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; - fn entry() -> CacheEntry { - CacheEntry { - timestamp: 123.0, - response: json!({"choices": [{"text": "cached"}]}), - } + fn entry() -> serde_json::Value { + json!({"deployment": "model-a", "cooldown_seconds": 30}) } #[test] fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { assert_eq!( - RedisCache::>::ttl_seconds(Duration::ZERO), + RedisCache::>::ttl_seconds(Duration::ZERO), 1 ); assert_eq!( - RedisCache::>::ttl_seconds(Duration::from_millis(1500)), + RedisCache::>::ttl_seconds(Duration::from_millis(1500)), 2 ); assert_eq!( - RedisCache::>::ttl_seconds(Duration::from_secs(15)), + RedisCache::>::ttl_seconds(Duration::from_secs(15)), 15 ); } @@ -269,7 +266,9 @@ mod tests { #[test] fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { let value = entry(); - let payload = JsonCodec::::new().encode(&value).unwrap(); + let payload = JsonCodec::::new() + .encode(&value) + .unwrap(); let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") @@ -282,8 +281,9 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -308,8 +308,9 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache.flush_cache().unwrap(); } @@ -318,8 +319,9 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); assert_eq!( cache.test_connection().await.unwrap().status, diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index a0c4a1f74ef..04affb9872d 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -7,13 +7,14 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true py_literal = "0.4.0" -redis = "1.7.0" serde.workspace = true serde_json.workspace = true +sha2.workspace = true [dev-dependencies] +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true +redis = "1.7.0" redis-test = "1.0.4" tokio.workspace = true diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md new file mode 100644 index 00000000000..309296f7773 --- /dev/null +++ b/litellm-rust/crates/cache-response/README.md @@ -0,0 +1,55 @@ +# Response cache foundation + +`ResponseCache` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache` + +## Ownership + +`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations + +`litellm-cache-response` owns response keys, controls, entries, and the Python-compatible response codec. It has no runtime dependency on a specific cache backend or Python + +The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host + +## Native Rust use + +```rust +use std::{sync::Arc, time::Duration}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest}; +use serde_json::json; + +let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); +let request = ResponseCacheRequest::new(CacheKeyInput { + preset: Some("example:key".into()), + ..Default::default() +}); +let now = Duration::from_secs(100); +cache.store(&request, json!({"answer": 7}), now)?; +assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); +``` + +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Each Redis constructor currently opens its own connection; shared connection pools remain follow-up work + +Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it + +## Python integration boundary + +The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support synchronous and asynchronous response lookup and storage + +The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution + +Explicit facade registration checks object identity, method overrides, and configuration changes before selecting native execution. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python + +Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy + +## Adding another backend + +Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python + +Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade + +## Follow-up scope + +Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial batches, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths + +Redis cluster, disk, cloud stores, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub 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/caching.rs b/litellm-rust/crates/cache-response/src/caching.rs new file mode 100644 index 00000000000..53b34025ccd --- /dev/null +++ b/litellm-rust/crates/cache-response/src/caching.rs @@ -0,0 +1,143 @@ +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +pub enum CacheMode { + #[default] + #[serde(rename = "default_on")] + DefaultOn, + #[serde(rename = "default_off")] + DefaultOff, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct CacheKeyField { + pub name: String, + pub value: Option, + pub api_parameter: bool, + pub internal_parameter: bool, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(default)] +pub struct CacheKeyInput { + pub fields: Vec, + pub preset: Option, + pub namespace: Option, + pub include_provider_parameters: bool, +} + +#[derive(Default)] +pub struct CacheKeyContext { + pub model_group: Option, + pub caching_groups: Vec<(Vec, String)>, + pub file_checksum: Option, + pub file_object_name: Option, + pub metadata_file_name: Option, + pub parameters_file_name: Option, +} + +impl CacheKeyContext { + pub fn apply(self, input: &mut CacheKeyInput) { + let group = self.model_group.as_ref().and_then(|model| { + self.caching_groups + .iter() + .find(|(models, _)| models.contains(model)) + }); + for field in &mut input.fields { + match field.name.as_str() { + "model" => { + field.value = group + .map(|(_, formatted)| formatted.clone()) + .or_else(|| self.model_group.clone()) + .or_else(|| field.value.take()) + } + "file" => { + field.value = self + .file_checksum + .clone() + .or_else(|| self.file_object_name.clone()) + .or_else(|| self.metadata_file_name.clone()) + .or_else(|| self.parameters_file_name.clone()) + } + _ => {} + } + } + } +} + +pub fn get_cache_key(input: &CacheKeyInput) -> String { + cache_key(input) +} + +pub fn cache_key(input: &CacheKeyInput) -> String { + if let Some(preset) = &input.preset { + return preset.clone(); + } + let mut digest = Sha256::new(); + for field in &input.fields { + if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter)) + && let Some(value) = &field.value + { + digest.update(field.name.as_bytes()); + digest.update(b": "); + digest.update(value.as_bytes()); + } + } + let hash = format!("{:x}", digest.finalize()); + input + .namespace + .as_deref() + .filter(|namespace| !namespace.is_empty()) + .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}")) +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] +pub struct CacheControls { + pub supported_call_type: bool, + pub configured: bool, + pub native_backend: bool, + pub default_on: bool, + pub caching: Option, + pub no_cache: bool, + pub no_store: bool, + #[serde(default)] + pub use_cache: bool, +} + +impl CacheControls { + pub fn reads(self) -> bool { + self.supported_call_type + && self.configured + && self.caching.unwrap_or(true) + && !self.no_cache + && (self.default_on || self.use_cache) + } + + pub fn writes(self) -> bool { + self.supported_call_type + && self.configured + && !self.no_store + && (self.default_on || self.use_cache) + } +} + +pub fn should_use_cache(controls: CacheControls) -> bool { + controls.reads() || controls.writes() +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CacheEntry { + pub timestamp: f64, + pub response: Value, +} + +impl CacheEntry { + pub fn fresh(&self, now: Duration, max_age: Option) -> bool { + self.timestamp.is_finite() + && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) + } +} diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index 137cf61267a..f1d55ddefe8 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -1,4 +1,6 @@ -use litellm_cache::{CacheCodec, CacheEntry, Error}; +use litellm_cache::{CacheCodec, Error}; + +use crate::CacheEntry; use serde_json::Value; pub struct ResponseCacheCodec; diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 454a82e76de..efa0b04b9f7 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,7 +1,10 @@ +mod caching; mod codec; -mod native; mod response; +pub use caching::{ + CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, + get_cache_key, should_use_cache, +}; pub use codec::ResponseCacheCodec; -pub use native::NativeResponseCache; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 3e9807b6d1b..987a47f0554 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{ - BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key, -}; +use litellm_cache::{BaseCache, CacheKwargs, Error}; + +use crate::{CacheControls, CacheEntry, CacheKeyInput, cache_key}; use serde_json::Value; #[derive(Clone)] diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs new file mode 100644 index 00000000000..d403791e421 --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -0,0 +1,83 @@ +use litellm_cache_response::{ + CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key, +}; +use sha2::{Digest, Sha256}; + +#[test] +fn keys_match_python_order_groups_files_presets_and_namespaces() { + let mut input = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".into(), + value: Some("deployment".into()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "file".into(), + value: None, + api_parameter: true, + internal_parameter: false, + }, + ], + namespace: Some("team".into()), + ..Default::default() + }; + CacheKeyContext { + model_group: Some("group".into()), + caching_groups: vec![(vec!["group".into()], "['group']".into())], + file_checksum: Some("checksum".into()), + ..Default::default() + } + .apply(&mut input); + assert_eq!( + cache_key(&input), + format!( + "team:{:x}", + Sha256::digest(b"model: ['group']file: checksum") + ) + ); + input.preset = Some("preset".into()); + assert_eq!(get_cache_key(&input), "preset"); +} + +#[test] +fn cache_controls_honor_default_modes_and_directives() { + let enabled = CacheControls { + supported_call_type: true, + configured: true, + default_on: true, + ..Default::default() + }; + assert!(enabled.reads()); + assert!(enabled.writes()); + assert!( + !CacheControls { + default_on: false, + ..enabled + } + .reads() + ); + assert!( + CacheControls { + default_on: false, + use_cache: true, + ..enabled + } + .reads() + ); + assert!( + !CacheControls { + no_cache: true, + ..enabled + } + .reads() + ); + assert!( + !CacheControls { + no_store: true, + ..enabled + } + .writes() + ); +} diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index a4beda9fa4c..88b53fdfe86 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -6,15 +6,23 @@ use std::{ time::Duration, }; -use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error}; +use litellm_cache::{BaseCache, CacheCodec, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; +fn memory() -> Arc>> { + Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new( + Some(8), + Some(Duration::from_secs(600)), + )))) +} + fn request() -> ResponseCacheRequest { ResponseCacheRequest::new(CacheKeyInput { preset: Some("tenant:key".into()), @@ -87,7 +95,7 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { #[tokio::test] async fn directives_skip_io_and_keep_reads_and_writes_independent() { - let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let cache = memory(); let mut request = request(); let now = Duration::from_secs(100); request.controls.no_store = true; @@ -112,7 +120,7 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent() { } #[tokio::test] -async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { +async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("GET").arg("tenant:key"), @@ -133,7 +141,7 @@ async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_ .assert_all_commands_consumed(); let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec) .with_namespace(Some("tenant".into())); - let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))); + let cache = ResponseCache::new(Arc::new(backend)); let request = request(); let expected = json!({"ok": true, "text": "cached"}); assert_eq!( @@ -154,10 +162,10 @@ async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_ } #[tokio::test] -async fn captured_enum_keeps_the_selected_backend_for_background_writes() { - let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); +async fn captured_service_keeps_the_selected_backend_for_background_writes() { + let original = memory(); let captured = original.clone(); - let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let replacement = memory(); let request = request(); let writer = tokio::spawn({ let request = request.clone(); @@ -186,7 +194,7 @@ async fn captured_enum_keeps_the_selected_backend_for_background_writes() { #[test] fn generated_keys_preserve_namespace_and_explicit_keys() { - let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let cache = memory(); let key = CacheKeyInput { fields: vec![CacheKeyField { name: "model".into(), @@ -199,7 +207,7 @@ fn generated_keys_preserve_namespace_and_explicit_keys() { }; let generated = ResponseCacheRequest::new(key.clone()); let explicit = ResponseCacheRequest::new(CacheKeyInput { - preset: Some(litellm_cache::cache_key(&key)), + preset: Some(litellm_cache_response::cache_key(&key)), ..Default::default() }); cache @@ -288,3 +296,15 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { Error::InvalidEntry ); } + +#[test] +fn response_entries_preserve_the_existing_json_representation() { + let codec = ResponseCacheCodec; + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + let bytes = codec.encode(&entry).unwrap(); + assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); + assert_eq!(codec.decode(&bytes).unwrap(), entry); +} diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 350db4b1adb..0c504ab727a 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -8,7 +8,6 @@ repository.workspace = true [dependencies] serde.workspace = true serde_json.workspace = true -sha2.workspace = true thiserror.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 1d1df966ba2..39479694f3b 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -1,152 +1,9 @@ use std::sync::Arc; -use std::time::Duration; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use sha2::{Digest, Sha256}; use crate::{BaseCache, CacheKwargs, Error}; pub use crate::BaseCache as Cache; -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] -pub enum CacheMode { - #[default] - #[serde(rename = "default_on")] - DefaultOn, - #[serde(rename = "default_off")] - DefaultOff, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct CacheKeyField { - pub name: String, - pub value: Option, - pub api_parameter: bool, - pub internal_parameter: bool, -} - -#[derive(Clone, Debug, Default, Deserialize, Serialize)] -#[serde(default)] -pub struct CacheKeyInput { - pub fields: Vec, - pub preset: Option, - pub namespace: Option, - pub include_provider_parameters: bool, -} - -#[derive(Default)] -pub struct CacheKeyContext { - pub model_group: Option, - pub caching_groups: Vec<(Vec, String)>, - pub file_checksum: Option, - pub file_object_name: Option, - pub metadata_file_name: Option, - pub parameters_file_name: Option, -} - -impl CacheKeyContext { - pub fn apply(self, input: &mut CacheKeyInput) { - let group = self.model_group.as_ref().and_then(|model| { - self.caching_groups - .iter() - .find(|(models, _)| models.contains(model)) - }); - for field in &mut input.fields { - match field.name.as_str() { - "model" => { - field.value = group - .map(|(_, formatted)| formatted.clone()) - .or_else(|| self.model_group.clone()) - .or_else(|| field.value.take()) - } - "file" => { - field.value = self - .file_checksum - .clone() - .or_else(|| self.file_object_name.clone()) - .or_else(|| self.metadata_file_name.clone()) - .or_else(|| self.parameters_file_name.clone()) - } - _ => {} - } - } - } -} - -pub fn get_cache_key(input: &CacheKeyInput) -> String { - cache_key(input) -} - -pub fn cache_key(input: &CacheKeyInput) -> String { - if let Some(preset) = &input.preset { - return preset.clone(); - } - let mut digest = Sha256::new(); - for field in &input.fields { - if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter)) - && let Some(value) = &field.value - { - digest.update(field.name.as_bytes()); - digest.update(b": "); - digest.update(value.as_bytes()); - } - } - let hash = format!("{:x}", digest.finalize()); - input - .namespace - .as_deref() - .filter(|namespace| !namespace.is_empty()) - .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}")) -} - -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] -pub struct CacheControls { - pub supported_call_type: bool, - pub configured: bool, - pub native_backend: bool, - pub default_on: bool, - pub caching: Option, - pub no_cache: bool, - pub no_store: bool, - #[serde(default)] - pub use_cache: bool, -} - -impl CacheControls { - pub fn reads(self) -> bool { - self.supported_call_type - && self.configured - && self.caching.unwrap_or(true) - && !self.no_cache - && (self.default_on || self.use_cache) - } - - pub fn writes(self) -> bool { - self.supported_call_type - && self.configured - && !self.no_store - && (self.default_on || self.use_cache) - } -} - -pub fn should_use_cache(controls: CacheControls) -> bool { - controls.reads() || controls.writes() -} - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct CacheEntry { - pub timestamp: f64, - pub response: Value, -} - -impl CacheEntry { - pub fn fresh(&self, now: Duration, max_age: Option) -> bool { - self.timestamp.is_finite() - && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) - } -} - pub fn get_cache( cache: &B, key: &str, diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index a1d9d1402bb..4ff02319bdc 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -4,9 +4,6 @@ mod codec; mod error; pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; -pub use caching::{ - Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, - CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache, -}; +pub use caching::{Cache, CacheBackend, get_cache, set_cache}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 5c250c6b3c9..824de00bdf4 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,17 +1,13 @@ -use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, - CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, -}; -use sha2::{Digest, Sha256}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; use std::{sync::Mutex, time::Duration}; struct TestCache { default_ttl: Duration, - writes: Mutex>, + writes: Mutex>, } impl BaseCache for TestCache { - type Value = CacheEntry; + type Value = String; fn default_ttl(&self) -> Duration { self.default_ttl @@ -83,10 +79,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { default_ttl: Duration::from_secs(60), writes: Mutex::default(), }; - let entry = CacheEntry { - timestamp: 123.0, - response: serde_json::json!("cached"), - }; + let entry = String::from("cached"); let kwargs = CacheKwargs { ttl: Some(Duration::from_secs(5)), ..Default::default() @@ -116,82 +109,3 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { ] ); } - -#[test] -fn keys_match_python_order_groups_files_presets_and_namespaces() { - let mut input = CacheKeyInput { - fields: vec![ - CacheKeyField { - name: "model".into(), - value: Some("deployment".into()), - api_parameter: true, - internal_parameter: false, - }, - CacheKeyField { - name: "file".into(), - value: None, - api_parameter: true, - internal_parameter: false, - }, - ], - namespace: Some("team".into()), - ..Default::default() - }; - CacheKeyContext { - model_group: Some("group".into()), - caching_groups: vec![(vec!["group".into()], "['group']".into())], - file_checksum: Some("checksum".into()), - ..Default::default() - } - .apply(&mut input); - assert_eq!( - cache_key(&input), - format!( - "team:{:x}", - Sha256::digest(b"model: ['group']file: checksum") - ) - ); - input.preset = Some("preset".into()); - assert_eq!(get_cache_key(&input), "preset"); -} - -#[test] -fn cache_controls_honor_default_modes_and_directives() { - let enabled = CacheControls { - supported_call_type: true, - configured: true, - default_on: true, - ..Default::default() - }; - assert!(enabled.reads()); - assert!(enabled.writes()); - assert!( - !CacheControls { - default_on: false, - ..enabled - } - .reads() - ); - assert!( - CacheControls { - default_on: false, - use_cache: true, - ..enabled - } - .reads() - ); - assert!( - !CacheControls { - no_cache: true, - ..enabled - } - .reads() - ); - assert!( - !CacheControls { - no_store: true, - ..enabled - } - .writes() - ); -} diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs index dad5398a879..e24545caad6 100644 --- a/litellm-rust/crates/cache/tests/codec.rs +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use litellm_cache::{CacheCodec, CacheEntry, Error, JsonCodec}; +use litellm_cache::{CacheCodec, Error, JsonCodec}; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -25,18 +25,6 @@ fn json_codec_round_trips_typed_domain_values() { ); } -#[test] -fn response_entries_preserve_the_existing_json_representation() { - let codec = JsonCodec::::new(); - let entry = CacheEntry { - timestamp: 123.0, - response: json!({"choices": [{"text": "cached"}]}), - }; - let bytes = codec.encode(&entry).unwrap(); - assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); - assert_eq!(codec.decode(&bytes).unwrap(), entry); -} - #[test] fn json_codec_rejects_malformed_and_wrongly_typed_entries() { let codec = JsonCodec::::new(); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 308dfd2dd7a..1eb2ec28036 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,8 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 32360e72469..eb07118e964 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,4 +1,3 @@ -use litellm_cache_response::NativeResponseCache; use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -8,7 +7,7 @@ use pyo3::{ }; use serde_json::Value; -use super::NativeCacheHandle; +use super::{NativeCacheHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index f00cceeb86e..5918967009a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,9 +1,10 @@ mod facade; +mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{CacheControls, CacheKeyInput, Error}; -use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest}; +use litellm_cache::Error; +use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; use pyo3::{ PyTraverseError, PyVisit, @@ -15,6 +16,7 @@ use serde::Deserialize; use serde_json::Value; use facade::FacadeGuard; +use native::NativeResponseCache; #[derive(Deserialize)] #[serde(deny_unknown_fields)] diff --git a/litellm-rust/crates/cache-response/src/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs similarity index 75% rename from litellm-rust/crates/cache-response/src/native.rs rename to litellm-rust/crates/python-bridge/src/cache/native.rs index c25c61cee82..6af04bfe2b2 100644 --- a/litellm-rust/crates/cache-response/src/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,33 +1,30 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheEntry, Error}; +use litellm_cache::{CacheCodec, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use serde_json::Value; -use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_response::{CacheEntry, ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; -pub enum NativeResponseCache -where - C: redis::ConnectionLike + Send + 'static, -{ +#[derive(Clone)] +pub(super) enum NativeResponseCache { Memory(Arc>>), - Redis(Arc>>), -} - -impl Clone for NativeResponseCache { - fn clone(&self) -> Self { - match self { - Self::Memory(cache) => Self::Memory(Arc::clone(cache)), - Self::Redis(cache) => Self::Redis(Arc::clone(cache)), - } - } + Redis(Arc>>), } impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { Self::Memory(Arc::new(ResponseCache::new(Arc::new( - InMemoryCache::response_cache(capacity, ttl, max_entry_bytes), + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(ttl), + Some(max_entry_bytes), + Some(Arc::new(|entry| { + ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) + })), + super::now, + ), )))) } @@ -41,7 +38,7 @@ impl NativeResponseCache { } } -impl NativeResponseCache { +impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d1cea860cb0..493baac228a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -224,8 +224,24 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): _native.NativeCacheHandle.memory(ttl_seconds=-1) + + +async def test_memory_size_policy_is_applied_by_the_native_host() -> None: + handle: Final = _native.NativeCacheHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _native.CacheResolver(SimpleNamespace(cache=handle)).resolve() + small: Final = {"answer": "ok"} + binding.store(request("small"), small) + assert await binding.async_lookup(request("small")) == small + await binding.async_store(request("large"), {"answer": "x" * 256}) + assert binding.lookup(request("large")) is None + assert binding.lookup(request("small")) == small + disabled: Final = _native.CacheResolver( + SimpleNamespace(cache=_native.NativeCacheHandle.memory(capacity=0)) + ).resolve() + await disabled.async_store(request(), small) + assert await disabled.async_lookup(request()) is None From cc8cadebb46d22286166bb9d4152b58bac11d91b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 07:56:20 -0700 Subject: [PATCH 22/56] fix(cache): close native parity gaps --- litellm-rust/Cargo.lock | 22 ++ litellm-rust/crates/cache-memory/src/cache.rs | 54 ++- .../crates/cache-memory/tests/cache.rs | 51 ++- litellm-rust/crates/cache-redis/Cargo.toml | 3 +- litellm-rust/crates/cache-redis/src/cache.rs | 356 +++++++++++++++--- .../crates/cache-redis/tests/cache.rs | 79 +++- litellm-rust/crates/cache-response/README.md | 8 +- .../crates/cache-response/src/caching.rs | 10 +- .../crates/cache-response/src/codec.rs | 24 +- .../crates/cache-response/src/embedding.rs | 22 ++ litellm-rust/crates/cache-response/src/lib.rs | 2 + .../crates/cache-response/src/response.rs | 164 +++++++- .../crates/cache-response/tests/caching.rs | 7 + .../crates/cache-response/tests/response.rs | 87 ++++- litellm-rust/crates/cache/src/base_cache.rs | 45 +++ litellm-rust/crates/cache/src/capabilities.rs | 39 ++ litellm-rust/crates/cache/src/dual.rs | 105 ++++++ litellm-rust/crates/cache/src/lib.rs | 7 +- litellm-rust/crates/cache/tests/dual.rs | 130 +++++++ .../crates/python-bridge/python_settings.json | 3 + .../crates/python-bridge/src/cache/facade.rs | 14 +- .../crates/python-bridge/src/cache/mod.rs | 178 ++++++++- .../crates/python-bridge/src/cache/native.rs | 129 ++++++- .../python-bridge/src/python_settings.rs | 5 +- litellm/caching/dual_cache.py | 79 ++-- litellm/rust_bridge/_native.pyi | 21 ++ litellm/rust_bridge/settings.py | 11 + tests/test_litellm/caching/test_dual_cache.py | 79 ++-- tests/test_litellm_rust/test_cache.py | 65 ++++ 29 files changed, 1615 insertions(+), 184 deletions(-) create mode 100644 litellm-rust/crates/cache-response/src/embedding.rs create mode 100644 litellm-rust/crates/cache/src/capabilities.rs create mode 100644 litellm-rust/crates/cache/src/dual.rs create mode 100644 litellm-rust/crates/cache/tests/dual.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8b299f5455b..725d2cfef41 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2479,6 +2479,7 @@ name = "litellm-cache-redis" version = "0.1.0" dependencies = [ "litellm-cache", + "r2d2", "redis", "redis-test", "serde_json", @@ -3564,6 +3565,17 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "rand" version = "0.8.7" @@ -3700,6 +3712,7 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", + "r2d2", "ryu", "sha1_smol", "socket2 0.6.5", @@ -4096,6 +4109,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + [[package]] name = "schemars" version = "0.9.0" diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 43186faf3f7..074c8dcb0fd 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -3,7 +3,10 @@ use std::collections::{BinaryHeap, HashMap}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error}; +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, + Error, +}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); @@ -168,6 +171,55 @@ impl InMemoryCache { } } +impl ClaimCache for InMemoryCache +where + V: Clone + PartialEq + Send + Sync + 'static, +{ + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + kwargs: CacheKwargs, + ) -> Result { + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now); + let winner = match state.values.get(key) { + Some(existing) if eligible.is_empty() => existing.clone(), + Some(existing) if eligible.contains(existing) => existing.clone(), + _ => candidate, + }; + let expiration = now + self.get_ttl(&kwargs); + state.values.insert(key.into(), winner.clone()); + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + Ok(winner) + } +} + +impl CounterCache for InMemoryCache { + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now); + let value = state.values.get(key).copied().unwrap_or_default() + amount; + let expiration = state + .expirations + .get(key) + .copied() + .unwrap_or_else(|| now + self.get_ttl(&kwargs)); + state.values.insert(key.into(), value); + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + Ok(value) + } +} + impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 370145bebef..e5831dfd6d5 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -3,7 +3,8 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, + get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -183,3 +184,51 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() { None ); } + +#[test] +fn claims_are_atomic_and_refresh_eligible_winners() { + let clock = clock(); + let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), { + let clock = clock.clone(); + move || Duration::from_secs(clock.load(Ordering::SeqCst)) + }); + let kwargs = CacheKwargs { + ttl: Some(Duration::from_secs(10)), + ..Default::default() + }; + assert_eq!( + cache + .claim_cache("affinity", "first".to_string(), &[], kwargs.clone()) + .unwrap(), + "first" + ); + clock.store(105, Ordering::SeqCst); + assert_eq!( + cache + .claim_cache( + "affinity", + "second".to_string(), + &["first".to_string(), "second".to_string()], + kwargs, + ) + .unwrap(), + "first" + ); + assert_eq!( + cache.expires_at("affinity").unwrap(), + Some(Duration::from_secs(115)) + ); +} + +#[test] +fn counters_increment_under_one_lock() { + let cache = InMemoryCache::::default(); + assert_eq!( + CounterCache::increment_cache(&cache, "counter", 1.5, CacheKwargs::default()).unwrap(), + 1.5 + ); + assert_eq!( + CounterCache::increment_cache(&cache, "counter", 2.0, CacheKwargs::default()).unwrap(), + 3.5 + ); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index a60813b6260..1a1bb505a7c 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = "1.7.0" +redis = { version = "1.7.0", features = ["r2d2"] } +r2d2 = "0.8.10" tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index d4ca0cf0522..281ffe17534 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,15 +1,83 @@ -use std::sync::{Arc, Mutex, MutexGuard}; +use std::sync::{Arc, Mutex}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error, + BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, + ClaimCache, CounterCache, Error, }; use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; + +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 connection = pool.get().map_err(|_| Error::Unavailable)?; + connection + .set_read_timeout(Some(REDIS_TIMEOUT)) + .map_err(|_| Error::Unavailable)?; + connection + .set_write_timeout(Some(REDIS_TIMEOUT)) + .map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + Self::Fixed(connection) => { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + } + } +} pub struct RedisCache { - connection: Arc>, + connections: Arc>, default_ttl: Duration, codec: S, namespace: Option, @@ -18,8 +86,18 @@ pub struct RedisCache { impl RedisCache { pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; - let connection = client.get_connection().map_err(|_| Error::Unavailable)?; - Ok(Self::with_connection(connection, default_ttl, codec)) + let pool = r2d2::Pool::builder() + .max_size(REDIS_POOL_SIZE) + .min_idle(Some(0)) + .connection_timeout(REDIS_TIMEOUT) + .build(client) + .map_err(|_| Error::Unavailable)?; + Ok(Self { + connections: Arc::new(Connections::Pool(pool)), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + codec, + namespace: None, + }) } } @@ -30,17 +108,13 @@ where { pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { - connection: Arc::new(Mutex::new(connection)), + connections: Arc::new(Connections::Fixed(Mutex::new(connection))), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, } } - fn connection(&self) -> Result, Error> { - self.connection.lock().map_err(|_| Error::Unavailable) - } - pub fn with_namespace(self, namespace: Option) -> Self { Self { namespace: namespace.filter(|value| !value.is_empty()), @@ -72,6 +146,29 @@ where Ok(format!("{escaped}:*")) } + fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> { + let mut cursor = 0u64; + loop { + let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") + .cursor_arg(cursor) + .arg("MATCH") + .arg(pattern) + .arg("COUNT") + .arg(1000) + .query(connection) + .map_err(|_| Error::Unavailable)?; + if !keys.is_empty() { + connection + .del::<_, usize>(keys) + .map_err(|_| Error::Unavailable)?; + } + if next_cursor == 0 { + return Ok(()); + } + cursor = next_cursor; + } + } + fn decode_response(&self, value: redis::Value) -> Result, Error> { match value { redis::Value::Nil => Ok(None), @@ -81,23 +178,29 @@ where } } + fn decode_batch_response(&self, value: redis::Value) -> Result, Error> { + match self.decode_response(value) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + } + } + fn ttl_seconds(ttl: Duration) -> u64 { ttl.as_secs() .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - async fn run_blocking(connection: Arc>, operation: F) -> Result + async fn run_blocking(connections: Arc>, operation: F) -> Result where T: Send + 'static, - F: FnOnce(&mut C) -> Result + Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, { - tokio::task::spawn_blocking(move || { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; - operation(&mut connection) - }) - .await - .map_err(|_| Error::Unavailable)? + tokio::task::spawn_blocking(move || connections.execute(operation)) + .await + .map_err(|_| Error::Unavailable)? } } @@ -115,40 +218,55 @@ where fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { let payload = self.codec.encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - self.connection()? - .set_ex::<_, _, ()>(self.namespaced_key(key), payload, ttl) - .map_err(|_| Error::Unavailable) + let key = self.namespaced_key(key); + self.connections.execute(|connection| { + connection + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable) + }) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - let value = self - .connection()? - .get::<_, redis::Value>(self.namespaced_key(key)) - .map_err(|_| Error::Unavailable)?; + let key = self.namespaced_key(key); + let value = self.connections.execute(|connection| { + connection + .get::<_, redis::Value>(key) + .map_err(|_| Error::Unavailable) + })?; self.decode_response(value) } + fn get_cache_batch( + &self, + keys: &[String], + _: &CacheKwargs, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = self.connections.execute(|connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + })?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } + fn delete_cache(&self, key: &str) -> Result<(), Error> { - self.connection()? - .del::<_, ()>(self.namespaced_key(key)) - .map_err(|_| Error::Unavailable) + let key = self.namespaced_key(key); + self.connections + .execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) } fn flush_cache(&self) -> Result<(), Error> { let pattern = self.namespaced_pattern()?; - let mut connection = self.connection()?; - let keys = connection - .scan_match(pattern) - .map_err(|_| Error::Unavailable)? - .collect::>>() - .map_err(|_| Error::Unavailable)?; - if keys.is_empty() { - return Ok(()); - } - connection - .del::<_, usize>(keys) - .map(|_| ()) - .map_err(|_| Error::Unavailable) + self.connections + .execute(|connection| Self::flush_matching(connection, &pattern)) } async fn async_set_cache( @@ -160,7 +278,7 @@ where let payload = self.codec.encode(&value)?; let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) @@ -174,7 +292,7 @@ where _: &CacheKwargs, ) -> Result, Error> { let key = self.namespaced_key(key); - let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| { + let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) @@ -183,6 +301,28 @@ where self.decode_response(value) } + async fn async_get_cache_batch( + &self, + keys: Vec, + _: CacheKwargs, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } + async fn async_set_cache_pipeline( &self, cache_list: Vec<(String, Self::Value)>, @@ -197,41 +337,137 @@ where }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); for (key, payload) in entries { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable)?; + pipeline + .cmd("SETEX") + .arg(key) + .arg(ttl) + .arg(payload) + .ignore(); } - Ok(()) + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) }) .await } async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) .await } + async fn async_flush_cache(&self) -> Result<(), Error> { + let pattern = self.namespaced_pattern()?; + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Self::flush_matching(connection, &pattern) + }) + .await + } + async fn disconnect(&self) -> Result<(), Error> { Ok(()) } async fn test_connection(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connection), |connection| { - redis::cmd("PING") - .query::(connection) + match Self::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +impl CounterCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + const SCRIPT: &str = concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" + ); + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + self.connections.execute(|connection| { + redis::cmd("EVAL") + .arg(SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) .map_err(|_| Error::Unavailable) }) - .await?; - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "Redis cache connection test successful".into(), - error: None, - }) + } +} + +impl ClaimCache for RedisCache +where + S: CacheCodec, + S::Value: PartialEq, + C: redis::ConnectionLike + Send + 'static, +{ + fn claim_cache( + &self, + key: &str, + candidate: S::Value, + eligible: &[S::Value], + kwargs: CacheKwargs, + ) -> Result { + const SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false then redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", + "return ARGV[1]; end; if #ARGV > 2 then for index = 3, #ARGV do ", + "if current == ARGV[index] then redis.call('EXPIRE', KEYS[1], ARGV[2]); ", + "return current; end; end; redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", + "return ARGV[1]; end; if current == ARGV[1] then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return current" + ); + let key = self.namespaced_key(key); + let candidate = self.codec.encode(&candidate)?; + let eligible = eligible + .iter() + .map(|value| self.codec.encode(value)) + .collect::, _>>()?; + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let value = self.connections.execute(|connection| { + redis::cmd("EVAL") + .arg(SCRIPT) + .arg(1) + .arg(key) + .arg(candidate) + .arg(ttl) + .arg(eligible) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + self.decode_response(value)?.ok_or(Error::Unavailable) } } @@ -302,7 +538,9 @@ mod tests { redis::cmd("SCAN") .cursor_arg(0) .arg("MATCH") - .arg("litellm-cache:*"), + .arg("litellm-cache:*") + .arg("COUNT") + .arg(1000), Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])), ), MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index d5bba19a8bd..6aef9bb36bf 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,6 +1,9 @@ use std::time::Duration; -use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, Error, JsonCodec, get_cache, set_cache}; +use litellm_cache::{ + BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, Error, JsonCodec, + get_cache, set_cache, +}; use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; @@ -174,7 +177,9 @@ fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { redis::cmd("SCAN") .cursor_arg(0) .arg("MATCH") - .arg("team\\*:*"), + .arg("team\\*:*") + .arg("COUNT") + .arg(1000), Ok(redis_test::redis_value!(["0", ["team*:key"]])), ), MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)), @@ -184,3 +189,73 @@ fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { .with_namespace(Some("team*".into())); scoped.flush_cache().unwrap(); } + +#[tokio::test] +async fn connection_failures_use_the_python_result_contract() { + let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused")); + let connection = + MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::(error))]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Failed); + assert!(result.message.starts_with("Redis connection failed:")); + assert!(result.error.is_some()); +} + +#[tokio::test] +async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() { + let connection = MockRedisConnection::new([MockCmd::new( + redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"), + Ok(vec![ + redis::Value::BulkString(vec![42, 7]), + redis::Value::Nil, + redis::Value::BulkString(vec![99, 7]), + ]), + )]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + + assert_eq!( + cache + .async_get_cache_batch( + vec!["hit".into(), "miss".into(), "invalid".into()], + CacheKwargs::default(), + ) + .await + .unwrap(), + vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid] + ); +} + +#[tokio::test] +async fn async_flush_deletes_each_scan_page_separately() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team:*") + .arg("COUNT") + .arg(1000), + Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(7) + .arg("MATCH") + .arg("team:*") + .arg("COUNT") + .arg(1000), + Ok(redis_test::redis_value!(["0", ["team:c"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + cache.async_flush_cache().await.unwrap(); +} diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 309296f7773..4e6694c191f 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -28,17 +28,17 @@ cache.store(&request, json!({"answer": 7}), now)?; assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); ``` -For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Each Redis constructor currently opens its own connection; shared connection pools remain follow-up work +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers move that blocking work off the executor Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it ## Python integration boundary -The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support synchronous and asynchronous response lookup and storage +The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution -Explicit facade registration checks object identity, method overrides, and configuration changes before selecting native execution. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python +Explicit facade registration checks object identity, method overrides, effective TTL, and configuration changes before selecting native execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy @@ -50,6 +50,6 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na ## Follow-up scope -Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial batches, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths +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, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub 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/caching.rs b/litellm-rust/crates/cache-response/src/caching.rs index 53b34025ccd..afae4dfe4a5 100644 --- a/litellm-rust/crates/cache-response/src/caching.rs +++ b/litellm-rust/crates/cache-response/src/caching.rs @@ -120,6 +120,7 @@ impl CacheControls { pub fn writes(self) -> bool { self.supported_call_type && self.configured + && self.caching.unwrap_or(true) && !self.no_store && (self.default_on || self.use_cache) } @@ -131,13 +132,16 @@ pub fn should_use_cache(controls: CacheControls) -> bool { #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct CacheEntry { - pub timestamp: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub timestamp: Option, pub response: Value, } impl CacheEntry { pub fn fresh(&self, now: Duration, max_age: Option) -> bool { - self.timestamp.is_finite() - && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) + self.timestamp.is_none_or(|timestamp| { + timestamp.is_finite() + && max_age.is_none_or(|age| now.as_secs_f64() - timestamp <= age.as_secs_f64()) + }) } } diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index f1d55ddefe8..eaba3d0c349 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -9,7 +9,10 @@ impl CacheCodec for ResponseCacheCodec { type Value = CacheEntry; fn encode(&self, value: &CacheEntry) -> Result, Error> { - if !value.timestamp.is_finite() { + if value + .timestamp + .is_some_and(|timestamp| !timestamp.is_finite()) + { return Err(Error::InvalidEntry); } serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) @@ -17,12 +20,21 @@ impl CacheCodec for ResponseCacheCodec { fn decode(&self, bytes: &[u8]) -> Result { let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; - let entry: CacheEntry = - serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?; - if !entry.timestamp.is_finite() { + let value = decode_value(text)?; + let Some(timestamp) = value.get("timestamp") else { + return Ok(CacheEntry { + timestamp: None, + response: value, + }); + }; + let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { return Err(Error::InvalidEntry); - } - Ok(entry) + }; + let response = value.get("response").cloned().ok_or(Error::InvalidEntry)?; + Ok(CacheEntry { + timestamp: Some(timestamp), + response, + }) } } diff --git a/litellm-rust/crates/cache-response/src/embedding.rs b/litellm-rust/crates/cache-response/src/embedding.rs new file mode 100644 index 00000000000..d1f8a2bc0a6 --- /dev/null +++ b/litellm-rust/crates/cache-response/src/embedding.rs @@ -0,0 +1,22 @@ +use serde::Serialize; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct PartialHits { + pub values: Vec>, + pub missing_indices: Vec, +} + +impl PartialHits { + pub fn new(values: Vec>) -> Self { + let missing_indices = values + .iter() + .enumerate() + .filter_map(|(index, value)| value.is_none().then_some(index)) + .collect(); + Self { + values, + missing_indices, + } + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index efa0b04b9f7..72a507f8ee9 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,5 +1,6 @@ mod caching; mod codec; +mod embedding; mod response; pub use caching::{ @@ -7,4 +8,5 @@ pub use caching::{ get_cache_key, should_use_cache, }; pub use codec::ResponseCacheCodec; +pub use embedding::PartialHits; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 987a47f0554..8f0fe953df6 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{BaseCache, CacheKwargs, Error}; +use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error}; -use crate::{CacheControls, CacheEntry, CacheKeyInput, cache_key}; +use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; use serde_json::Value; #[derive(Clone)] @@ -39,6 +39,18 @@ impl> ResponseCache { Self { backend } } + pub fn default_ttl(&self) -> Duration { + self.backend.default_ttl() + } + + pub async fn async_flush(&self) -> Result<(), Error> { + self.backend.async_flush_cache().await + } + + pub async fn test_connection(&self) -> Result { + self.backend.test_connection().await + } + pub fn lookup( &self, request: &ResponseCacheRequest, @@ -47,10 +59,15 @@ impl> ResponseCache { if !request.controls.reads() { return Ok(None); } - let entry = self + let entry = match self .backend - .get_cache(&cache_key(&request.key), &request.kwargs)?; - Self::fresh_response(entry, now, request.max_age) + .get_cache(&cache_key(&request.key), &request.kwargs) + { + Ok(entry) => entry, + Err(Error::InvalidEntry) => None, + Err(error) => return Err(error), + }; + Self::fresh_or_miss(entry, now, request.max_age) } pub async fn async_lookup( @@ -61,11 +78,62 @@ impl> ResponseCache { if !request.controls.reads() { return Ok(None); } - let entry = self + let entry = match self .backend .async_get_cache(&cache_key(&request.key), &request.kwargs) - .await?; - Self::fresh_response(entry, now, request.max_age) + .await + { + Ok(entry) => entry, + Err(Error::InvalidEntry) => None, + Err(error) => return Err(error), + }; + Self::fresh_or_miss(entry, now, request.max_age) + } + + pub fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + let readable = requests + .iter() + .enumerate() + .filter(|(_, request)| request.controls.reads()) + .collect::>(); + let keys = readable + .iter() + .map(|(_, request)| cache_key(&request.key)) + .collect::>(); + let entries = if let Some((_, request)) = readable.first() { + self.backend.get_cache_batch(&keys, &request.kwargs)? + } else { + Vec::new() + }; + Self::partial_hits(requests, readable, entries, now) + } + + pub async fn async_lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + let readable = requests + .iter() + .enumerate() + .filter(|(_, request)| request.controls.reads()) + .collect::>(); + let keys = readable + .iter() + .map(|(_, request)| cache_key(&request.key)) + .collect::>(); + let entries = if let Some((_, request)) = readable.first() { + self.backend + .async_get_cache_batch(keys, request.kwargs.clone()) + .await? + } else { + Vec::new() + }; + Self::partial_hits(requests, readable, entries, now) } pub fn store( @@ -80,7 +148,7 @@ impl> ResponseCache { self.backend.set_cache( &cache_key(&request.key), CacheEntry { - timestamp: now.as_secs_f64(), + timestamp: Some(now.as_secs_f64()), response, }, request.kwargs.clone(), @@ -100,7 +168,7 @@ impl> ResponseCache { .async_set_cache( &cache_key(&request.key), CacheEntry { - timestamp: now.as_secs_f64(), + timestamp: Some(now.as_secs_f64()), response, }, request.kwargs.clone(), @@ -108,6 +176,76 @@ impl> ResponseCache { .await } + pub async fn async_store_batch( + &self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> Result<(), Error> { + let writable = entries + .into_iter() + .filter(|(request, _)| request.controls.writes()) + .map(|(request, response)| { + ( + cache_key(&request.key), + CacheEntry { + timestamp: Some(now.as_secs_f64()), + response, + }, + request.kwargs, + ) + }) + .collect::>(); + let Some((_, _, first_kwargs)) = writable.first() else { + return Ok(()); + }; + if writable.iter().all(|(_, _, kwargs)| kwargs == first_kwargs) { + let kwargs = first_kwargs.clone(); + let cache_list = writable + .into_iter() + .map(|(key, entry, _)| (key, entry)) + .collect(); + return self + .backend + .async_set_cache_pipeline(cache_list, kwargs) + .await; + } + for (key, entry, kwargs) in writable { + self.backend.async_set_cache(&key, entry, kwargs).await?; + } + Ok(()) + } + + fn partial_hits( + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, + entries: Vec>, + now: Duration, + ) -> Result { + if readable.len() != entries.len() { + return Err(Error::Unavailable); + } + let mut values = vec![None; requests.len()]; + for ((index, request), entry) in readable.into_iter().zip(entries) { + let response = match entry { + BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age)?, + BatchEntry::Miss | BatchEntry::Invalid => None, + }; + values[index] = response; + } + Ok(PartialHits::new(values)) + } + + fn fresh_or_miss( + entry: Option, + now: Duration, + max_age: Option, + ) -> Result, Error> { + match Self::fresh_response(entry, now, max_age) { + Err(Error::InvalidEntry) => Ok(None), + result => result, + } + } + fn fresh_response( entry: Option, now: Duration, @@ -115,9 +253,9 @@ impl> ResponseCache { ) -> Result, Error> { entry .filter(|entry| entry.fresh(now, max_age)) - .map(|entry| match entry.response { - Value::String(text) => crate::codec::decode_value(&text), - value => Ok(value), + .map(|entry| match (entry.timestamp, entry.response) { + (Some(_), Value::String(text)) => crate::codec::decode_value(&text), + (_, value) => Ok(value), }) .transpose() } diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs index d403791e421..0e8ce9b3b1d 100644 --- a/litellm-rust/crates/cache-response/tests/caching.rs +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -80,4 +80,11 @@ fn cache_controls_honor_default_modes_and_directives() { } .writes() ); + assert!( + !CacheControls { + caching: Some(false), + ..enabled + } + .writes() + ); } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 88b53fdfe86..e4a04b3dec5 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -245,7 +245,7 @@ fn response_codec_accepts_python_literals_without_executing_code() { assert_eq!( ResponseCacheCodec .encode(&CacheEntry { - timestamp: f64::NAN, + timestamp: Some(f64::NAN), response: json!({}) }) .unwrap_err(), @@ -254,7 +254,7 @@ fn response_codec_accepts_python_literals_without_executing_code() { } #[tokio::test] -async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() { +async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() { let connection = MockRedisConnection::new([MockCmd::new( redis::cmd("GET").arg("tenant:key"), Ok(b"invalid".to_vec()), @@ -267,22 +267,19 @@ async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redi assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); request.controls.no_cache = false; assert_eq!( - cache - .async_lookup(&request, Duration::ZERO) - .await - .unwrap_err(), - Error::InvalidEntry + cache.async_lookup(&request, Duration::ZERO).await.unwrap(), + None ); } #[test] -fn malformed_memory_entries_are_rejected_by_the_response_consumer() { +fn malformed_memory_entries_are_treated_as_misses() { let backend = Arc::new(InMemoryCache::default()); BaseCache::set_cache( backend.as_ref(), "tenant:key", CacheEntry { - timestamp: 100.0, + timestamp: Some(100.0), response: json!("not a serialized response"), }, Default::default(), @@ -290,10 +287,8 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { .unwrap(); let cache = ResponseCache::new(backend); assert_eq!( - cache - .lookup(&request(), Duration::from_secs(100)) - .unwrap_err(), - Error::InvalidEntry + cache.lookup(&request(), Duration::from_secs(100)).unwrap(), + None ); } @@ -301,10 +296,74 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { fn response_entries_preserve_the_existing_json_representation() { let codec = ResponseCacheCodec; let entry = CacheEntry { - timestamp: 123.0, + timestamp: Some(123.0), response: json!({"choices": [{"text": "cached"}]}), }; let bytes = codec.encode(&entry).unwrap(); assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); assert_eq!(codec.decode(&bytes).unwrap(), entry); } + +#[test] +fn response_codec_preserves_values_without_timestamps() { + let codec = ResponseCacheCodec; + let raw = json!({"choices": [{"text": "legacy"}]}); + let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap(); + assert_eq!(entry.timestamp, None); + assert_eq!(entry.response, raw); + + let backend = Arc::new(InMemoryCache::default()); + BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, Default::default()).unwrap(); + let cache = ResponseCache::new(backend); + assert_eq!( + cache.lookup(&request(), Duration::from_secs(100)).unwrap(), + Some(json!({"choices": [{"text": "legacy"}]})) + ); +} + +#[tokio::test] +async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() { + let cache = memory(); + let requests = ["hit", "miss", "disabled"].map(|key| { + ResponseCacheRequest::new(CacheKeyInput { + preset: Some(key.into()), + ..Default::default() + }) + }); + cache + .store(&requests[0], json!({"value": 1}), Duration::from_secs(100)) + .unwrap(); + let mut requests = requests.to_vec(); + requests[2].controls.caching = Some(false); + + let partial = cache + .async_lookup_batch(&requests, Duration::from_secs(100)) + .await + .unwrap(); + assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]); + assert_eq!(partial.missing_indices, vec![1, 2]); + + cache + .async_store_batch( + vec![ + (requests[1].clone(), json!({"value": 2})), + (requests[2].clone(), json!({"value": 3})), + ], + Duration::from_secs(100), + ) + .await + .unwrap(); + assert_eq!( + cache + .lookup(&requests[1], Duration::from_secs(100)) + .unwrap(), + Some(json!({"value": 2})) + ); + requests[2].controls.caching = None; + assert_eq!( + cache + .lookup(&requests[2], Duration::from_secs(100)) + .unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 1891e417cb0..5bc1ebc2945 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -6,6 +6,13 @@ use serde_json::{Map, Value}; use crate::Error; +#[derive(Clone, Debug, PartialEq)] +pub enum BatchEntry { + Hit(V), + Miss, + Invalid, +} + #[derive(Clone, Debug, Default, PartialEq)] pub struct CacheKwargs { pub ttl: Option, @@ -42,6 +49,21 @@ pub trait BaseCache: Send + Sync { fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>; + fn get_cache_batch( + &self, + keys: &[String], + kwargs: &CacheKwargs, + ) -> Result>, Error> { + keys.iter() + .map(|key| match self.get_cache(key, kwargs) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + }) + .collect() + } + fn async_set_cache( &self, key: &str, @@ -59,6 +81,25 @@ pub trait BaseCache: Send + Sync { async move { self.get_cache(key, kwargs) } } + fn async_get_cache_batch( + &self, + keys: Vec, + kwargs: CacheKwargs, + ) -> impl Future>, Error>> + Send { + async move { + let mut entries = Vec::with_capacity(keys.len()); + for key in keys { + entries.push(match self.async_get_cache(&key, &kwargs).await { + Ok(Some(value)) => BatchEntry::Hit(value), + Ok(None) => BatchEntry::Miss, + Err(Error::InvalidEntry) => BatchEntry::Invalid, + Err(error) => return Err(error), + }); + } + Ok(entries) + } + } + fn async_set_cache_pipeline( &self, cache_list: Vec<(String, Self::Value)>, @@ -89,6 +130,10 @@ pub trait BaseCache: Send + Sync { fn flush_cache(&self) -> Result<(), Error>; + fn async_flush_cache(&self) -> impl Future> + Send { + async move { self.flush_cache() } + } + fn disconnect(&self) -> impl Future> + Send; fn test_connection(&self) -> impl Future> + Send; diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs new file mode 100644 index 00000000000..0b9deab1f5a --- /dev/null +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -0,0 +1,39 @@ +use std::future::Future; + +use crate::{BaseCache, CacheKwargs, Error}; + +pub trait CounterCache: BaseCache { + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; + + fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> impl Future> + Send { + async move { self.increment_cache(key, amount, kwargs) } + } +} + +pub trait ClaimCache: BaseCache +where + Self::Value: PartialEq, +{ + fn claim_cache( + &self, + key: &str, + candidate: Self::Value, + eligible: &[Self::Value], + kwargs: CacheKwargs, + ) -> Result; + + fn async_claim_cache( + &self, + key: &str, + candidate: Self::Value, + eligible: Vec, + kwargs: CacheKwargs, + ) -> impl Future> + Send { + async move { self.claim_cache(key, candidate, &eligible, kwargs) } + } +} diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs new file mode 100644 index 00000000000..17a2c430ddd --- /dev/null +++ b/litellm-rust/crates/cache/src/dual.rs @@ -0,0 +1,105 @@ +use std::sync::Arc; + +use crate::{BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error}; + +pub struct DualCache { + l1: Arc, + l2: Arc, +} + +impl DualCache { + pub fn new(l1: Arc, l2: Arc) -> Self { + Self { l1, l2 } + } +} + +impl BaseCache for DualCache +where + V: Clone + Send + Sync + 'static, + L1: BaseCache, + L2: BaseCache, +{ + type Value = V; + + fn default_ttl(&self) -> std::time::Duration { + self.l2.default_ttl() + } + + fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { + self.l2.set_cache(key, value.clone(), kwargs.clone())?; + self.l1.set_cache(key, value, kwargs) + } + + fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { + if let Some(value) = self.l1.get_cache(key, kwargs)? { + return Ok(Some(value)); + } + let value = self.l2.get_cache(key, kwargs)?; + if let Some(value) = &value { + self.l1.set_cache(key, value.clone(), kwargs.clone())?; + } + Ok(value) + } + + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.l2.delete_cache(key)?; + self.l1.delete_cache(key) + } + + fn flush_cache(&self) -> Result<(), Error> { + self.l2.flush_cache()?; + self.l1.flush_cache() + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + self.l2.async_flush_cache().await?; + self.l1.async_flush_cache().await + } + + async fn disconnect(&self) -> Result<(), Error> { + self.l2.disconnect().await?; + self.l1.disconnect().await + } + + async fn test_connection(&self) -> Result { + self.l2.test_connection().await + } +} + +impl CounterCache for DualCache +where + L1: BaseCache, + L2: CounterCache, +{ + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + let value = self.l2.increment_cache(key, amount, kwargs.clone())?; + self.l1.set_cache(key, value, kwargs)?; + Ok(value) + } +} + +impl ClaimCache for DualCache +where + V: Clone + PartialEq + Send + Sync + 'static, + L1: ClaimCache, + L2: ClaimCache, +{ + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + kwargs: CacheKwargs, + ) -> Result { + match self + .l2 + .claim_cache(key, candidate.clone(), eligible, kwargs.clone()) + { + Ok(winner) => { + self.l1.set_cache(key, winner.clone(), kwargs)?; + Ok(winner) + } + Err(_) => self.l1.claim_cache(key, candidate, eligible, kwargs), + } + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 4ff02319bdc..ed67eb2fe15 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -1,9 +1,14 @@ mod base_cache; mod caching; +mod capabilities; mod codec; +pub mod dual; mod error; -pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; +pub use base_cache::{ + BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, +}; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; +pub use capabilities::{ClaimCache, CounterCache}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs new file mode 100644 index 00000000000..1be1556734c --- /dev/null +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -0,0 +1,130 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, dual::DualCache, +}; + +struct TestCache { + value: Mutex>, + fail: bool, +} + +impl TestCache { + fn new(value: Option, fail: bool) -> Self { + Self { + value: Mutex::new(value), + fail, + } + } +} + +impl BaseCache for TestCache +where + V: Clone + Send + Sync + 'static, +{ + type Value = V; + + fn set_cache(&self, _: &str, value: V, _: CacheKwargs) -> Result<(), Error> { + *self.value.lock().unwrap() = Some(value); + Ok(()) + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + Ok(self.value.lock().unwrap().clone()) + } + + fn delete_cache(&self, _: &str) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } + + fn flush_cache(&self) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +impl CounterCache for TestCache { + fn increment_cache(&self, _: &str, amount: f64, _: CacheKwargs) -> Result { + if self.fail { + return Err(Error::Unavailable); + } + let mut value = self.value.lock().unwrap(); + let incremented = value.unwrap_or_default() + amount; + *value = Some(incremented); + Ok(incremented) + } +} + +impl ClaimCache for TestCache +where + V: Clone + PartialEq + Send + Sync + 'static, +{ + fn claim_cache( + &self, + _: &str, + candidate: V, + eligible: &[V], + _: CacheKwargs, + ) -> Result { + if self.fail { + return Err(Error::Unavailable); + } + let mut value = self.value.lock().unwrap(); + let winner = match value.as_ref() { + Some(existing) if eligible.is_empty() || eligible.contains(existing) => { + existing.clone() + } + _ => candidate, + }; + *value = Some(winner.clone()); + Ok(winner) + } +} + +#[test] +fn failed_l2_increment_leaves_l1_unchanged() { + let l1 = Arc::new(TestCache::new(Some(10.0), false)); + let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true))); + + assert_eq!( + cache.increment_cache("counter", 2.0, CacheKwargs::default()), + Err(Error::Unavailable) + ); + assert_eq!( + l1.get_cache("counter", &CacheKwargs::default()).unwrap(), + Some(10.0) + ); +} + +#[test] +fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { + let l1 = Arc::new(TestCache::new(Some("first".to_string()), false)); + let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))); + + assert_eq!( + cache + .claim_cache( + "affinity", + "second".into(), + &["first".into(), "second".into()], + CacheKwargs { + ttl: Some(Duration::from_secs(60)), + ..Default::default() + }, + ) + .unwrap(), + "first" + ); +} diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 0af55083bef..15d0d603ac7 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -22,5 +22,8 @@ ], "secret_manager": [ "readable" + ], + "cache_settings": [ + "default_redis_ttl" ] } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index eb07118e964..58550c2987d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -6,6 +6,7 @@ use pyo3::{ types::{PyDict, PyTuple, PyType}, }; use serde_json::Value; +use std::time::Duration; use super::{NativeCacheHandle, native::NativeResponseCache}; @@ -126,7 +127,12 @@ impl ObjectGuard { } impl FacadeGuard { - pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult { + pub(super) fn capture( + py: Python<'_>, + facade: &Bound<'_, PyAny>, + kind: &str, + native_default_ttl: Duration, + ) -> PyResult { let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -146,6 +152,12 @@ impl FacadeGuard { "facade and native backend types must match", )); } + let python_default_ttl = backend.getattr("default_ttl")?.extract::()?; + if python_default_ttl != native_default_ttl.as_secs_f64() { + return Err(PyTypeError::new_err( + "facade and native backend default TTLs must match", + )); + } Ok(Self { outer: ObjectGuard::capture( py, diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 5918967009a..4bab03fe243 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -4,7 +4,7 @@ mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use litellm_cache::Error; -use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; use pyo3::{ PyTraverseError, PyVisit, @@ -15,9 +15,26 @@ use pyo3::{ use serde::Deserialize; use serde_json::Value; +use crate::python_settings::PythonSettings; use facade::FacadeGuard; use native::NativeResponseCache; +const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(FromPyObject)] +struct PythonCacheSettings { + default_redis_ttl: Option, +} + +fn redis_default_ttl(py: Python<'_>) -> PyResult { + let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; + settings + .default_redis_ttl + .map(duration) + .transpose() + .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct RequestInput { @@ -29,6 +46,10 @@ struct RequestInput { 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; @@ -38,6 +59,13 @@ fn request(value: &Bound<'_, PyAny>) -> PyResult { Ok(request) } +fn requests(value: &Bound<'_, PyAny>) -> PyResult> { + from_py::>(value)? + .into_iter() + .map(request_input) + .collect() +} + fn duration(seconds: f64) -> PyResult { Duration::try_from_secs_f64(seconds) .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) @@ -94,7 +122,10 @@ impl NativeCacheHandle { ttl_seconds: Option, namespace: Option, ) -> PyResult { - let ttl = ttl_seconds.map(duration).transpose()?; + let ttl = Some(match ttl_seconds { + Some(seconds) => duration(seconds)?, + None => redis_default_ttl(py)?, + }); let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) .map_err(cache_error)?; Ok(Self { @@ -111,7 +142,12 @@ impl NativeCacheHandle { fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, self.backend())?; + let guard = FacadeGuard::capture(py, facade, self.backend(), service.default_ttl())?; + let service = service.with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); let handle = Py::new( py, Self { @@ -249,6 +285,37 @@ impl ResolvedCache { } } + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn lookup_batch( + &self, + py: Python<'_>, + requests: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + to_py(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup_batch(&requests, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "batch_get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(Bound::unbind), + } + } + #[pyo3(signature = (request, *, callback_kwargs=None))] fn async_lookup<'py>( &self, @@ -292,6 +359,102 @@ impl ResolvedCache { } } + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + ready_value(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup_batch(&requests, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_batch_get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + #[pyo3(signature = (requests, responses, *, callback_kwargs=None))] + fn async_store_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + responses: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let responses: Vec = from_py(responses)?; + if requests.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let entries = requests.into_iter().zip(responses).collect(); + let service = service.clone(); + run_async( + py, + async move { service.async_store_batch(entries, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_set_cache_pipeline", + (responses,), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async(py, async move { service.async_flush().await }, cache_error) + } + CacheBinding::PythonCallback(object) => { + object.bind(py).call_method0("flush_cache")?; + ready_none(py) + } + } + } + + fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async( + py, + async move { service.test_connection().await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method0("test_connection"), + } + } + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { if let CacheBinding::PythonCallback(object) = &self.binding { visit.call(object)?; @@ -309,11 +472,18 @@ fn callback_kwargs<'a, 'py>( } fn ready_none(py: Python<'_>) -> PyResult> { + ready_value(py, &()) +} + +fn ready_value<'py, T: serde::Serialize>( + py: Python<'py>, + value: &T, +) -> PyResult> { let future = py .import("asyncio")? .call_method0("get_running_loop")? .call_method0("create_future")?; - future.call_method1("set_result", (py.None(),))?; + future.call_method1("set_result", (to_py(py, value)?,))?; Ok(future) } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 6af04bfe2b2..20891719550 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,16 +1,27 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use serde_json::Value; +use tokio::sync::Mutex; -use litellm_cache_response::{CacheEntry, ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_response::{ + CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), - Redis(Arc>>), + Redis { + cache: Arc>>, + buffer: Option>, + }, +} + +pub(super) struct RedisWriteBuffer { + flush_size: usize, + entries: Mutex>, } impl NativeResponseCache { @@ -34,7 +45,10 @@ impl NativeResponseCache { namespace: Option, ) -> Result { let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace); - Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend))))) + Ok(Self::Redis { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + buffer: None, + }) } } @@ -42,7 +56,29 @@ impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", - Self::Redis(_) => "redis", + Self::Redis { .. } => "redis", + } + } + + pub fn default_ttl(&self) -> Duration { + match self { + Self::Memory(cache) => cache.default_ttl(), + Self::Redis { cache, .. } => cache.default_ttl(), + } + } + + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { + match self { + Self::Redis { cache, .. } => Self::Redis { + cache, + buffer: flush_size.map(|flush_size| { + Arc::new(RedisWriteBuffer { + flush_size: flush_size.max(1), + entries: Mutex::new(Vec::new()), + }) + }), + }, + memory => memory, } } @@ -53,7 +89,7 @@ impl NativeResponseCache { ) -> Result, Error> { match self { Self::Memory(cache) => cache.lookup(request, now), - Self::Redis(cache) => cache.lookup(request, now), + Self::Redis { cache, .. } => cache.lookup(request, now), } } @@ -65,7 +101,18 @@ impl NativeResponseCache { ) -> Result<(), Error> { match self { Self::Memory(cache) => cache.store(request, response, now), - Self::Redis(cache) => cache.store(request, response, now), + Self::Redis { cache, .. } => cache.store(request, response, now), + } + } + + pub fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + match self { + Self::Memory(cache) => cache.lookup_batch(requests, now), + Self::Redis { cache, .. } => cache.lookup_batch(requests, now), } } @@ -76,7 +123,7 @@ impl NativeResponseCache { ) -> Result, Error> { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis(cache) => cache.async_lookup(request, now).await, + Self::Redis { cache, .. } => cache.async_lookup(request, now).await, } } @@ -88,7 +135,71 @@ impl NativeResponseCache { ) -> Result<(), Error> { match self { Self::Memory(cache) => cache.async_store(request, response, now).await, - Self::Redis(cache) => cache.async_store(request, response, now).await, + Self::Redis { + cache, + buffer: None, + } => cache.async_store(request, response, now).await, + Self::Redis { + cache, + buffer: Some(buffer), + } => { + let pending = { + let mut entries = buffer.entries.lock().await; + entries.push((request.clone(), response)); + (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) + }; + let Some(pending) = pending else { + return Ok(()); + }; + if let Err(error) = cache.async_store_batch(pending.clone(), now).await { + let mut entries = buffer.entries.lock().await; + let current = std::mem::take(&mut *entries); + *entries = pending.into_iter().chain(current).collect(); + return Err(error); + } + Ok(()) + } + } + } + + pub async fn async_lookup_batch( + &self, + requests: &[ResponseCacheRequest], + 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, + } + } + + pub async fn async_store_batch( + &self, + entries: Vec<(ResponseCacheRequest, 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, + } + } + + pub async fn async_flush(&self) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_flush().await, + Self::Redis { cache, buffer } => { + if let Some(buffer) = buffer { + buffer.entries.lock().await.clear(); + } + cache.async_flush().await + } + } + } + + pub async fn test_connection(&self) -> Result { + match self { + Self::Memory(cache) => cache.test_connection().await, + Self::Redis { cache, .. } => cache.test_connection().await, } } } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 7ac23a05542..90819c5b3fc 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -8,15 +8,17 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, + Cache, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 4] = [ + pub(crate) const ALL: [Self; 5] = [ Self::Http, Self::UrlPolicy, Self::ProviderDefaults, Self::SecretManager, + Self::Cache, ]; pub(crate) fn name(self) -> &'static str { @@ -25,6 +27,7 @@ impl PythonSettings { Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", + Self::Cache => "cache_settings", } } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 66be77dbb40..64a5af2c618 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -12,7 +12,7 @@ import logging import time from collections.abc import Sequence from threading import Lock -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeVar if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -34,14 +34,20 @@ else: from collections import OrderedDict +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") -class LimitedSizeOrderedDict(OrderedDict): - def __init__(self, *args, max_size=100, **kwargs): - super().__init__(*args, **kwargs) + +class LimitedSizeOrderedDict(OrderedDict[_KeyT, _ValueT]): + def __init__(self, *, max_size: int = 100) -> None: + super().__init__() self.max_size = max_size - def __setitem__(self, key, value): - # If inserting a new key exceeds max size, remove the oldest item + def __setitem__(self, key: _KeyT, value: _ValueT) -> None: + if key in self: + super().__setitem__(key, value) + self.move_to_end(key) + return if len(self) >= self.max_size: self.popitem(last=False) super().__setitem__(key, value) @@ -68,7 +74,9 @@ class DualCache(BaseCache): self.in_memory_cache = in_memory_cache or InMemoryCache() # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache - self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size) + self.last_redis_batch_access_time: LimitedSizeOrderedDict[str, float] = LimitedSizeOrderedDict( + max_size=default_max_redis_batch_cache_size + ) self._last_redis_batch_access_time_lock = Lock() self.redis_batch_cache_expiry = ( default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10 @@ -131,7 +139,7 @@ class DualCache(BaseCache): except Exception as e: print_verbose(e) - def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> int: + def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> float: """ Key - the key in cache @@ -140,14 +148,15 @@ class DualCache(BaseCache): Returns - int - the incremented value """ try: - result: int = value - if self.in_memory_cache is not None: - result = self.in_memory_cache.increment_cache(key, value, **kwargs) - if self.redis_cache is not None and local_only is False: - result = self.redis_cache.increment_cache(key, value, **kwargs) + result: Final = self.redis_cache.increment_cache(key, value, **kwargs) + if self.in_memory_cache is not None: + self.in_memory_cache.set_cache(key, result, **kwargs) + return result - return result + if self.in_memory_cache is not None: + return self.in_memory_cache.increment_cache(key, value, **kwargs) + return value except Exception as e: verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e) raise e @@ -421,29 +430,30 @@ class DualCache(BaseCache): Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ - result: float | None = None try: - if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment(key, value, **kwargs) - if self.redis_cache is not None and local_only is False: - result = await self.redis_cache.async_increment( + result: Final = await self.redis_cache.async_increment( key, value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), refresh_ttl=refresh_ttl, ) + if self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache(key, result, **kwargs) + return result - return result + if self.in_memory_cache is not None: + return await self.in_memory_cache.async_increment(key, value, **kwargs) + return None except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache failed, falling back to in-memory result", + "Redis async_increment_cache failed; local counter unchanged", e, ) - return result + return None async def async_increment_cache_pipeline( self, @@ -452,29 +462,32 @@ class DualCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> list[float] | None: - result: list[float] | None = None try: - if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment_pipeline( - increment_list=increment_list, - parent_otel_span=parent_otel_span, - ) - if self.redis_cache is not None and local_only is False: - result = await self.redis_cache.async_increment_pipeline( + result: Final = await self.redis_cache.async_increment_pipeline( increment_list=increment_list, parent_otel_span=parent_otel_span, ) + if result is not None and self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache_pipeline( + cache_list=tuple((increment["key"], value) for increment, value in zip(increment_list, result)) + ) + return result - return result + if self.in_memory_cache is not None: + return await self.in_memory_cache.async_increment_pipeline( + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + return None except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result", + "Redis async_increment_cache_pipeline failed; local counters unchanged", e, ) - return result + return None async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None: """ diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 4fd2f0829a3..68b1742dc41 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -126,6 +126,12 @@ class CacheBinding: *, callback_kwargs: dict[str, object] | None = None, ) -> None: ... + def lookup_batch( + self, + requests: Sequence[Mapping[str, object]], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> object: ... def async_lookup( self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None ) -> Awaitable[object]: ... @@ -136,6 +142,21 @@ class CacheBinding: *, callback_kwargs: dict[str, object] | None = None, ) -> Awaitable[object]: ... + def async_lookup_batch( + self, + requests: Sequence[Mapping[str, object]], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + def async_store_batch( + self, + requests: Sequence[Mapping[str, object]], + responses: Sequence[object], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + def async_flush(self) -> Awaitable[None]: ... + def ping(self) -> Awaitable[dict[str, object] | None]: ... @final class TokenCounter: diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 3aa2d742862..862a116496d 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -36,6 +36,11 @@ class SecretManager: readable: bool +@dataclass(frozen=True, slots=True) +class CacheSettings: + default_redis_ttl: float | None + + def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -50,6 +55,12 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) +def cache_settings() -> CacheSettings: + import litellm + + return CacheSettings(default_redis_ttl=litellm.default_redis_ttl) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 5f59de9cca5..149c9b34bd5 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -6,18 +6,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE -from litellm.caching.dual_cache import DualCache +from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.types.caching import RedisPipelineIncrementOperation @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads(): - dual_cache = DualCache( - redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) keys = ["shared_a", "shared_b"] start_gate = asyncio.Event() @@ -44,9 +42,7 @@ async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_error(): - dual_cache = DualCache( - redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) keys = ["shared_a", "shared_b"] with patch.object( @@ -116,9 +112,7 @@ def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) first = dual_cache.batch_get_cache(keys=["absent_key"]) second = dual_cache.batch_get_cache(keys=["absent_key"]) @@ -131,9 +125,7 @@ def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): mock_redis = MagicMock(spec=RedisCache) mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) first_result = dual_cache.batch_get_cache(keys=["shared_a"]) second_result = dual_cache.batch_get_cache(keys=["shared_a"]) @@ -146,9 +138,7 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() result = dual_cache.batch_get_cache(keys=["throttled_key"]) @@ -257,9 +247,7 @@ async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl(): default_in_memory_ttl, same as the single-key path.""" in_memory_cache = InMemoryCache(default_ttl=600) mock_redis = MagicMock(spec=RedisCache) - mock_redis.async_batch_get_cache = AsyncMock( - return_value={"batch_backfill_key": "redis_value"} - ) + mock_redis.async_batch_get_cache = AsyncMock(return_value={"batch_backfill_key": "redis_value"}) dual_cache = DualCache( in_memory_cache=in_memory_cache, redis_cache=mock_redis, @@ -371,9 +359,7 @@ async def test_circuit_breaker_open_skips_redis(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=3, recovery_timeout=60 - ) + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) self._circuit_breaker._state = "open" self._circuit_breaker._opened_at = time.time() self.call_count = 0 @@ -426,9 +412,7 @@ def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): # All subsequent concurrent callers: HALF_OPEN → fast-fail (return True) for _ in range(10): - assert ( - cb.is_open() is True - ), "concurrent callers should be fast-failed in HALF_OPEN" + assert cb.is_open() is True, "concurrent callers should be fast-failed in HALF_OPEN" def test_circuit_breaker_disabled_never_opens(): @@ -472,9 +456,7 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=1, recovery_timeout=60, enabled=False - ) + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) self.call_count = 0 @_redis_circuit_breaker_guard @@ -512,6 +494,30 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re ) +@pytest.mark.asyncio +async def test_failed_redis_increment_does_not_change_the_local_counter(): + memory = InMemoryCache() + memory.set_cache("counter", 10) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) + cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) + + assert await cache.async_increment_cache("counter", 2) is None + assert memory.get_cache("counter") == 10 + + +@pytest.mark.asyncio +async def test_successful_redis_increment_replaces_the_local_counter_with_the_authoritative_value(): + memory = InMemoryCache() + memory.set_cache("counter", 10) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment = AsyncMock(return_value=42.0) + cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) + + assert await cache.async_increment_cache("counter", 2) == 42.0 + assert memory.get_cache("counter") == 42.0 + + def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): """ Typical lazy startup (sync): DualCache runs with in-memory only, then Redis @@ -742,7 +748,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in visible] == [ ( logging.WARNING, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" + "Redis async_increment_cache_pipeline failed; local counters unchanged:" " Timeout reading from 127.0.0.1:6379", ) ] @@ -756,7 +762,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ ( logging.WARNING, - "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" + "Redis async_increment_cache failed; local counter unchanged: Timeout reading from 127.0.0.1:6379" " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", ) ] @@ -791,3 +797,14 @@ async def test_async_delete_cache_keys_on_empty_list_touches_no_backend(): await dual_cache.async_delete_cache_keys([]) redis_cache.delete_cache_keys.assert_not_awaited() + + +def test_limited_ordered_dict_refreshes_recency_without_evicting_another_key(): + tracker = LimitedSizeOrderedDict(max_size=2) + tracker["hot"] = 1 + tracker["cold"] = 2 + + tracker["hot"] = 3 + tracker["new"] = 4 + + assert list(tracker.items()) == [("hot", 3), ("new", 4)] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 493baac228a..7456a1499f7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -8,6 +8,7 @@ import weakref from collections.abc import Generator from types import SimpleNamespace from typing import Final, Protocol, cast +from urllib.parse import urlparse import fakeredis import pytest @@ -209,8 +210,12 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + client.set("team:raw", json.dumps(response)) + client.set("team:invalid", "not a cache entry") assert binding.lookup(request("sync")) == response assert await binding.async_lookup(request("team:async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) stored: Final = client.get("team:native") assert isinstance(stored, bytes) @@ -245,3 +250,63 @@ async def test_memory_size_policy_is_applied_by_the_native_host() -> None: ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None + + +async def test_native_batch_lookup_and_store_report_partial_hits() -> None: + binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } + + +async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + with rebound(litellm, "default_redis_ttl", 7): + binding: Final = _native.CacheResolver( + SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url)) + ).resolve() + await binding.async_store(request("native-default"), {"value": 1}) + + assert 0 < client.ttl("native-default") <= 7 + client.close() + + +async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: + parsed: Final = urlparse(redis_url) + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache( + type=LiteLLMCacheType.REDIS, + host=parsed.hostname, + port=str(parsed.port), + redis_flush_size=2, + ) + with pytest.raises(TypeError, match="default TTLs must match"): + _native.NativeCacheHandle.redis(redis_url, ttl_seconds=61).bind_facade(facade) + _native.NativeCacheHandle.redis(redis_url).bind_facade(facade) + binding: Final = _native.CacheResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(redis_url) + + await binding.async_store(request("first"), {"value": 1}) + assert client.get("first") is None + await binding.async_store(request("second"), {"value": 2}) + + assert client.get("first") is not None + assert client.get("second") is not None + await facade.cache.disconnect() + client.close() From 9783b7a377009982abd9541b305d466590770252 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 07:56:58 -0700 Subject: [PATCH 23/56] docs(cache): align native cache follow-up scope --- litellm-rust/crates/cache-response/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 4e6694c191f..d8ffd6d6a50 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -52,4 +52,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, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache now provides L2-first counters and atomic affinity claims with local fallback, but public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees From da402b8aeee24203fad87a646ce3b5fec7b9ff51 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 08:27:17 -0700 Subject: [PATCH 24/56] fix(cache): preserve batch callback contracts --- .../crates/python-bridge/src/cache/mod.rs | 41 ++++++++++--- litellm/caching/dual_cache.py | 61 +++++++++---------- tests/test_litellm/caching/test_dual_cache.py | 28 +-------- tests/test_litellm_rust/test_cache.py | 39 ++++++++++++ 4 files changed, 102 insertions(+), 67 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4bab03fe243..1e2e42600ff 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 @@ use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::PyDict, + types::{PyDict, PyList}, }; use serde::Deserialize; use serde_json::Value; @@ -309,7 +309,7 @@ impl ResolvedCache { .bind(py) .call_method( "batch_get_cache", - (), + (callback_keys(py, requests)?,), Some(self::callback_kwargs(callback_kwargs)?), ) .map(Bound::unbind), @@ -383,7 +383,7 @@ impl ResolvedCache { } CacheBinding::PythonCallback(object) => object.bind(py).call_method( "async_batch_get_cache", - (), + (callback_keys(py, requests)?,), Some(self::callback_kwargs(callback_kwargs)?), ), } @@ -416,11 +416,24 @@ impl ResolvedCache { cache_error, ) } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_set_cache_pipeline", - (responses,), - Some(self::callback_kwargs(callback_kwargs)?), - ), + CacheBinding::PythonCallback(object) => { + let keys = callback_keys(py, requests)?; + let responses = responses.try_iter()?.collect::>>()?; + if keys.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let cache_list = PyList::empty(py); + for (key, response) in keys.iter().zip(responses) { + cache_list.append((key, response))?; + } + object.bind(py).call_method( + "async_set_cache_pipeline", + (cache_list,), + Some(self::callback_kwargs(callback_kwargs)?), + ) + } } } @@ -471,6 +484,18 @@ fn callback_kwargs<'a, 'py>( }) } +fn callback_keys<'py>( + py: Python<'py>, + requests: &Bound<'py, PyAny>, +) -> PyResult> { + PyList::new( + py, + self::requests(requests)? + .into_iter() + .map(|request| litellm_cache_response::cache_key(&request.key)), + ) +} + fn ready_none(py: Python<'_>) -> PyResult> { ready_value(py, &()) } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 64a5af2c618..04c82232784 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -139,7 +139,7 @@ class DualCache(BaseCache): except Exception as e: print_verbose(e) - def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> float: + def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> int: """ Key - the key in cache @@ -148,15 +148,14 @@ class DualCache(BaseCache): Returns - int - the incremented value """ try: - if self.redis_cache is not None and local_only is False: - result: Final = self.redis_cache.increment_cache(key, value, **kwargs) - if self.in_memory_cache is not None: - self.in_memory_cache.set_cache(key, result, **kwargs) - return result - + result: int = value if self.in_memory_cache is not None: - return self.in_memory_cache.increment_cache(key, value, **kwargs) - return value + result = self.in_memory_cache.increment_cache(key, value, **kwargs) + + if self.redis_cache is not None and local_only is False: + result = self.redis_cache.increment_cache(key, value, **kwargs) + + return result except Exception as e: verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e) raise e @@ -430,30 +429,29 @@ class DualCache(BaseCache): Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ + result: float | None = None try: + if self.in_memory_cache is not None: + result = await self.in_memory_cache.async_increment(key, value, **kwargs) + if self.redis_cache is not None and local_only is False: - result: Final = await self.redis_cache.async_increment( + result = await self.redis_cache.async_increment( key, value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), refresh_ttl=refresh_ttl, ) - if self.in_memory_cache is not None: - await self.in_memory_cache.async_set_cache(key, result, **kwargs) - return result - if self.in_memory_cache is not None: - return await self.in_memory_cache.async_increment(key, value, **kwargs) - return None + return result except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache failed; local counter unchanged", + "Redis async_increment_cache failed, falling back to in-memory result", e, ) - return None + return result async def async_increment_cache_pipeline( self, @@ -462,32 +460,29 @@ class DualCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> list[float] | None: + result: list[float] | None = None try: - if self.redis_cache is not None and local_only is False: - result: Final = await self.redis_cache.async_increment_pipeline( - increment_list=increment_list, - parent_otel_span=parent_otel_span, - ) - if result is not None and self.in_memory_cache is not None: - await self.in_memory_cache.async_set_cache_pipeline( - cache_list=tuple((increment["key"], value) for increment, value in zip(increment_list, result)) - ) - return result - if self.in_memory_cache is not None: - return await self.in_memory_cache.async_increment_pipeline( + result = await self.in_memory_cache.async_increment_pipeline( increment_list=increment_list, parent_otel_span=parent_otel_span, ) - return None + + if self.redis_cache is not None and local_only is False: + result = await self.redis_cache.async_increment_pipeline( + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + + return result except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache_pipeline failed; local counters unchanged", + "Redis async_increment_cache_pipeline failed, falling back to in-memory result", e, ) - return None + return result async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None: """ diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 149c9b34bd5..05a4ef68e0b 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -494,30 +494,6 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re ) -@pytest.mark.asyncio -async def test_failed_redis_increment_does_not_change_the_local_counter(): - memory = InMemoryCache() - memory.set_cache("counter", 10) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) - cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) - - assert await cache.async_increment_cache("counter", 2) is None - assert memory.get_cache("counter") == 10 - - -@pytest.mark.asyncio -async def test_successful_redis_increment_replaces_the_local_counter_with_the_authoritative_value(): - memory = InMemoryCache() - memory.set_cache("counter", 10) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_increment = AsyncMock(return_value=42.0) - cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) - - assert await cache.async_increment_cache("counter", 2) == 42.0 - assert memory.get_cache("counter") == 42.0 - - def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): """ Typical lazy startup (sync): DualCache runs with in-memory only, then Redis @@ -748,7 +724,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in visible] == [ ( logging.WARNING, - "Redis async_increment_cache_pipeline failed; local counters unchanged:" + "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" " Timeout reading from 127.0.0.1:6379", ) ] @@ -762,7 +738,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ ( logging.WARNING, - "Redis async_increment_cache failed; local counter unchanged: Timeout reading from 127.0.0.1:6379" + "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", ) ] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 7456a1499f7..d38583dad0a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -275,6 +275,45 @@ async def test_native_batch_lookup_and_store_report_partial_hits() -> None: } +async def test_python_batch_callbacks_receive_keys_and_key_value_pairs() -> None: + first: Final = object() + second: Final = object() + + class CustomCache: + def batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: + return keys, marker + + async def async_batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: + return keys, marker + + async def async_set_cache_pipeline( + self, cache_list: list[tuple[str, object]], *, marker: object + ) -> tuple[list[tuple[str, object]], object]: + return cache_list, marker + + marker: Final = object() + binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() + requests: Final = [request("first"), request("second")] + + assert binding.lookup_batch(requests, callback_kwargs={"marker": marker}) == (["first", "second"], marker) + assert await binding.async_lookup_batch(requests, callback_kwargs={"marker": marker}) == ( + ["first", "second"], + marker, + ) + stored: Final = cast( + tuple[list[tuple[str, object]], object], + await binding.async_store_batch( + requests, + [first, second], + callback_kwargs={"marker": marker}, + ), + ) + assert [key for key, _ in stored[0]] == ["first", "second"] + assert stored[1] is marker + assert stored[0][0][1] is first + assert stored[0][1][1] is second + + async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) with rebound(litellm, "default_redis_ttl", 7): From ef14fdaf339006a1f2bcb443ecc14225554d1568 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 09:53:41 -0700 Subject: [PATCH 25/56] fix(cache): harden native foundation parity --- litellm-rust/Cargo.lock | 1 - litellm-rust/crates/cache-memory/src/cache.rs | 98 ++++-- .../crates/cache-memory/tests/cache.rs | 79 ++++- litellm-rust/crates/cache-redis/Cargo.toml | 2 +- litellm-rust/crates/cache-redis/src/cache.rs | 236 ++++++++++---- .../crates/cache-redis/tests/cache.rs | 127 +++++++- litellm-rust/crates/cache-response/README.md | 12 +- .../crates/cache-response/src/codec.rs | 25 +- .../crates/cache-response/src/response.rs | 50 +-- .../crates/cache-response/tests/response.rs | 75 ++++- litellm-rust/crates/cache/src/base_cache.rs | 3 +- litellm-rust/crates/cache/src/caching.rs | 3 +- litellm-rust/crates/cache/src/codec.rs | 8 + litellm-rust/crates/cache/src/dual.rs | 302 +++++++++++++++++- litellm-rust/crates/cache/tests/caching.rs | 3 +- litellm-rust/crates/cache/tests/dual.rs | 216 ++++++++++++- .../crates/python-bridge/src/cache/facade.rs | 32 +- .../crates/python-bridge/src/cache/mod.rs | 136 ++++---- .../crates/python-bridge/src/cache/native.rs | 37 ++- litellm-rust/crates/python-bridge/src/lib.rs | 16 +- litellm/rust_bridge/_native.pyi | 26 +- tests/test_litellm/caching/test_dual_cache.py | 2 +- tests/test_litellm_rust/test_cache.py | 176 ++++++---- 23 files changed, 1347 insertions(+), 318 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 725d2cfef41..2d6fb6c3082 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3712,7 +3712,6 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", - "r2d2", "ryu", "sha1_smol", "socket2 0.6.5", diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 074c8dcb0fd..a9814ff6fd6 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -1,7 +1,9 @@ -use std::cmp::Reverse; -use std::collections::{BinaryHeap, HashMap}; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::{ + cmp::Reverse, + collections::{BinaryHeap, HashMap}, + sync::{Arc, Mutex}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; use litellm_cache::{ BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, @@ -95,15 +97,13 @@ impl InMemoryCache { } let now = (self.now)(); let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; - Self::evict(&mut state, self.max_size_in_memory, now); let key = key.into(); - state.values.insert(key.clone(), value); + Self::evict(&mut state, self.max_size_in_memory, now, &key); let expiration = state.expirations.get(&key).copied(); if expiration.is_none_or(|expiration| expiration < now) { - let expiration = now + ttl.unwrap_or(self.default_ttl); - state.expirations.insert(key.clone(), expiration); - state.expiration_heap.push(Reverse((expiration, key))); + Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl)); } + state.values.insert(key, value); Ok(CacheWrite::Stored) } @@ -120,6 +120,10 @@ impl InMemoryCache { Ok(state.values.get(key).cloned()) } + pub fn max_size_in_memory(&self) -> usize { + self.max_size_in_memory + } + pub fn expires_at(&self, key: &str) -> Result, Error> { Ok(self .state @@ -144,7 +148,7 @@ impl InMemoryCache { Ok(()) } - fn evict(state: &mut CacheState, capacity: usize, now: Duration) { + fn evict(state: &mut CacheState, capacity: usize, now: Duration, key: &str) { while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() { if state.expirations.get(&key).copied() != Some(expiration) { state.expiration_heap.pop(); @@ -155,6 +159,9 @@ impl InMemoryCache { break; } } + if state.values.contains_key(key) { + return; + } while state.values.len() >= capacity { let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else { break; @@ -165,6 +172,15 @@ impl InMemoryCache { } } + fn set_expiration(state: &mut CacheState, key: &str, expiration: Duration) { + if state.expirations.get(key).copied() != Some(expiration) { + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + } + } + fn remove(state: &mut CacheState, key: &str) { state.values.remove(key); state.expirations.remove(key); @@ -182,40 +198,44 @@ where eligible: &[V], kwargs: CacheKwargs, ) -> Result { + if self.max_size_in_memory == 0 { + return Ok(candidate); + } let now = (self.now)(); let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; - Self::evict(&mut state, self.max_size_in_memory, now); - let winner = match state.values.get(key) { - Some(existing) if eligible.is_empty() => existing.clone(), - Some(existing) if eligible.contains(existing) => existing.clone(), - _ => candidate, - }; - let expiration = now + self.get_ttl(&kwargs); + Self::evict(&mut state, self.max_size_in_memory, now, key); + let existing = state + .values + .get(key) + .filter(|existing| eligible.is_empty() || eligible.contains(existing)) + .cloned(); + // Matches the Redis claim: an unconditional claim only extends its own winner. + if let Some(existing) = &existing + && eligible.is_empty() + && *existing != candidate + { + return Ok(existing.clone()); + } + let winner = existing.unwrap_or(candidate); + Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs)); state.values.insert(key.into(), winner.clone()); - state.expirations.insert(key.into(), expiration); - state - .expiration_heap - .push(Reverse((expiration, key.into()))); Ok(winner) } } impl CounterCache for InMemoryCache { fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + if self.max_size_in_memory == 0 { + return Ok(amount); + } let now = (self.now)(); let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; - Self::evict(&mut state, self.max_size_in_memory, now); + Self::evict(&mut state, self.max_size_in_memory, now, key); let value = state.values.get(key).copied().unwrap_or_default() + amount; - let expiration = state - .expirations - .get(key) - .copied() - .unwrap_or_else(|| now + self.get_ttl(&kwargs)); + if !state.expirations.contains_key(key) { + Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs)); + } state.values.insert(key.into(), value); - state.expirations.insert(key.into(), expiration); - state - .expiration_heap - .push(Reverse((expiration, key.into()))); Ok(value) } } @@ -256,3 +276,19 @@ impl BaseCache for InMemoryCache { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn repeated_increments_keep_one_heap_entry_per_expiration() { + let cache = InMemoryCache::::new(Some(4), None); + for _ in 0..100 { + cache + .increment_cache("counter", 1.0, CacheKwargs::default()) + .unwrap(); + } + assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1); + } +} diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index e5831dfd6d5..22c5595da52 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -1,6 +1,10 @@ -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; use litellm_cache::{ BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, @@ -202,6 +206,17 @@ fn claims_are_atomic_and_refresh_eligible_winners() { .unwrap(), "first" ); + clock.store(103, Ordering::SeqCst); + assert_eq!( + cache + .claim_cache("affinity", "second".to_string(), &[], kwargs.clone()) + .unwrap(), + "first" + ); + assert_eq!( + cache.expires_at("affinity").unwrap(), + Some(Duration::from_secs(110)) + ); clock.store(105, Ordering::SeqCst); assert_eq!( cache @@ -232,3 +247,61 @@ fn counters_increment_under_one_lock() { 3.5 ); } + +#[rstest] +fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc) { + let cache = cache(clock, 2); + cache + .set_cache("hot", "1".into(), Some(Duration::from_secs(10))) + .unwrap(); + cache + .set_cache("cold", "2".into(), Some(Duration::from_secs(20))) + .unwrap(); + + cache.set_cache("cold", "3".into(), None).unwrap(); + assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into())); + assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into())); + + cache + .claim_cache("cold", "4".into(), &[], CacheKwargs::default()) + .unwrap(); + assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into())); + + cache.set_cache("new", "5".into(), None).unwrap(); + assert_eq!(cache.get_cache("hot").unwrap(), None); + assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into())); + assert_eq!(cache.get_cache("new").unwrap(), Some("5".into())); +} + +#[test] +fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() { + let cache = InMemoryCache::::new(Some(2), None); + for key in ["a", "b", "a", "b"] { + cache + .increment_cache(key, 1.0, CacheKwargs::default()) + .unwrap(); + } + assert_eq!(cache.get_cache("a").unwrap(), Some(2.0)); + assert_eq!(cache.get_cache("b").unwrap(), Some(2.0)); +} + +#[test] +fn disabled_cache_does_not_retain_claims_or_counters() { + let claims = InMemoryCache::::new(Some(0), None); + assert_eq!( + claims + .claim_cache("key", "first".into(), &[], CacheKwargs::default()) + .unwrap(), + "first" + ); + assert_eq!(claims.get_cache("key").unwrap(), None); + + let counters = InMemoryCache::::new(Some(0), None); + assert_eq!( + counters + .increment_cache("key", 2.0, CacheKwargs::default()) + .unwrap(), + 2.0 + ); + assert_eq!(counters.get_cache("key").unwrap(), None); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 1a1bb505a7c..f123e774158 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = { version = "1.7.0", features = ["r2d2"] } +redis = "1.7.0" r2d2 = "0.8.10" tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 281ffe17534..b8f0d3857e2 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,5 +1,7 @@ -use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use litellm_cache::{ BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, @@ -11,8 +13,59 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; +struct PooledConnection { + connection: redis::Connection, + failed: bool, +} + +/// Pools connections without a checkout PING, which would double every operation's round trips. +/// A timed-out command leaves its reply on the socket while redis still reports the connection +/// open, so any connection whose operation failed is discarded instead of being reused. +struct ConnectionManager(redis::Client); + +impl r2d2::ManageConnection for ConnectionManager { + type Connection = PooledConnection; + type Error = redis::RedisError; + + fn connect(&self) -> Result { + let connection = self.0.get_connection()?; + connection.set_read_timeout(Some(REDIS_TIMEOUT))?; + connection.set_write_timeout(Some(REDIS_TIMEOUT))?; + Ok(PooledConnection { + connection, + failed: false, + }) + } + + fn is_valid(&self, connection: &mut PooledConnection) -> Result<(), redis::RedisError> { + redis::cmd("PING").query::(&mut connection.connection)?; + Ok(()) + } + + fn has_broken(&self, connection: &mut PooledConnection) -> bool { + connection.failed || !redis::ConnectionLike::is_open(&connection.connection) + } +} + +const INCREMENT_SCRIPT: &str = concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" +); + +// Compare-and-set against the exact bytes the claim decision was made on. +// ARGV: [1] expected payload or "" when absent, [2] ttl, [3] new payload, [4] refresh ttl. +const CLAIM_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", + "elseif current ~= ARGV[1] then return 0; end; ", + "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", + "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" +); +const CLAIM_ATTEMPTS: usize = 8; + enum Connections { - Pool(r2d2::Pool), + Pool(r2d2::Pool), Fixed(Mutex), } @@ -59,14 +112,10 @@ where ) -> Result { match self { Self::Pool(pool) => { - let mut connection = pool.get().map_err(|_| Error::Unavailable)?; - connection - .set_read_timeout(Some(REDIS_TIMEOUT)) - .map_err(|_| Error::Unavailable)?; - connection - .set_write_timeout(Some(REDIS_TIMEOUT)) - .map_err(|_| Error::Unavailable)?; - operation(&mut ConnectionRef(&mut *connection)) + 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)?; @@ -90,7 +139,8 @@ impl RedisCache { .max_size(REDIS_POOL_SIZE) .min_idle(Some(0)) .connection_timeout(REDIS_TIMEOUT) - .build(client) + .test_on_check_out(false) + .build(ConnectionManager(client)) .map_err(|_| Error::Unavailable)?; Ok(Self { connections: Arc::new(Connections::Pool(pool)), @@ -122,6 +172,10 @@ where } } + pub fn namespace(&self) -> Option<&str> { + self.namespace.as_deref() + } + fn namespaced_key(&self, key: &str) -> String { match &self.namespace { Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { @@ -407,29 +461,105 @@ where C: redis::ConnectionLike + Send + 'static, { fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { - const SCRIPT: &str = concat!( - "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", - "if redis.call('TTL', KEYS[1]) == -1 then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" - ); let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - self.connections.execute(|connection| { - redis::cmd("EVAL") - .arg(SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(|_| Error::Unavailable) - }) + self.connections + .execute(|connection| increment(connection, key, amount, ttl)) } + + async fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + increment(connection, key, amount, ttl) + }) + .await + } +} + +fn increment( + connection: &mut ConnectionRef<'_>, + key: String, + amount: f64, + ttl: u64, +) -> Result { + redis::cmd("EVAL") + .arg(INCREMENT_SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) +} + +fn stored_bytes(value: redis::Value) -> Result>, Error> { + match value { + redis::Value::Nil => Ok(None), + redis::Value::BulkString(bytes) => Ok(Some(bytes)), + redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())), + _ => Err(Error::InvalidEntry), + } +} + +/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's +/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the +/// bytes that decision was made on, retried when another claimant wins the race. +fn claim( + connection: &mut ConnectionRef<'_>, + codec: &S, + key: &str, + candidate: S::Value, + eligible: &[S::Value], + ttl: u64, +) -> Result +where + S::Value: PartialEq, +{ + let payload = codec.encode(&candidate)?; + if payload.is_empty() { + return Err(Error::InvalidEntry); + } + for _ in 0..CLAIM_ATTEMPTS { + let current = stored_bytes( + connection + .get::<_, redis::Value>(key) + .map_err(|_| Error::Unavailable)?, + )? + .filter(|bytes| !bytes.is_empty()); + let existing = current + .as_deref() + .and_then(|bytes| codec.decode(bytes).ok()) + .filter(|existing| eligible.is_empty() || eligible.contains(existing)); + let refresh = existing + .as_ref() + .is_some_and(|existing| !eligible.is_empty() || *existing == candidate); + let write: &[u8] = if existing.is_some() { b"" } else { &payload }; + let applied = redis::cmd("EVAL") + .arg(CLAIM_SCRIPT) + .arg(1) + .arg(key) + .arg(current.as_deref().unwrap_or_default()) + .arg(ttl) + .arg(write) + .arg(u8::from(refresh)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + if applied { + return Ok(existing.unwrap_or(candidate)); + } + } + Err(Error::Unavailable) } impl ClaimCache for RedisCache where - S: CacheCodec, + S: CacheCodec + Clone + 'static, S::Value: PartialEq, C: redis::ConnectionLike + Send + 'static, { @@ -440,44 +570,38 @@ where eligible: &[S::Value], kwargs: CacheKwargs, ) -> Result { - const SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if current == false then redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", - "return ARGV[1]; end; if #ARGV > 2 then for index = 3, #ARGV do ", - "if current == ARGV[index] then redis.call('EXPIRE', KEYS[1], ARGV[2]); ", - "return current; end; end; redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", - "return ARGV[1]; end; if current == ARGV[1] then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return current" - ); let key = self.namespaced_key(key); - let candidate = self.codec.encode(&candidate)?; - let eligible = eligible - .iter() - .map(|value| self.codec.encode(value)) - .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - let value = self.connections.execute(|connection| { - redis::cmd("EVAL") - .arg(SCRIPT) - .arg(1) - .arg(key) - .arg(candidate) - .arg(ttl) - .arg(eligible) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; - self.decode_response(value)?.ok_or(Error::Unavailable) + self.connections + .execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl)) + } + + async fn async_claim_cache( + &self, + key: &str, + candidate: S::Value, + eligible: Vec, + kwargs: CacheKwargs, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let codec = self.codec.clone(); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + claim(connection, &codec, &key, candidate, &eligible, ttl) + }) + .await } } #[cfg(test)] mod tests { - use super::RedisCache; + use std::time::Duration; + use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; - use std::time::Duration; + + use super::RedisCache; fn entry() -> serde_json::Value { json!({"deployment": "model-a", "cooldown_seconds": 30}) diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 6aef9bb36bf..6a61b80b84c 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,8 +1,8 @@ use std::time::Duration; use litellm_cache::{ - BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, Error, JsonCodec, - get_cache, set_cache, + BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache, + CounterCache, Error, JsonCodec, get_cache, set_cache, }; use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; @@ -259,3 +259,126 @@ async fn async_flush_deletes_each_scan_page_separately() { cache.async_flush_cache().await.unwrap(); } + +const CLAIM_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", + "elseif current ~= ARGV[1] then return 0; end; ", + "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", + "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" +); + +fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd { + let mut cmd = redis::cmd("EVAL"); + cmd.arg(CLAIM_SCRIPT) + .arg(1) + .arg("pin") + .arg(expected) + .arg(600) + .arg(write) + .arg(u8::from(refresh)); + cmd +} + +#[tokio::test] +async fn claims_match_eligible_values_written_by_another_encoder() { + let python_payload = r#"{"model_id": "a", "deployment": "east"}"#; + let stored = serde_json::json!({"deployment": "east", "model_id": "a"}); + let candidate = serde_json::json!({"model_id": "b"}); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)), + MockCmd::new(claim_eval(python_payload, "", true), Ok(1)), + ]) + .assert_all_commands_consumed(); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .async_claim_cache( + "pin", + candidate, + vec![stored.clone()], + CacheKwargs::default() + ) + .await + .unwrap(), + stored + ); +} + +#[test] +fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() { + let candidate = serde_json::json!({"model_id": "b"}); + let payload = r#"{"model_id":"b"}"#; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)), + MockCmd::new(claim_eval("", payload, false), Ok(0)), + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)), + MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)), + ]) + .assert_all_commands_consumed(); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .claim_cache( + "pin", + candidate.clone(), + &[serde_json::json!({"model_id": "a"})], + CacheKwargs::default() + ) + .unwrap(), + candidate + ); +} + +#[test] +fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() { + let stored = r#"{"model_id": "a"}"#; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)), + MockCmd::new(claim_eval(stored, "", false), Ok(1)), + ]) + .assert_all_commands_consumed(); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .claim_cache( + "pin", + serde_json::json!({"model_id": "b"}), + &[], + CacheKwargs::default() + ) + .unwrap(), + serde_json::json!({"model_id": "a"}) + ); +} + +#[tokio::test] +async fn async_increment_runs_the_atomic_script() { + let mut eval = redis::cmd("EVAL"); + eval.arg(concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" + )) + .arg(1) + .arg("counter") + .arg(2.5f64) + .arg(600); + let connection = + MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .async_increment_cache("counter", 2.5, CacheKwargs::default()) + .await + .unwrap(), + 4.5 + ); +} diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index d8ffd6d6a50..013cf4b7baf 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -28,17 +28,21 @@ cache.store(&request, json!({"answer": 7}), now)?; assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); ``` -For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers move that blocking work off the executor +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it ## Python integration boundary -The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring +The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API + +Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution -Explicit facade registration checks object identity, method overrides, effective TTL, and configuration changes before selecting native execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python +Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend + +The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and configuration changes before selecting native execution. It does not compare Redis connection settings. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy @@ -52,4 +56,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 now provides L2-first counters and atomic affinity claims with local fallback, but public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +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, queues, 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/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index eaba3d0c349..6b0f29e0a58 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -1,8 +1,9 @@ use litellm_cache::{CacheCodec, Error}; - -use crate::CacheEntry; use serde_json::Value; +use crate::CacheEntry; + +#[derive(Clone, Copy, Debug, Default)] pub struct ResponseCacheCodec; impl CacheCodec for ResponseCacheCodec { @@ -15,7 +16,18 @@ impl CacheCodec for ResponseCacheCodec { { return Err(Error::InvalidEntry); } - serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + // Python reads a `response` that is either a dict or a serialized string, so every + // other shape is written serialized. A string on the wire is therefore always a + // serialized response, which keeps string-valued responses unambiguous. + if value.timestamp.is_none() || value.response.is_object() { + return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry); + } + let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?; + serde_json::to_vec(&CacheEntry { + timestamp: value.timestamp, + response: Value::String(response), + }) + .map_err(|_| Error::InvalidEntry) } fn decode(&self, bytes: &[u8]) -> Result { @@ -30,7 +42,10 @@ impl CacheCodec for ResponseCacheCodec { let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { return Err(Error::InvalidEntry); }; - let response = value.get("response").cloned().ok_or(Error::InvalidEntry)?; + let response = match value.get("response").ok_or(Error::InvalidEntry)? { + Value::String(text) => decode_value(text)?, + response => response.clone(), + }; Ok(CacheEntry { timestamp: Some(timestamp), response, @@ -38,7 +53,7 @@ impl CacheCodec for ResponseCacheCodec { } } -pub(crate) fn decode_value(text: &str) -> Result { +fn decode_value(text: &str) -> Result { if let Ok(value) = serde_json::from_str(text) { return Ok(value); } diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 8f0fe953df6..f18a48863c5 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,9 +1,9 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error}; +use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; -use serde_json::Value; #[derive(Clone)] pub struct ResponseCacheRequest { @@ -39,6 +39,10 @@ impl> ResponseCache { Self { backend } } + pub fn backend(&self) -> &B { + &self.backend + } + pub fn default_ttl(&self) -> Duration { self.backend.default_ttl() } @@ -67,7 +71,7 @@ impl> ResponseCache { Err(Error::InvalidEntry) => None, Err(error) => return Err(error), }; - Self::fresh_or_miss(entry, now, request.max_age) + Ok(Self::fresh_or_miss(entry, now, request.max_age)) } pub async fn async_lookup( @@ -87,7 +91,7 @@ impl> ResponseCache { Err(Error::InvalidEntry) => None, Err(error) => return Err(error), }; - Self::fresh_or_miss(entry, now, request.max_age) + Ok(Self::fresh_or_miss(entry, now, request.max_age)) } pub fn lookup_batch( @@ -180,11 +184,26 @@ impl> ResponseCache { &self, entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, + ) -> Result<(), Error> { + self.async_store_entries( + entries + .into_iter() + .map(|(request, response)| (request, response, now)) + .collect(), + ) + .await + } + + /// Stores entries that each carry the time they were produced, so a deferred write keeps + /// the freshness of its original response. + pub async fn async_store_entries( + &self, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() - .filter(|(request, _)| request.controls.writes()) - .map(|(request, response)| { + .filter(|(request, _, _)| request.controls.writes()) + .map(|(request, response, now)| { ( cache_key(&request.key), CacheEntry { @@ -227,7 +246,7 @@ impl> ResponseCache { let mut values = vec![None; requests.len()]; for ((index, request), entry) in readable.into_iter().zip(entries) { let response = match entry { - BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age)?, + BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age), BatchEntry::Miss | BatchEntry::Invalid => None, }; values[index] = response; @@ -239,24 +258,9 @@ impl> ResponseCache { entry: Option, now: Duration, max_age: Option, - ) -> Result, Error> { - match Self::fresh_response(entry, now, max_age) { - Err(Error::InvalidEntry) => Ok(None), - result => result, - } - } - - fn fresh_response( - entry: Option, - now: Duration, - max_age: Option, - ) -> Result, Error> { + ) -> Option { entry .filter(|entry| entry.fresh(now, max_age)) - .map(|entry| match (entry.timestamp, entry.response) { - (Some(_), Value::String(text)) => crate::codec::decode_value(&text), - (_, value) => Ok(value), - }) - .transpose() + .map(|entry| entry.response) } } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index e4a04b3dec5..7c69e5a1d55 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -273,22 +273,45 @@ async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() { } #[test] -fn malformed_memory_entries_are_treated_as_misses() { - let backend = Arc::new(InMemoryCache::default()); - BaseCache::set_cache( - backend.as_ref(), - "tenant:key", - CacheEntry { +fn string_responses_round_trip_through_typed_and_wire_backends() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let now = Duration::from_secs(100); + for response in [json!("hello world"), json!("123"), json!("null")] { + cache.store(&request(), response.clone(), now).unwrap(); + assert_eq!( + cache.lookup(&request(), now).unwrap(), + Some(response.clone()) + ); + + let wire = ResponseCacheCodec + .encode(&CacheEntry { + timestamp: Some(100.0), + response: response.clone(), + }) + .unwrap(); + assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response); + } +} + +#[test] +fn non_object_responses_are_written_as_python_readable_serialized_strings() { + let wire = ResponseCacheCodec + .encode(&CacheEntry { timestamp: Some(100.0), - response: json!("not a serialized response"), - }, - Default::default(), - ) - .unwrap(); - let cache = ResponseCache::new(backend); + response: json!([1, 2]), + }) + .unwrap(); assert_eq!( - cache.lookup(&request(), Duration::from_secs(100)).unwrap(), - None + serde_json::from_slice::(&wire).unwrap(), + json!({"timestamp": 100.0, "response": "[1,2]"}) + ); + assert_eq!( + ResponseCacheCodec.decode(&wire).unwrap().response, + json!([1, 2]) + ); + assert_eq!( + ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#), + Err(Error::InvalidEntry) ); } @@ -367,3 +390,27 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() { None ); } + +#[tokio::test] +async fn deferred_entries_keep_the_time_they_were_produced() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let mut request = request(); + request.max_age = Some(Duration::from_secs(10)); + cache + .async_store_entries(vec![( + request.clone(), + json!({"answer": 7}), + Duration::from_secs(100), + )]) + .await + .unwrap(); + + assert_eq!( + cache.lookup(&request, Duration::from_secs(110)).unwrap(), + Some(json!({"answer": 7})) + ); + assert_eq!( + cache.lookup(&request, Duration::from_secs(111)).unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 5bc1ebc2945..d6ef8052c4f 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -1,5 +1,4 @@ -use std::future::Future; -use std::time::Duration; +use std::{future::Future, time::Duration}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 39479694f3b..4ec94244ac5 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -1,8 +1,7 @@ use std::sync::Arc; -use crate::{BaseCache, CacheKwargs, Error}; - pub use crate::BaseCache as Cache; +use crate::{BaseCache, CacheKwargs, Error}; pub fn get_cache( cache: &B, diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs index 09bee6032f6..6d47c682406 100644 --- a/litellm-rust/crates/cache/src/codec.rs +++ b/litellm-rust/crates/cache/src/codec.rs @@ -14,6 +14,14 @@ pub trait CacheCodec: Send + Sync { pub struct JsonCodec(PhantomData V>); +impl Clone for JsonCodec { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for JsonCodec {} + impl Default for JsonCodec { fn default() -> Self { Self::new() diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index 17a2c430ddd..9858c5e2748 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -1,15 +1,140 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; -use crate::{BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error}; +use crate::{ + BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, +}; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ReadPolicy { + #[default] + LocalThenRemote, + LocalOnly, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum WritePolicy { + #[default] + Both, + LocalOnly, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RemoteFailurePolicy { + #[default] + Propagate, + UseLocal, +} pub struct DualCache { l1: Arc, l2: Arc, + read_policy: ReadPolicy, + write_policy: WritePolicy, + remote_failure_policy: RemoteFailurePolicy, + promotion_ttl: Option, } impl DualCache { pub fn new(l1: Arc, l2: Arc) -> Self { - Self { l1, l2 } + Self { + l1, + l2, + read_policy: ReadPolicy::default(), + write_policy: WritePolicy::default(), + remote_failure_policy: RemoteFailurePolicy::default(), + promotion_ttl: None, + } + } + + pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self { + Self { + read_policy, + ..self + } + } + + pub fn with_write_policy(self, write_policy: WritePolicy) -> Self { + Self { + write_policy, + ..self + } + } + + pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self { + Self { + remote_failure_policy, + ..self + } + } + + pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self { + Self { + promotion_ttl: Some(promotion_ttl), + ..self + } + } + + fn reads_remote(&self) -> bool { + self.read_policy == ReadPolicy::LocalThenRemote + } + + fn writes_remote(&self) -> bool { + self.write_policy == WritePolicy::Both + } + + fn remote(&self, result: Result) -> Result, Error> { + match result { + Ok(value) => Ok(Some(value)), + Err(Error::Unavailable) + if self.remote_failure_policy == RemoteFailurePolicy::UseLocal => + { + Ok(None) + } + Err(error) => Err(error), + } + } + + fn promotion_kwargs(&self, kwargs: &CacheKwargs) -> CacheKwargs { + CacheKwargs { + ttl: self.promotion_ttl.or(kwargs.ttl), + extras: kwargs.extras.clone(), + } + } +} + +impl DualCache +where + V: Clone + Send + Sync + 'static, + L1: BaseCache, + L2: BaseCache, +{ + fn missing(entries: &[BatchEntry]) -> Vec { + entries + .iter() + .enumerate() + .filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index)) + .collect() + } + + fn merge_batch( + &self, + keys: &[String], + kwargs: &CacheKwargs, + mut entries: Vec>, + missing: Vec, + remote: Vec>, + ) -> Result>, Error> { + if missing.len() != remote.len() { + return Err(Error::Unavailable); + } + for (index, entry) in missing.into_iter().zip(remote) { + if let BatchEntry::Hit(value) = &entry { + self.l1 + .set_cache(&keys[index], value.clone(), self.promotion_kwargs(kwargs))?; + } + entries[index] = entry; + } + Ok(entries) } } @@ -21,12 +146,14 @@ where { type Value = V; - fn default_ttl(&self) -> std::time::Duration { + fn default_ttl(&self) -> Duration { self.l2.default_ttl() } fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { - self.l2.set_cache(key, value.clone(), kwargs.clone())?; + if self.writes_remote() { + self.remote(self.l2.set_cache(key, value.clone(), kwargs.clone()))?; + } self.l1.set_cache(key, value, kwargs) } @@ -34,25 +161,130 @@ where if let Some(value) = self.l1.get_cache(key, kwargs)? { return Ok(Some(value)); } - let value = self.l2.get_cache(key, kwargs)?; + if !self.reads_remote() { + return Ok(None); + } + let value = self.remote(self.l2.get_cache(key, kwargs))?.flatten(); if let Some(value) = &value { - self.l1.set_cache(key, value.clone(), kwargs.clone())?; + self.l1 + .set_cache(key, value.clone(), self.promotion_kwargs(kwargs))?; } Ok(value) } + fn get_cache_batch( + &self, + keys: &[String], + kwargs: &CacheKwargs, + ) -> Result>, Error> { + let entries = self.l1.get_cache_batch(keys, kwargs)?; + let missing = Self::missing(&entries); + if missing.is_empty() || !self.reads_remote() { + return Ok(entries); + } + let remote_keys = missing + .iter() + .map(|index| keys[*index].clone()) + .collect::>(); + match self.remote(self.l2.get_cache_batch(&remote_keys, kwargs))? { + Some(remote) => self.merge_batch(keys, kwargs, entries, missing, remote), + None => Ok(entries), + } + } + + async fn async_set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { + if self.writes_remote() { + self.remote( + self.l2 + .async_set_cache(key, value.clone(), kwargs.clone()) + .await, + )?; + } + self.l1.async_set_cache(key, value, kwargs).await + } + + async fn async_get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { + if let Some(value) = self.l1.async_get_cache(key, kwargs).await? { + return Ok(Some(value)); + } + if !self.reads_remote() { + return Ok(None); + } + let value = self + .remote(self.l2.async_get_cache(key, kwargs).await)? + .flatten(); + if let Some(value) = &value { + self.l1 + .async_set_cache(key, value.clone(), self.promotion_kwargs(kwargs)) + .await?; + } + Ok(value) + } + + async fn async_get_cache_batch( + &self, + keys: Vec, + kwargs: CacheKwargs, + ) -> Result>, Error> { + let entries = self + .l1 + .async_get_cache_batch(keys.clone(), kwargs.clone()) + .await?; + let missing = Self::missing(&entries); + if missing.is_empty() || !self.reads_remote() { + return Ok(entries); + } + let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect(); + match self.remote( + self.l2 + .async_get_cache_batch(remote_keys, kwargs.clone()) + .await, + )? { + Some(remote) => self.merge_batch(&keys, &kwargs, entries, missing, remote), + None => Ok(entries), + } + } + + async fn async_set_cache_pipeline( + &self, + cache_list: Vec<(String, V)>, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + if self.writes_remote() { + self.remote( + self.l2 + .async_set_cache_pipeline(cache_list.clone(), kwargs.clone()) + .await, + )?; + } + self.l1.async_set_cache_pipeline(cache_list, kwargs).await + } + fn delete_cache(&self, key: &str) -> Result<(), Error> { - self.l2.delete_cache(key)?; + if self.writes_remote() { + self.remote(self.l2.delete_cache(key))?; + } self.l1.delete_cache(key) } + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + if self.writes_remote() { + self.remote(self.l2.async_delete_cache(key).await)?; + } + self.l1.async_delete_cache(key).await + } + fn flush_cache(&self) -> Result<(), Error> { - self.l2.flush_cache()?; + if self.writes_remote() { + self.remote(self.l2.flush_cache())?; + } self.l1.flush_cache() } async fn async_flush_cache(&self) -> Result<(), Error> { - self.l2.async_flush_cache().await?; + if self.writes_remote() { + self.remote(self.l2.async_flush_cache().await)?; + } self.l1.async_flush_cache().await } @@ -76,6 +308,20 @@ where self.l1.set_cache(key, value, kwargs)?; Ok(value) } + + async fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> Result { + let value = self + .l2 + .async_increment_cache(key, amount, kwargs.clone()) + .await?; + self.l1.async_set_cache(key, value, kwargs).await?; + Ok(value) + } } impl ClaimCache for DualCache @@ -91,15 +337,39 @@ where eligible: &[V], kwargs: CacheKwargs, ) -> Result { - match self - .l2 - .claim_cache(key, candidate.clone(), eligible, kwargs.clone()) - { - Ok(winner) => { + match self.remote( + self.l2 + .claim_cache(key, candidate.clone(), eligible, kwargs.clone()), + )? { + Some(winner) => { self.l1.set_cache(key, winner.clone(), kwargs)?; Ok(winner) } - Err(_) => self.l1.claim_cache(key, candidate, eligible, kwargs), + None => self.l1.claim_cache(key, candidate, eligible, kwargs), + } + } + + async fn async_claim_cache( + &self, + key: &str, + candidate: V, + eligible: Vec, + kwargs: CacheKwargs, + ) -> Result { + match self.remote( + self.l2 + .async_claim_cache(key, candidate.clone(), eligible.clone(), kwargs.clone()) + .await, + )? { + Some(winner) => { + self.l1.async_set_cache(key, winner.clone(), kwargs).await?; + Ok(winner) + } + None => { + self.l1 + .async_claim_cache(key, candidate, eligible, kwargs) + .await + } } } } diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 824de00bdf4..5c27aa9ff4e 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,6 +1,7 @@ -use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; use std::{sync::Mutex, time::Duration}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; + struct TestCache { default_ttl: Duration, writes: Mutex>, diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index 1be1556734c..2e1e72c7119 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -4,7 +4,8 @@ use std::{ }; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, dual::DualCache, + BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, + dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy}, }; struct TestCache { @@ -111,7 +112,8 @@ fn failed_l2_increment_leaves_l1_unchanged() { #[test] fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { let l1 = Arc::new(TestCache::new(Some("first".to_string()), false)); - let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))); + let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); assert_eq!( cache @@ -128,3 +130,213 @@ fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { "first" ); } + +struct SyncPanics(TestCache); + +impl BaseCache for SyncPanics { + type Value = String; + + fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> { + panic!("sync L2 write on an async path") + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + panic!("sync L2 read on an async path") + } + + async fn async_set_cache( + &self, + key: &str, + value: String, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + self.0.set_cache(key, value, kwargs) + } + + async fn async_get_cache( + &self, + key: &str, + kwargs: &CacheKwargs, + ) -> Result, Error> { + self.0.get_cache(key, kwargs) + } + + async fn async_get_cache_batch( + &self, + keys: Vec, + kwargs: CacheKwargs, + ) -> Result>, Error> { + assert_eq!(keys, ["missing"]); + Ok(vec![match self.0.get_cache("missing", &kwargs)? { + Some(value) => litellm_cache::BatchEntry::Hit(value), + None => litellm_cache::BatchEntry::Miss, + }]) + } + + async fn async_set_cache_pipeline( + &self, + cache_list: Vec<(String, String)>, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + for (key, value) in cache_list { + self.0.set_cache(&key, value, kwargs.clone())?; + } + Ok(()) + } + + fn delete_cache(&self, _: &str) -> Result<(), Error> { + panic!("sync L2 delete on an async path") + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + self.0.delete_cache(key) + } + + fn flush_cache(&self) -> Result<(), Error> { + panic!("sync L2 flush on an async path") + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +#[tokio::test] +async fn async_operations_use_the_async_l2_methods() { + let l1 = Arc::new(TestCache::new(None, false)); + let cache = DualCache::new( + l1.clone(), + Arc::new(SyncPanics(TestCache::new( + Some("remote".to_string()), + false, + ))), + ); + let kwargs = CacheKwargs::default(); + + assert_eq!( + cache.async_get_cache("missing", &kwargs).await.unwrap(), + Some("remote".into()) + ); + assert_eq!( + l1.get_cache("missing", &kwargs).unwrap(), + Some("remote".into()) + ); + + l1.delete_cache("missing").unwrap(); + assert_eq!( + cache + .async_get_cache_batch(vec!["missing".into()], kwargs.clone()) + .await + .unwrap(), + [litellm_cache::BatchEntry::Hit("remote".to_string())] + ); + cache + .async_set_cache("missing", "written".into(), kwargs.clone()) + .await + .unwrap(); + cache + .async_set_cache_pipeline(vec![("missing".into(), "piped".into())], kwargs.clone()) + .await + .unwrap(); + cache.async_delete_cache("missing").await.unwrap(); + assert_eq!( + cache.async_get_cache("missing", &kwargs).await.unwrap(), + None + ); +} + +struct Unavailable; + +impl BaseCache for Unavailable { + type Value = String; + + fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> { + Err(Error::Unavailable) + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + Err(Error::Unavailable) + } + + fn delete_cache(&self, _: &str) -> Result<(), Error> { + Err(Error::Unavailable) + } + + fn flush_cache(&self) -> Result<(), Error> { + Err(Error::Unavailable) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +impl ClaimCache for Unavailable { + fn claim_cache( + &self, + _: &str, + _: String, + _: &[String], + _: CacheKwargs, + ) -> Result { + Err(Error::InvalidEntry) + } +} + +#[test] +fn remote_failure_policy_selects_propagation_or_the_local_tier() { + let kwargs = CacheKwargs::default(); + let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); + assert_eq!( + strict.set_cache("key", "value".into(), kwargs.clone()), + Err(Error::Unavailable) + ); + assert_eq!(strict.get_cache("key", &kwargs), Err(Error::Unavailable)); + + let l1 = Arc::new(TestCache::new(None, false)); + let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable)) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); + assert_eq!(degraded.get_cache("key", &kwargs), Ok(None)); + degraded + .set_cache("key", "value".into(), kwargs.clone()) + .unwrap(); + assert_eq!(degraded.get_cache("key", &kwargs), Ok(Some("value".into()))); + degraded.delete_cache("key").unwrap(); + assert_eq!(l1.get_cache("key", &kwargs), Ok(None)); +} + +#[test] +fn claim_fallback_does_not_hide_non_availability_errors() { + let cache = DualCache::new( + Arc::new(TestCache::new(Some("first".to_string()), false)), + Arc::new(Unavailable), + ) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); + assert_eq!( + cache.claim_cache("affinity", "second".into(), &[], CacheKwargs::default()), + Err(Error::InvalidEntry) + ); +} + +#[test] +fn local_only_policies_never_touch_l2() { + let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false)); + let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone()) + .with_read_policy(ReadPolicy::LocalOnly) + .with_write_policy(WritePolicy::LocalOnly); + let kwargs = CacheKwargs::default(); + + assert_eq!(cache.get_cache("key", &kwargs), Ok(None)); + cache + .set_cache("key", "local".into(), kwargs.clone()) + .unwrap(); + assert_eq!(l2.get_cache("key", &kwargs), Ok(Some("remote".into()))); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 58550c2987d..2ad13ce200f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -6,9 +8,8 @@ use pyo3::{ types::{PyDict, PyTuple, PyType}, }; use serde_json::Value; -use std::time::Duration; -use super::{NativeCacheHandle, native::NativeResponseCache}; +use super::{CacheTestHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, @@ -130,9 +131,10 @@ impl FacadeGuard { pub(super) fn capture( py: Python<'_>, facade: &Bound<'_, PyAny>, - kind: &str, - native_default_ttl: Duration, + service: &NativeResponseCache, ) -> PyResult { + let kind = service.kind(); + let native_default_ttl: Duration = service.default_ttl(); let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -158,6 +160,23 @@ impl FacadeGuard { "facade and native backend default TTLs must match", )); } + let namespace = match backend.getattr_opt("namespace")? { + Some(namespace) => namespace.extract::>()?, + None => None, + } + .filter(|namespace| !namespace.is_empty()); + if kind == "redis" && namespace.as_deref() != service.namespace() { + return Err(PyTypeError::new_err( + "facade and native backend namespaces must match", + )); + } + if let Some(capacity) = service.capacity() + && backend.getattr("max_size_in_memory")?.extract::()? != capacity + { + return Err(PyTypeError::new_err( + "facade and native backend capacities must match", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -169,6 +188,7 @@ impl FacadeGuard { "namespace", "supported_call_types", "redis_flush_size", + "semantic_cache_scope", ], )?, backend: ObjectGuard::capture( @@ -179,6 +199,8 @@ impl FacadeGuard { "default_ttl", "max_size_in_memory", "max_size_per_item", + "redis_kwargs", + "redis_flush_size", ], )?, }) @@ -208,7 +230,7 @@ pub(super) fn resolve( let Some(handle) = dict.get_item("_native_cache_handle")? else { return Ok(None); }; - let Ok(handle) = handle.extract::>() else { + let Ok(handle) = handle.extract::>() else { return Ok(None); }; let Some(guard) = &handle.guard else { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 1e2e42600ff..f83788f6036 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -3,21 +3,21 @@ mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use facade::FacadeGuard; use litellm_cache::Error; use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use native::NativeResponseCache; use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::{PyDict, PyList}, + types::{PyDict, PyList, PyTuple}, }; use serde::Deserialize; use serde_json::Value; use crate::python_settings::PythonSettings; -use facade::FacadeGuard; -use native::NativeResponseCache; const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); @@ -84,14 +84,14 @@ fn cache_error(error: Error) -> PyErr { } } -#[pyclass(frozen)] -pub(crate) struct NativeCacheHandle { +#[pyclass(frozen, name = "_CacheTestHandle")] +pub(crate) struct CacheTestHandle { service: NativeResponseCache, guard: Option, pid: u32, } -impl NativeCacheHandle { +impl CacheTestHandle { fn service(&self) -> PyResult { if self.pid != std::process::id() { return Err(PyRuntimeError::new_err( @@ -103,7 +103,7 @@ impl NativeCacheHandle { } #[pymethods] -impl NativeCacheHandle { +impl CacheTestHandle { #[staticmethod] #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { @@ -140,9 +140,9 @@ impl NativeCacheHandle { self.service.kind() } - fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { + fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, self.backend(), service.default_ttl())?; + let guard = FacadeGuard::capture(py, facade, &service)?; let service = service.with_redis_flush_size( facade .getattr("redis_flush_size")? @@ -173,7 +173,7 @@ enum CacheBinding { PythonCallback(Py), } -#[pyclass(frozen, name = "CacheBinding")] +#[pyclass(frozen, name = "_CacheTestBinding")] pub(crate) struct ResolvedCache { binding: CacheBinding, pid: u32, @@ -285,12 +285,15 @@ impl ResolvedCache { } } + /// Native bindings return `{values, missing_indices}`. The built-in `Cache` API has no batch + /// read, so a Python callback receives one `get_cache(**kwargs)` call per request, in order, + /// and the results come back as a list. #[pyo3(signature = (requests, *, callback_kwargs=None))] fn lookup_batch( &self, py: Python<'_>, requests: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, + callback_kwargs: Option<&Bound<'_, PyAny>>, ) -> PyResult> { self.check_process()?; match &self.binding { @@ -305,14 +308,17 @@ impl ResolvedCache { .map_err(cache_error)?; to_py(py, &response) } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "batch_get_cache", - (callback_keys(py, requests)?,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(Bound::unbind), + CacheBinding::PythonCallback(object) => { + let results = PyList::empty(py); + for kwargs in batch_callback_kwargs(requests, callback_kwargs)? { + results.append(object.bind(py).call_method( + "get_cache", + (), + Some(&kwargs), + )?)?; + } + Ok(results.into_any().unbind()) + } } } @@ -364,7 +370,7 @@ impl ResolvedCache { &self, py: Python<'py>, requests: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, + callback_kwargs: Option<&Bound<'py, PyAny>>, ) -> PyResult> { self.check_process()?; match &self.binding { @@ -381,20 +387,30 @@ impl ResolvedCache { cache_error, ) } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_batch_get_cache", - (callback_keys(py, requests)?,), - Some(self::callback_kwargs(callback_kwargs)?), - ), + CacheBinding::PythonCallback(object) => { + let awaitables = batch_callback_kwargs(requests, callback_kwargs)? + .iter() + .map(|kwargs| { + object + .bind(py) + .call_method("async_get_cache", (), Some(kwargs)) + }) + .collect::>>()?; + py.import("asyncio")? + .call_method1("gather", PyTuple::new(py, awaitables)?) + } } } - #[pyo3(signature = (requests, responses, *, callback_kwargs=None))] + /// A Python callback receives the caller's original result through `callback_result`, because + /// the built-in `Cache.async_add_cache_pipeline` splits the batch itself. + #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] fn async_store_batch<'py>( &self, py: Python<'py>, requests: &Bound<'py, PyAny>, responses: &Bound<'py, PyAny>, + callback_result: Option<&Bound<'py, PyAny>>, callback_kwargs: Option<&Bound<'py, PyDict>>, ) -> PyResult> { self.check_process()?; @@ -417,20 +433,14 @@ impl ResolvedCache { ) } CacheBinding::PythonCallback(object) => { - let keys = callback_keys(py, requests)?; - let responses = responses.try_iter()?.collect::>>()?; - if keys.len() != responses.len() { - return Err(PyValueError::new_err( - "batch cache requests and responses must have equal lengths", - )); - } - let cache_list = PyList::empty(py); - for (key, response) in keys.iter().zip(responses) { - cache_list.append((key, response))?; - } + let result = callback_result.ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require their original callback_result", + ) + })?; object.bind(py).call_method( - "async_set_cache_pipeline", - (cache_list,), + "async_add_cache_pipeline", + (result,), Some(self::callback_kwargs(callback_kwargs)?), ) } @@ -445,8 +455,17 @@ impl ResolvedCache { let service = service.clone(); run_async(py, async move { service.async_flush().await }, cache_error) } + // The built-in `Cache` facade has no flush of its own; its backend does. CacheBinding::PythonCallback(object) => { - object.bind(py).call_method0("flush_cache")?; + let object = object.bind(py); + let backend = match object.getattr_opt("cache")? { + Some(backend) if !backend.is_none() => backend, + _ => object.clone(), + }; + if backend.hasattr("async_flush_cache")? { + return backend.call_method0("async_flush_cache"); + } + backend.call_method0("flush_cache")?; ready_none(py) } } @@ -464,7 +483,7 @@ impl ResolvedCache { cache_error, ) } - CacheBinding::PythonCallback(object) => object.bind(py).call_method0("test_connection"), + CacheBinding::PythonCallback(object) => object.bind(py).call_method0("ping"), } } @@ -484,16 +503,25 @@ fn callback_kwargs<'a, 'py>( }) } -fn callback_keys<'py>( - py: Python<'py>, +fn batch_callback_kwargs<'py>( requests: &Bound<'py, PyAny>, -) -> PyResult> { - PyList::new( - py, - self::requests(requests)? - .into_iter() - .map(|request| litellm_cache_response::cache_key(&request.key)), - ) + kwargs: Option<&Bound<'py, PyAny>>, +) -> PyResult>> { + let kwargs = kwargs + .ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require one original callback_kwargs mapping per request", + ) + })? + .try_iter()? + .map(|item| Ok(item?.cast_into::()?)) + .collect::>>()?; + if kwargs.len() != requests.len()? { + return Err(PyValueError::new_err( + "batch cache requests and callback_kwargs must have equal lengths", + )); + } + Ok(kwargs) } fn ready_none(py: Python<'_>) -> PyResult> { @@ -512,13 +540,13 @@ fn ready_value<'py, T: serde::Serialize>( Ok(future) } -#[pyclass(frozen)] -pub(crate) struct CacheResolver { +#[pyclass(frozen, name = "_CacheTestResolver")] +pub(crate) struct CacheTestResolver { namespace: Py, } #[pymethods] -impl CacheResolver { +impl CacheTestResolver { #[new] fn new(namespace: Py) -> Self { Self { namespace } @@ -528,7 +556,7 @@ impl CacheResolver { let object = self.namespace.bind(py).getattr("cache")?; let binding = if object.is_none() { CacheBinding::Disabled - } else if let Ok(handle) = object.extract::>() { + } else if let Ok(handle) = object.extract::>() { CacheBinding::Native(handle.service()?) } else if let Some(service) = facade::resolve(py, &object)? { CacheBinding::Native(service) diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 20891719550..3fc8f61dff6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -3,12 +3,11 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; -use serde_json::Value; -use tokio::sync::Mutex; - use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, }; +use serde_json::Value; +use tokio::sync::Mutex; #[derive(Clone)] pub(super) enum NativeResponseCache { @@ -21,7 +20,7 @@ pub(super) enum NativeResponseCache { pub(super) struct RedisWriteBuffer { flush_size: usize, - entries: Mutex>, + entries: Mutex>, } impl NativeResponseCache { @@ -67,6 +66,20 @@ impl NativeResponseCache { } } + 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 with_redis_flush_size(self, flush_size: Option) -> Self { match self { Self::Redis { cache, .. } => Self::Redis { @@ -145,19 +158,15 @@ impl NativeResponseCache { } => { let pending = { let mut entries = buffer.entries.lock().await; - entries.push((request.clone(), response)); + entries.push((request.clone(), response, now)); (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) }; - let Some(pending) = pending else { - return Ok(()); - }; - if let Err(error) = cache.async_store_batch(pending.clone(), now).await { - let mut entries = buffer.entries.lock().await; - let current = std::mem::take(&mut *entries); - *entries = pending.into_iter().chain(current).collect(); - return Err(error); + // A failed flush drops its batch, as Python does. Requeueing would grow the + // buffer and re-send an ever larger pipeline on every write during an outage. + match pending { + Some(pending) => cache.async_store_entries(pending).await, + None => Ok(()), } - Ok(()) } } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 621c111a35b..bd62c5aadf1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -10,8 +10,7 @@ mod token_counter; #[pymodule(gil_used = true)] mod _native { - #[pymodule_export] - use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache}; + use crate::cache::{CacheTestHandle, CacheTestResolver, ResolvedCache}; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -35,6 +34,16 @@ mod _native { use crate::token_counter::TokenCounter; #[pymodule_export] use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; + use pyo3::{prelude::*, types::PyModule}; + + #[pymodule_init] + fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + let py = module.py(); + let dict = module.dict(); + dict.set_item("_CacheTestHandle", py.get_type::())?; + dict.set_item("_CacheTestResolver", py.get_type::())?; + dict.set_item("_CacheTestBinding", py.get_type::()) + } } use pyo3::prelude::*; @@ -68,9 +77,6 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "TokenCounter", - "CacheResolver", - "NativeCacheHandle", - "CacheBinding", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 68b1742dc41..ab4639bc876 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -94,25 +94,25 @@ class ResponsesWebSocketConnection: def close(self) -> Future[None]: ... @final -class NativeCacheHandle: +class _CacheTestHandle: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @staticmethod def memory( *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576 - ) -> NativeCacheHandle: ... + ) -> _CacheTestHandle: ... @staticmethod - def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ... + def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> _CacheTestHandle: ... @property def backend(self) -> str: ... - def bind_facade(self, facade: object) -> None: ... + def _bind_facade(self, facade: object) -> None: ... @final -class CacheResolver: - def __new__(cls, namespace: object) -> CacheResolver: ... - def resolve(self) -> CacheBinding: ... +class _CacheTestResolver: + def __new__(cls, namespace: object) -> _CacheTestResolver: ... + def resolve(self) -> _CacheTestBinding: ... @final -class CacheBinding: +class _CacheTestBinding: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @property def kind(self) -> str: ... @@ -130,7 +130,7 @@ class CacheBinding: self, requests: Sequence[Mapping[str, object]], *, - callback_kwargs: dict[str, object] | None = None, + callback_kwargs: Sequence[dict[str, object]] | None = None, ) -> object: ... def async_lookup( self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None @@ -146,17 +146,18 @@ class CacheBinding: self, requests: Sequence[Mapping[str, object]], *, - callback_kwargs: dict[str, object] | None = None, + callback_kwargs: Sequence[dict[str, object]] | None = None, ) -> Awaitable[object]: ... def async_store_batch( self, requests: Sequence[Mapping[str, object]], responses: Sequence[object], *, + callback_result: object = None, callback_kwargs: dict[str, object] | None = None, ) -> Awaitable[object]: ... def async_flush(self) -> Awaitable[None]: ... - def ping(self) -> Awaitable[dict[str, object] | None]: ... + def ping(self) -> Awaitable[object]: ... @final class TokenCounter: @@ -174,10 +175,7 @@ def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... __all__ = [ - "CacheBinding", - "CacheResolver", "ForkedAfterNativeRuntimeStarted", - "NativeCacheHandle", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 05a4ef68e0b..d34d23ca1d7 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -6,10 +6,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync -from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.types.caching import RedisPipelineIncrementOperation diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d38583dad0a..cf46a0566f7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -15,7 +15,7 @@ import pytest import redis import litellm -from litellm.caching.caching import Cache +from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType @@ -50,20 +50,43 @@ def test_existing_constructor_and_global_are_unchanged() -> None: assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) with rebound(litellm, "cache", facade): - resolver: Final = _native.CacheResolver(litellm) + resolver: Final = _native._CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} +def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: + resolver: Final = _native._CacheTestResolver(litellm) + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) + enabled: Final = litellm.cache + assert isinstance(enabled, Cache) + assert enabled.ttl == 30 + assert resolver.resolve().kind == "python_callback" + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + assert litellm.cache is enabled + + update_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + updated: Final = litellm.cache + assert isinstance(updated, Cache) + assert updated is not enabled + assert updated.ttl == 60 + + disable_cache() + assert litellm.cache is None + assert resolver.resolve().kind == "disabled" + + async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory()) - resolver: Final = _native.CacheResolver(namespace) + namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) + resolver: Final = _native._CacheTestResolver(namespace) selected: Final = resolver.resolve() assert selected.kind == "native" selected.store(request(), {"answer": 1}) assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _native.NativeCacheHandle.memory()): + with rebound(namespace, "cache", _native._CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -96,7 +119,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native.CacheResolver(namespace).resolve() + binding: Final = _native._CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -117,7 +140,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -132,9 +155,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _native.NativeCacheHandle.memory() - handle.bind_facade(facade) - resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + handle: Final = _native._CacheTestHandle.memory() + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -165,16 +188,18 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native.NativeCacheHandle.memory() + handle: Final = _native._CacheTestHandle.memory() with pytest.raises(TypeError): - handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) + handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle.bind_facade(facade) - resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" def custom_key(**_kwargs: object) -> str: return "custom" @@ -193,7 +218,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: def cyclic_reference() -> weakref.ReferenceType[CustomCache]: callback: Final = CustomCache() namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _native.CacheResolver(namespace).resolve() + binding: Final = _native._CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -204,8 +229,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team")) - binding: Final = _native.CacheResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _native._CacheTestResolver(namespace).resolve() response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) @@ -227,33 +252,33 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _native.NativeCacheHandle.memory(ttl_seconds=-1) + _native._CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native.NativeCacheHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native.CacheResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() small: Final = {"answer": "ok"} binding.store(request("small"), small) assert await binding.async_lookup(request("small")) == small await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _native.CacheResolver( - SimpleNamespace(cache=_native.NativeCacheHandle.memory(capacity=0)) + disabled: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -275,50 +300,72 @@ async def test_native_batch_lookup_and_store_report_partial_hits() -> None: } -async def test_python_batch_callbacks_receive_keys_and_key_value_pairs() -> None: - first: Final = object() - second: Final = object() - - class CustomCache: - def batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: - return keys, marker - - async def async_batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: - return keys, marker - - async def async_set_cache_pipeline( - self, cache_list: list[tuple[str, object]], *, marker: object - ) -> tuple[list[tuple[str, object]], object]: - return cache_list, marker - +async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: + result: Final = object() marker: Final = object() - binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() - requests: Final = [request("first"), request("second")] - assert binding.lookup_batch(requests, callback_kwargs={"marker": marker}) == (["first", "second"], marker) - assert await binding.async_lookup_batch(requests, callback_kwargs={"marker": marker}) == ( - ["first", "second"], - marker, - ) + class CustomCache(Cache): + def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("sync", kwargs) + + async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("async", kwargs) + + async def async_add_cache_pipeline( + self, result: object, dynamic_cache_object: object = None, **kwargs: object + ) -> object: + return result, kwargs + + binding: Final = _native._CacheTestResolver( + SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) + ).resolve() + assert binding.kind == "python_callback" + requests: Final = [request("first"), request("second")] + kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] + + assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])] + assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [ + ("async", kwargs[0]), + ("async", kwargs[1]), + ] + with pytest.raises(ValueError, match="equal lengths"): + binding.lookup_batch(requests, callback_kwargs=kwargs[:1]) + with pytest.raises(TypeError, match="callback_result"): + await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker}) stored: Final = cast( - tuple[list[tuple[str, object]], object], - await binding.async_store_batch( - requests, - [first, second], - callback_kwargs={"marker": marker}, - ), + tuple[object, dict[str, object]], + await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}), ) - assert [key for key, _ in stored[0]] == ["first", "second"] - assert stored[1] is marker - assert stored[0][0][1] is first - assert stored[0][1][1] is second + assert stored[0] is result + assert stored[1] == {"marker": marker} + + +async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: + async def ping() -> str: + return "pong" + + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache.cache.set_cache("key", "value") + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + assert binding.kind == "python_callback" + + setattr(cache.cache, "ping", ping) + assert await binding.ping() == "pong" + await binding.async_flush() + assert cache.cache.get_cache("key") is None + + +def test_facade_registration_rejects_mismatched_capacity() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + with pytest.raises(TypeError, match="capacities must match"): + _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) with rebound(litellm, "default_redis_ttl", 7): - binding: Final = _native.CacheResolver( - SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url)) + binding: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url)) ).resolve() await binding.async_store(request("native-default"), {"value": 1}) @@ -336,11 +383,16 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: redis_flush_size=2, ) with pytest.raises(TypeError, match="default TTLs must match"): - _native.NativeCacheHandle.redis(redis_url, ttl_seconds=61).bind_facade(facade) - _native.NativeCacheHandle.redis(redis_url).bind_facade(facade) - binding: Final = _native.CacheResolver(SimpleNamespace(cache=facade)).resolve() + _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + with pytest.raises(TypeError, match="namespaces must match"): + _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _native._CacheTestHandle.redis(redis_url)._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None await binding.async_store(request("second"), {"value": 2}) From dc5f0c58a4decfdd6227fbf3af661bf6458ab04e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 10:21:54 -0700 Subject: [PATCH 26/56] feat(cache): add native backend operation primitives --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/cache-memory/src/cache.rs | 42 +- .../crates/cache-memory/tests/cache.rs | 45 +- litellm-rust/crates/cache-redis/Cargo.toml | 2 +- litellm-rust/crates/cache-redis/src/cache.rs | 4 + .../cache-redis/src/cache/operations.rs | 500 ++++++++++++++++++ litellm-rust/crates/cache-redis/src/lib.rs | 2 +- .../crates/cache-redis/tests/cache.rs | 307 ++++++++++- litellm-rust/crates/cache-response/README.md | 6 +- litellm-rust/crates/cache/src/capabilities.rs | 7 + litellm-rust/crates/cache/src/lib.rs | 2 +- 11 files changed, 910 insertions(+), 9 deletions(-) create mode 100644 litellm-rust/crates/cache-redis/src/cache/operations.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2d6fb6c3082..ed4ae4e3353 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3712,6 +3712,8 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", + "rustls 0.23.42", + "rustls-native-certs", "ryu", "sha1_smol", "socket2 0.6.5", diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index a9814ff6fd6..77635893640 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -7,7 +7,7 @@ use std::{ use litellm_cache::{ BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, - Error, + Error, IncrementOperation, }; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; @@ -134,6 +134,25 @@ impl InMemoryCache { .copied()) } + pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { + self.expires_at(key) + } + + pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result, Error> { + let state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut expirations = state + .expirations + .iter() + .map(|(key, expiration)| (key.clone(), *expiration)) + .collect::>(); + expirations.sort_unstable_by_key(|(_, expiration)| *expiration); + Ok(expirations + .into_iter() + .take(count) + .map(|(key, _)| key) + .collect()) + } + pub fn delete_cache(&self, key: &str) -> Result<(), Error> { let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; Self::remove(&mut state, key); @@ -240,6 +259,27 @@ impl CounterCache for InMemoryCache { } } +impl InMemoryCache { + pub async fn async_increment_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + operations + .into_iter() + .map(|operation| { + self.increment_cache( + &operation.key, + operation.amount, + CacheKwargs { + ttl: operation.ttl, + ..CacheKwargs::default() + }, + ) + }) + .collect() + } +} + impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 22c5595da52..c44590b63ca 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -8,7 +8,7 @@ use std::{ use litellm_cache::{ BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, - get_cache, set_cache, + IncrementOperation, get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -305,3 +305,46 @@ fn disabled_cache_does_not_retain_claims_or_counters() { ); assert_eq!(counters.get_cache("key").unwrap(), None); } + +#[tokio::test] +async fn ttl_and_oldest_key_operations_use_the_stored_expirations() { + let clock = Arc::new(AtomicU64::new(100)); + let cache = cache(clock, 3); + cache + .set_cache("later", "2".into(), Some(Duration::from_secs(20))) + .unwrap(); + cache + .set_cache("first", "1".into(), Some(Duration::from_secs(10))) + .unwrap(); + + assert_eq!( + cache.async_get_ttl("first").await.unwrap(), + Some(Duration::from_secs(110)) + ); + assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]); + assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); +} + +#[tokio::test] +async fn increment_pipeline_preserves_operation_order() { + let cache = InMemoryCache::::new(Some(3), None); + assert_eq!( + cache + .async_increment_pipeline(vec![ + IncrementOperation { + key: "a".into(), + amount: 1.0, + ttl: Some(Duration::from_secs(10)), + }, + IncrementOperation { + key: "a".into(), + amount: 2.0, + ttl: Some(Duration::from_secs(20)), + }, + ]) + .await + .unwrap(), + [1.0, 3.0] + ); + assert_eq!(cache.get_cache("a").unwrap(), Some(3.0)); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index f123e774158..5818f75ff3d 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = "1.7.0" +redis = { version = "1.7.0", features = ["tls-rustls"] } r2d2 = "0.8.10" tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index b8f0d3857e2..23b1fabbab4 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -9,6 +9,10 @@ use litellm_cache::{ }; use redis::Commands; +mod operations; + +pub use operations::{RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; + const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs new file mode 100644 index 00000000000..f8c2bf4078c --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -0,0 +1,500 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{CacheCodec, Error, IncrementOperation}; +use redis::Commands; + +use super::{ConnectionRef, RedisCache}; + +const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return count" +); +const SET_MAX_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", + "redis.call('SET', KEYS[1], ARGV[1]); ", + "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return ARGV[1]; end; return current" +); + +#[derive(Clone, Debug, PartialEq)] +pub enum RedisArg { + Bytes(Vec), + Integer(i64), + Float(f64), +} + +impl From<&str> for RedisArg { + fn from(value: &str) -> Self { + Self::Bytes(value.as_bytes().to_vec()) + } +} + +impl From for RedisArg { + fn from(value: String) -> Self { + Self::Bytes(value.into_bytes()) + } +} + +impl From> for RedisArg { + fn from(value: Vec) -> Self { + Self::Bytes(value) + } +} + +impl From for RedisArg { + fn from(value: i64) -> Self { + Self::Integer(value) + } +} + +impl From for RedisArg { + fn from(value: f64) -> Self { + Self::Float(value) + } +} + +impl redis::ToRedisArgs for RedisArg { + fn write_redis_args(&self, out: &mut W) + where + W: ?Sized + redis::RedisWrite, + { + match self { + Self::Bytes(value) => value.write_redis_args(out), + Self::Integer(value) => value.write_redis_args(out), + Self::Float(value) => value.write_redis_args(out), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RedisRpushOperation { + pub key: String, + pub values: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RedisLpopOperation { + pub key: String, + pub count: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RedisLpopResult { + Missing, + Value(Vec), + Values(Vec>), +} + +impl RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub async fn delete_cache_keys(&self, keys: Vec) -> Result { + if keys.is_empty() { + return Ok(0); + } + let keys = keys + .into_iter() + .map(|key| self.namespaced_key(&key)) + .collect::>(); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + connection.del(keys).map_err(|_| Error::Unavailable) + }) + .await + } + + pub fn batch_get_counts(&self, keys: &[String]) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = self.connections.execute(|connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + })?; + values.into_iter().map(count).collect() + } + + pub async fn async_batch_get_counts( + &self, + keys: Vec, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values.into_iter().map(count).collect() + } + + pub fn sync_ping(&self) -> Result { + self.connections.execute(|connection| { + redis::cmd("PING") + .query::(connection) + .map(|response| response == "PONG") + .map_err(|_| Error::Unavailable) + }) + } + + pub async fn ping(&self) -> Result { + Self::run_blocking(Arc::clone(&self.connections), |connection| { + redis::cmd("PING") + .query::(connection) + .map(|response| response == "PONG") + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { + let key = self.namespaced_key(key); + let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("TTL") + .arg(key) + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + Ok((ttl >= 0).then_some(ttl)) + } + + pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { + let pattern = format!("{}*", self.namespaced_key(pattern)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut cursor = 0u64; + let mut matches = Vec::new(); + loop { + let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") + .cursor_arg(cursor) + .arg("MATCH") + .arg(&pattern) + .arg("COUNT") + .arg(count) + .query(connection) + .map_err(|_| Error::Unavailable)?; + matches.extend(keys); + if matches.len() >= count || next_cursor == 0 { + matches.truncate(count); + return Ok(matches); + } + cursor = next_cursor; + } + }) + .await + } + + pub async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result { + if values.is_empty() { + return Err(Error::InvalidEntry); + } + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + pipeline.cmd("SADD").arg(&key).arg(values); + pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore(); + pipeline + .query::<(usize,)>(connection) + .map(|(added,)| added) + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_rpush(&self, key: &str, values: Vec) -> Result { + if values.is_empty() { + return Err(Error::InvalidEntry); + } + let key = self.namespaced_key(key); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("RPUSH") + .arg(key) + .arg(values) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_rpush_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let operations = operations + .into_iter() + .map(|operation| { + if operation.values.is_empty() { + return Err(Error::InvalidEntry); + } + Ok((self.namespaced_key(&operation.key), operation.values)) + }) + .collect::, _>>()?; + if operations.is_empty() { + return Ok(Vec::new()); + } + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + for (key, values) in operations { + pipeline.cmd("RPUSH").arg(key).arg(values); + } + pipeline.query(connection).map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_lpop( + &self, + key: &str, + count: Option, + ) -> Result { + let key = self.namespaced_key(key); + let multiple = count.is_some(); + let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut command = redis::cmd("LPOP"); + command.arg(key); + if let Some(count) = count { + command.arg(count); + } + command + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + lpop_result(value, multiple) + } + + pub async fn async_lpop_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let operations = operations + .into_iter() + .map(|operation| (self.namespaced_key(&operation.key), operation.count)) + .collect::>(); + if operations.is_empty() { + return Ok(Vec::new()); + } + let multiple = operations + .iter() + .map(|(_, count)| count.is_some()) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + for (key, count) in operations { + let command = pipeline.cmd("LPOP").arg(key); + if let Some(count) = count { + command.arg(count); + } + } + pipeline + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values + .into_iter() + .zip(multiple) + .map(|(value, multiple)| lpop_result(value, multiple)) + .collect() + } + + pub async fn async_eval( + &self, + script: String, + keys: Vec, + arguments: Vec, + ) -> Result { + let keys = keys + .into_iter() + .map(|key| self.namespaced_key(&key)) + .collect::>(); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("EVAL") + .arg(script) + .arg(keys.len()) + .arg(keys) + .arg(arguments) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub fn client_list(&self) -> Result { + self.connections.execute(|connection| { + redis::cmd("CLIENT") + .arg("LIST") + .query(connection) + .map_err(|_| Error::Unavailable) + }) + } + + pub fn info(&self) -> Result { + self.connections.execute(|connection| { + redis::cmd("INFO") + .query(connection) + .map_err(|_| Error::Unavailable) + }) + } + + pub fn flushall(&self) -> Result<(), Error> { + self.connections.execute(|connection| { + redis::cmd("FLUSHALL") + .query(connection) + .map_err(|_| Error::Unavailable) + }) + } +} + +impl RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub fn increment_with_floor( + &self, + key: &str, + amount: i64, + ttl: Duration, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl); + self.connections + .execute(|connection| increment_with_floor(connection, key, amount, ttl)) + } + + pub async fn async_increment_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let operations = operations + .into_iter() + .map(|operation| { + ( + self.namespaced_key(&operation.key), + operation.amount, + operation.ttl.map(Self::ttl_seconds), + ) + }) + .collect::>(); + if operations.is_empty() { + return Ok(Vec::new()); + } + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + for (key, amount, ttl) in operations { + pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount); + if let Some(ttl) = ttl { + pipeline.cmd("EXPIRE").arg(key).arg(ttl).ignore(); + } + } + pipeline.query(connection).map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_increment_with_floor( + &self, + key: &str, + amount: i64, + ttl: Duration, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + increment_with_floor(connection, key, amount, ttl) + }) + .await + } + + pub async fn async_set_max( + &self, + key: &str, + value: f64, + ttl: Option, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("EVAL") + .arg(SET_MAX_SCRIPT) + .arg(1) + .arg(key) + .arg(value) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + .await + } +} + +fn redis_bytes(value: redis::Value) -> Result, Error> { + match value { + redis::Value::BulkString(bytes) => Ok(bytes), + redis::Value::SimpleString(text) => Ok(text.into_bytes()), + _ => Err(Error::InvalidEntry), + } +} + +fn lpop_result(value: redis::Value, multiple: bool) -> Result { + match value { + redis::Value::Nil => Ok(RedisLpopResult::Missing), + redis::Value::Array(values) if multiple => values + .into_iter() + .map(redis_bytes) + .collect::, _>>() + .map(RedisLpopResult::Values), + value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), + _ => Err(Error::InvalidEntry), + } +} + +fn count(value: redis::Value) -> Result, Error> { + match value { + redis::Value::Nil => Ok(None), + redis::Value::Int(value) => Ok(Some(value)), + redis::Value::BulkString(value) => std::str::from_utf8(&value) + .ok() + .and_then(|value| value.parse().ok()) + .map(Some) + .ok_or(Error::InvalidEntry), + redis::Value::SimpleString(value) => { + value.parse().map(Some).map_err(|_| Error::InvalidEntry) + } + _ => Err(Error::InvalidEntry), + } +} + +fn increment_with_floor( + connection: &mut ConnectionRef<'_>, + key: String, + amount: i64, + ttl: u64, +) -> Result { + redis::cmd("EVAL") + .arg(INCREMENT_WITH_FLOOR_SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 37b35c5ea4a..2548c7ac3c6 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -1,3 +1,3 @@ mod cache; -pub use cache::RedisCache; +pub use cache::{RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 6a61b80b84c..3d755d0f44c 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -2,9 +2,11 @@ use std::time::Duration; use litellm_cache::{ BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache, - CounterCache, Error, JsonCodec, get_cache, set_cache, + CounterCache, Error, IncrementOperation, JsonCodec, get_cache, set_cache, +}; +use litellm_cache_redis::{ + RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, }; -use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; struct TaggedByteCodec(u8); @@ -260,6 +262,307 @@ async fn async_flush_deletes_each_scan_page_separately() { cache.async_flush_cache().await.unwrap(); } +#[tokio::test] +async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { + let mut sadd_pipeline = redis::pipe(); + sadd_pipeline + .cmd("SADD") + .arg("team:members") + .arg("a") + .arg("b") + .cmd("EXPIRE") + .arg("team:members") + .arg(600u64) + .ignore(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("MGET").arg("team:count").arg("team:missing"), + Ok(redis_test::redis_value!(["7", nil])), + ), + MockCmd::new( + redis::cmd("MGET").arg("team:count").arg("team:missing"), + Ok(redis_test::redis_value!(["7", nil])), + ), + MockCmd::new(redis::cmd("PING"), Ok("PONG")), + MockCmd::new(redis::cmd("PING"), Ok("PONG")), + MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team:job-*") + .arg("COUNT") + .arg(25), + Ok(redis_test::redis_value!(["4", ["team:job-a"]])), + ), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(4) + .arg("MATCH") + .arg("team:job-*") + .arg("COUNT") + .arg(25), + Ok(redis_test::redis_value!(["0", ["team:job-b"]])), + ), + MockCmd::new( + redis::cmd("DEL").arg("team:job-a").arg("team:job-b"), + Ok(2u32), + ), + MockCmd::with_values( + sadd_pipeline, + Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]), + ), + MockCmd::new( + redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"), + Ok(2u32), + ), + MockCmd::new( + redis::cmd("LPOP").arg("team:queue").arg(2usize), + Ok(redis_test::redis_value!(["a", "b"])), + ), + MockCmd::new( + redis::cmd("EVAL") + .arg("return KEYS[1]") + .arg(1usize) + .arg("team:key"), + Ok("team:key"), + ), + MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")), + MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")), + MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + assert_eq!( + cache + .batch_get_counts(&["count".into(), "missing".into()]) + .unwrap(), + [Some(7), None] + ); + assert_eq!( + cache + .async_batch_get_counts(vec!["count".into(), "missing".into()]) + .await + .unwrap(), + [Some(7), None] + ); + assert!(cache.sync_ping().unwrap()); + assert!(cache.ping().await.unwrap()); + assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); + assert_eq!( + cache.async_scan_iter("job-", 25).await.unwrap(), + ["team:job-a", "team:job-b"] + ); + assert_eq!( + cache + .delete_cache_keys(vec!["job-a".into(), "job-b".into()]) + .await + .unwrap(), + 2 + ); + assert_eq!( + cache + .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) + .await + .unwrap(), + 2 + ); + assert_eq!( + cache + .async_rpush("queue", vec!["a".into(), "b".into()]) + .await + .unwrap(), + 2 + ); + assert_eq!( + cache.async_lpop("queue", Some(2)).await.unwrap(), + RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) + ); + assert_eq!( + cache + .async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new()) + .await + .unwrap(), + redis::Value::BulkString(b"team:key".to_vec()) + ); + assert_eq!(cache.client_list().unwrap(), "id=1"); + assert_eq!(cache.info().unwrap(), "redis_version:7"); + cache.flushall().unwrap(); +} + +#[tokio::test] +async fn direct_redis_pipelines_preserve_operation_order() { + let mut rpush_pipeline = redis::pipe(); + rpush_pipeline + .cmd("RPUSH") + .arg("team:a") + .arg("one") + .cmd("RPUSH") + .arg("team:b") + .arg("two"); + let mut lpop_pipeline = redis::pipe(); + lpop_pipeline + .cmd("LPOP") + .arg("team:a") + .arg(2usize) + .cmd("LPOP") + .arg("team:b"); + let connection = MockRedisConnection::new([ + MockCmd::with_values( + rpush_pipeline, + Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]), + ), + MockCmd::with_values( + lpop_pipeline, + Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]), + ), + ]) + .assert_all_commands_consumed(); + let queue = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + assert_eq!( + queue + .async_rpush_pipeline(vec![ + RedisRpushOperation { + key: "a".into(), + values: vec![RedisArg::from("one")], + }, + RedisRpushOperation { + key: "b".into(), + values: vec![RedisArg::from("two")], + }, + ]) + .await + .unwrap(), + [1, 2] + ); + assert_eq!( + queue + .async_lpop_pipeline(vec![ + RedisLpopOperation { + key: "a".into(), + count: Some(2), + }, + RedisLpopOperation { + key: "b".into(), + count: None, + }, + ]) + .await + .unwrap(), + [ + RedisLpopResult::Values(vec![b"one".to_vec()]), + RedisLpopResult::Missing, + ] + ); + + let mut increment_pipeline = redis::pipe(); + increment_pipeline + .cmd("INCRBYFLOAT") + .arg("team:counter") + .arg(1.5f64) + .cmd("EXPIRE") + .arg("team:counter") + .arg(10u64) + .ignore() + .cmd("INCRBYFLOAT") + .arg("team:counter") + .arg(2.0f64); + let connection = MockRedisConnection::new([MockCmd::with_values( + increment_pipeline, + Ok(vec![ + redis::Value::BulkString(b"1.5".to_vec()), + redis::Value::Int(1), + redis::Value::BulkString(b"3.5".to_vec()), + ]), + )]) + .assert_all_commands_consumed(); + let counters = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + assert_eq!( + counters + .async_increment_pipeline(vec![ + IncrementOperation { + key: "counter".into(), + amount: 1.5, + ttl: Some(Duration::from_secs(10)), + }, + IncrementOperation { + key: "counter".into(), + amount: 2.0, + ttl: None, + }, + ]) + .await + .unwrap(), + [1.5, 3.5] + ); +} + +const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return count" +); +const SET_MAX_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", + "redis.call('SET', KEYS[1], ARGV[1]); ", + "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return ARGV[1]; end; return current" +); + +#[tokio::test] +async fn counter_repairs_are_atomic_and_use_default_ttl() { + let floor = || { + redis::cmd("EVAL") + .arg(INCREMENT_WITH_FLOOR_SCRIPT) + .arg(1) + .arg("team:counter") + .arg(-2i64) + .arg(30u64) + .clone() + }; + let connection = MockRedisConnection::new([ + MockCmd::new(floor(), Ok(0i64)), + MockCmd::new(floor(), Ok(0i64)), + MockCmd::new( + redis::cmd("EVAL") + .arg(SET_MAX_SCRIPT) + .arg(1) + .arg("team:counter") + .arg(4.5f64) + .arg(600u64), + Ok("4.5"), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + assert_eq!( + cache + .increment_with_floor("counter", -2, Duration::from_secs(30)) + .unwrap(), + 0 + ); + assert_eq!( + cache + .async_increment_with_floor("counter", -2, Duration::from_secs(30)) + .await + .unwrap(), + 0 + ); + assert_eq!( + cache.async_set_max("counter", 4.5, None).await.unwrap(), + 4.5 + ); +} + const CLAIM_SCRIPT: &str = concat!( "local current = redis.call('GET', KEYS[1]); ", "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 013cf4b7baf..f9e7a148706 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -42,10 +42,12 @@ The resolver reads the namespace's `cache` attribute each time it resolves. A ca Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend -The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and configuration changes before selecting native execution. It does not compare Redis connection settings. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python +The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy +The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations + ## Adding another backend Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python @@ -56,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, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +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 diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index 0b9deab1f5a..1613a5786f4 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -2,6 +2,13 @@ use std::future::Future; use crate::{BaseCache, CacheKwargs, Error}; +#[derive(Clone, Debug, PartialEq)] +pub struct IncrementOperation { + pub key: String, + pub amount: f64, + pub ttl: Option, +} + pub trait CounterCache: BaseCache { fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ed67eb2fe15..cdd5589aa39 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -9,6 +9,6 @@ pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, }; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; -pub use capabilities::{ClaimCache, CounterCache}; +pub use capabilities::{ClaimCache, CounterCache, IncrementOperation}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; From d56479767e51dc58b4a5a4db278827d85461bf37 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 17:27:12 +0000 Subject: [PATCH 27/56] fix(bedrock): send s3BucketOwner on batch input and output data config Resolve s3_bucket_owner from litellm_params, then optional_params, then AWS_S3_BUCKET_OWNER and emit it on both S3 data configs so cross-account batch buckets pass Bedrock ownership validation. Omitted when unset Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/bedrock/batches/transformation.py | 51 +++++++---- litellm/llms/bedrock/common_utils.py | 26 +++++- litellm/types/llms/bedrock.py | 4 +- litellm/types/router.py | 1 + litellm/types/utils.py | 1 + .../bedrock/batches/test_transformation.py | 85 +++++++++++++++++++ tests/test_litellm/test_router.py | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 8 files changed, 155 insertions(+), 19 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 973388ca5bd..dd6c1bd9302 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -33,6 +33,7 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( CommonBatchFilesUtils, merge_bedrock_aws_request_params, + resolve_s3_bucket_owner, resolve_s3_encryption_key_id, ) @@ -51,6 +52,26 @@ _S3_BATCH_FILE_UUID_SUFFIX_PATTERN: Final = re.compile( _BEDROCK_TAGS_ADAPTER: Final[TypeAdapter[list[BedrockTag]]] = TypeAdapter(list[BedrockTag]) +def _build_s3_input_config(s3_uri: str, s3_bucket_owner: str | None) -> BedrockS3InputDataConfig: + if s3_bucket_owner is None: + return BedrockS3InputDataConfig(s3Uri=s3_uri) + return BedrockS3InputDataConfig(s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner) + + +def _build_s3_output_config( + s3_uri: str, s3_bucket_owner: str | None, s3_encryption_key_id: str | None +) -> BedrockS3OutputDataConfig: + match (s3_bucket_owner, s3_encryption_key_id): + case (None, None): + return BedrockS3OutputDataConfig(s3Uri=s3_uri) + case (str() as owner, None): + return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=owner) + case (None, str() as key_id): + return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3EncryptionKeyId=key_id) + case (str() as owner, str() as key_id): + return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=owner, s3EncryptionKeyId=key_id) + + def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: try: return _BEDROCK_TAGS_ADAPTER.validate_python(raw_tags, strict=True) @@ -214,25 +235,23 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): job_name: Final = self.common_utils.generate_unique_job_name(model, prefix="litellm") output_key: Final = f"litellm-batch-outputs/{job_name}/" - # Build input data config - input_data_config: Final[BedrockInputDataConfig] = { - "s3InputDataConfig": BedrockS3InputDataConfig(s3Uri=f"s3://{input_bucket}/{input_key}") - } - - # Build output data config - s3_output_config: Final[BedrockS3OutputDataConfig] = BedrockS3OutputDataConfig( - s3Uri=f"s3://{output_bucket}/{output_key}" - ) - - # Add optional KMS encryption key ID if provided - s3_encryption_key_id = resolve_s3_encryption_key_id( + s3_bucket_owner: Final = resolve_s3_bucket_owner(litellm_params=litellm_params, optional_params=optional_params) + s3_encryption_key_id: Final = resolve_s3_encryption_key_id( litellm_params=litellm_params, optional_params=optional_params, ) - if s3_encryption_key_id: - s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id - - output_data_config: Final[BedrockOutputDataConfig] = {"s3OutputDataConfig": s3_output_config} + input_data_config: Final[BedrockInputDataConfig] = { + "s3InputDataConfig": _build_s3_input_config( + s3_uri=f"s3://{input_bucket}/{input_key}", s3_bucket_owner=s3_bucket_owner + ) + } + output_data_config: Final[BedrockOutputDataConfig] = { + "s3OutputDataConfig": _build_s3_output_config( + s3_uri=f"s3://{output_bucket}/{output_key}", + s3_bucket_owner=s3_bucket_owner, + s3_encryption_key_id=s3_encryption_key_id, + ) + } # Create Bedrock batch request with proper typing bedrock_request: Final[BedrockCreateBatchRequest] = { diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 7e24292a87e..f1066643874 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1555,11 +1555,33 @@ def resolve_s3_encryption_key_id( Precedence: `s3_encryption_key_id` in litellm_params, then optional_params (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. """ + return _resolve_s3_setting("s3_encryption_key_id", "AWS_S3_ENCRYPTION_KEY_ID", litellm_params, optional_params) + + +def resolve_s3_bucket_owner( + litellm_params: Mapping[str, object], + optional_params: Mapping[str, object] | None = None, +) -> str | None: + """ + Resolve the AWS account id that owns the S3 buckets used by Bedrock batch jobs. + + Precedence: `s3_bucket_owner` in litellm_params, then optional_params + (client-side / request params), then the AWS_S3_BUCKET_OWNER env var. + """ + return _resolve_s3_setting("s3_bucket_owner", "AWS_S3_BUCKET_OWNER", litellm_params, optional_params) + + +def _resolve_s3_setting( + param_name: str, + env_var: str, + litellm_params: Mapping[str, object], + optional_params: Mapping[str, object] | None, +) -> str | None: candidates: Final = tuple( - source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None + source.get(param_name) for source in (litellm_params, optional_params) if source is not None ) explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None) - return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + return explicit or get_secret_str(env_var) class CommonBatchFilesUtils: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 10082cf2373..f9eaef5e891 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -4,7 +4,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict -from typing_extensions import ReadOnly, Required, TypedDict, override +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, override from .openai import ChatCompletionToolCallChunk @@ -1082,6 +1082,7 @@ class BedrockS3InputDataConfig(TypedDict): """S3 input data configuration for Bedrock batch jobs.""" s3Uri: str + s3BucketOwner: NotRequired[ReadOnly[str]] class BedrockInputDataConfig(TypedDict): @@ -1095,6 +1096,7 @@ class BedrockS3OutputDataConfig(TypedDict, total=False): s3Uri: str s3EncryptionKeyId: str | None + s3BucketOwner: ReadOnly[str] class BedrockOutputDataConfig(TypedDict): diff --git a/litellm/types/router.py b/litellm/types/router.py index a75b4654cab..6b43573d0e2 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -306,6 +306,7 @@ class CredentialLiteLLMParams(BaseModel): s3_endpoint_url: str | None = None s3_region_name: str | None = None s3_encryption_key_id: str | None = None + s3_bucket_owner: str | None = None aws_batch_role_arn: str | None = None s3_output_bucket_name: str | None = None bedrock_tags: list | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5a80644347e..45556acdfb5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3830,6 +3830,7 @@ bedrock_batch_litellm_params: Final = ( "s3_region_name", "s3_endpoint_url", "s3_output_bucket_name", + "s3_bucket_owner", "bedrock_tags", ) diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 7e5716a7495..eb08c19cbdf 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -185,6 +185,91 @@ def test_create_request_omits_kms_key_when_absent(config): assert "s3EncryptionKeyId" not in s3out +def _signed_batch_request(config, litellm_params: dict, optional_params: dict) -> dict: + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://in-bucket/in.jsonl"}, + optional_params=optional_params, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r", **litellm_params}, + ) + return mock_sign.call_args.kwargs["data"] + + +@pytest.mark.parametrize( + ("litellm_params", "optional_params", "env_owner", "expected_owner"), + [ + pytest.param({"s3_bucket_owner": "111111111111"}, {}, None, "111111111111", id="litellm_params"), + pytest.param({}, {"s3_bucket_owner": "222222222222"}, None, "222222222222", id="optional_params"), + pytest.param({}, {}, "333333333333", "333333333333", id="env"), + pytest.param( + {"s3_bucket_owner": "111111111111"}, + {"s3_bucket_owner": "222222222222"}, + "333333333333", + "111111111111", + id="litellm_params_wins", + ), + pytest.param( + {}, {"s3_bucket_owner": "222222222222"}, "333333333333", "222222222222", id="optional_params_beats_env" + ), + ], +) +def test_create_request_sets_s3_bucket_owner_on_input_and_output( + config, monkeypatch, litellm_params, optional_params, env_owner, expected_owner +): + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + if env_owner is None: + monkeypatch.delenv("AWS_S3_BUCKET_OWNER", raising=False) + else: + monkeypatch.setenv("AWS_S3_BUCKET_OWNER", env_owner) + + bedrock_request = _signed_batch_request(config, litellm_params, optional_params) + + assert bedrock_request["inputDataConfig"] == { + "s3InputDataConfig": {"s3Uri": "s3://in-bucket/in.jsonl", "s3BucketOwner": expected_owner} + } + assert bedrock_request["outputDataConfig"] == { + "s3OutputDataConfig": { + "s3Uri": "s3://in-bucket/litellm-batch-outputs/litellm-batch-1/", + "s3BucketOwner": expected_owner, + } + } + + +def test_create_request_omits_s3_bucket_owner_when_unset(config, monkeypatch): + monkeypatch.delenv("AWS_S3_BUCKET_OWNER", raising=False) + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + bedrock_request = _signed_batch_request(config, {}, {}) + + assert bedrock_request["inputDataConfig"] == {"s3InputDataConfig": {"s3Uri": "s3://in-bucket/in.jsonl"}} + assert bedrock_request["outputDataConfig"] == { + "s3OutputDataConfig": {"s3Uri": "s3://in-bucket/litellm-batch-outputs/litellm-batch-1/"} + } + + +def test_create_request_keeps_kms_key_alongside_s3_bucket_owner(config, monkeypatch): + monkeypatch.delenv("AWS_S3_BUCKET_OWNER", raising=False) + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + bedrock_request = _signed_batch_request( + config, {"s3_bucket_owner": "111111111111", "s3_encryption_key_id": "kms-key-123"}, {} + ) + + assert bedrock_request["outputDataConfig"] == { + "s3OutputDataConfig": { + "s3Uri": "s3://in-bucket/litellm-batch-outputs/litellm-batch-1/", + "s3BucketOwner": "111111111111", + "s3EncryptionKeyId": "kms-key-123", + } + } + + def test_create_request_missing_input_file_id_raises(config): with pytest.raises(ValueError, match="input_file_id is required"): config.transform_create_batch_request( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8310d30d90e..d20fdff894a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6294,6 +6294,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): "s3_bucket_name": "my-batch-bucket", "s3_region_name": "us-east-1", "s3_encryption_key_id": "arn:aws:kms:us-west-2:123:key/abc", + "s3_bucket_owner": "111111111111", "aws_batch_role_arn": "arn:aws:iam::123:role/batch-role", }, } @@ -6311,6 +6312,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): assert credentials["s3_bucket_name"] == "my-batch-bucket" assert credentials["s3_region_name"] == "us-east-1" assert credentials["s3_encryption_key_id"] == "arn:aws:kms:us-west-2:123:key/abc" + assert credentials["s3_bucket_owner"] == "111111111111" assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index aa71adfad42..f5f1567d401 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31136,6 +31136,8 @@ export interface components { rpm?: number | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; + /** S3 Bucket Owner */ + s3_bucket_owner?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; /** S3 Endpoint Url */ @@ -41860,6 +41862,8 @@ export interface components { rpm?: number | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; + /** S3 Bucket Owner */ + s3_bucket_owner?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; /** S3 Endpoint Url */ From 595711829aa2b0a00b2573070d2fc97a850bbc5a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 10:30:42 -0700 Subject: [PATCH 28/56] wip --- litellm-rust/crates/cache-response/README.md | 4 +- .../crates/cache-response/src/buffer.rs | 46 ++ litellm-rust/crates/cache-response/src/lib.rs | 2 + .../crates/cache-response/tests/response.rs | 70 ++- .../crates/python-bridge/src/cache/binding.rs | 294 +++++++++ .../python-bridge/src/cache/callback.rs | 169 ++++++ .../crates/python-bridge/src/cache/facade.rs | 2 +- .../crates/python-bridge/src/cache/future.rs | 18 + .../crates/python-bridge/src/cache/handle.rs | 106 ++++ .../crates/python-bridge/src/cache/mod.rs | 570 +----------------- .../crates/python-bridge/src/cache/native.rs | 35 +- .../crates/python-bridge/src/cache/request.rs | 48 ++ .../python-bridge/src/cache/resolver.rs | 39 ++ 13 files changed, 810 insertions(+), 593 deletions(-) create mode 100644 litellm-rust/crates/cache-response/src/buffer.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/binding.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/callback.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/future.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/handle.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/request.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/resolver.rs diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index f9e7a148706..9863c46783f 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -6,9 +6,9 @@ `litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations -`litellm-cache-response` owns response keys, controls, entries, and the Python-compatible response codec. It has no runtime dependency on a specific cache backend or Python +`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python -The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host +The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host ## Native Rust use diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs new file mode 100644 index 00000000000..68af5278c8c --- /dev/null +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -0,0 +1,46 @@ +use std::{sync::Mutex, time::Duration}; + +use litellm_cache::{BaseCache, Error}; +use serde_json::Value; + +use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; + +/// Defers async writes until `flush_size` entries are pending, then stores them as one batch. +pub struct WriteBuffer { + flush_size: usize, + entries: Mutex>, +} + +impl WriteBuffer { + pub fn new(flush_size: usize) -> Self { + Self { + flush_size: flush_size.max(1), + entries: Mutex::new(Vec::new()), + } + } + + pub async fn async_store>( + &self, + cache: &ResponseCache, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + let pending = { + let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?; + entries.push((request.clone(), response, now)); + (entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries)) + }; + // A failed flush drops its batch, as Python does. Requeueing would grow the + // buffer and re-send an ever larger pipeline on every write during an outage. + match pending { + Some(pending) => cache.async_store_entries(pending).await, + None => Ok(()), + } + } + + pub fn clear(&self) -> Result<(), Error> { + self.entries.lock().map_err(|_| Error::Unavailable)?.clear(); + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 72a507f8ee9..91b36ebe24b 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,8 +1,10 @@ +mod buffer; mod caching; mod codec; mod embedding; mod response; +pub use buffer::WriteBuffer; pub use caching::{ CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, get_cache_key, should_use_cache, diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 7c69e5a1d55..94b25626f86 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -11,7 +11,7 @@ use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, - ResponseCacheRequest, + ResponseCacheRequest, WriteBuffer, }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; @@ -414,3 +414,71 @@ async fn deferred_entries_keep_the_time_they_were_produced() { None ); } + +#[tokio::test] +async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let buffer = WriteBuffer::new(2); + let mut first = request(); + first.max_age = Some(Duration::from_secs(10)); + let mut second = request(); + second.key.preset = Some("tenant:other".into()); + + buffer + .async_store( + &cache, + &first, + json!({"answer": 7}), + Duration::from_secs(100), + ) + .await + .unwrap(); + assert_eq!( + cache.lookup(&first, Duration::from_secs(100)).unwrap(), + None + ); + + buffer + .async_store( + &cache, + &second, + json!({"answer": 8}), + Duration::from_secs(200), + ) + .await + .unwrap(); + assert_eq!( + cache.lookup(&first, Duration::from_secs(110)).unwrap(), + Some(json!({"answer": 7})) + ); + assert_eq!( + cache.lookup(&first, Duration::from_secs(111)).unwrap(), + None + ); + assert_eq!( + cache.lookup(&second, Duration::from_secs(200)).unwrap(), + Some(json!({"answer": 8})) + ); +} + +#[tokio::test] +async fn write_buffer_clear_drops_pending_entries() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let buffer = WriteBuffer::new(2); + let mut other = request(); + other.key.preset = Some("tenant:other".into()); + let now = Duration::from_secs(100); + + buffer + .async_store(&cache, &request(), json!({"answer": 7}), now) + .await + .unwrap(); + buffer.clear().unwrap(); + buffer + .async_store(&cache, &other, json!({"answer": 8}), now) + .await + .unwrap(); + + assert_eq!(cache.lookup(&request(), now).unwrap(), None); + assert_eq!(cache.lookup(&other, now).unwrap(), None); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs new file mode 100644 index 00000000000..44c133d4611 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -0,0 +1,294 @@ +use litellm_cache_response::PartialHits; +use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::PyDict, +}; +use serde_json::Value; + +use super::{ + cache_error, + callback::PythonCallback, + future::{ready_none, ready_value}, + native::NativeResponseCache, + request::{now, request, requests}, +}; + +pub(super) enum CacheBinding { + Disabled, + Native(NativeResponseCache), + PythonCallback(PythonCallback), +} + +#[pyclass(frozen, name = "_CacheTestBinding")] +pub(crate) struct ResolvedCache { + binding: CacheBinding, + pid: u32, +} + +impl ResolvedCache { + pub(super) fn new(binding: CacheBinding) -> Self { + Self { + binding, + pid: std::process::id(), + } + } + + fn check_process(&self) -> PyResult<()> { + if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache bindings must be resolved again after fork", + )); + } + Ok(()) + } + + pub(crate) fn lookup_step( + &self, + py: Python<'_>, + input: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.check_process()?; + let awaitable = match &self.binding { + CacheBinding::Disabled => ready_none(py)?, + CacheBinding::Native(service) => { + let request = request(input)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup(&request, now()).await }, + cache_error, + )? + } + CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?, + }; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +#[pymethods] +impl ResolvedCache { + #[getter] + fn kind(&self) -> &'static str { + match self.binding { + CacheBinding::Disabled => "disabled", + CacheBinding::Native(_) => "native", + CacheBinding::PythonCallback(_) => "python_callback", + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn lookup( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(py.None()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup(&request, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(callback) => { + callback.lookup(py, callback_kwargs).map(Bound::unbind) + } + } + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn store( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + response: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + release_gil(py, move || service.store(&request, response, now())) + .map_err(cache_error) + } + CacheBinding::PythonCallback(callback) => callback.store(py, response, callback_kwargs), + } + } + + /// Native bindings return `{values, missing_indices}`, while a Python callback returns the + /// list of its per-request results. + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn lookup_batch( + &self, + py: Python<'_>, + requests: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyAny>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + to_py(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup_batch(&requests, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(callback) => callback + .lookup_batch(py, requests, callback_kwargs) + .map(Bound::unbind), + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn async_lookup<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? + else { + unreachable!() + }; + Ok(awaitable.into_bound(py)) + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn async_store<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + response: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + run_async( + py, + async move { service.async_store(&request, response, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_store(py, response, callback_kwargs) + } + } + } + + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + ready_value(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup_batch(&requests, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_lookup_batch(py, requests, callback_kwargs) + } + } + } + + /// A Python callback receives the caller's original result through `callback_result`. + #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] + fn async_store_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + responses: &Bound<'py, PyAny>, + callback_result: Option<&Bound<'py, PyAny>>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let responses: Vec = from_py(responses)?; + if requests.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let entries = requests.into_iter().zip(responses).collect(); + let service = service.clone(); + run_async( + py, + async move { service.async_store_batch(entries, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_store_batch(py, callback_result, callback_kwargs) + } + } + } + + fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async(py, async move { service.async_flush().await }, cache_error) + } + CacheBinding::PythonCallback(callback) => callback.async_flush(py), + } + } + + fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async( + py, + async move { service.test_connection().await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => callback.ping(py), + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let CacheBinding::PythonCallback(callback) = &self.binding { + callback.traverse(&visit)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/callback.rs b/litellm-rust/crates/python-bridge/src/cache/callback.rs new file mode 100644 index 00000000000..318f9d02080 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/callback.rs @@ -0,0 +1,169 @@ +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{PyDict, PyList, PyTuple}, +}; + +use super::future::ready_none; + +/// A custom Python cache object, driven through the built-in `Cache` API so a `Cache` subclass +/// works unchanged. +pub(super) struct PythonCallback(Py); + +impl PythonCallback { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + + pub(super) fn lookup<'py>( + &self, + py: Python<'py>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0 + .bind(py) + .call_method("get_cache", (), Some(callback_kwargs(kwargs)?)) + } + + pub(super) fn async_lookup<'py>( + &self, + py: Python<'py>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0 + .bind(py) + .call_method("async_get_cache", (), Some(callback_kwargs(kwargs)?)) + } + + pub(super) fn store( + &self, + py: Python<'_>, + response: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.0 + .bind(py) + .call_method("add_cache", (response,), Some(callback_kwargs(kwargs)?)) + .map(|_| ()) + } + + pub(super) fn async_store<'py>( + &self, + py: Python<'py>, + response: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0.bind(py).call_method( + "async_add_cache", + (response,), + Some(callback_kwargs(kwargs)?), + ) + } + + /// The built-in `Cache` API has no batch read, so the callback receives one + /// `get_cache(**kwargs)` call per request, in order, and the results come back as a list. + pub(super) fn lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let results = PyList::empty(py); + for kwargs in batch_callback_kwargs(requests, kwargs)? { + results.append( + self.0 + .bind(py) + .call_method("get_cache", (), Some(&kwargs))?, + )?; + } + Ok(results.into_any()) + } + + pub(super) fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let awaitables = batch_callback_kwargs(requests, kwargs)? + .iter() + .map(|kwargs| { + self.0 + .bind(py) + .call_method("async_get_cache", (), Some(kwargs)) + }) + .collect::>>()?; + py.import("asyncio")? + .call_method1("gather", PyTuple::new(py, awaitables)?) + } + + /// Receives the caller's original result, because the built-in + /// `Cache.async_add_cache_pipeline` splits the batch itself. + pub(super) fn async_store_batch<'py>( + &self, + py: Python<'py>, + result: Option<&Bound<'py, PyAny>>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let result = result.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_result") + })?; + self.0.bind(py).call_method( + "async_add_cache_pipeline", + (result,), + Some(callback_kwargs(kwargs)?), + ) + } + + /// The built-in `Cache` facade has no flush of its own; its backend does. + pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + let object = self.0.bind(py); + let backend = match object.getattr_opt("cache")? { + Some(backend) if !backend.is_none() => backend, + _ => object.clone(), + }; + if backend.hasattr("async_flush_cache")? { + return backend.call_method0("async_flush_cache"); + } + backend.call_method0("flush_cache")?; + ready_none(py) + } + + pub(super) fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.0.bind(py).call_method0("ping") + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } +} + +fn callback_kwargs<'a, 'py>( + kwargs: Option<&'a Bound<'py, PyDict>>, +) -> PyResult<&'a Bound<'py, PyDict>> { + kwargs.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") + }) +} + +fn batch_callback_kwargs<'py>( + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, +) -> PyResult>> { + let kwargs = kwargs + .ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require one original callback_kwargs mapping per request", + ) + })? + .try_iter()? + .map(|item| Ok(item?.cast_into::()?)) + .collect::>>()?; + if kwargs.len() != requests.len()? { + return Err(PyValueError::new_err( + "batch cache requests and callback_kwargs must have equal lengths", + )); + } + Ok(kwargs) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 2ad13ce200f..19220bf868f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -9,7 +9,7 @@ use pyo3::{ }; use serde_json::Value; -use super::{CacheTestHandle, native::NativeResponseCache}; +use super::{handle::CacheTestHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, diff --git a/litellm-rust/crates/python-bridge/src/cache/future.rs b/litellm-rust/crates/python-bridge/src/cache/future.rs new file mode 100644 index 00000000000..42593eee1f4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/future.rs @@ -0,0 +1,18 @@ +use litellm_host_python::to_py; +use pyo3::prelude::*; + +pub(super) fn ready_none(py: Python<'_>) -> PyResult> { + ready_value(py, &()) +} + +pub(super) fn ready_value<'py, T: serde::Serialize>( + py: Python<'py>, + value: &T, +) -> PyResult> { + let future = py + .import("asyncio")? + .call_method0("get_running_loop")? + .call_method0("create_future")?; + future.call_method1("set_result", (to_py(py, value)?,))?; + Ok(future) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs new file mode 100644 index 00000000000..42d7f2c2f3d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -0,0 +1,106 @@ +use std::time::Duration; + +use litellm_host_python::release_gil; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; + +use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use crate::python_settings::PythonSettings; + +const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(FromPyObject)] +struct PythonCacheSettings { + default_redis_ttl: Option, +} + +fn redis_default_ttl(py: Python<'_>) -> PyResult { + let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; + settings + .default_redis_ttl + .map(duration) + .transpose() + .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) +} + +#[pyclass(frozen, name = "_CacheTestHandle")] +pub(crate) struct CacheTestHandle { + service: NativeResponseCache, + pub(super) guard: Option, + pid: u32, +} + +impl CacheTestHandle { + pub(super) fn service(&self) -> PyResult { + if self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache handles must be recreated after fork", + )); + } + Ok(self.service.clone()) + } +} + +#[pymethods] +impl CacheTestHandle { + #[staticmethod] + #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] + fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { + Ok(Self { + service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), + guard: None, + pid: std::process::id(), + }) + } + + #[staticmethod] + #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + fn redis( + py: Python<'_>, + url: String, + ttl_seconds: Option, + namespace: Option, + ) -> PyResult { + let ttl = Some(match ttl_seconds { + Some(seconds) => duration(seconds)?, + None => redis_default_ttl(py)?, + }); + let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + + #[getter] + fn backend(&self) -> &'static str { + self.service.kind() + } + + 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 handle = Py::new( + py, + Self { + service, + guard: Some(guard), + pid: self.pid, + }, + )?; + facade.setattr("_native_cache_handle", handle) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(guard) = &self.guard { + guard.traverse(visit)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index f83788f6036..7955ed934b2 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,81 +1,21 @@ +mod binding; +mod callback; mod facade; +mod future; +mod handle; mod native; +mod request; +mod resolver; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use facade::FacadeGuard; use litellm_cache::Error; -use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; -use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; -use native::NativeResponseCache; use pyo3::{ - PyTraverseError, PyVisit, - exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + exceptions::{PyRuntimeError, PyValueError}, prelude::*, - types::{PyDict, PyList, PyTuple}, }; -use serde::Deserialize; -use serde_json::Value; -use crate::python_settings::PythonSettings; - -const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); - -#[derive(FromPyObject)] -struct PythonCacheSettings { - default_redis_ttl: Option, -} - -fn redis_default_ttl(py: Python<'_>) -> PyResult { - let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; - settings - .default_redis_ttl - .map(duration) - .transpose() - .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) -} - -#[derive(Deserialize)] -#[serde(deny_unknown_fields)] -struct RequestInput { - key: CacheKeyInput, - controls: Option, - ttl_seconds: Option, - max_age_seconds: Option, -} - -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.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) -} - -fn requests(value: &Bound<'_, PyAny>) -> PyResult> { - from_py::>(value)? - .into_iter() - .map(request_input) - .collect() -} - -fn duration(seconds: f64) -> PyResult { - Duration::try_from_secs_f64(seconds) - .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) -} - -fn now() -> Duration { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() -} +pub(crate) use self::{ + binding::ResolvedCache, handle::CacheTestHandle, resolver::CacheTestResolver, +}; fn cache_error(error: Error) -> PyErr { match error { @@ -83,493 +23,3 @@ fn cache_error(error: Error) -> PyErr { _ => PyRuntimeError::new_err(error.to_string()), } } - -#[pyclass(frozen, name = "_CacheTestHandle")] -pub(crate) struct CacheTestHandle { - service: NativeResponseCache, - guard: Option, - pid: u32, -} - -impl CacheTestHandle { - fn service(&self) -> PyResult { - if self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native cache handles must be recreated after fork", - )); - } - Ok(self.service.clone()) - } -} - -#[pymethods] -impl CacheTestHandle { - #[staticmethod] - #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] - fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { - Ok(Self { - service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), - guard: None, - pid: std::process::id(), - }) - } - - #[staticmethod] - #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] - fn redis( - py: Python<'_>, - url: String, - ttl_seconds: Option, - namespace: Option, - ) -> PyResult { - let ttl = Some(match ttl_seconds { - Some(seconds) => duration(seconds)?, - None => redis_default_ttl(py)?, - }); - let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) - .map_err(cache_error)?; - Ok(Self { - service, - guard: None, - pid: std::process::id(), - }) - } - - #[getter] - fn backend(&self) -> &'static str { - self.service.kind() - } - - 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 handle = Py::new( - py, - Self { - service, - guard: Some(guard), - pid: self.pid, - }, - )?; - facade.setattr("_native_cache_handle", handle) - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let Some(guard) = &self.guard { - guard.traverse(visit)?; - } - Ok(()) - } -} - -enum CacheBinding { - Disabled, - Native(NativeResponseCache), - PythonCallback(Py), -} - -#[pyclass(frozen, name = "_CacheTestBinding")] -pub(crate) struct ResolvedCache { - binding: CacheBinding, - pid: u32, -} - -impl ResolvedCache { - fn check_process(&self) -> PyResult<()> { - if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native cache bindings must be resolved again after fork", - )); - } - Ok(()) - } - - pub(crate) fn lookup_step( - &self, - py: Python<'_>, - input: &Bound<'_, PyAny>, - kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult { - self.check_process()?; - let awaitable = match &self.binding { - CacheBinding::Disabled => ready_none(py)?, - CacheBinding::Native(service) => { - let request = request(input)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup(&request, now()).await }, - cache_error, - )? - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_get_cache", - (), - Some(callback_kwargs(kwargs)?), - )?, - }; - Ok(ExecutionStep::Await(awaitable.unbind())) - } -} - -#[pymethods] -impl ResolvedCache { - #[getter] - fn kind(&self) -> &'static str { - match self.binding { - CacheBinding::Disabled => "disabled", - CacheBinding::Native(_) => "native", - CacheBinding::PythonCallback(_) => "python_callback", - } - } - - #[pyo3(signature = (request, *, callback_kwargs=None))] - fn lookup( - &self, - py: Python<'_>, - request: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => Ok(py.None()), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let service = service.clone(); - let response = release_gil(py, move || service.lookup(&request, now())) - .map_err(cache_error)?; - to_py(py, &response) - } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "get_cache", - (), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(Bound::unbind), - } - } - - #[pyo3(signature = (request, response, *, callback_kwargs=None))] - fn store( - &self, - py: Python<'_>, - request: &Bound<'_, PyAny>, - response: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult<()> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => Ok(()), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let response: Value = from_py(response)?; - let service = service.clone(); - release_gil(py, move || service.store(&request, response, now())) - .map_err(cache_error) - } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "add_cache", - (response,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(|_| ()), - } - } - - /// Native bindings return `{values, missing_indices}`. The built-in `Cache` API has no batch - /// read, so a Python callback receives one `get_cache(**kwargs)` call per request, in order, - /// and the results come back as a list. - #[pyo3(signature = (requests, *, callback_kwargs=None))] - fn lookup_batch( - &self, - py: Python<'_>, - requests: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyAny>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => { - let requests = self::requests(requests)?; - to_py(py, &PartialHits::new(vec![None; requests.len()])) - } - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let service = service.clone(); - let response = release_gil(py, move || service.lookup_batch(&requests, now())) - .map_err(cache_error)?; - to_py(py, &response) - } - CacheBinding::PythonCallback(object) => { - let results = PyList::empty(py); - for kwargs in batch_callback_kwargs(requests, callback_kwargs)? { - results.append(object.bind(py).call_method( - "get_cache", - (), - Some(&kwargs), - )?)?; - } - Ok(results.into_any().unbind()) - } - } - } - - #[pyo3(signature = (request, *, callback_kwargs=None))] - fn async_lookup<'py>( - &self, - py: Python<'py>, - request: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? - else { - unreachable!() - }; - Ok(awaitable.into_bound(py)) - } - - #[pyo3(signature = (request, response, *, callback_kwargs=None))] - fn async_store<'py>( - &self, - py: Python<'py>, - request: &Bound<'py, PyAny>, - response: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let response: Value = from_py(response)?; - let service = service.clone(); - run_async( - py, - async move { service.async_store(&request, response, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_add_cache", - (response,), - Some(self::callback_kwargs(callback_kwargs)?), - ), - } - } - - #[pyo3(signature = (requests, *, callback_kwargs=None))] - fn async_lookup_batch<'py>( - &self, - py: Python<'py>, - requests: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyAny>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => { - let requests = self::requests(requests)?; - ready_value(py, &PartialHits::new(vec![None; requests.len()])) - } - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup_batch(&requests, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => { - let awaitables = batch_callback_kwargs(requests, callback_kwargs)? - .iter() - .map(|kwargs| { - object - .bind(py) - .call_method("async_get_cache", (), Some(kwargs)) - }) - .collect::>>()?; - py.import("asyncio")? - .call_method1("gather", PyTuple::new(py, awaitables)?) - } - } - } - - /// A Python callback receives the caller's original result through `callback_result`, because - /// the built-in `Cache.async_add_cache_pipeline` splits the batch itself. - #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] - fn async_store_batch<'py>( - &self, - py: Python<'py>, - requests: &Bound<'py, PyAny>, - responses: &Bound<'py, PyAny>, - callback_result: Option<&Bound<'py, PyAny>>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let responses: Vec = from_py(responses)?; - if requests.len() != responses.len() { - return Err(PyValueError::new_err( - "batch cache requests and responses must have equal lengths", - )); - } - let entries = requests.into_iter().zip(responses).collect(); - let service = service.clone(); - run_async( - py, - async move { service.async_store_batch(entries, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => { - let result = callback_result.ok_or_else(|| { - PyTypeError::new_err( - "Python cache callbacks require their original callback_result", - ) - })?; - object.bind(py).call_method( - "async_add_cache_pipeline", - (result,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - } - } - } - - fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let service = service.clone(); - run_async(py, async move { service.async_flush().await }, cache_error) - } - // The built-in `Cache` facade has no flush of its own; its backend does. - CacheBinding::PythonCallback(object) => { - let object = object.bind(py); - let backend = match object.getattr_opt("cache")? { - Some(backend) if !backend.is_none() => backend, - _ => object.clone(), - }; - if backend.hasattr("async_flush_cache")? { - return backend.call_method0("async_flush_cache"); - } - backend.call_method0("flush_cache")?; - ready_none(py) - } - } - } - - fn ping<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let service = service.clone(); - run_async( - py, - async move { service.test_connection().await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method0("ping"), - } - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let CacheBinding::PythonCallback(object) = &self.binding { - visit.call(object)?; - } - Ok(()) - } -} - -fn callback_kwargs<'a, 'py>( - kwargs: Option<&'a Bound<'py, PyDict>>, -) -> PyResult<&'a Bound<'py, PyDict>> { - kwargs.ok_or_else(|| { - PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") - }) -} - -fn batch_callback_kwargs<'py>( - requests: &Bound<'py, PyAny>, - kwargs: Option<&Bound<'py, PyAny>>, -) -> PyResult>> { - let kwargs = kwargs - .ok_or_else(|| { - PyTypeError::new_err( - "Python cache callbacks require one original callback_kwargs mapping per request", - ) - })? - .try_iter()? - .map(|item| Ok(item?.cast_into::()?)) - .collect::>>()?; - if kwargs.len() != requests.len()? { - return Err(PyValueError::new_err( - "batch cache requests and callback_kwargs must have equal lengths", - )); - } - Ok(kwargs) -} - -fn ready_none(py: Python<'_>) -> PyResult> { - ready_value(py, &()) -} - -fn ready_value<'py, T: serde::Serialize>( - py: Python<'py>, - value: &T, -) -> PyResult> { - let future = py - .import("asyncio")? - .call_method0("get_running_loop")? - .call_method0("create_future")?; - future.call_method1("set_result", (to_py(py, value)?,))?; - Ok(future) -} - -#[pyclass(frozen, name = "_CacheTestResolver")] -pub(crate) struct CacheTestResolver { - namespace: Py, -} - -#[pymethods] -impl CacheTestResolver { - #[new] - fn new(namespace: Py) -> Self { - Self { namespace } - } - - pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { - let object = self.namespace.bind(py).getattr("cache")?; - let binding = if object.is_none() { - CacheBinding::Disabled - } else if let Ok(handle) = object.extract::>() { - CacheBinding::Native(handle.service()?) - } else if let Some(service) = facade::resolve(py, &object)? { - CacheBinding::Native(service) - } else { - CacheBinding::PythonCallback(object.unbind()) - }; - Ok(ResolvedCache { - binding, - pid: std::process::id(), - }) - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.namespace) - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 3fc8f61dff6..a718d07b286 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -4,25 +4,19 @@ use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; use serde_json::Value; -use tokio::sync::Mutex; #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), Redis { cache: Arc>>, - buffer: Option>, + buffer: Option>, }, } -pub(super) struct RedisWriteBuffer { - flush_size: usize, - entries: Mutex>, -} - impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { Self::Memory(Arc::new(ResponseCache::new(Arc::new( @@ -33,7 +27,7 @@ impl NativeResponseCache { Some(Arc::new(|entry| { ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) })), - super::now, + super::request::now, ), )))) } @@ -84,12 +78,7 @@ impl NativeResponseCache { match self { Self::Redis { cache, .. } => Self::Redis { cache, - buffer: flush_size.map(|flush_size| { - Arc::new(RedisWriteBuffer { - flush_size: flush_size.max(1), - entries: Mutex::new(Vec::new()), - }) - }), + buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, memory => memory, } @@ -155,19 +144,7 @@ impl NativeResponseCache { Self::Redis { cache, buffer: Some(buffer), - } => { - let pending = { - let mut entries = buffer.entries.lock().await; - entries.push((request.clone(), response, now)); - (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) - }; - // A failed flush drops its batch, as Python does. Requeueing would grow the - // buffer and re-send an ever larger pipeline on every write during an outage. - match pending { - Some(pending) => cache.async_store_entries(pending).await, - None => Ok(()), - } - } + } => buffer.async_store(cache, request, response, now).await, } } @@ -198,7 +175,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.async_flush().await, Self::Redis { cache, buffer } => { if let Some(buffer) = buffer { - buffer.entries.lock().await.clear(); + buffer.clear()?; } cache.async_flush().await } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs new file mode 100644 index 00000000000..d8793abd115 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -0,0 +1,48 @@ +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_host_python::from_py; +use pyo3::{exceptions::PyValueError, prelude::*}; +use serde::Deserialize; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestInput { + key: CacheKeyInput, + controls: Option, + ttl_seconds: Option, + max_age_seconds: 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.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; + request.max_age = input.max_age_seconds.map(duration).transpose()?; + Ok(request) +} + +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { + from_py::>(value)? + .into_iter() + .map(request_input) + .collect() +} + +pub(super) fn duration(seconds: f64) -> PyResult { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) +} + +pub(super) fn now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() +} diff --git a/litellm-rust/crates/python-bridge/src/cache/resolver.rs b/litellm-rust/crates/python-bridge/src/cache/resolver.rs new file mode 100644 index 00000000000..ef6f142e0a1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/resolver.rs @@ -0,0 +1,39 @@ +use pyo3::{PyTraverseError, PyVisit, prelude::*}; + +use super::{ + binding::{CacheBinding, ResolvedCache}, + callback::PythonCallback, + facade, + handle::CacheTestHandle, +}; + +#[pyclass(frozen, name = "_CacheTestResolver")] +pub(crate) struct CacheTestResolver { + namespace: Py, +} + +#[pymethods] +impl CacheTestResolver { + #[new] + fn new(namespace: Py) -> Self { + Self { namespace } + } + + pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { + let object = self.namespace.bind(py).getattr("cache")?; + let binding = if object.is_none() { + CacheBinding::Disabled + } else if let Ok(handle) = object.extract::>() { + CacheBinding::Native(handle.service()?) + } else if let Some(service) = facade::resolve(py, &object)? { + CacheBinding::Native(service) + } else { + CacheBinding::PythonCallback(PythonCallback::new(object.unbind())) + }; + Ok(ResolvedCache::new(binding)) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.namespace) + } +} From e647255909c00c7ebf54152c4c869b9f841498fa Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 17:36:03 +0000 Subject: [PATCH 29/56] fix(policy_engine): resolve policies once and apply fallback semantics in get_matching_policies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/policy_engine/policy_matcher.py | 21 ++++---- .../policy_engine/test_policy_matcher.py | 48 +++++++++++++++++++ 2 files changed, 60 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index 2b54b5dbe41..e0f558b5085 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -114,7 +114,7 @@ class PolicyMatcher: verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list") return [] - return registry.get_attached_policies(context) + return registry.get_attached_policies(context, PolicyMatcher.policy_applies(context)) @staticmethod def get_matching_policies_from_registry( @@ -137,14 +137,22 @@ class PolicyMatcher: policies: dict[str, Policy] | None = None, ) -> Callable[[str], bool]: """Predicate telling whether a policy exists and its condition matches the context.""" + resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies() return lambda policy_name: bool( PolicyMatcher.get_policies_with_matching_conditions( policy_names=(policy_name,), context=context, - policies=policies, + policies=resolved, ) ) + @staticmethod + def _registry_policies() -> dict[str, Policy]: + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + registry: Final = get_policy_registry() + return registry.get_all_policies() if registry.is_initialized() else {} + @staticmethod def get_policies_with_matching_conditions( policy_names: Sequence[str], @@ -167,17 +175,12 @@ class PolicyMatcher: List of policy names whose conditions match the context """ from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator - from litellm.proxy.policy_engine.policy_registry import get_policy_registry - if policies is None: - registry: Final = get_policy_registry() - if not registry.is_initialized(): - return [] - policies = registry.get_all_policies() + resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies() matching_policies: Final = [] for policy_name in policy_names: - policy = policies.get(policy_name) + policy = resolved.get(policy_name) if policy is None: continue # Policy matches if it has no condition OR condition evaluates to True diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py index 6143898ccbe..b07137893ec 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py @@ -8,8 +8,11 @@ Tests: import pytest +import litellm.proxy.policy_engine.attachment_registry as attachment_registry_module +import litellm.proxy.policy_engine.policy_registry as policy_registry_module from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry from litellm.types.proxy.policy_engine import ( PolicyMatchContext, PolicyScope, @@ -196,3 +199,48 @@ class TestPolicyMatcherWithAttachments: attached = registry.get_attached_policies(context) assert "healthcare-policy" not in attached + + +def _global_registries(monkeypatch): + policies = PolicyRegistry() + policies.load_policies( + { + "guardrail-y": {"guardrails": {"add": ["y"]}}, + "guardrail-x": {"guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}}, + } + ) + attachments = AttachmentRegistry() + attachments.load_attachments( + [ + {"policy": "guardrail-x", "tags": ["opt-in"]}, + {"policy": "guardrail-y", "scope": "*", "default": True}, + ] + ) + monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies) + monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments) + return policies + + +class TestGetMatchingPoliciesFallback: + def test_condition_failing_opt_in_falls_back_to_default(self, monkeypatch): + _global_registries(monkeypatch) + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"]) + + assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-y"] + + def test_condition_passing_opt_in_suppresses_default(self, monkeypatch): + _global_registries(monkeypatch) + context = PolicyMatchContext(team_alias="t", key_alias="k", model="claude-haiku", tags=["opt-in"]) + + assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-x"] + + def test_policy_applies_reads_registry_once(self, monkeypatch): + policies = _global_registries(monkeypatch) + calls = [] + original = policies.get_all_policies + monkeypatch.setattr(policies, "get_all_policies", lambda: calls.append(1) or original()) + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"]) + + PolicyMatcher.get_matching_policies(context=context) + + assert len(calls) == 1 From 3e2d03297fdbb8243a9c5a9e2cc3d8ef4b10abf9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 18:04:34 +0000 Subject: [PATCH 30/56] refactor(bedrock): build batch output config with explicit returns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/batches/transformation.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index dd6c1bd9302..e4001566b8c 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -61,15 +61,15 @@ def _build_s3_input_config(s3_uri: str, s3_bucket_owner: str | None) -> BedrockS def _build_s3_output_config( s3_uri: str, s3_bucket_owner: str | None, s3_encryption_key_id: str | None ) -> BedrockS3OutputDataConfig: - match (s3_bucket_owner, s3_encryption_key_id): - case (None, None): + if s3_bucket_owner is None: + if s3_encryption_key_id is None: return BedrockS3OutputDataConfig(s3Uri=s3_uri) - case (str() as owner, None): - return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=owner) - case (None, str() as key_id): - return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3EncryptionKeyId=key_id) - case (str() as owner, str() as key_id): - return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=owner, s3EncryptionKeyId=key_id) + return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3EncryptionKeyId=s3_encryption_key_id) + if s3_encryption_key_id is None: + return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner) + return BedrockS3OutputDataConfig( + s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner, s3EncryptionKeyId=s3_encryption_key_id + ) def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: From 1b6b704ddd639898895b1add55cdb9579a278c61 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:15:44 -0700 Subject: [PATCH 31/56] refactor(cache): keep native foundation isolated --- litellm-rust/crates/cache-response/README.md | 2 +- .../crates/python-bridge/python_settings.json | 3 - .../crates/python-bridge/src/cache/handle.rs | 28 +------- .../python-bridge/src/python_settings.rs | 5 +- litellm/caching/dual_cache.py | 22 ++---- litellm/rust_bridge/_native.pyi | 68 +------------------ litellm/rust_bridge/settings.py | 11 --- tests/test_litellm/caching/test_dual_cache.py | 49 +++++++------ tests/test_litellm_rust/test_cache.py | 14 +--- 9 files changed, 42 insertions(+), 160 deletions(-) diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 9863c46783f..56c1646d343 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -42,7 +42,7 @@ The resolver reads the namespace's `cache` attribute each time it resolves. A ca Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend -The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python +The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 15d0d603ac7..0af55083bef 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -22,8 +22,5 @@ ], "secret_manager": [ "readable" - ], - "cache_settings": [ - "default_redis_ttl" ] } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 42d7f2c2f3d..8251b3df06c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,26 +1,7 @@ -use std::time::Duration; - use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; -use crate::python_settings::PythonSettings; - -const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); - -#[derive(FromPyObject)] -struct PythonCacheSettings { - default_redis_ttl: Option, -} - -fn redis_default_ttl(py: Python<'_>) -> PyResult { - let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; - settings - .default_redis_ttl - .map(duration) - .transpose() - .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) -} #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -53,17 +34,14 @@ impl CacheTestHandle { } #[staticmethod] - #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + #[pyo3(signature = (url, *, ttl_seconds=60.0, namespace=None))] fn redis( py: Python<'_>, url: String, - ttl_seconds: Option, + ttl_seconds: f64, namespace: Option, ) -> PyResult { - let ttl = Some(match ttl_seconds { - Some(seconds) => duration(seconds)?, - None => redis_default_ttl(py)?, - }); + let ttl = Some(duration(ttl_seconds)?); let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) .map_err(cache_error)?; Ok(Self { diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 90819c5b3fc..7ac23a05542 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -8,17 +8,15 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, - Cache, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 5] = [ + pub(crate) const ALL: [Self; 4] = [ Self::Http, Self::UrlPolicy, Self::ProviderDefaults, Self::SecretManager, - Self::Cache, ]; pub(crate) fn name(self) -> &'static str { @@ -27,7 +25,6 @@ impl PythonSettings { Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", - Self::Cache => "cache_settings", } } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 04c82232784..66be77dbb40 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -12,7 +12,7 @@ import logging import time from collections.abc import Sequence from threading import Lock -from typing import TYPE_CHECKING, Any, Final, TypeVar +from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -34,20 +34,14 @@ else: from collections import OrderedDict -_KeyT = TypeVar("_KeyT") -_ValueT = TypeVar("_ValueT") - -class LimitedSizeOrderedDict(OrderedDict[_KeyT, _ValueT]): - def __init__(self, *, max_size: int = 100) -> None: - super().__init__() +class LimitedSizeOrderedDict(OrderedDict): + def __init__(self, *args, max_size=100, **kwargs): + super().__init__(*args, **kwargs) self.max_size = max_size - def __setitem__(self, key: _KeyT, value: _ValueT) -> None: - if key in self: - super().__setitem__(key, value) - self.move_to_end(key) - return + def __setitem__(self, key, value): + # If inserting a new key exceeds max size, remove the oldest item if len(self) >= self.max_size: self.popitem(last=False) super().__setitem__(key, value) @@ -74,9 +68,7 @@ class DualCache(BaseCache): self.in_memory_cache = in_memory_cache or InMemoryCache() # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache - self.last_redis_batch_access_time: LimitedSizeOrderedDict[str, float] = LimitedSizeOrderedDict( - max_size=default_max_redis_batch_cache_size - ) + self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size) self._last_redis_batch_access_time_lock = Lock() self.redis_batch_cache_expiry = ( default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index ab4639bc876..05a6df6d5af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,5 +1,5 @@ from asyncio import Future -from collections.abc import AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -93,72 +93,6 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... -@final -class _CacheTestHandle: - def __new__(cls, _uninstantiable: Never, /) -> Never: ... - @staticmethod - def memory( - *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576 - ) -> _CacheTestHandle: ... - @staticmethod - def redis(url: str, *, ttl_seconds: float | None = None, namespace: 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 _CacheTestBinding: - def __new__(cls, _uninstantiable: Never, /) -> Never: ... - @property - def kind(self) -> str: ... - def lookup( - self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None - ) -> object: ... - def store( - self, - request: Mapping[str, object] | None, - response: object, - *, - callback_kwargs: dict[str, object] | None = None, - ) -> None: ... - def lookup_batch( - self, - requests: Sequence[Mapping[str, object]], - *, - callback_kwargs: Sequence[dict[str, object]] | None = None, - ) -> object: ... - def async_lookup( - self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None - ) -> Awaitable[object]: ... - def async_store( - self, - request: Mapping[str, object] | None, - response: object, - *, - callback_kwargs: dict[str, object] | None = None, - ) -> Awaitable[object]: ... - def async_lookup_batch( - self, - requests: Sequence[Mapping[str, object]], - *, - callback_kwargs: Sequence[dict[str, object]] | None = None, - ) -> Awaitable[object]: ... - def async_store_batch( - self, - requests: Sequence[Mapping[str, object]], - responses: Sequence[object], - *, - callback_result: object = None, - callback_kwargs: dict[str, object] | None = None, - ) -> Awaitable[object]: ... - def async_flush(self) -> Awaitable[None]: ... - def ping(self) -> Awaitable[object]: ... - @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 862a116496d..3aa2d742862 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -36,11 +36,6 @@ class SecretManager: readable: bool -@dataclass(frozen=True, slots=True) -class CacheSettings: - default_redis_ttl: float | None - - def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -55,12 +50,6 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) -def cache_settings() -> CacheSettings: - import litellm - - return CacheSettings(default_redis_ttl=litellm.default_redis_ttl) - - def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index d34d23ca1d7..5f59de9cca5 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE -from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync from litellm.types.caching import RedisPipelineIncrementOperation @@ -15,7 +15,9 @@ from litellm.types.caching import RedisPipelineIncrementOperation @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads(): - dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 + ) keys = ["shared_a", "shared_b"] start_gate = asyncio.Event() @@ -42,7 +44,9 @@ async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_error(): - dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 + ) keys = ["shared_a", "shared_b"] with patch.object( @@ -112,7 +116,9 @@ def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) - dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) first = dual_cache.batch_get_cache(keys=["absent_key"]) second = dual_cache.batch_get_cache(keys=["absent_key"]) @@ -125,7 +131,9 @@ def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): mock_redis = MagicMock(spec=RedisCache) mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") - dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) first_result = dual_cache.batch_get_cache(keys=["shared_a"]) second_result = dual_cache.batch_get_cache(keys=["shared_a"]) @@ -138,7 +146,9 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) - dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() result = dual_cache.batch_get_cache(keys=["throttled_key"]) @@ -247,7 +257,9 @@ async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl(): default_in_memory_ttl, same as the single-key path.""" in_memory_cache = InMemoryCache(default_ttl=600) mock_redis = MagicMock(spec=RedisCache) - mock_redis.async_batch_get_cache = AsyncMock(return_value={"batch_backfill_key": "redis_value"}) + mock_redis.async_batch_get_cache = AsyncMock( + return_value={"batch_backfill_key": "redis_value"} + ) dual_cache = DualCache( in_memory_cache=in_memory_cache, redis_cache=mock_redis, @@ -359,7 +371,9 @@ async def test_circuit_breaker_open_skips_redis(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + self._circuit_breaker = RedisCircuitBreaker( + failure_threshold=3, recovery_timeout=60 + ) self._circuit_breaker._state = "open" self._circuit_breaker._opened_at = time.time() self.call_count = 0 @@ -412,7 +426,9 @@ def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): # All subsequent concurrent callers: HALF_OPEN → fast-fail (return True) for _ in range(10): - assert cb.is_open() is True, "concurrent callers should be fast-failed in HALF_OPEN" + assert ( + cb.is_open() is True + ), "concurrent callers should be fast-failed in HALF_OPEN" def test_circuit_breaker_disabled_never_opens(): @@ -456,7 +472,9 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) + self._circuit_breaker = RedisCircuitBreaker( + failure_threshold=1, recovery_timeout=60, enabled=False + ) self.call_count = 0 @_redis_circuit_breaker_guard @@ -773,14 +791,3 @@ async def test_async_delete_cache_keys_on_empty_list_touches_no_backend(): await dual_cache.async_delete_cache_keys([]) redis_cache.delete_cache_keys.assert_not_awaited() - - -def test_limited_ordered_dict_refreshes_recency_without_evicting_another_key(): - tracker = LimitedSizeOrderedDict(max_size=2) - tracker["hot"] = 1 - tracker["cold"] = 2 - - tracker["hot"] = 3 - tracker["new"] = 4 - - assert list(tracker.items()) == [("hot", 3), ("new", 4)] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index cf46a0566f7..796a0ec36ac 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -361,18 +361,6 @@ def test_facade_registration_rejects_mismatched_capacity() -> None: _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) -async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: - client: Final = redis.Redis.from_url(redis_url) - with rebound(litellm, "default_redis_ttl", 7): - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url)) - ).resolve() - await binding.async_store(request("native-default"), {"value": 1}) - - assert 0 < client.ttl("native-default") <= 7 - client.close() - - async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: parsed: Final = urlparse(redis_url) with rebound(litellm, "default_redis_ttl", 60): @@ -386,7 +374,7 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) with pytest.raises(TypeError, match="namespaces must match"): _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _native._CacheTestHandle.redis(redis_url)._bind_facade(facade) + _native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) From 22995d1575ec7b6e9ee2712efaf24ef45a59e07a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:23:47 -0700 Subject: [PATCH 32/56] fix(cache): remove redundant source comments --- litellm-rust/crates/cache-memory/src/cache.rs | 1 - litellm-rust/crates/cache-redis/src/cache.rs | 2 -- litellm-rust/crates/cache-response/src/buffer.rs | 1 - litellm-rust/crates/python-bridge/src/cache/binding.rs | 3 --- litellm-rust/crates/python-bridge/src/cache/callback.rs | 7 ------- 5 files changed, 14 deletions(-) diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 77635893640..45c638f5178 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -228,7 +228,6 @@ where .get(key) .filter(|existing| eligible.is_empty() || eligible.contains(existing)) .cloned(); - // Matches the Redis claim: an unconditional claim only extends its own winner. if let Some(existing) = &existing && eligible.is_empty() && *existing != candidate diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 23b1fabbab4..2249966dd79 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -57,8 +57,6 @@ const INCREMENT_SCRIPT: &str = concat!( "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" ); -// Compare-and-set against the exact bytes the claim decision was made on. -// ARGV: [1] expected payload or "" when absent, [2] ttl, [3] new payload, [4] refresh ttl. const CLAIM_SCRIPT: &str = concat!( "local current = redis.call('GET', KEYS[1]); ", "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 68af5278c8c..1fd2bb809de 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -5,7 +5,6 @@ use serde_json::Value; use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; -/// Defers async writes until `flush_size` entries are pending, then stores them as one batch. pub struct WriteBuffer { flush_size: usize, entries: Mutex>, diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 44c133d4611..ad64b24d3c1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -125,8 +125,6 @@ impl ResolvedCache { } } - /// Native bindings return `{values, missing_indices}`, while a Python callback returns the - /// list of its per-request results. #[pyo3(signature = (requests, *, callback_kwargs=None))] fn lookup_batch( &self, @@ -222,7 +220,6 @@ impl ResolvedCache { } } - /// A Python callback receives the caller's original result through `callback_result`. #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] fn async_store_batch<'py>( &self, diff --git a/litellm-rust/crates/python-bridge/src/cache/callback.rs b/litellm-rust/crates/python-bridge/src/cache/callback.rs index 318f9d02080..492e0329672 100644 --- a/litellm-rust/crates/python-bridge/src/cache/callback.rs +++ b/litellm-rust/crates/python-bridge/src/cache/callback.rs @@ -7,8 +7,6 @@ use pyo3::{ use super::future::ready_none; -/// A custom Python cache object, driven through the built-in `Cache` API so a `Cache` subclass -/// works unchanged. pub(super) struct PythonCallback(Py); impl PythonCallback { @@ -61,8 +59,6 @@ impl PythonCallback { ) } - /// The built-in `Cache` API has no batch read, so the callback receives one - /// `get_cache(**kwargs)` call per request, in order, and the results come back as a list. pub(super) fn lookup_batch<'py>( &self, py: Python<'py>, @@ -98,8 +94,6 @@ impl PythonCallback { .call_method1("gather", PyTuple::new(py, awaitables)?) } - /// Receives the caller's original result, because the built-in - /// `Cache.async_add_cache_pipeline` splits the batch itself. pub(super) fn async_store_batch<'py>( &self, py: Python<'py>, @@ -116,7 +110,6 @@ impl PythonCallback { ) } - /// The built-in `Cache` facade has no flush of its own; its backend does. pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { let object = self.0.bind(py); let backend = match object.getattr_opt("cache")? { From 5ea4fe620fdefe92d7674c7d3a985919fa3dbc37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:55:58 -0700 Subject: [PATCH 33/56] test(router): give each prompt caching check test a fresh callback registry --- .../test_prompt_caching_deployment_check.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 3a3ed2c45f4..a87b24656f3 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1,7 +1,7 @@ import asyncio import copy import functools -from typing import cast +from typing import Final, cast import pytest @@ -20,6 +20,23 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p MODEL_GROUP_ALIAS = "my-claude-group" OPUS_4_6_MIN_TOKENS = 4096 +CALLBACK_REGISTRIES: Final = ( + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + "callbacks", +) + + +@pytest.fixture(autouse=True) +def _fresh_callback_registries(monkeypatch): + """`litellm.logging_callback_manager` keeps one callback per class, so a + `PromptCachingDeploymentCheck` or `_SentMessagesCapture` left behind by an + earlier test would swallow the next test's success events.""" + for registry in CALLBACK_REGISTRIES: + monkeypatch.setattr(litellm, registry, []) @pytest.fixture From a7bc8e373e0d5b582003b2aafe62292139c61e12 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:56:04 -0700 Subject: [PATCH 34/56] feat(cache): project Python backend configuration --- litellm-rust/crates/cache-memory/src/cache.rs | 4 + .../crates/python-bridge/src/cache/config.rs | 586 ++++++++++++++++++ .../crates/python-bridge/src/cache/facade.rs | 45 +- .../crates/python-bridge/src/cache/mod.rs | 1 + .../crates/python-bridge/src/cache/native.rs | 7 + 5 files changed, 616 insertions(+), 27 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/config.rs diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 45c638f5178..54d831378e4 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -124,6 +124,10 @@ impl InMemoryCache { self.max_size_in_memory } + pub fn max_entry_bytes(&self) -> Option { + self.max_entry_bytes + } + pub fn expires_at(&self, key: &str) -> Result, Error> { Ok(self .state diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs new file mode 100644 index 00000000000..637bdab4055 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -0,0 +1,586 @@ +use std::time::Duration; + +use pyo3::{ + exceptions::{PyOverflowError, PyTypeError, PyValueError}, + prelude::*, + types::{PyAny, PyDict}, +}; + +use super::{native::NativeResponseCache, request::duration}; + +#[derive(PartialEq)] +pub(super) struct CachePolicy { + pub(super) mode: String, + pub(super) ttl: Option, + pub(super) namespace: Option, + pub(super) supported_call_types: Option>, + pub(super) redis_flush_size: Option, + pub(super) semantic_cache_scope: String, +} + +#[derive(PartialEq)] +pub(super) struct MemoryCacheConfig { + pub(super) default_ttl: Duration, + pub(super) capacity: usize, + pub(super) max_entry_bytes: usize, +} + +#[derive(Debug, PartialEq)] +pub(super) enum RedisProtocol { + Resp2, + Resp3, +} + +#[derive(Debug, PartialEq)] +pub(super) enum CertificateRequirement { + None, + Optional, + Required, +} + +#[derive(PartialEq)] +pub(super) struct RedisTlsConfig { + pub(super) certificate_requirement: CertificateRequirement, + pub(super) check_hostname: bool, + pub(super) ca_certificate: Option, + pub(super) ca_data: Option>, + pub(super) client_certificate: Option, + pub(super) client_key: Option, +} + +#[derive(PartialEq)] +pub(super) struct RedisConnectionConfig { + pub(super) host: String, + pub(super) port: u16, + pub(super) database: i64, + pub(super) username: Option, + pub(super) password: Option, + pub(super) protocol: RedisProtocol, + pub(super) pool_size: usize, + pub(super) read_timeout: Option, + pub(super) connect_timeout: Option, + pub(super) socket_keepalive: Option, + pub(super) health_check_interval: Duration, + pub(super) client_name: Option, + pub(super) tls: Option, +} + +#[derive(PartialEq)] +pub(super) struct RedisCacheConfig { + pub(super) default_ttl: Duration, + pub(super) namespace: Option, + pub(super) flush_size: usize, + pub(super) connection: RedisConnectionConfig, +} + +#[derive(PartialEq)] +pub(super) enum CacheBackendConfig { + Memory(MemoryCacheConfig), + Redis(Box), +} + +#[derive(PartialEq)] +pub(super) struct NativeCacheConfig { + pub(super) policy: CachePolicy, + pub(super) backend: CacheBackendConfig, +} + +pub(super) enum UnsupportedCacheConfig { + Backend(String), + RedisMode(&'static str), + RedisOption(String), +} + +impl UnsupportedCacheConfig { + pub(super) fn message(&self) -> String { + match self { + Self::Backend(backend) => { + format!("native cache backend {backend:?} is not implemented") + } + Self::RedisMode(mode) => format!("native Redis {mode} mode is not implemented"), + Self::RedisOption(option) => { + format!("native Redis option {option:?} is not implemented") + } + } + } +} + +pub(super) enum CacheConfigProjection { + Native(Box), + Unsupported(UnsupportedCacheConfig), +} + +impl NativeCacheConfig { + pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult { + let backend_name = facade.getattr("type")?.extract::()?; + let policy = CachePolicy { + mode: facade.getattr("mode")?.extract::()?, + ttl: optional_duration(facade.getattr("ttl")?)?, + namespace: optional_string(facade.getattr("namespace")?)?, + supported_call_types: facade + .getattr("supported_call_types")? + .extract::>>()?, + redis_flush_size: facade + .getattr("redis_flush_size")? + .extract::>()?, + semantic_cache_scope: facade + .getattr("semantic_cache_scope")? + .extract::()?, + }; + let backend = facade.getattr("cache")?; + match backend_name.as_str() { + "local" => project_memory(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Memory(backend), + })) + }), + "redis" => match project_redis(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Redis(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, + _ => Ok(CacheConfigProjection::Unsupported( + UnsupportedCacheConfig::Backend(backend_name), + )), + } + } + + pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { + if service.default_ttl() + != match &self.backend { + CacheBackendConfig::Memory(config) => config.default_ttl, + CacheBackendConfig::Redis(config) => config.default_ttl, + } + { + return Some("facade and native backend default TTLs must match"); + } + match &self.backend { + CacheBackendConfig::Memory(config) if service.kind() != "memory" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => { + Some("facade and native backend capacities must match") + } + CacheBackendConfig::Memory(config) + if service.max_entry_bytes() != Some(config.max_entry_bytes) => + { + Some("facade and native backend item limits must match") + } + CacheBackendConfig::Memory(_) => None, + CacheBackendConfig::Redis(_) if service.kind() != "redis" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Redis(config) => (service.namespace() + != config.namespace.as_deref()) + .then_some("facade and native backend namespaces must match"), + } + } +} + +fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { + let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; + Ok(MemoryCacheConfig { + default_ttl: duration(backend.getattr("default_ttl")?.extract::()?)?, + capacity: backend.getattr("max_size_in_memory")?.extract::()?, + max_entry_bytes: max_size_kib + .checked_mul(1024) + .ok_or_else(|| PyOverflowError::new_err("memory cache item limit is too large"))?, + }) +} + +fn project_redis( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let source = backend.getattr("redis_kwargs")?.cast_into::()?; + if has_value(&source, "startup_nodes")? { + return Ok(Err(UnsupportedCacheConfig::RedisMode("cluster"))); + } + if has_value(&source, "sentinel_nodes")? { + return Ok(Err(UnsupportedCacheConfig::RedisMode("sentinel"))); + } + for key in [ + "credential_provider", + "redis_connect_func", + "connection_pool", + ] { + if has_value(&source, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + for key in [ + "retry", + "retry_on_error", + "socket_keepalive_options", + "unix_socket_path", + "cache", + "cache_config", + "event_dispatcher", + "ssl_ca_path", + "ssl_password", + "ssl_min_version", + "ssl_ciphers", + "ssl_validate_ocsp", + "ssl_validate_ocsp_stapled", + "ssl_ocsp_context", + "ssl_ocsp_expected_cert", + ] { + if has_value(&source, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + for key in ["retry_on_timeout", "single_connection_client"] { + if optional_coerced_bool(&source, key)?.unwrap_or(false) { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + + let client = backend.getattr("redis_client")?; + let pool = client.getattr("connection_pool")?; + let pool_class = class_identity(&pool)?; + if pool_class != ("redis.connection".to_owned(), "ConnectionPool".to_owned()) { + return Ok(Err(UnsupportedCacheConfig::RedisMode("custom pool"))); + } + 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::RedisOption(key.to_owned()))); + } + } + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + let connection_class = ( + connection_class + .getattr("__module__")? + .extract::()?, + connection_class + .getattr("__qualname__")? + .extract::()?, + ); + let tls = match connection_class { + (module, name) if module == "redis.connection" && name == "Connection" => None, + (module, name) if module == "redis.connection" && name == "SSLConnection" => { + Some(project_tls(&resolved)?) + } + _ => return Ok(Err(UnsupportedCacheConfig::RedisMode("custom connection"))), + }; + + let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { + 2 => RedisProtocol::Resp2, + 3 => RedisProtocol::Resp3, + value => { + return Err(PyValueError::new_err(format!( + "unsupported Redis protocol version {value}" + ))); + } + }; + let health_check_interval = + duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?; + Ok(Ok(RedisCacheConfig { + default_ttl: duration(backend.getattr("default_ttl")?.extract::()?)?, + namespace: optional_attribute_string(backend, "namespace")?, + flush_size: backend.getattr("redis_flush_size")?.extract::()?, + connection: RedisConnectionConfig { + host: required_string(&resolved, "host")?, + port: required_u16(&resolved, "port")?, + database: optional_i64(&resolved, "db")?.unwrap_or(0), + username: optional_dict_string(&resolved, "username")?, + password: optional_dict_string(&resolved, "password")?, + protocol, + pool_size: pool.getattr("max_connections")?.extract::()?, + read_timeout: optional_dict_duration(&resolved, "socket_timeout")?, + connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?, + socket_keepalive: optional_bool(&resolved, "socket_keepalive")?, + health_check_interval, + client_name: optional_dict_string(&resolved, "client_name")?, + tls, + }, + })) +} + +fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { + Ok(RedisTlsConfig { + certificate_requirement: certificate_requirement(values)?, + check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false), + ca_certificate: optional_dict_string(values, "ssl_ca_certs")?, + ca_data: optional_bytes(values, "ssl_ca_data")?, + client_certificate: optional_dict_string(values, "ssl_certfile")?, + client_key: optional_dict_string(values, "ssl_keyfile")?, + }) +} + +fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult { + let Some(value) = values.get_item("ssl_cert_reqs")? else { + return Ok(CertificateRequirement::Required); + }; + if value.is_none() { + return Ok(CertificateRequirement::Required); + } + if let Ok(number) = value.extract::() { + return match number { + 0 => Ok(CertificateRequirement::None), + 1 => Ok(CertificateRequirement::Optional), + 2 => Ok(CertificateRequirement::Required), + _ => Err(PyValueError::new_err( + "invalid Redis TLS certificate requirement", + )), + }; + } + match value.str()?.to_str()?.to_ascii_lowercase().as_str() { + "none" | "cert_none" => Ok(CertificateRequirement::None), + "optional" | "cert_optional" => Ok(CertificateRequirement::Optional), + "required" | "cert_required" => Ok(CertificateRequirement::Required), + _ => Err(PyValueError::new_err( + "invalid Redis TLS certificate requirement", + )), + } +} + +fn class_identity(value: &Bound<'_, PyAny>) -> PyResult<(String, String)> { + let class = value.get_type(); + Ok(( + class.getattr("__module__")?.extract::()?, + class.getattr("__qualname__")?.extract::()?, + )) +} + +fn optional_duration(value: Bound<'_, PyAny>) -> PyResult> { + value.extract::>()?.map(duration).transpose() +} + +fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + match value.getattr(name) { + Ok(value) => optional_string(value), + Err(error) if error.is_instance_of::(value.py()) => { + Ok(None) + } + Err(error) => Err(error), + } +} + +fn optional_string(value: Bound<'_, PyAny>) -> PyResult> { + Ok(value + .extract::>()? + .filter(|value| !value.is_empty())) +} + +fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult { + Ok(values.get_item(key)?.is_some_and(|value| !value.is_none())) +} + +fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { + values + .get_item(key)? + .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .extract::() +} + +fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { + values + .get_item(key)? + .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .extract::() +} + +fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) if !value.is_none() => optional_string(value), + _ => Ok(None), + } +} + +fn optional_bytes(values: &Bound<'_, PyDict>, key: &str) -> PyResult>> { + let Some(value) = values.get_item(key)? else { + return Ok(None); + }; + if value.is_none() { + return Ok(None); + } + if let Ok(bytes) = value.extract::>() { + return Ok(Some(bytes)); + } + Ok(Some(value.extract::()?.into_bytes())) +} + +fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + let Some(value) = values.get_item(key)? else { + return Ok(None); + }; + if value.is_none() { + return Ok(None); + } + if let Ok(text) = value.extract::() { + return Ok(Some(matches!( + text.to_ascii_lowercase().as_str(), + "true" | "1" | "yes" + ))); + } + value.extract::().map(Some) +} + +fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + optional_f64(values, key)?.map(duration).transpose() +} + +#[cfg(test)] +mod tests { + use std::ffi::CString; + + use pyo3::{prelude::*, types::PyDict}; + + use super::{ + CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, + RedisProtocol, + }; + use crate::cache::native::NativeResponseCache; + + fn 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\ + ConnectionPool = type('ConnectionPool', (), {{'__module__': 'redis.connection'}})\n\ + Connection = type('Connection', (), {{'__module__': 'redis.connection'}})\n\ + SSLConnection = type('SSLConnection', (), {{'__module__': 'redis.connection'}})\n\ + {body}" + )) + .unwrap(), + None, + Some(&locals), + ) + .unwrap(); + locals.get_item("facade").unwrap().unwrap() + } + + #[test] + fn projects_effective_memory_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(default_ttl=913, max_size_in_memory=37, max_size_per_item=8)\n\ + facade = SimpleNamespace(type='local', mode='default-on', ttl=11.5, namespace=None, supported_call_types=['completion'], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("memory cache should be supported"); + }; + assert_eq!( + config.policy.ttl.unwrap(), + std::time::Duration::from_secs_f64(11.5) + ); + let CacheBackendConfig::Memory(memory) = config.backend else { + panic!("expected memory configuration"); + }; + assert_eq!(memory.default_ttl, std::time::Duration::from_secs(913)); + assert_eq!(memory.capacity, 37); + assert_eq!(memory.max_entry_bytes, 8192); + let matching = + NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8192); + let mismatched = + NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8191); + let matching_config = NativeCacheConfig { + policy: config.policy, + backend: CacheBackendConfig::Memory(memory), + }; + assert_eq!(matching_config.service_mismatch(&matching), None); + assert_eq!( + matching_config.service_mismatch(&mismatched), + Some("facade and native backend item limits must match") + ); + }); + } + + #[test] + fn projects_resolved_redis_tls_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = SSLConnection\n\ + pool.max_connections = 29\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\ + facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Redis cache should be supported"); + }; + let CacheBackendConfig::Redis(redis) = config.backend else { + panic!("expected Redis configuration"); + }; + assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777)); + assert_eq!(redis.namespace.as_deref(), Some("team")); + assert_eq!(redis.flush_size, 31); + assert_eq!(redis.connection.host, "cache.internal"); + assert_eq!(redis.connection.port, 6380); + assert_eq!(redis.connection.database, 4); + assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); + assert_eq!(redis.connection.pool_size, 29); + let tls = redis.connection.tls.unwrap(); + assert_eq!( + tls.certificate_requirement, + CertificateRequirement::Optional + ); + assert!(tls.check_hostname); + assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem")); + assert_eq!(tls.ca_data.as_deref(), Some(b"CA DATA".as_slice())); + assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); + assert_eq!(tls.client_key.as_deref(), Some("/client.key")); + }); + } + + #[test] + fn dynamic_redis_auth_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\ + facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("dynamic authentication must stay on Python"); + }; + assert!(reason.message().contains("credential_provider")); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 19220bf868f..6f54d6a121a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,5 +1,3 @@ -use std::time::Duration; - use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -9,7 +7,11 @@ use pyo3::{ }; use serde_json::Value; -use super::{handle::CacheTestHandle, native::NativeResponseCache}; +use super::{ + config::{CacheConfigProjection, NativeCacheConfig}, + handle::CacheTestHandle, + native::NativeResponseCache, +}; struct ClassGuard { class: Py, @@ -26,6 +28,7 @@ struct ObjectGuard { pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, + config: NativeCacheConfig, } impl ObjectGuard { @@ -134,7 +137,6 @@ impl FacadeGuard { service: &NativeResponseCache, ) -> PyResult { let kind = service.kind(); - let native_default_ttl: Duration = service.default_ttl(); let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -154,28 +156,14 @@ impl FacadeGuard { "facade and native backend types must match", )); } - let python_default_ttl = backend.getattr("default_ttl")?.extract::()?; - if python_default_ttl != native_default_ttl.as_secs_f64() { - return Err(PyTypeError::new_err( - "facade and native backend default TTLs must match", - )); - } - let namespace = match backend.getattr_opt("namespace")? { - Some(namespace) => namespace.extract::>()?, - None => None, - } - .filter(|namespace| !namespace.is_empty()); - if kind == "redis" && namespace.as_deref() != service.namespace() { - return Err(PyTypeError::new_err( - "facade and native backend namespaces must match", - )); - } - if let Some(capacity) = service.capacity() - && backend.getattr("max_size_in_memory")?.extract::()? != capacity - { - return Err(PyTypeError::new_err( - "facade and native backend capacities must match", - )); + let config = match NativeCacheConfig::project(facade)? { + CacheConfigProjection::Native(config) => *config, + CacheConfigProjection::Unsupported(reason) => { + return Err(PyTypeError::new_err(reason.message())); + } + }; + if let Some(message) = config.service_mismatch(service) { + return Err(PyTypeError::new_err(message)); } Ok(Self { outer: ObjectGuard::capture( @@ -203,12 +191,15 @@ impl FacadeGuard { "redis_flush_size", ], )?, + config, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { + let projected = NativeCacheConfig::project(facade)?; Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)?) + && self.backend.matches(py, &facade.getattr("cache")?)? + && matches!(projected, CacheConfigProjection::Native(config) if *config == self.config)) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 7955ed934b2..aec08610f6e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,5 +1,6 @@ mod binding; mod callback; +mod config; mod facade; mod future; mod handle; diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a718d07b286..6a3835ac84d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -74,6 +74,13 @@ impl NativeResponseCache { } } + pub fn max_entry_bytes(&self) -> Option { + match self { + Self::Memory(cache) => cache.backend().max_entry_bytes(), + Self::Redis { .. } => None, + } + } + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { match self { Self::Redis { cache, .. } => Self::Redis { From 2fe5c8990ef744add2e74c2d73dce985df959898 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:00:38 -0700 Subject: [PATCH 35/56] fix(bedrock_mantle): bill Mantle's un-versioned Claude ids from a Mantle cost row Mantle serves anthropic.claude-haiku-4-5 without the dated -20251001-v1:0 suffix the Bedrock row carries, so the native route billed it at 0. Add a bedrock_mantle/anthropic.claude-haiku-4-5 row and let a bedrock_mantle// name fall back to the region-free bedrock_mantle/ row before the provider-prefixed lookup. Also satisfy the mutable-collection gate in the native messages transformation. --- .../bedrock_mantle/messages/transformation.py | 11 ++++++-- ...odel_prices_and_context_window_backup.json | 28 +++++++++++++++++++ litellm/utils.py | 7 ++++- model_prices_and_context_window.json | 28 +++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 27 ++++++++++++++++++ tests/test_litellm/test_utils.py | 15 ++++++++++ 6 files changed, 112 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py index 480fe82ef4c..6e975d072ed 100644 --- a/litellm/llms/bedrock_mantle/messages/transformation.py +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -33,7 +33,7 @@ _MANTLE_REQUEST: Final = TypeAdapter(dict[str, object]) def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str: - region: Final = resolve_mantle_region({**litellm_params, "api_base": api_base}) + region: Final = resolve_mantle_region(MappingProxyType({**litellm_params, "api_base": api_base})) configured: Final = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") or f"https://bedrock-mantle.{region}.api.aws" ).rstrip("/") @@ -96,7 +96,10 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM ) if any(name.lower() == "anthropic-version" for name in merged_headers): return merged_headers, resolved_api_base - return {**merged_headers, "anthropic-version": DEFAULT_ANTHROPIC_API_VERSION}, resolved_api_base + return { # mutable-ok: the base class contract returns a dict the handler signs into in place + **merged_headers, + "anthropic-version": DEFAULT_ANTHROPIC_API_VERSION, + }, resolved_api_base def transform_anthropic_messages_request( self, @@ -119,4 +122,6 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM if betas is not None: header_betas: Final = ",".join(_ANTHROPIC_BETAS.validate_python(betas)) headers["anthropic-beta"] = header_betas # rebind-ok: the handler signs and sends this same dict - return {key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS} + return { # mutable-ok: the base class contract returns the dict the handler serializes as the body + key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS + } diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5755c1e7f9b..b24d5ba6916 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59210,6 +59210,34 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/anthropic.claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "bedrock_mantle", + "supports_tool_search": true, + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://aws.amazon.com/about-aws/whats-new/2025/10/claude-4-5-haiku-anthropic-amazon-bedrock", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 + }, "us.xai.grok-4.6": { "input_cost_per_token": 2.2e-06, "output_cost_per_token": 6.6e-06, diff --git a/litellm/utils.py b/litellm/utils.py index 20b8461066c..bd6d6a336da 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5665,6 +5665,11 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P region_free_split_model: Final = ( _strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model ) + region_free_combined_stripped_model_name: Final = ( + f"bedrock_mantle/{_strip_model_name(model=region_free_split_model, custom_llm_provider=custom_llm_provider)}" + if custom_llm_provider == "bedrock_mantle" + else combined_stripped_model_name + ) provider_model_info: Final = ( ProviderConfigManager.get_provider_model_info( model=region_free_split_model, provider=LlmProviders(custom_llm_provider) @@ -5680,7 +5685,7 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P split_model=region_free_split_model, combined_model_name=combined_model_name, stripped_model_name=stripped_model_name, - combined_stripped_model_name=combined_stripped_model_name, + combined_stripped_model_name=region_free_combined_stripped_model_name, provider_prefixed_model_name=provider_cost_key or provider_prefixed_model_name, custom_llm_provider=cast(str, custom_llm_provider), ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5755c1e7f9b..b24d5ba6916 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59210,6 +59210,34 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/anthropic.claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "bedrock_mantle", + "supports_tool_search": true, + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://aws.amazon.com/about-aws/whats-new/2025/10/claude-4-5-haiku-anthropic-amazon-bedrock", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 + }, "us.xai.grok-4.6": { "input_cost_per_token": 2.2e-06, "output_cost_per_token": 6.6e-06, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 2d5798d1561..1d6c229f9ce 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3547,6 +3547,33 @@ def test_completion_cost_mantle_native_messages_prices_claude_from_the_bedrock_r ) == pytest.approx(expected) +def test_completion_cost_mantle_native_messages_prices_haiku_from_the_mantle_row(_local_model_cost_map): + """Mantle serves Anthropic's un-versioned haiku id, which has no bare Bedrock row (Bedrock's carries + the -20251001-v1:0 suffix), and Claude Code sends every small-fast-model call to it. Both the plain + and the region-prefixed deployment names must price from bedrock_mantle/anthropic.claude-haiku-4-5 + instead of billing $0.""" + + response = litellm.ModelResponse( + id="msg_x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="claude-haiku-4-5", + usage={"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110}, + ) + row = litellm.model_cost["bedrock_mantle/anthropic.claude-haiku-4-5"] + expected = 100 * row["input_cost_per_token"] + 10 * row["output_cost_per_token"] + assert expected > 0 + + for model in ( + "bedrock_mantle/anthropic.claude-haiku-4-5", + "bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5", + ): + assert litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock_mantle", + ) == pytest.approx(expected), model + + def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map): """An explicit base_model keeps pricing on that model's own key even when the request carries a region with different regional rates, so the private provider model never widens region pricing.""" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f24fc284cd6..2ccb88b29db 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1163,6 +1163,21 @@ def test_get_model_info_bedrock_regional_inference_profile_pricing(local_model_c assert control["key"] == "au.anthropic.claude-opus-4-8" +def test_get_model_info_bedrock_mantle_region_prefix_falls_back_to_the_mantle_row(local_model_cost_map): + """A Mantle deployment name may carry the region as a prefix (bedrock_mantle/us-east-2/). + That name has no cost row of its own, so pricing must fall through to the region-free + bedrock_mantle/ row instead of raising, while a region that has its own row keeps it.""" + for model, expected_key in ( + ("bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5", "bedrock_mantle/anthropic.claude-haiku-4-5"), + ("bedrock_mantle/us-east-2/openai.gpt-5.6-sol", "bedrock_mantle/openai.gpt-5.6-sol"), + ("bedrock_mantle/us-gov-west-1/openai.gpt-5.4", "bedrock_mantle/us-gov-west-1/openai.gpt-5.4"), + ): + info = litellm.get_model_info(model=model, custom_llm_provider="bedrock_mantle") + assert info["key"] == expected_key, model + assert info["input_cost_per_token"] == litellm.model_cost[expected_key]["input_cost_per_token"], model + assert info["input_cost_per_token"] > 0, model + + def test_openai_models_in_model_info(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") From 5db2a97829fde8c04f019edaf1b96a9c54da4b13 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:01:26 -0700 Subject: [PATCH 36/56] chore: keep main's lazy OpenAPI snapshot The snapshot check runs on Python 3.12, which keeps the indentation of a route docstring that Python 3.13+ strips at compile time, so regenerating it locally on 3.14 produces a file CI rejects. --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 391f0042ed0..06e157498aa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19632,7 +19632,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 8d73ce756a10a298c0aaf26310e3e92a3a035ca9 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:01:54 +0000 Subject: [PATCH 37/56] feat(fal_ai): add MiniMax H3 text-to-video and reference-to-video Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 88 ++++++++++++++----- ...odel_prices_and_context_window_backup.json | 39 ++++++++ litellm/types/router.py | 2 + litellm/types/utils.py | 4 + litellm/utils.py | 2 + model_prices_and_context_window.json | 39 ++++++++ model_prices_and_context_window.schema.json | 8 ++ tests/integration/contracts.json | 3 + .../providers/test_fal_ai_video_wire.py | 54 ++++++++++++ .../test_fal_ai_video_transformation.py | 63 ++++++++++++- tests/test_litellm/test_utils.py | 4 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++ 12 files changed, 289 insertions(+), 25 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 98528c82f6b..8a355b3d226 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -1,6 +1,8 @@ import math +import sys import time from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeAlias @@ -36,11 +38,33 @@ class FalAIVideoError(BaseLLMException): _ALLOWED_ASPECT_RATIOS: Final[frozenset[str]] = frozenset({"auto", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}) -_ALLOWED_RESOLUTIONS: Final[frozenset[str]] = frozenset({"480p", "720p", "1080p", "4k"}) -_RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( - (480, "480p"), - (720, "720p"), - (1080, "1080p"), + + +@dataclass(frozen=True, slots=True) +class _ModelProfile: + resolutions: frozenset[str] + resolution_tiers: tuple[tuple[int, str], ...] + default_resolution: str + integer_duration: bool + reference_key: str + reference_as_list: bool + + +_SEEDANCE_PROFILE: Final[_ModelProfile] = _ModelProfile( + resolutions=frozenset({"480p", "720p", "1080p", "4k"}), + resolution_tiers=((480, "480p"), (720, "720p"), (1080, "1080p"), (sys.maxsize, "4k")), + default_resolution="720p", + integer_duration=False, + reference_key="image_url", + reference_as_list=False, +) +_H3_PROFILE: Final[_ModelProfile] = _ModelProfile( + resolutions=frozenset({"480P", "768P", "2K", "4K"}), + resolution_tiers=((480, "480P"), (768, "768P"), (1440, "2K"), (sys.maxsize, "4K")), + default_resolution="2K", + integer_duration=True, + reference_key="reference_image_urls", + reference_as_list=True, ) _QUEUE_NAMESPACES: Final[frozenset[str]] = frozenset(("workflows", "comfy")) _STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType( @@ -75,8 +99,12 @@ def _duration_value(value: object) -> str | None: return None -def _resolution_for_short_side(short_side: int) -> str: - return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") +def _profile_for_model(model: str) -> _ModelProfile: + return _H3_PROFILE if model.startswith("minimax/h3/") else _SEEDANCE_PROFILE + + +def _resolution_for_short_side(short_side: int, profile: _ModelProfile) -> str: + return next(resolution for threshold, resolution in profile.resolution_tiers if short_side <= threshold) def _model_path_from_request_url(raw_response: httpx.Response) -> str | None: @@ -97,14 +125,19 @@ def _request_id_from_request_url(raw_response: httpx.Response) -> str | None: return segments[request_id_index] if len(segments) > request_id_index else None -def _size_params(size: object) -> Mapping[str, str]: +def _size_params(size: object, profile: _ModelProfile) -> Mapping[str, str]: if not isinstance(size, str): return MappingProxyType({}) - if size in _ALLOWED_RESOLUTIONS: - return MappingProxyType({"resolution": size}) - if size.count("x") != 1: + normalized_size: Final[str] = size.lower() + canonical_resolution: Final[str | None] = next( + (resolution for resolution in profile.resolutions if resolution.lower() == normalized_size), + None, + ) + if canonical_resolution is not None: + return MappingProxyType({"resolution": canonical_resolution}) + if normalized_size.count("x") != 1: return MappingProxyType({}) - width_text, height_text = size.split("x") + width_text, height_text = normalized_size.split("x") if not (width_text.isdigit() and height_text.isdigit()): return MappingProxyType({}) width: Final[int] = int(width_text) @@ -113,7 +146,7 @@ def _size_params(size: object) -> Mapping[str, str]: return MappingProxyType({}) reduced_gcd: Final[int] = math.gcd(width, height) aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" - resolution: Final[str] = _resolution_for_short_side(min(width, height)) + resolution: Final[str] = _resolution_for_short_side(min(width, height), profile) if aspect_ratio in _ALLOWED_ASPECT_RATIOS: return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) return MappingProxyType({"resolution": resolution}) @@ -158,18 +191,27 @@ class FalAIVideoConfig(BaseVideoConfig): input_reference: Final[object] = video_create_optional_params.get("input_reference") if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): raise ValueError("fal.ai needs a public image URL for input_reference") - input_reference_params: Final[Mapping[str, str]] = ( + profile: Final[_ModelProfile] = _profile_for_model(model) + input_reference_params: Final[Mapping[str, object]] = ( MappingProxyType({}) if not isinstance(input_reference, str) - else MappingProxyType({"image_url": input_reference}) + else MappingProxyType( + { + profile.reference_key: ( + [input_reference] # mutable-ok: fal.ai expects a list for H3 references + if profile.reference_as_list + else input_reference + ), + } + ) ) - duration_params: Final[Mapping[str, str]] = ( + duration_params: Final[Mapping[str, object]] = ( MappingProxyType({}) if "seconds" not in video_create_optional_params - else self._duration_params(video_create_optional_params["seconds"]) + else self._duration_params(video_create_optional_params["seconds"], profile) ) size_params: Final[Mapping[str, str]] = ( - _size_params(video_create_optional_params["size"]) + _size_params(video_create_optional_params["size"], profile) if "size" in video_create_optional_params else MappingProxyType({}) ) @@ -190,11 +232,11 @@ class FalAIVideoConfig(BaseVideoConfig): return mapped_params @staticmethod - def _duration_params(seconds: object) -> Mapping[str, str]: + def _duration_params(seconds: object, profile: _ModelProfile) -> Mapping[str, object]: duration: Final[str | None] = _duration_value(seconds) if duration is None: raise ValueError("fal.ai seconds must be a numeric value") - return MappingProxyType({"duration": duration}) + return MappingProxyType({"duration": int(duration) if profile.integer_duration else duration}) def validate_environment( self, @@ -251,6 +293,7 @@ class FalAIVideoConfig(BaseVideoConfig): request_data: Mapping[str, object] | None = None, ) -> VideoObject: response_data: Final[Mapping[str, object]] = _response_data(raw_response) + profile: Final[_ModelProfile] = _profile_for_model(model) request_params: Final[Mapping[str, object]] = request_data or MappingProxyType({}) request_id: Final[str] = _response_string(response_data, "request_id") provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER @@ -262,7 +305,10 @@ class FalAIVideoConfig(BaseVideoConfig): key: value for key, value in ( ("duration_seconds", duration), - ("video_resolution", resolution if isinstance(resolution, str) else "720p"), + ( + "video_resolution", + resolution if isinstance(resolution, str) else profile.default_resolution, + ), ) if value is not None } diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 87edd1544ca..e39ede803ed 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22889,6 +22889,45 @@ "video" ] }, + "fal_ai/minimax/h3/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/minimax/h3/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/bytedance/seedance-2.0/text-to-video": { "litellm_provider": "fal_ai", "mode": "video_generation", diff --git a/litellm/types/router.py b/litellm/types/router.py index a75b4654cab..d5cbb3df409 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -539,6 +539,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): output_cost_per_second: float | None output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_720p: ReadOnly[float | None] + output_cost_per_second_768p: ReadOnly[float | None] + output_cost_per_second_2k: ReadOnly[float | None] output_cost_per_second_1080p: float | None output_cost_per_second_4k: ReadOnly[float | None] num_retries: int | None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5a80644347e..ae2df17b94a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -329,6 +329,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_720p: ReadOnly[float | None] + output_cost_per_second_768p: ReadOnly[float | None] + output_cost_per_second_2k: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_page_batches: ReadOnly[float | None] @@ -3610,6 +3612,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_second_1080p: float | None = None output_cost_per_second_480p: float | None = None output_cost_per_second_720p: float | None = None + output_cost_per_second_768p: float | None = None + output_cost_per_second_2k: float | None = None output_cost_per_second_4k: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 9a80b115d4b..3b922500961 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6084,6 +6084,8 @@ def _get_model_info_helper( output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None), + output_cost_per_second_768p=_model_info.get("output_cost_per_second_768p", None), + output_cost_per_second_2k=_model_info.get("output_cost_per_second_2k", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 87edd1544ca..e39ede803ed 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22889,6 +22889,45 @@ "video" ] }, + "fal_ai/minimax/h3/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/minimax/h3/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.13, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_768p": 0.06, + "output_cost_per_second_2k": 0.13, + "output_cost_per_second_4k": 0.16, + "source": "https://fal.ai/models/minimax/h3/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/bytedance/seedance-2.0/text-to-video": { "litellm_provider": "fal_ai", "mode": "video_generation", diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 5b0a23adfea..0509516ac32 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -607,6 +607,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_2k": { + "type": "number", + "minimum": 0 + }, "output_cost_per_second_480p": { "type": "number", "minimum": 0 @@ -619,6 +623,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_768p": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token": { "type": "number", "minimum": 0, diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index cd7e84f81b6..365456c0cec 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -175,6 +175,9 @@ "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row": [ "other.provider_wire.fal_ai.image_edit_json_data_urls_and_keyed_pricing" ], + "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_h3_video_create_uses_canonical_body_and_status_path": [ + "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" + ], "tests/integration/mcp/test_mcp_lifecycle.py::test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit": [ "mcp.call_tool.saved_headers.reach_actual_transport" ], diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py index 8c72810ffb6..ceb53c77c83 100644 --- a/tests/integration/providers/test_fal_ai_video_wire.py +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -7,6 +7,7 @@ from integration._support.client import Gateway from integration._support.wire import Reply, Request, wire_server _MODEL: Final = "bytedance/seedance-2.5/text-to-video" +_H3_MODEL: Final = "minimax/h3/text-to-video" _MP4: Final = b"\x00\x00\x00\x18ftypmp42" + uuid.uuid4().bytes * 4 @@ -67,3 +68,56 @@ def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: ("GET", f"/bytedance/seedance-2.5/requests/{request_id}"), ("GET", f"/files/{request_id}.mp4"), ] + + +@pytest.mark.covers("other.provider_wire.fal_ai.video_queue_create_status_and_content_download") +def test_fal_h3_video_create_uses_canonical_body_and_status_path(gateway: Gateway) -> None: + request_id: Final = "fal-h3-req-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_H3_MODEL}" + assert json.loads(request.body) == { + "prompt": "a cat playing volleyball on a beach", + "duration": 6, + "resolution": "2K", + } + return Reply( + body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode() + ) + assert request.method == "GET" + if request.target == f"/minimax/h3/requests/{request_id}/status": + return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode()) + assert request.target == f"/minimax/h3/requests/{request_id}" + return Reply(body=json.dumps({"video": {"url": f"{wire_url}/files/{request_id}.mp4"}}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model=f"fal_ai/{_H3_MODEL}", + api_base=wire.url, + api_key="synthetic-fal-key", + ) + created: Final = gateway.post( + "/v1/videos", + { + "model": model, + "prompt": "a cat playing volleyball on a beach", + "seconds": 6, + "size": "2k", + }, + ) + assert created["status"] == "queued" + video_id: Final = created["id"] + status: Final = gateway.get(f"/v1/videos/{video_id}") + assert status["status"] == "completed" + content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.content == _MP4 + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"/{_H3_MODEL}"), + ("GET", f"/minimax/h3/requests/{request_id}/status"), + ("GET", f"/minimax/h3/requests/{request_id}"), + ("GET", f"/files/{request_id}.mp4"), + ] diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index 5e2e4532265..963ed7eac47 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -11,12 +11,15 @@ from litellm.llms.fal_ai.videos.transformation import ( FalAIVideoError, _queue_request_base_path, ) +from litellm.llms.openai.cost_calculation import video_generation_cost from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.types.videos.utils import decode_video_id_with_provider from litellm.utils import ProviderConfigManager MODEL = "bytedance/seedance-2.5/text-to-video" +H3_TEXT_MODEL = "minimax/h3/text-to-video" +H3_REFERENCE_MODEL = "minimax/h3/reference-to-video" class TestFalAIVideoTransformation: @@ -64,6 +67,24 @@ class TestFalAIVideoTransformation: with pytest.raises(ValueError, match="public image URL"): self.config.map_openai_params({"input_reference": b"image"}, MODEL, False) + def test_map_openai_params_supports_h3_profiles(self): + url = "https://example.com/image.png" + + assert self.config.map_openai_params({"size": "2k"}, H3_TEXT_MODEL, False) == {"resolution": "2K"} + assert self.config.map_openai_params({"size": "1024x768"}, H3_TEXT_MODEL, False) == { + "resolution": "768P", + "aspect_ratio": "4:3", + } + mapped = self.config.map_openai_params( + {"seconds": 6, "input_reference": url}, + H3_REFERENCE_MODEL, + False, + ) + assert mapped["duration"] == 6 + assert isinstance(mapped["duration"], int) + assert mapped["reference_image_urls"] == [url] + assert "image_url" not in mapped + def test_transform_video_create_request(self): body, files, url = self.config.transform_video_create_request( model=MODEL, @@ -141,6 +162,20 @@ class TestFalAIVideoTransformation: assert auto_video.seconds is None assert auto_video.size is None + def test_transform_video_create_response_uses_h3_default_resolution(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + + video = self.config.transform_video_create_response( + model=H3_TEXT_MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": 5}, + ) + + assert video.usage == {"duration_seconds": 5.0, "video_resolution": "2K"} + def test_status_request_uses_queue_base_path(self): response = Mock(spec=httpx.Response) response.json.return_value = {"request_id": "abc"} @@ -290,9 +325,29 @@ class TestFalAIVideoTransformation: } assert rows for model, row in rows.items(): - assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="480p") == ( - 5 * row["output_cost_per_second_480p"] - ) - assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="720p") == ( + for key, value in row.items(): + if key.startswith("output_cost_per_second_") and value is not None: + tier = key.removeprefix("output_cost_per_second_") + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution=tier) == 5 * value + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="9999p") == ( 5 * row["output_cost_per_second"] ) + + def test_h3_video_cost_uses_model_info_tiers(self, local_model_cost_map): + row = litellm.model_cost[f"fal_ai/{H3_TEXT_MODEL}"] + model_info = litellm.get_model_info(model=H3_TEXT_MODEL, custom_llm_provider="fal_ai") + + assert video_generation_cost( + model=H3_TEXT_MODEL, + duration_seconds=5, + custom_llm_provider="fal_ai", + model_info=model_info, + video_resolution="2K", + ) == 5 * row["output_cost_per_second_2k"] + assert video_generation_cost( + model=H3_TEXT_MODEL, + duration_seconds=5, + custom_llm_provider="fal_ai", + model_info=model_info, + video_resolution="768p", + ) == 5 * row["output_cost_per_second_768p"] diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8bdda0490c0..e9720b83ad2 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -619,6 +619,8 @@ def validate_model_cost_values(model_data, exceptions=None): "output_cost_per_second", "output_cost_per_second_480p", "output_cost_per_second_720p", + "output_cost_per_second_768p", + "output_cost_per_second_2k", "output_cost_per_second_1080p", "output_cost_per_second_4k", "input_cost_per_query", @@ -838,6 +840,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_second": {"type": "number"}, "output_cost_per_second_480p": {"type": "number"}, "output_cost_per_second_720p": {"type": "number"}, + "output_cost_per_second_768p": {"type": "number"}, + "output_cost_per_second_2k": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, "output_cost_per_second_4k": {"type": "number"}, "output_cost_per_token": {"type": "number"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6211eeaf962..4932f4aa7b2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31138,12 +31138,16 @@ export interface components { output_cost_per_second?: number | null; /** Output Cost Per Second 1080P */ output_cost_per_second_1080p?: number | null; + /** Output Cost Per Second 2K */ + output_cost_per_second_2k?: number | null; /** Output Cost Per Second 480P */ output_cost_per_second_480p?: number | null; /** Output Cost Per Second 4K */ output_cost_per_second_4k?: number | null; /** Output Cost Per Second 720P */ output_cost_per_second_720p?: number | null; + /** Output Cost Per Second 768P */ + output_cost_per_second_768p?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Output Cost Per Token Above 128K Tokens */ @@ -41878,12 +41882,16 @@ export interface components { output_cost_per_second?: number | null; /** Output Cost Per Second 1080P */ output_cost_per_second_1080p?: number | null; + /** Output Cost Per Second 2K */ + output_cost_per_second_2k?: number | null; /** Output Cost Per Second 480P */ output_cost_per_second_480p?: number | null; /** Output Cost Per Second 4K */ output_cost_per_second_4k?: number | null; /** Output Cost Per Second 720P */ output_cost_per_second_720p?: number | null; + /** Output Cost Per Second 768P */ + output_cost_per_second_768p?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Output Cost Per Token Above 128K Tokens */ From cda890438297ac133df56a93d69e0e3eac65340b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:03:13 -0700 Subject: [PATCH 38/56] fix(cache): keep native wheel within size budget --- .../crates/python-bridge/src/cache/config.rs | 84 ++++++++----------- 1 file changed, 33 insertions(+), 51 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 637bdab4055..37eee2048e0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,7 +1,7 @@ use std::time::Duration; use pyo3::{ - exceptions::{PyOverflowError, PyTypeError, PyValueError}, + exceptions::{PyTypeError, PyValueError}, prelude::*, types::{PyAny, PyDict}, }; @@ -43,7 +43,7 @@ pub(super) struct RedisTlsConfig { pub(super) certificate_requirement: CertificateRequirement, pub(super) check_hostname: bool, pub(super) ca_certificate: Option, - pub(super) ca_data: Option>, + pub(super) ca_data: Option, pub(super) client_certificate: Option, pub(super) client_key: Option, } @@ -86,21 +86,21 @@ pub(super) struct NativeCacheConfig { } pub(super) enum UnsupportedCacheConfig { - Backend(String), - RedisMode(&'static str), - RedisOption(String), + Backend, + RedisTopology, + RedisCredentials, + RedisConnection, + RedisOption, } impl UnsupportedCacheConfig { - pub(super) fn message(&self) -> String { + pub(super) fn message(&self) -> &'static str { match self { - Self::Backend(backend) => { - format!("native cache backend {backend:?} is not implemented") - } - Self::RedisMode(mode) => format!("native Redis {mode} mode is not implemented"), - Self::RedisOption(option) => { - format!("native Redis option {option:?} is not implemented") - } + Self::Backend => "native cache backend is not implemented", + Self::RedisTopology => "native Redis topology is not implemented", + Self::RedisCredentials => "native Redis credentials require Python", + Self::RedisConnection => "native Redis connection type is not implemented", + Self::RedisOption => "native Redis configuration requires Python", } } } @@ -143,7 +143,7 @@ impl NativeCacheConfig { Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, _ => Ok(CacheConfigProjection::Unsupported( - UnsupportedCacheConfig::Backend(backend_name), + UnsupportedCacheConfig::Backend, )), } } @@ -187,7 +187,7 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { capacity: backend.getattr("max_size_in_memory")?.extract::()?, max_entry_bytes: max_size_kib .checked_mul(1024) - .ok_or_else(|| PyOverflowError::new_err("memory cache item limit is too large"))?, + .ok_or_else(|| PyValueError::new_err("memory cache item limit is too large"))?, }) } @@ -196,20 +196,19 @@ fn project_redis( ) -> PyResult> { let source = backend.getattr("redis_kwargs")?.cast_into::()?; if has_value(&source, "startup_nodes")? { - return Ok(Err(UnsupportedCacheConfig::RedisMode("cluster"))); + return Ok(Err(UnsupportedCacheConfig::RedisTopology)); } if has_value(&source, "sentinel_nodes")? { - return Ok(Err(UnsupportedCacheConfig::RedisMode("sentinel"))); + return Ok(Err(UnsupportedCacheConfig::RedisTopology)); } - for key in [ - "credential_provider", - "redis_connect_func", - "connection_pool", - ] { + for key in ["credential_provider", "redis_connect_func"] { if has_value(&source, key)? { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } } + if has_value(&source, "connection_pool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } for key in [ "retry", "retry_on_error", @@ -228,12 +227,12 @@ fn project_redis( "ssl_ocsp_expected_cert", ] { if has_value(&source, key)? { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisOption)); } } for key in ["retry_on_timeout", "single_connection_client"] { if optional_coerced_bool(&source, key)?.unwrap_or(false) { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisOption)); } } @@ -241,12 +240,12 @@ fn project_redis( let pool = client.getattr("connection_pool")?; let pool_class = class_identity(&pool)?; if pool_class != ("redis.connection".to_owned(), "ConnectionPool".to_owned()) { - return Ok(Err(UnsupportedCacheConfig::RedisMode("custom pool"))); + 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::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } } let connection_class = resolved @@ -265,17 +264,13 @@ fn project_redis( (module, name) if module == "redis.connection" && name == "SSLConnection" => { Some(project_tls(&resolved)?) } - _ => return Ok(Err(UnsupportedCacheConfig::RedisMode("custom connection"))), + _ => return Ok(Err(UnsupportedCacheConfig::RedisConnection)), }; let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, 3 => RedisProtocol::Resp3, - value => { - return Err(PyValueError::new_err(format!( - "unsupported Redis protocol version {value}" - ))); - } + _ => return Err(PyValueError::new_err("unsupported Redis protocol version")), }; let health_check_interval = duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?; @@ -306,7 +301,7 @@ fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { certificate_requirement: certificate_requirement(values)?, check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false), ca_certificate: optional_dict_string(values, "ssl_ca_certs")?, - ca_data: optional_bytes(values, "ssl_ca_data")?, + ca_data: optional_dict_string(values, "ssl_ca_data")?, client_certificate: optional_dict_string(values, "ssl_certfile")?, client_key: optional_dict_string(values, "ssl_keyfile")?, }) @@ -374,14 +369,14 @@ fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult { fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? - .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))? .extract::() } fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? - .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))? .extract::() } @@ -392,19 +387,6 @@ fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult, key: &str) -> PyResult>> { - let Some(value) = values.get_item(key)? else { - return Ok(None); - }; - if value.is_none() { - return Ok(None); - } - if let Ok(bytes) = value.extract::>() { - return Ok(Some(bytes)); - } - Ok(Some(value.extract::()?.into_bytes())) -} - fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -560,7 +542,7 @@ mod tests { ); assert!(tls.check_hostname); assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem")); - assert_eq!(tls.ca_data.as_deref(), Some(b"CA DATA".as_slice())); + assert_eq!(tls.ca_data.as_deref(), Some("CA DATA")); assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); assert_eq!(tls.client_key.as_deref(), Some("/client.key")); }); @@ -580,7 +562,7 @@ mod tests { else { panic!("dynamic authentication must stay on Python"); }; - assert!(reason.message().contains("credential_provider")); + assert_eq!(reason.message(), "native Redis credentials require Python"); }); } } From 59dbbe5ce71f382c64790ac5823c3f93a14fc8b5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:10:06 -0700 Subject: [PATCH 39/56] fix(cache): outline Python configuration extraction --- .../crates/python-bridge/src/cache/config.rs | 66 ++++++++++++------- 1 file changed, 43 insertions(+), 23 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 37eee2048e0..bcacf7b4e36 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -3,7 +3,7 @@ use std::time::Duration; use pyo3::{ exceptions::{PyTypeError, PyValueError}, prelude::*, - types::{PyAny, PyDict}, + types::{PyAny, PyDict, PyString}, }; use super::{native::NativeResponseCache, request::duration}; @@ -111,6 +111,7 @@ pub(super) enum CacheConfigProjection { } impl NativeCacheConfig { + #[inline(never)] pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult { let backend_name = facade.getattr("type")?.extract::()?; let policy = CachePolicy { @@ -180,6 +181,7 @@ impl NativeCacheConfig { } } +#[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; Ok(MemoryCacheConfig { @@ -191,6 +193,7 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { }) } +#[inline(never)] fn project_redis( backend: &Bound<'_, PyAny>, ) -> PyResult> { @@ -238,8 +241,7 @@ fn project_redis( let client = backend.getattr("redis_client")?; let pool = client.getattr("connection_pool")?; - let pool_class = class_identity(&pool)?; - if pool_class != ("redis.connection".to_owned(), "ConnectionPool".to_owned()) { + if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); } let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; @@ -251,20 +253,12 @@ fn project_redis( let connection_class = resolved .get_item("connection_class")? .unwrap_or(pool.getattr("connection_class")?); - let connection_class = ( - connection_class - .getattr("__module__")? - .extract::()?, - connection_class - .getattr("__qualname__")? - .extract::()?, - ); - let tls = match connection_class { - (module, name) if module == "redis.connection" && name == "Connection" => None, - (module, name) if module == "redis.connection" && name == "SSLConnection" => { - Some(project_tls(&resolved)?) - } - _ => return Ok(Err(UnsupportedCacheConfig::RedisConnection)), + 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 protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { @@ -296,6 +290,7 @@ fn project_redis( })) } +#[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { certificate_requirement: certificate_requirement(values)?, @@ -307,6 +302,7 @@ fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { }) } +#[inline(never)] fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult { let Some(value) = values.get_item("ssl_cert_reqs")? else { return Ok(CertificateRequirement::Required); @@ -334,18 +330,31 @@ fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult) -> PyResult<(String, String)> { - let class = value.get_type(); - Ok(( - class.getattr("__module__")?.extract::()?, - class.getattr("__qualname__")?.extract::()?, - )) +#[inline(never)] +fn instance_class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult { + class_is(value.get_type().as_any(), module, name) } +#[inline(never)] +fn class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult { + Ok(value + .getattr("__module__")? + .cast_into::()? + .to_str()? + == module + && value + .getattr("__qualname__")? + .cast_into::()? + .to_str()? + == name) +} + +#[inline(never)] fn optional_duration(value: Bound<'_, PyAny>) -> PyResult> { value.extract::>()?.map(duration).transpose() } +#[inline(never)] fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { match value.getattr(name) { Ok(value) => optional_string(value), @@ -356,16 +365,19 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult) -> PyResult> { Ok(value .extract::>()? .filter(|value| !value.is_empty())) } +#[inline(never)] fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult { Ok(values.get_item(key)?.is_some_and(|value| !value.is_none())) } +#[inline(never)] fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? @@ -373,6 +385,7 @@ fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { .extract::() } +#[inline(never)] fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? @@ -380,6 +393,7 @@ fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { .extract::() } +#[inline(never)] fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) if !value.is_none() => optional_string(value), @@ -387,6 +401,7 @@ fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -394,6 +409,7 @@ fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> } } +#[inline(never)] fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -401,6 +417,7 @@ fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> } } +#[inline(never)] fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -408,6 +425,7 @@ fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { } } +#[inline(never)] fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -415,6 +433,7 @@ fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult } } +#[inline(never)] fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { let Some(value) = values.get_item(key)? else { return Ok(None); @@ -431,6 +450,7 @@ fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult().map(Some) } +#[inline(never)] fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { optional_f64(values, key)?.map(duration).transpose() } From adc4e6a13284678ce322359e09aa88b3fe473a42 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:00:12 +0000 Subject: [PATCH 40/56] fix(fal_ai): price images from the dimensions fal returns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../amazon_nova_canvas_transformation.py | 2 +- .../amazon_stability1_transformation.py | 2 +- .../amazon_stability3_transformation.py | 2 +- .../amazon_titan_transformation.py | 2 +- litellm/llms/fal_ai/cost_calculator.py | 138 ++++++++++++++---- .../flux_pro_v11_ultra_transformation.py | 23 +-- .../gpt_image_2_transformation.py | 2 +- .../fal_ai/image_generation/transformation.py | 41 +++++- .../llms/gemini/image_edit/transformation.py | 2 +- .../vertex_gemini_transformation.py | 2 +- .../vertex_imagen_transformation.py | 2 +- .../image_generation_handler.py | 2 +- ...odel_prices_and_context_window_backup.json | 3 +- litellm/types/utils.py | 1 + model_prices_and_context_window.json | 3 +- .../providers/test_fal_ai_image_wire.py | 67 ++++++--- .../test_fal_ai_image_edit_transformation.py | 19 ++- .../test_fal_ai_flux_dev_transformation.py | 29 +++- .../llms/fal_ai/test_cost_calculator.py | 45 ++++++ 19 files changed, 304 insertions(+), 83 deletions(-) diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index ce61a6253f6..32f069be6c3 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -181,7 +181,7 @@ class AmazonNovaCanvasConfig: for _img in nova_response.get("images", []): openai_images.append(Image(b64_json=_img)) - model_response.data = openai_images + model_response.data = openai_images # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response @classmethod diff --git a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py index df91e736239..a8ae5f6e1eb 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py @@ -125,7 +125,7 @@ class AmazonStabilityConfig: _image = Image(b64_json=artifact["base64"]) image_list.append(_image) - model_response.data = image_list + model_response.data = image_list # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index 98e4cbbfd4d..489440a7430 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -101,7 +101,7 @@ class AmazonStability3Config: for _img in stability_3_response.get("images", []): openai_images.append(Image(b64_json=_img)) - model_response.data = openai_images + model_response.data = openai_images # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response @classmethod diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index e1b06791c9d..93b23672377 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -134,7 +134,7 @@ class AmazonTitanImageGenerationConfig: _image = Image(b64_json=image) image_list.append(_image) - model_response.data = image_list + model_response.data = image_list # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index f23bd1b46bc..497a792fd8a 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -1,9 +1,12 @@ from collections.abc import Mapping +from math import ceil from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + import litellm -from litellm.types.utils import ImageResponse +from litellm.types.utils import ImageObject, ImageResponse FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high" FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768" @@ -18,14 +21,20 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( } ) +_MODEL_COST_MAP: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter( + Mapping[str, Mapping[str, object]] +) +_OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + def _keyed_size(optional_params: Mapping[str, object]) -> str | None: image_size: Final = optional_params.get("image_size") if image_size is None or image_size == "auto": return FAL_TEXT_TO_IMAGE_DEFAULT_SIZE if isinstance(image_size, Mapping): - width: Final = image_size.get("width") - height: Final = image_size.get("height") + image_size_map: Final = _OBJECT_MAP.validate_python(image_size) + width: Final = image_size_map.get("width") + height: Final = image_size_map.get("height") if isinstance(width, int) and isinstance(height, int): return f"{width}-x-{height}" return None @@ -34,21 +43,71 @@ def _keyed_size(optional_params: Mapping[str, object]) -> str | None: return None -def _keyed_cost_per_image(model: str, optional_params: Mapping[str, object] | None) -> float | None: - if optional_params is None: +def _response_size(image: object) -> str | None: + if not isinstance(image, ImageObject): return None - size: Final = _keyed_size(optional_params) - if size is None: + raw_provider_specific_fields: Final = image.provider_specific_fields + if not isinstance(raw_provider_specific_fields, Mapping): return None + provider_specific_fields: Final = _OBJECT_MAP.validate_python(raw_provider_specific_fields) + width: Final = provider_specific_fields.get("width") + height: Final = provider_specific_fields.get("height") + if not isinstance(width, int) or not isinstance(height, int): + return None + return f"{width}-x-{height}" + + +def _keyed_quality(optional_params: Mapping[str, object]) -> str: raw_quality: Final = optional_params.get("quality") - quality: Final = ( - raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY - ) - keyed_entry: Final = litellm.model_cost.get(f"fal_ai/{quality}/{size}/{model}") - if keyed_entry is None: + return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY + + +def _keyed_cost_per_image( + model: str, + image: object, + optional_params: Mapping[str, object], + model_cost_map: Mapping[str, Mapping[str, object]], +) -> float | None: + quality: Final = _keyed_quality(optional_params) + request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE + sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE) + for size in sizes: + if size is None: + continue + keyed_entry = model_cost_map.get(f"fal_ai/{quality}/{size}/{model}") + if keyed_entry is None: + continue + keyed_cost = keyed_entry.get("output_cost_per_image") + if isinstance(keyed_cost, (int, float)): + return float(keyed_cost) + return None + + +def _image_dimensions(image: object) -> tuple[int, int] | None: + if not isinstance(image, ImageObject): return None - keyed_cost: Final = keyed_entry.get("output_cost_per_image") - return float(keyed_cost) if isinstance(keyed_cost, (int, float)) else None + raw_provider_specific_fields: Final = image.provider_specific_fields + if not isinstance(raw_provider_specific_fields, Mapping): + return None + provider_specific_fields: Final = _OBJECT_MAP.validate_python(raw_provider_specific_fields) + width: Final = provider_specific_fields.get("width") + height: Final = provider_specific_fields.get("height") + if not isinstance(width, int) or not isinstance(height, int): + return None + return width, height + + +def _flat_cost_per_image( + image: object, + output_cost_per_image: float, + output_cost_per_pixel: float | None, +) -> float: + dimensions: Final = _image_dimensions(image) + if dimensions is None or output_cost_per_pixel is None: + return output_cost_per_image + width, height = dimensions + megapixels: Final = 1 if (width, height) == (1024, 1024) else ceil(width * height / 1_000_000) + return output_cost_per_pixel * 1_000_000 * megapixels def cost_calculator( @@ -61,15 +120,44 @@ def cost_calculator( """ if not isinstance(image_response, ImageResponse): raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") - # the proxy cost path passes the provider-prefixed model name - model = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/") - num_images: Final[int] = len(image_response.data) if image_response.data else 0 - keyed_cost_per_image: Final = _keyed_cost_per_image(model=model, optional_params=optional_params) - if keyed_cost_per_image is not None: - return keyed_cost_per_image * num_images - _model_info: Final = litellm.get_model_info( - model=model, - custom_llm_provider=litellm.LlmProviders.FAL_AI.value, + normalized_model: Final = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/") + params: Final[Mapping[str, object]] = optional_params or MappingProxyType({}) + images: Final = tuple(image_response.data or ()) + raw_model_cost: Final[object] = litellm.model_cost # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # global catalog is untyped + model_cost_map: Final = _MODEL_COST_MAP.validate_python(raw_model_cost) + keyed_costs: Final = tuple( + _keyed_cost_per_image( + model=normalized_model, + image=image, + optional_params=params, + model_cost_map=model_cost_map, + ) + for image in images + ) + if all(cost is not None for cost in keyed_costs): + return sum(cost for cost in keyed_costs if cost is not None) + model_info_entry: Final = next( + ( + entry + for key in (f"{litellm.LlmProviders.FAL_AI.value}/{normalized_model}", normalized_model) + if (entry := model_cost_map.get(key)) is not None + ), + None, + ) + model_info: Final = _OBJECT_MAP.validate_python(model_info_entry or MappingProxyType({})) + raw_output_cost_per_image: Final = model_info.get("output_cost_per_image") + output_cost_per_image: Final = ( + float(raw_output_cost_per_image) if isinstance(raw_output_cost_per_image, (int, float)) else 0.0 + ) + raw_output_cost_per_pixel: Final = model_info.get("output_cost_per_pixel") + output_cost_per_pixel: Final = ( + float(raw_output_cost_per_pixel) if isinstance(raw_output_cost_per_pixel, (int, float)) else None + ) + return sum( + _flat_cost_per_image( + image=image, + output_cost_per_image=output_cost_per_image, + output_cost_per_pixel=output_cost_per_pixel, + ) + for image in images ) - output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0 - return output_cost_per_image * num_images diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 228dd9257ce..6b8558b8124 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -3,9 +3,9 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageObject, ImageResponse +from litellm.types.utils import ImageResponse -from .transformation import FalAIBaseConfig +from .transformation import FalAIBaseConfig, fal_images_to_image_objects if TYPE_CHECKING: import tiktoken @@ -229,25 +229,8 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): if not model_response.data: model_response.data = [] - # Handle Flux Pro v1.1-ultra response format images: Final = response_data.get("images", []) - if isinstance(images, list): - for image_data in images: - if isinstance(image_data, dict): - model_response.data.append( - ImageObject( - url=image_data.get("url", None), - b64_json=None, # Flux Pro returns URLs only - ) - ) - elif isinstance(image_data, str): - # If images is just a list of URLs - model_response.data.append( - ImageObject( - url=image_data, - b64_json=None, - ) - ) + model_response.data.extend(fal_images_to_image_objects(images)) # Add additional metadata from Flux Pro response if hasattr(model_response, "_hidden_params"): diff --git a/litellm/llms/fal_ai/image_generation/gpt_image_2_transformation.py b/litellm/llms/fal_ai/image_generation/gpt_image_2_transformation.py index 3dfc26f8f46..ca301662cf8 100644 --- a/litellm/llms/fal_ai/image_generation/gpt_image_2_transformation.py +++ b/litellm/llms/fal_ai/image_generation/gpt_image_2_transformation.py @@ -51,7 +51,7 @@ def supported_gpt_image_qualities( and "-x-" in parts[2] and "/".join(parts[3:]) == qualified_endpoint ) - return qualities | {"auto"} if qualities else frozenset() + return qualities | frozenset({"auto"}) if qualities else frozenset() def map_gpt_image_quality( diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 7f6a417e8a1..a5afdc41e9a 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -1,6 +1,9 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -22,16 +25,40 @@ else: LiteLLMLoggingObj = Any +class FalImageProviderSpecificFields(TypedDict, total=False): + width: ReadOnly[int] + height: ReadOnly[int] + content_type: ReadOnly[str] + + +_FAL_IMAGE_DATA: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]: if not isinstance(images, list): return () - return tuple( - ImageObject(url=image_data.get("url", None), b64_json=image_data.get("b64_json", None)) - if isinstance(image_data, dict) - else ImageObject(url=image_data, b64_json=None) - for image_data in images - if isinstance(image_data, (dict, str)) - ) + + def to_image_object(image_data: object) -> ImageObject: + if isinstance(image_data, Mapping): + image_map: Final = _FAL_IMAGE_DATA.validate_python(image_data) + url: Final = image_map.get("url") + b64_json: Final = image_map.get("b64_json") + width: Final = image_map.get("width") + height: Final = image_map.get("height") + content_type: Final = image_map.get("content_type") + provider_specific_fields: Final[FalImageProviderSpecificFields] = { + **({"width": width} if isinstance(width, int) else {}), + **({"height": height} if isinstance(height, int) else {}), + **({"content_type": content_type} if isinstance(content_type, str) else {}), + } + return ImageObject( + url=url if isinstance(url, str) else None, + b64_json=b64_json if isinstance(b64_json, str) else None, + provider_specific_fields=provider_specific_fields or None, + ) + return ImageObject(url=image_data if isinstance(image_data, str) else None, b64_json=None) + + return tuple(to_image_object(image_data) for image_data in images if isinstance(image_data, (Mapping, str))) class FalAIBaseConfig(BaseImageGenerationConfig): diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index e6c22dc60b4..594b4dfd4ce 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -148,7 +148,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): ) ) - model_response.data = cast(list[OpenAIImage], data_list) + model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type if "usageMetadata" in response_json: model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) return model_response diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 725a7f39917..6cf32733550 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -219,7 +219,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): if (inline_data := part.get("inlineData")) and (b64_json := inline_data.get("data")) ] - model_response.data = cast(list[OpenAIImage], data_list) + model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index c6ad5928b74..d0485275bf3 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -220,7 +220,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) ) - model_response.data = cast(list[OpenAIImage], data_list) + model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index 6a5bb484540..d0fb6a2f937 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -40,7 +40,7 @@ class VertexImageGeneration(VertexLLM): image_object = Image(b64_json=bytes_base64_encoded) response_data.append(image_object) - model_response.data = response_data + model_response.data = response_data # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type return model_response def transform_optional_params(self, optional_params: dict | None) -> dict: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 87edd1544ca..855ef1fb04d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24917,10 +24917,11 @@ "fal_ai/fal-ai/flux/dev": { "litellm_provider": "fal_ai", "metadata": { - "notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. Every named fal image_size (including the landscape_4_3 default) rounds up to 1 megapixel, so this flat per-image price is exact for them" + "notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. The per-pixel rate is used when Fal reports the output size, and the flat per-image price is the fallback when dimensions are unavailable" }, "mode": "image_generation", "output_cost_per_image": 0.025, + "output_cost_per_pixel": 2.5e-08, "source": "https://fal.ai/models/fal-ai/flux/dev", "supported_endpoints": [ "/v1/images/generations" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5a80644347e..7c33579be5a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2543,6 +2543,7 @@ from openai.types.images_response import ImagesResponse as OpenAIImageResponse class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = {} + data: list[ImageObject] usage: ImageUsage | None = None """ Users might use litellm with older python versions, we don't want this to break for them. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 87edd1544ca..855ef1fb04d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24917,10 +24917,11 @@ "fal_ai/fal-ai/flux/dev": { "litellm_provider": "fal_ai", "metadata": { - "notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. Every named fal image_size (including the landscape_4_3 default) rounds up to 1 megapixel, so this flat per-image price is exact for them" + "notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. The per-pixel rate is used when Fal reports the output size, and the flat per-image price is the fallback when dimensions are unavailable" }, "mode": "image_generation", "output_cost_per_image": 0.025, + "output_cost_per_pixel": 2.5e-08, "source": "https://fal.ai/models/fal-ai/flux/dev", "supported_endpoints": [ "/v1/images/generations" diff --git a/tests/integration/providers/test_fal_ai_image_wire.py b/tests/integration/providers/test_fal_ai_image_wire.py index 23ab7e08c16..e0c05e82b33 100644 --- a/tests/integration/providers/test_fal_ai_image_wire.py +++ b/tests/integration/providers/test_fal_ai_image_wire.py @@ -23,14 +23,14 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) -def _catalog_cost(key: str) -> float: +def _catalog_cost(key: str, field: str = "output_cost_per_image") -> float: cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) - cost_value: Final = cost_map[key]["output_cost_per_image"] + cost_value: Final = cost_map[key][field] assert isinstance(cost_value, (int, float)) return float(cost_value) -def _image_response(urls: tuple[str, ...], prompt: str) -> bytes: +def _image_response(images: tuple[tuple[str, int, int], ...], prompt: str) -> bytes: return json.dumps( { "images": [ @@ -39,10 +39,10 @@ def _image_response(urls: tuple[str, ...], prompt: str) -> bytes: "content_type": "image/png", "file_name": url.rsplit("/", 1)[-1], "file_size": 123456, - "width": 1024, - "height": 768, + "width": width, + "height": height, } - for url in urls + for url, width, height in images ], "timings": {"inference": 2.1}, "seed": 1234567, @@ -69,9 +69,9 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro body: Final = _JSON_OBJECT.validate_json(request.body) if body.get("quality") == "high": assert body == {"prompt": _PROMPT, "quality": "high", "image_size": {"width": 1024, "height": 1536}} - return Reply(body=_image_response((f"{wire_url}/files/high.png",), _PROMPT)) + return Reply(body=_image_response(((f"{wire_url}/files/high.png", 1024, 1536),), _PROMPT)) assert body == {"prompt": _PROMPT, "quality": "low"} - return Reply(body=_image_response((f"{wire_url}/files/low.png",), _PROMPT)) + return Reply(body=_image_response(((f"{wire_url}/files/low.png", 1024, 1536),), _PROMPT)) with wire_server(respond) as wire, gateway.scenario() as scenario: wire_url: Final = wire.url @@ -85,7 +85,14 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro ) assert high_response.status_code == 200, high_response.text high_payload: Final = _JSON_OBJECT.validate_json(high_response.content) - assert high_payload["data"] == [{"url": f"{wire.url}/files/high.png", "b64_json": None, "revised_prompt": None}] + assert high_payload["data"] == [ + { + "url": f"{wire.url}/files/high.png", + "b64_json": None, + "revised_prompt": None, + "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"}, + } + ] high_cost: Final = _response_cost(high_response) assert high_cost == _approx(_catalog_cost("fal_ai/high/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image")) @@ -96,9 +103,16 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro ) assert low_response.status_code == 200, low_response.text low_payload: Final = _JSON_OBJECT.validate_json(low_response.content) - assert low_payload["data"] == [{"url": f"{wire.url}/files/low.png", "b64_json": None, "revised_prompt": None}] + assert low_payload["data"] == [ + { + "url": f"{wire.url}/files/low.png", + "b64_json": None, + "revised_prompt": None, + "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"}, + } + ] low_cost: Final = _response_cost(low_response) - assert low_cost == _approx(_catalog_cost("fal_ai/low/1024-x-768/openai/gpt-image-2.5/flare/text-to-image")) + assert low_cost == _approx(_catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image")) assert high_cost != low_cost assert [(request.method, request.target) for request in wire.drain()] == [ ("POST", "/openai/gpt-image-2.5/flare/text-to-image"), @@ -119,7 +133,7 @@ def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gate } return Reply( body=_image_response( - (f"{wire_url}/files/flux-1.png", f"{wire_url}/files/flux-2.png"), + ((f"{wire_url}/files/flux-1.png", 1024, 1024), (f"{wire_url}/files/flux-2.png", 1920, 1080)), _PROMPT, ) ) @@ -135,11 +149,21 @@ def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gate assert response.status_code == 200, response.text payload: Final = _JSON_OBJECT.validate_json(response.content) assert payload["data"] == [ - {"url": f"{wire.url}/files/flux-1.png", "b64_json": None, "revised_prompt": None}, - {"url": f"{wire.url}/files/flux-2.png", "b64_json": None, "revised_prompt": None}, + { + "url": f"{wire.url}/files/flux-1.png", + "b64_json": None, + "revised_prompt": None, + "provider_specific_fields": {"width": 1024, "height": 1024, "content_type": "image/png"}, + }, + { + "url": f"{wire.url}/files/flux-2.png", + "b64_json": None, + "revised_prompt": None, + "provider_specific_fields": {"width": 1920, "height": 1080, "content_type": "image/png"}, + }, ] cost: Final = _response_cost(response) - assert cost == _approx(2 * _catalog_cost("fal_ai/fal-ai/flux/dev")) + assert cost == _approx(4 * _catalog_cost("fal_ai/fal-ai/flux/dev", "output_cost_per_pixel") * 1_000_000) assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/fal-ai/flux/dev")] @@ -155,7 +179,7 @@ def test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row( "image_urls": ["data:image/png;base64," + base64.b64encode(_PNG_BYTES).decode()], "quality": "low", } - return Reply(body=_image_response((f"{wire_url}/files/edit.png",), _PROMPT)) + return Reply(body=_image_response(((f"{wire_url}/files/edit.png", 1024, 1536),), _PROMPT)) with wire_server(respond) as wire, gateway.scenario() as scenario: wire_url: Final = wire.url @@ -168,9 +192,16 @@ def test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row( ) assert response.status_code == 200, response.text payload: Final = _JSON_OBJECT.validate_json(response.content) - assert payload["data"] == [{"url": f"{wire.url}/files/edit.png", "b64_json": None, "revised_prompt": None}] + assert payload["data"] == [ + { + "url": f"{wire.url}/files/edit.png", + "b64_json": None, + "revised_prompt": None, + "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"}, + } + ] cost: Final = _response_cost(response) - assert cost == _approx(_catalog_cost("fal_ai/low/1024-x-768/openai/gpt-image-2.5/flare/edit")) + assert cost == _approx(_catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/edit")) assert [(request.method, request.target) for request in wire.drain()] == [ ("POST", "/openai/gpt-image-2.5/flare/edit") ] diff --git a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py b/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py index 65b04e1f1b8..6c55760b625 100644 --- a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py @@ -120,12 +120,29 @@ def test_transform_request_reads_every_file_types_input(tmp_path, image_factory) def test_transform_response_maps_fal_images(): - raw = httpx.Response(200, json={"images": [{"url": "https://fal.media/out.png"}]}) + raw = httpx.Response( + 200, + json={ + "images": [ + { + "url": "https://fal.media/out.png", + "width": 1024, + "height": 1536, + "content_type": "image/png", + } + ] + }, + ) response = FalAIImageEditConfig().transform_image_edit_response( model="openai/gpt-image-2.5/flare/edit", raw_response=raw, logging_obj=None ) assert isinstance(response, ImageResponse) assert [image.url for image in response.data] == ["https://fal.media/out.png"] + assert response.data[0].provider_specific_fields == { + "width": 1024, + "height": 1536, + "content_type": "image/png", + } @pytest.mark.parametrize("image", [None, []]) diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py index 09c9bc4b5f7..273b9151a0a 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py @@ -47,7 +47,15 @@ def test_flux_dev_maps_openai_params_and_builds_request(): def test_flux_dev_response_yields_one_image_object_per_fal_image(): - raw = httpx.Response(200, json={"images": [{"url": "https://fal.media/a.png"}, {"url": "https://fal.media/b.png"}]}) + raw = httpx.Response( + 200, + json={ + "images": [ + {"url": "https://fal.media/a.png", "width": 1024, "height": 768, "content_type": "image/png"}, + {"url": "https://fal.media/b.png", "width": 512, "height": 512, "content_type": "image/webp"}, + ] + }, + ) response = FalAIFluxDevConfig().transform_image_generation_response( model="fal-ai/flux/dev", raw_response=raw, @@ -59,3 +67,22 @@ def test_flux_dev_response_yields_one_image_object_per_fal_image(): encoding=None, ) assert [image.url for image in response.data] == ["https://fal.media/a.png", "https://fal.media/b.png"] + assert [image.provider_specific_fields for image in response.data] == [ + {"width": 1024, "height": 768, "content_type": "image/png"}, + {"width": 512, "height": 512, "content_type": "image/webp"}, + ] + + +def test_flux_dev_response_omits_provider_specific_fields_when_fal_omits_metadata(): + raw = httpx.Response(200, json={"images": [{"url": "https://fal.media/a.png"}]}) + response = FalAIFluxDevConfig().transform_image_generation_response( + model="fal-ai/flux/dev", + raw_response=raw, + model_response=ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert response.data[0].provider_specific_fields is None diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py index 989b5855803..05b6014c8e2 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -19,6 +19,18 @@ def _image_response(num_images: int = 1) -> ImageResponse: return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)]) +def _image_response_with_dimensions(dimensions: tuple[tuple[int, int], ...]) -> ImageResponse: + return ImageResponse( + data=[ + ImageObject( + url=f"https://example.com/img-{index}.png", + provider_specific_fields={"width": width, "height": height}, + ) + for index, (width, height) in enumerate(dimensions) + ] + ) + + GPT_IMAGE_25_MODELS = ( "openai/gpt-image-2.5/flare/text-to-image", "openai/gpt-image-2.5/flare/edit", @@ -55,6 +67,28 @@ def test_gpt_image_25_edit_auto_size_still_honors_quality(): assert 0 < low < high +def test_gpt_image_response_dimensions_override_request_size(): + model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image" + cost = cost_calculator( + model=model, + image_response=_image_response_with_dimensions(((1024, 1536),)), + optional_params={"quality": "low", "image_size": {"width": 1024, "height": 768}}, + ) + expected = litellm.model_cost[f"fal_ai/low/1024-x-1536/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"] + assert cost == expected + + +def test_gpt_image_response_dimensions_fall_back_to_request_size_when_unpriced(): + model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image" + cost = cost_calculator( + model=model, + image_response=_image_response_with_dimensions(((777, 888),)), + optional_params={"quality": "low", "image_size": {"width": 1024, "height": 1536}}, + ) + expected = litellm.model_cost[f"fal_ai/low/1024-x-1536/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"] + assert cost == expected + + def test_gpt_image_25_quality_tiers_are_monotonic(): costs = tuple( cost_calculator( @@ -78,6 +112,17 @@ def test_flux_dev_cost_is_nonzero_and_distinct_from_schnell(): assert dev == 3 * litellm.model_cost["fal_ai/fal-ai/flux/dev"]["output_cost_per_image"] +def test_flux_dev_cost_uses_response_megapixels_per_image(): + model = "fal_ai/fal-ai/flux/dev" + cost = cost_calculator( + model=model, + image_response=_image_response_with_dimensions(((1024, 1024), (1920, 1080), (512, 512))), + optional_params={}, + ) + output_cost_per_pixel = litellm.model_cost[model]["output_cost_per_pixel"] + assert cost == pytest.approx(output_cost_per_pixel * 1_000_000 * (1 + 3 + 1)) + + def test_image_edit_call_type_routes_to_fal_keyed_pricing(): model = "openai/gpt-image-2.5/flare/edit" cost = CostCalculatorUtils.route_image_generation_cost_calculator( From 9b54c4b0779a83584f8c38fe002e8d8c9ffd5259 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:08:16 +0000 Subject: [PATCH 41/56] refactor(fal_ai): bill flux dev per 1024x1024 megapixel and drop ImageResponse retyping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../amazon_nova_canvas_transformation.py | 2 +- .../amazon_stability1_transformation.py | 2 +- .../amazon_stability3_transformation.py | 2 +- .../amazon_titan_transformation.py | 2 +- litellm/llms/fal_ai/cost_calculator.py | 99 +++++++++---------- .../llms/gemini/image_edit/transformation.py | 2 +- .../vertex_gemini_transformation.py | 2 +- .../vertex_imagen_transformation.py | 2 +- .../image_generation_handler.py | 2 +- ...odel_prices_and_context_window_backup.json | 2 +- litellm/types/utils.py | 5 +- model_prices_and_context_window.json | 2 +- .../providers/test_fal_ai_image_wire.py | 2 +- .../llms/fal_ai/test_cost_calculator.py | 2 +- 14 files changed, 63 insertions(+), 65 deletions(-) diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index 32f069be6c3..ce61a6253f6 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -181,7 +181,7 @@ class AmazonNovaCanvasConfig: for _img in nova_response.get("images", []): openai_images.append(Image(b64_json=_img)) - model_response.data = openai_images # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = openai_images return model_response @classmethod diff --git a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py index a8ae5f6e1eb..df91e736239 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py @@ -125,7 +125,7 @@ class AmazonStabilityConfig: _image = Image(b64_json=artifact["base64"]) image_list.append(_image) - model_response.data = image_list # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = image_list return model_response diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index 489440a7430..98e4cbbfd4d 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -101,7 +101,7 @@ class AmazonStability3Config: for _img in stability_3_response.get("images", []): openai_images.append(Image(b64_json=_img)) - model_response.data = openai_images # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = openai_images return model_response @classmethod diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index 93b23672377..e1b06791c9d 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -134,7 +134,7 @@ class AmazonTitanImageGenerationConfig: _image = Image(b64_json=image) image_list.append(_image) - model_response.data = image_list # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = image_list return model_response diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 497a792fd8a..114bbabe4e6 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -10,6 +10,7 @@ from litellm.types.utils import ImageObject, ImageResponse FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high" FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768" +FAL_PIXELS_PER_MEGAPIXEL: Final[int] = 1_048_576 FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( { "square_hd": "1024-x-1024", @@ -21,10 +22,8 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( } ) -_MODEL_COST_MAP: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter( - Mapping[str, Mapping[str, object]] -) _OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) def _keyed_size(optional_params: Mapping[str, object]) -> str | None: @@ -43,46 +42,6 @@ def _keyed_size(optional_params: Mapping[str, object]) -> str | None: return None -def _response_size(image: object) -> str | None: - if not isinstance(image, ImageObject): - return None - raw_provider_specific_fields: Final = image.provider_specific_fields - if not isinstance(raw_provider_specific_fields, Mapping): - return None - provider_specific_fields: Final = _OBJECT_MAP.validate_python(raw_provider_specific_fields) - width: Final = provider_specific_fields.get("width") - height: Final = provider_specific_fields.get("height") - if not isinstance(width, int) or not isinstance(height, int): - return None - return f"{width}-x-{height}" - - -def _keyed_quality(optional_params: Mapping[str, object]) -> str: - raw_quality: Final = optional_params.get("quality") - return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY - - -def _keyed_cost_per_image( - model: str, - image: object, - optional_params: Mapping[str, object], - model_cost_map: Mapping[str, Mapping[str, object]], -) -> float | None: - quality: Final = _keyed_quality(optional_params) - request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE - sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE) - for size in sizes: - if size is None: - continue - keyed_entry = model_cost_map.get(f"fal_ai/{quality}/{size}/{model}") - if keyed_entry is None: - continue - keyed_cost = keyed_entry.get("output_cost_per_image") - if isinstance(keyed_cost, (int, float)): - return float(keyed_cost) - return None - - def _image_dimensions(image: object) -> tuple[int, int] | None: if not isinstance(image, ImageObject): return None @@ -97,6 +56,39 @@ def _image_dimensions(image: object) -> tuple[int, int] | None: return width, height +def _response_size(image: object) -> str | None: + dimensions: Final = _image_dimensions(image) + if dimensions is None: + return None + width, height = dimensions + return f"{width}-x-{height}" + + +def _keyed_quality(optional_params: Mapping[str, object]) -> str: + raw_quality: Final = optional_params.get("quality") + return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY + + +def _keyed_cost_per_image( + model: str, + image: object, + optional_params: Mapping[str, object], +) -> float | None: + quality: Final = _keyed_quality(optional_params) + request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE + sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE) + for size in sizes: + if size is None: + continue + keyed_entry = _entry(f"fal_ai/{quality}/{size}/{model}") + if keyed_entry is None: + continue + keyed_cost = keyed_entry.get("output_cost_per_image") + if isinstance(keyed_cost, (int, float)): + return float(keyed_cost) + return None + + def _flat_cost_per_image( image: object, output_cost_per_image: float, @@ -106,8 +98,15 @@ def _flat_cost_per_image( if dimensions is None or output_cost_per_pixel is None: return output_cost_per_image width, height = dimensions - megapixels: Final = 1 if (width, height) == (1024, 1024) else ceil(width * height / 1_000_000) - return output_cost_per_pixel * 1_000_000 * megapixels + megapixels: Final = ceil(width * height / FAL_PIXELS_PER_MEGAPIXEL) + return output_cost_per_pixel * FAL_PIXELS_PER_MEGAPIXEL * megapixels + + +def _entry(key: str) -> Mapping[str, object] | None: + raw_entry: Final[object] = litellm.model_cost.get(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # global catalog is untyped + if not isinstance(raw_entry, Mapping): + return None + return _OBJECT_MAP.validate_python(raw_entry) def cost_calculator( @@ -123,28 +122,24 @@ def cost_calculator( normalized_model: Final = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/") params: Final[Mapping[str, object]] = optional_params or MappingProxyType({}) images: Final = tuple(image_response.data or ()) - raw_model_cost: Final[object] = litellm.model_cost # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # global catalog is untyped - model_cost_map: Final = _MODEL_COST_MAP.validate_python(raw_model_cost) keyed_costs: Final = tuple( _keyed_cost_per_image( model=normalized_model, image=image, optional_params=params, - model_cost_map=model_cost_map, ) for image in images ) if all(cost is not None for cost in keyed_costs): return sum(cost for cost in keyed_costs if cost is not None) - model_info_entry: Final = next( + model_info: Final = next( ( entry for key in (f"{litellm.LlmProviders.FAL_AI.value}/{normalized_model}", normalized_model) - if (entry := model_cost_map.get(key)) is not None + if (entry := _entry(key)) is not None ), - None, + _EMPTY_ENTRY, ) - model_info: Final = _OBJECT_MAP.validate_python(model_info_entry or MappingProxyType({})) raw_output_cost_per_image: Final = model_info.get("output_cost_per_image") output_cost_per_image: Final = ( float(raw_output_cost_per_image) if isinstance(raw_output_cost_per_image, (int, float)) else 0.0 diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 594b4dfd4ce..e6c22dc60b4 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -148,7 +148,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): ) ) - model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = cast(list[OpenAIImage], data_list) if "usageMetadata" in response_json: model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) return model_response diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 6cf32733550..725a7f39917 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -219,7 +219,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): if (inline_data := part.get("inlineData")) and (b64_json := inline_data.get("data")) ] - model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = cast(list[OpenAIImage], data_list) return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index d0485275bf3..c6ad5928b74 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -220,7 +220,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) ) - model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = cast(list[OpenAIImage], data_list) return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index d0fb6a2f937..6a5bb484540 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -40,7 +40,7 @@ class VertexImageGeneration(VertexLLM): image_object = Image(b64_json=bytes_base64_encoded) response_data.append(image_object) - model_response.data = response_data # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = response_data return model_response def transform_optional_params(self, optional_params: dict | None) -> dict: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 855ef1fb04d..570fa3485ee 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24921,7 +24921,7 @@ }, "mode": "image_generation", "output_cost_per_image": 0.025, - "output_cost_per_pixel": 2.5e-08, + "output_cost_per_pixel": 2.384185791015625e-08, "source": "https://fal.ai/models/fal-ai/flux/dev", "supported_endpoints": [ "/v1/images/generations" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 7c33579be5a..a2d92491188 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2543,7 +2543,6 @@ from openai.types.images_response import ImagesResponse as OpenAIImageResponse class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = {} - data: list[ImageObject] usage: ImageUsage | None = None """ Users might use litellm with older python versions, we don't want this to break for them. @@ -2552,6 +2551,10 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): model_config = ConfigDict(extra="allow", protected_namespaces=()) + @field_serializer("data") + def _serialize_image_data(self, data: list[OpenAIImage]) -> list[dict[str, object]]: + return [image.model_dump() for image in data] + def __init__( self, created: int | None = None, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 855ef1fb04d..570fa3485ee 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24921,7 +24921,7 @@ }, "mode": "image_generation", "output_cost_per_image": 0.025, - "output_cost_per_pixel": 2.5e-08, + "output_cost_per_pixel": 2.384185791015625e-08, "source": "https://fal.ai/models/fal-ai/flux/dev", "supported_endpoints": [ "/v1/images/generations" diff --git a/tests/integration/providers/test_fal_ai_image_wire.py b/tests/integration/providers/test_fal_ai_image_wire.py index e0c05e82b33..f9ceac0b037 100644 --- a/tests/integration/providers/test_fal_ai_image_wire.py +++ b/tests/integration/providers/test_fal_ai_image_wire.py @@ -163,7 +163,7 @@ def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gate }, ] cost: Final = _response_cost(response) - assert cost == _approx(4 * _catalog_cost("fal_ai/fal-ai/flux/dev", "output_cost_per_pixel") * 1_000_000) + assert cost == _approx(3 * _catalog_cost("fal_ai/fal-ai/flux/dev", "output_cost_per_pixel") * 1_048_576) assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/fal-ai/flux/dev")] diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py index 05b6014c8e2..e387586e417 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -120,7 +120,7 @@ def test_flux_dev_cost_uses_response_megapixels_per_image(): optional_params={}, ) output_cost_per_pixel = litellm.model_cost[model]["output_cost_per_pixel"] - assert cost == pytest.approx(output_cost_per_pixel * 1_000_000 * (1 + 3 + 1)) + assert cost == pytest.approx(output_cost_per_pixel * 1_048_576 * (1 + 2 + 1)) def test_image_edit_call_type_routes_to_fal_keyed_pricing(): From c12c20083fbc4cc3f0107766d033a171981b65f7 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:09:10 +0000 Subject: [PATCH 42/56] fix(fal_ai): keep image data serializer none safe Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a2d92491188..8d327b337a7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2552,8 +2552,8 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): model_config = ConfigDict(extra="allow", protected_namespaces=()) @field_serializer("data") - def _serialize_image_data(self, data: list[OpenAIImage]) -> list[dict[str, object]]: - return [image.model_dump() for image in data] + def _serialize_image_data(self, data: list[OpenAIImage] | None) -> list[dict[str, object]] | None: + return None if data is None else [image.model_dump() for image in data] def __init__( self, From cf00ab1bf84b5aa35cbe42049466dcd00d008aca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:13:43 -0700 Subject: [PATCH 43/56] fix: rename the mainland China brand to Qianwen AI Platform --- README.md | 2 +- litellm/llms/dashscope/common_utils.py | 2 +- litellm/llms/dashscope/qwen_ai_platform.py | 2 +- .../provider_endpoints_support_backup.json | 2 +- .../provider_create_fields.json | 6 +-- provider_endpoints_support.json | 2 +- .../llms/dashscope/test_qwen_brand_aliases.py | 48 +++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 19 +++++++- .../components/provider_info_helpers.test.tsx | 7 +++ .../src/components/provider_info_helpers.tsx | 2 +- 10 files changed, 82 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 1624d408419..3eb475f121e 100644 --- a/README.md +++ b/README.md @@ -356,7 +356,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse | [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | | | [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | | | [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | | -| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | +| [Qianwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | | [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | | [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | | | [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py index 9ed9c276e43..952960e8207 100644 --- a/litellm/llms/dashscope/common_utils.py +++ b/litellm/llms/dashscope/common_utils.py @@ -103,7 +103,7 @@ def missing_dashscope_family_key_message(custom_llm_provider: str) -> str: ) if custom_llm_provider == "qwen_ai_platform": return ( - "Missing API key for Qwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or " + "Missing API key for Qianwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or " "DASHSCOPE_API_KEY environment variable or pass api_key parameter." ) return "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter." diff --git a/litellm/llms/dashscope/qwen_ai_platform.py b/litellm/llms/dashscope/qwen_ai_platform.py index 6998a57e2b7..864f0f459e4 100644 --- a/litellm/llms/dashscope/qwen_ai_platform.py +++ b/litellm/llms/dashscope/qwen_ai_platform.py @@ -23,7 +23,7 @@ def _require_qwen_ai_platform_api_key(api_key: str | None) -> str: resolved: Final = _resolve_qwen_ai_platform_api_key(api_key) if resolved is None: raise ValueError( - "Qwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var " + "Qianwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var " "or pass api_key explicitly." ) return resolved diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 30c1e0b894e..e87db5bf593 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -691,7 +691,7 @@ } }, "qwen_ai_platform": { - "display_name": "Qwen AI Platform (`qwen_ai_platform`)", + "display_name": "Qianwen AI Platform (`qwen_ai_platform`)", "url": "https://docs.litellm.ai/docs/providers/qwencloud", "endpoints": { "chat_completions": true, diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index d7e100dd630..cb2bc540b55 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1129,12 +1129,12 @@ }, { "provider": "Qwen_AI_Platform", - "provider_display_name": "Qwen AI Platform", + "provider_display_name": "Qianwen AI Platform", "litellm_provider": "qwen_ai_platform", "credential_fields": [ { "key": "api_key", - "label": "Qwen AI Platform API Key", + "label": "Qianwen AI Platform API Key", "placeholder": null, "tooltip": null, "required": true, @@ -1146,7 +1146,7 @@ "key": "api_base", "label": "API Base", "placeholder": "https://dashscope.aliyuncs.com/compatible-mode/v1", - "tooltip": "The base URL for Qwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.", + "tooltip": "The base URL for Qianwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.", "required": true, "field_type": "text", "options": null, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index af9b194bbee..c208d4edaaa 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -744,7 +744,7 @@ } }, "qwen_ai_platform": { - "display_name": "Qwen AI Platform (`qwen_ai_platform`)", + "display_name": "Qianwen AI Platform (`qwen_ai_platform`)", "url": "https://docs.litellm.ai/docs/providers/qwencloud", "endpoints": { "chat_completions": true, diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py index 7862297bcd5..6f264ea1424 100644 --- a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py +++ b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py @@ -1,10 +1,14 @@ +import json import math +import re +from pathlib import Path import pytest import litellm from litellm import completion, get_llm_provider from litellm.llms.dashscope.chat.transformation import DashScopeChatConfig +from litellm.llms.dashscope.common_utils import missing_dashscope_family_key_message from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, ) @@ -53,6 +57,7 @@ BRAND_CASES = [ pytest.param( { "provider": "qwencloud", + "display_name": "QwenCloud", "enum": LlmProviders.QWENCLOUD, "key_env": "QWENCLOUD_API_KEY", "base_env": "QWENCLOUD_API_BASE", @@ -69,6 +74,7 @@ BRAND_CASES = [ pytest.param( { "provider": "qwen_ai_platform", + "display_name": "Qianwen AI Platform", "enum": LlmProviders.QWEN_AI_PLATFORM, "key_env": "QWEN_AI_PLATFORM_API_KEY", "base_env": "QWEN_AI_PLATFORM_API_BASE", @@ -250,6 +256,48 @@ class TestQwenBrandDefaultUrls: ) +class TestQwenBrandUserFacingNames: + RETIRED_MAINLAND_NAME = "Qwen AI Platform" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_missing_key_message_names_brand(self, brand): + message = missing_dashscope_family_key_message(brand["provider"]) + assert brand["display_name"] in message + assert brand["key_env"] in message + assert self.RETIRED_MAINLAND_NAME not in message + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_embedding_without_key_names_brand(self, brand): + with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: + litellm.embedding(model=f"{brand['provider']}/text-embedding-v4", input=["hello"]) + assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_rerank_without_key_names_brand(self, brand): + with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: + litellm.rerank(model=f"{brand['provider']}/gte-rerank-v2", query="q", documents=["a", "b"]) + assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_image_generation_without_key_names_brand(self, brand): + with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: + litellm.image_generation(model=f"{brand['provider']}/qwen-image", prompt="a cup of coffee") + assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + + @pytest.mark.parametrize("brand", BRAND_CASES) + @pytest.mark.parametrize( + "matrix_path", + [ + Path(litellm.__file__).parent / "provider_endpoints_support_backup.json", + Path(litellm.__file__).parent.parent / "provider_endpoints_support.json", + ], + ids=["backup", "root"], + ) + def test_supported_endpoints_matrix_display_name(self, brand, matrix_path): + matrix = json.loads(matrix_path.read_text()) + assert matrix["providers"][brand["provider"]]["display_name"] == f"{brand['display_name']} (`{brand['provider']}`)" + + class TestQwenBrandCostParity: @pytest.fixture(autouse=True) def setup_model_cost_map(self, monkeypatch): diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 680dd4df0ae..ba88fbd3951 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1,3 +1,4 @@ +import json import re from datetime import datetime, timezone from typing import Final @@ -339,7 +340,23 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False -def test_chatgpt_provider_fields(): +def test_qwen_mainland_provider_fields_carry_the_qianwen_brand(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + providers = test_client.get("/public/providers/fields").json() + + mainland = next(p for p in providers if p["litellm_provider"] == "qwen_ai_platform") + international = next(p for p in providers if p["litellm_provider"] == "qwencloud") + + assert mainland["provider_display_name"] == "Qianwen AI Platform" + assert international["provider_display_name"] == "QwenCloud" + + mainland_fields = {f["key"]: f for f in mainland["credential_fields"]} + assert mainland_fields["api_key"]["label"] == "Qianwen AI Platform API Key" + assert "Qianwen AI Platform" in mainland_fields["api_base"]["tooltip"] + assert "Qwen AI Platform" not in json.dumps(mainland) app_instance = FastAPI() app_instance.include_router(router) test_client = TestClient(app_instance) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 4c68e302267..439a5958874 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -133,6 +133,13 @@ describe("provider_info_helpers", () => { expect(result.logo).toBeTruthy(); }); + it("should resolve the qwen_ai_platform slug and Qwen_AI_Platform enum key to the Qianwen AI Platform display name", () => { + expect(getProviderLogoAndName("qwen_ai_platform").displayName).toBe("Qianwen AI Platform"); + expect(getProviderLogoAndName("Qwen_AI_Platform").displayName).toBe("Qianwen AI Platform"); + expect(getProviderLogoAndName("qwencloud").displayName).toBe("QwenCloud"); + expect(getProviderLogoAndName("qwen_ai_platform").logo).toBe(providerLogoMap[Providers.Qwen_AI_Platform]); + }); + it("should return provider value as display name when no mapping exists", () => { const unknownProvider = "unknown_provider"; const result = getProviderLogoAndName(unknownProvider); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 1a1a2fd73aa..288c443a5ad 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -152,7 +152,7 @@ export enum Providers { PETALS = "Petals", PG_VECTOR = "Pg Vector", PREDIBASE = "Predibase", - Qwen_AI_Platform = "Qwen AI Platform", + Qwen_AI_Platform = "Qianwen AI Platform", QwenCloud = "QwenCloud", RECRAFT = "Recraft", REPLICATE = "Replicate", From 34293483059d74af7138b62aa55865454237c165 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:21:53 +0000 Subject: [PATCH 44/56] fix(types): avoid mutable image serializer annotations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8d327b337a7..c661aa0839a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2552,7 +2552,7 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): model_config = ConfigDict(extra="allow", protected_namespaces=()) @field_serializer("data") - def _serialize_image_data(self, data: list[OpenAIImage] | None) -> list[dict[str, object]] | None: + def _serialize_image_data(self, data: Sequence[OpenAIImage] | None) -> Sequence[Mapping[str, object]] | None: return None if data is None else [image.model_dump() for image in data] def __init__( From 14febcc8788d0c83094a096d612827d56925fb5f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:22:18 -0700 Subject: [PATCH 45/56] fix(rust): shrink cache configuration bridge --- .../crates/python-bridge/src/cache/config.rs | 56 +++++++++---------- .../crates/python-bridge/src/cache/facade.rs | 6 +- 2 files changed, 27 insertions(+), 35 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index bcacf7b4e36..6e694c20705 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -8,7 +8,7 @@ use pyo3::{ use super::{native::NativeResponseCache, request::duration}; -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct CachePolicy { pub(super) mode: String, pub(super) ttl: Option, @@ -18,7 +18,6 @@ pub(super) struct CachePolicy { pub(super) semantic_cache_scope: String, } -#[derive(PartialEq)] pub(super) struct MemoryCacheConfig { pub(super) default_ttl: Duration, pub(super) capacity: usize, @@ -38,7 +37,7 @@ pub(super) enum CertificateRequirement { Required, } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisTlsConfig { pub(super) certificate_requirement: CertificateRequirement, pub(super) check_hostname: bool, @@ -48,7 +47,7 @@ pub(super) struct RedisTlsConfig { pub(super) client_key: Option, } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisConnectionConfig { pub(super) host: String, pub(super) port: u16, @@ -65,7 +64,7 @@ pub(super) struct RedisConnectionConfig { pub(super) tls: Option, } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisCacheConfig { pub(super) default_ttl: Duration, pub(super) namespace: Option, @@ -73,13 +72,12 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } -#[derive(PartialEq)] pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct NativeCacheConfig { pub(super) policy: CachePolicy, pub(super) backend: CacheBackendConfig, @@ -261,7 +259,7 @@ fn project_redis( return Ok(Err(UnsupportedCacheConfig::RedisConnection)); }; - let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { + let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, 3 => RedisProtocol::Resp3, _ => return Err(PyValueError::new_err("unsupported Redis protocol version")), @@ -274,7 +272,8 @@ fn project_redis( flush_size: backend.getattr("redis_flush_size")?.extract::()?, connection: RedisConnectionConfig { host: required_string(&resolved, "host")?, - port: required_u16(&resolved, "port")?, + 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")?, @@ -320,14 +319,20 @@ fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult Ok(CertificateRequirement::None), - "optional" | "cert_optional" => Ok(CertificateRequirement::Optional), - "required" | "cert_required" => Ok(CertificateRequirement::Required), - _ => Err(PyValueError::new_err( - "invalid Redis TLS certificate requirement", - )), + let text = value.str()?; + let text = text.to_str()?; + if text.eq_ignore_ascii_case("none") || text.eq_ignore_ascii_case("cert_none") { + return Ok(CertificateRequirement::None); } + if text.eq_ignore_ascii_case("optional") || text.eq_ignore_ascii_case("cert_optional") { + return Ok(CertificateRequirement::Optional); + } + if text.eq_ignore_ascii_case("required") || text.eq_ignore_ascii_case("cert_required") { + return Ok(CertificateRequirement::Required); + } + Err(PyValueError::new_err( + "invalid Redis TLS certificate requirement", + )) } #[inline(never)] @@ -386,11 +391,11 @@ fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { } #[inline(never)] -fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { +fn required_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? .ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))? - .extract::() + .extract::() } #[inline(never)] @@ -417,14 +422,6 @@ fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> } } -#[inline(never)] -fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { - match values.get_item(key)? { - Some(value) => value.extract::>(), - None => Ok(None), - } -} - #[inline(never)] fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { @@ -442,10 +439,9 @@ fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult() { - return Ok(Some(matches!( - text.to_ascii_lowercase().as_str(), - "true" | "1" | "yes" - ))); + return Ok(Some( + text == "1" || text.eq_ignore_ascii_case("true") || text.eq_ignore_ascii_case("yes"), + )); } value.extract::().map(Some) } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 6f54d6a121a..83508356263 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -28,7 +28,6 @@ struct ObjectGuard { pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, - config: NativeCacheConfig, } impl ObjectGuard { @@ -191,15 +190,12 @@ impl FacadeGuard { "redis_flush_size", ], )?, - config, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { - let projected = NativeCacheConfig::project(facade)?; Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)? - && matches!(projected, CacheConfigProjection::Native(config) if *config == self.config)) + && self.backend.matches(py, &facade.getattr("cache")?)?) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { From a835e75620b1e3fb8c5e11d39c70fc73f17feb04 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:24:19 -0700 Subject: [PATCH 46/56] fix(mcp): handle split UTF-8 routing previews in place --- .../proxy/_experimental/mcp_server/server.py | 20 +- .../mcp_server/test_mcp_server.py | 173 ++++++++---------- 2 files changed, 78 insertions(+), 115 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e0fb7cede03..397a82cfa45 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -277,24 +277,6 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: return False -def _utf8_boundary_prefix(data: bytes) -> bytes: - """``data`` with any trailing incomplete UTF-8 sequence removed. - - Cutting a body at a fixed byte budget can land in the middle of a multibyte - character, and ``json.loads`` on such bytes raises ``UnicodeDecodeError`` - rather than ``JSONDecodeError``. Trimming to a character boundary keeps the - truncated peek decodable so callers only have to handle malformed JSON. - """ - for trailing in range(0, min(3, len(data)) + 1): - candidate = data[: len(data) - trailing] - try: - candidate.decode("utf-8") - except UnicodeDecodeError: - continue - return candidate - return data - - def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. @@ -4057,7 +4039,7 @@ if MCP_AVAILABLE: # directly from the original `receive` via wrapped_receive. break - return consumed_messages, _utf8_boundary_prefix(b"".join(body_chunks)) + return consumed_messages, b"".join(body_chunks) async def _handle_stale_mcp_session( scope: Scope, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index aff3a2f0483..b2eded67430 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1992,6 +1992,7 @@ async def test_streamable_http_session_manager_is_stateless(): ( ("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True), ("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False), + ("POST", b"", False), ("GET", b"", False), ("DELETE", b"", False), ), @@ -2466,106 +2467,65 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): @pytest.mark.asyncio -async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap(): - """ - A tool-call POST whose UTF-8 body is larger than the routing peek cap, with a - multibyte character straddling the cap boundary, must still be forwarded - intact instead of blowing up with a UnicodeDecodeError 500. +@pytest.mark.parametrize("method", ("initialize", "tools/call")) +@pytest.mark.parametrize("chunked", (False, True)) +@pytest.mark.parametrize( + ("character", "bytes_before_cap"), + (("é", 0), ("é", 1), ("中", 1), ("中", 2), ("😀", 1), ("😀", 2), ("😀", 3)), +) +async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap( + method: str, chunked: bool, character: str, bytes_before_cap: int +) -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_module - Regression test for https://github.com/BerriAI/litellm/issues/34917 - """ - try: - from litellm.proxy._experimental.mcp_server import server as mcp_server - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - session_manager_stateful, - session_manager_stateless, - ) - except ImportError: - pytest.skip("MCP server not available") + params: Final = ( + { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {}, + "clientInfo": {"name": "<>", "version": "1"}, + } + if method == "initialize" + else {"name": "update_full_document", "arguments": {"markdown": "<>"}} + ) + template: Final = json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode() + prefix, suffix = template.split(b"<>") + cap: Final = mcp_module._MCP_ROUTING_PEEK_MAX_BYTES + body: Final = prefix + b"x" * (cap - bytes_before_cap - len(prefix)) + character.encode() + b"tail" + suffix + chunks: Final = (body[: cap - 1], body[cap - 1 : cap], body[cap:]) if chunked else (body,) + messages: Final[tuple[Message, ...]] = tuple( + {"type": "http.request", "body": chunk, "more_body": index < len(chunks) - 1} + for index, chunk in enumerate(chunks) + ) + receive: Final = AsyncMock(side_effect=messages) + send: Final = AsyncMock() + received: Final[asyncio.Future[bytes]] = asyncio.get_running_loop().create_future() - peek_cap = mcp_server._MCP_ROUTING_PEEK_MAX_BYTES - - def _splits_multibyte_at_cap(candidate: bytes) -> bool: - try: - candidate[:peek_cap].decode("utf-8") - except UnicodeDecodeError: - return True - return False - - def _build_body() -> bytes: - for pad in range(4): - candidate = json.dumps( - { - "jsonrpc": "2.0", - "id": 1, - "method": "tools/call", - "params": { - "name": "update_full_document" + "x" * pad, - "arguments": {"markdown": "щ" * 3000}, - }, - }, - ensure_ascii=False, - ).encode("utf-8") - if len(candidate) > peek_cap and _splits_multibyte_at_cap(candidate): - return candidate - raise AssertionError("could not build a body splitting a multibyte char at the peek cap") - - body = _build_body() - - messages = [{"type": "http.request", "body": body, "more_body": False}] - receive_calls = {"count": 0} - - async def receive(): - idx = receive_calls["count"] - receive_calls["count"] += 1 - return messages[idx] - - scope = { - "type": "http", - "method": "POST", - "path": "/mcp/progress_test", - "headers": [ - (b"content-type", b"application/json"), - (b"authorization", b"Bearer test-key"), - ], - } - send = AsyncMock() - - streamed_chunks = [] - - async def stateless_handle(s, r, se): - while True: - msg = await r() - if msg.get("type") != "http.request": - break - streamed_chunks.append(msg.get("body", b"") or b"") - if not msg.get("more_body", False): - break - - async def stateful_handle(s, r, se): - raise AssertionError("non-initialize POST should not reach stateful manager") + async def handle_request(_: Scope, downstream_receive: Receive, outgoing: Send) -> None: + assert receive.await_count == (2 if chunked else 1) + received.set_result(await _drain_body(downstream_receive)) + await outgoing({"type": "http.response.start", "status": 200, "headers": []}) + await outgoing({"type": "http.response.body", "body": b"{}"}) + stateless_handle: Final = AsyncMock(side_effect=handle_request) + stateful_handle: Final = AsyncMock() + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} with ( + _client_allowlist_patches({}, None), patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), ), - patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), patch( - "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", - True, + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), ), - patch.object(session_manager_stateless, "handle_request", side_effect=stateless_handle), - patch.object(session_manager_stateful, "handle_request", side_effect=stateful_handle), - patch.object(session_manager_stateless, "_server_instances", {}), - patch.object(session_manager_stateful, "_server_instances", {}), ): - await handle_streamable_http_mcp(scope, receive, send) + await mcp_module.handle_streamable_http_mcp(scope, receive, send) - assert send.await_count == 0, f"unexpected response emitted by the proxy: {send.await_args_list}" - assert b"".join(streamed_chunks) == body + assert send.call_args_list[0].args[0]["status"] == 200 + assert received.result() == body + stateless_handle.assert_awaited_once() + stateful_handle.assert_not_awaited() @pytest.mark.asyncio @@ -4119,7 +4079,12 @@ def test_jsonrpc_text_has_top_level_method_ignores_nested_method(): @pytest.mark.asyncio -async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): +@pytest.mark.parametrize("response_field", ("result", "error")) +@pytest.mark.parametrize(("character", "bytes_before_cap"), (("", 0), ("x", 0), ("é", 1), ("中", 2), ("😀", 3))) +@pytest.mark.parametrize("cancel_request", (False, True)) +async def test_truncated_jsonrpc_response_with_nested_method_skips_lock( + response_field: str, character: str, bytes_before_cap: int, cancel_request: bool +) -> None: """Regression: a large JSON-RPC *response* POST whose ``result`` payload nests a ``method`` key must skip the per-session lock so it does not deadlock behind the in-flight request POST that is holding the lock while @@ -4147,7 +4112,7 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): async def handle(s, r, se): msg = await r() body = msg.get("body", b"") or b"" - if b'"result"' in body: + if body == response_body: response_handled.set() else: request_in_handle.set() @@ -4174,9 +4139,16 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): # A JSON-RPC response larger than the routing peek cap so it can't be fully # parsed, with a nested "method" key in the first bytes to trip a flat # substring heuristic. - response_body = ( - '{"jsonrpc":"2.0","id":99,"result":{"toolResult":{"method":"GET","payload":"' + ("x" * 5000) + '"}}}' + response_prefix: Final = ( + '{"jsonrpc":"2.0","id":99,"' + response_field + + '":{"code":-32000,"message":"test","data":{"method":"GET","payload":"' ).encode() + response_body: Final = ( + response_prefix + + b"x" * (mcp_server._MCP_ROUTING_PEEK_MAX_BYTES - bytes_before_cap - len(response_prefix) if character else 0) + + character.encode() + + b'tail"}}}' + ) try: with ( @@ -4204,8 +4176,17 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): # lock held by req_task and this wait would time out (deadlock). await asyncio.wait_for(response_handled.wait(), timeout=1.0) - gate.set() - await asyncio.gather(req_task, resp_task) + await resp_task + assert not req_task.done() + if cancel_request: + req_task.cancel() + with pytest.raises(asyncio.CancelledError): + await req_task + else: + gate.set() + await req_task + assert not mcp_server._stateful_session_locks[session_id].locked() + assert session_id not in mcp_server._stateful_session_active_request_counts finally: gate.set() mcp_server._stateful_session_auth_contexts.pop(session_id, None) From efcafa7f120341187991bd2655cdec3257459b3b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:30:17 -0700 Subject: [PATCH 47/56] fix(rust): invalidate changed Redis pool settings --- .../crates/python-bridge/src/cache/facade.rs | 69 ++++++++++++++++++- tests/test_litellm_rust/test_cache.py | 4 ++ 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 83508356263..f2f86c14b37 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -25,9 +25,17 @@ struct ObjectGuard { config: Vec, } +struct RedisPoolGuard { + reference: Py, + connection_class: Py, + connection_kwargs: Py, + max_connections: usize, +} + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, + redis_pool: Option, } impl ObjectGuard { @@ -129,6 +137,45 @@ impl ObjectGuard { } } +impl RedisPoolGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + let pool = backend + .getattr("redis_client")? + .getattr("connection_pool")?; + Ok(Self { + reference: pool.clone().unbind(), + connection_class: pool.getattr("connection_class")?.unbind(), + connection_kwargs: pool + .getattr("connection_kwargs")? + .call_method0("copy")? + .unbind(), + max_connections: pool.getattr("max_connections")?.extract::()?, + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + let pool = backend + .getattr("redis_client")? + .getattr("connection_pool")?; + Ok(self.reference.bind(py).is(&pool) + && self + .connection_class + .bind(py) + .is(&pool.getattr("connection_class")?) + && self.max_connections == pool.getattr("max_connections")?.extract::()? + && self + .connection_kwargs + .bind(py) + .eq(pool.getattr("connection_kwargs")?)?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference)?; + visit.call(&self.connection_class)?; + visit.call(&self.connection_kwargs) + } +} + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -190,17 +237,33 @@ impl FacadeGuard { "redis_flush_size", ], )?, + redis_pool: (kind == "redis") + .then(|| RedisPoolGuard::capture(&backend)) + .transpose()?, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { - Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)?) + if !self.outer.matches(py, facade)? { + return Ok(false); + } + let backend = facade.getattr("cache")?; + if !self.backend.matches(py, &backend)? { + return Ok(false); + } + match &self.redis_pool { + Some(guard) => guard.matches(py, &backend), + None => Ok(true), + } } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.outer.traverse(&visit)?; - self.backend.traverse(&visit) + self.backend.traverse(&visit)?; + if let Some(guard) = &self.redis_pool { + guard.traverse(&visit)?; + } + Ok(()) } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 796a0ec36ac..c35cb1a20fb 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -381,6 +381,10 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + pool: Final = facade.cache.redis_client.connection_pool + with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None await binding.async_store(request("second"), {"value": 2}) From 2da7e6dfe20febb05e7c00302f9fd29ccb24dc8f Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:31:18 -0700 Subject: [PATCH 48/56] ci: trigger MCP fix checks against main From cbf5bb5e719b58915989279c796b410b7668400e Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 19:31:29 +0000 Subject: [PATCH 49/56] chore(prices): sync OpenRouter prices: 1 model openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index aa67dc58d69..d100fe62e72 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -42971,21 +42971,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.22722e-07, + "input_cost_per_token": 9.19242e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.845444e-06, + "output_cost_per_token": 1.838484e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.68935e-08, + "cache_read_input_token_cost": 7.66035e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index aa67dc58d69..d100fe62e72 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -42971,21 +42971,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.22722e-07, + "input_cost_per_token": 9.19242e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.845444e-06, + "output_cost_per_token": 1.838484e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.68935e-08, + "cache_read_input_token_cost": 7.66035e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From f851e6ddb7f7e9d6009079dc1af1498f7878e382 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 19:31:34 +0000 Subject: [PATCH 50/56] chore(prices): sync AWS Bedrock prices: 3 models [enrichment failed: AWS Bedrock, 18 held] us.moonshotai.kimi-k3: zai.glm-4.7: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_audio_input, supports_response_schema zai.glm-4.7-flash: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_audio_input, supports_response_schema --- ...odel_prices_and_context_window_backup.json | 24 ++++++++++++------- model_prices_and_context_window.json | 24 ++++++++++++------- 2 files changed, 30 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index aa67dc58d69..91a4291171a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54441,16 +54441,19 @@ "zai.glm-4.7": { "input_cost_per_token": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 203000, + "max_output_tokens": 4000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 2.2e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "zai.glm-5": { "input_cost_per_token": 1e-06, @@ -54470,16 +54473,19 @@ "zai.glm-4.7-flash": { "input_cost_per_token": 7e-08, "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 203000, + "max_output_tokens": 4000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 4e-07, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "zai/glm-5": { "cache_creation_input_token_cost": 0, @@ -76803,7 +76809,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.65e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3.html", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_audio_input": false, "supports_function_calling": true, "supports_prompt_caching": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index aa67dc58d69..91a4291171a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54441,16 +54441,19 @@ "zai.glm-4.7": { "input_cost_per_token": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 203000, + "max_output_tokens": 4000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 2.2e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "zai.glm-5": { "input_cost_per_token": 1e-06, @@ -54470,16 +54473,19 @@ "zai.glm-4.7-flash": { "input_cost_per_token": 7e-08, "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 203000, + "max_output_tokens": 4000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 4e-07, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "zai/glm-5": { "cache_creation_input_token_cost": 0, @@ -76803,7 +76809,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.65e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3.html", + "source": "https://aws.amazon.com/bedrock/pricing/", "supports_audio_input": false, "supports_function_calling": true, "supports_prompt_caching": true, From 991108cc4865efa649d0871966c275dc2593125e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:50:21 -0700 Subject: [PATCH 51/56] refactor(rust): align cache foundation with Python backends --- litellm-rust/crates/cache-memory/src/cache.rs | 113 ++++++-- .../crates/cache-memory/tests/cache.rs | 64 +++-- litellm-rust/crates/cache-redis/src/cache.rs | 243 ++++++++++------- .../cache-redis/src/cache/operations.rs | 137 +++++++++- litellm-rust/crates/cache-redis/src/lib.rs | 6 +- .../crates/cache-redis/src/topology.rs | 14 + .../crates/cache-redis/tests/cache.rs | 68 +++-- .../crates/cache-response/src/buffer.rs | 4 +- .../crates/cache-response/src/response.rs | 58 ++-- .../crates/cache-response/tests/response.rs | 4 +- litellm-rust/crates/cache/src/base_cache.rs | 106 +++----- litellm-rust/crates/cache/src/cache_type.rs | 85 ++++++ litellm-rust/crates/cache/src/caching.rs | 10 +- litellm-rust/crates/cache/src/capabilities.rs | 143 +++++++++- litellm-rust/crates/cache/src/dual.rs | 255 +++++++++--------- litellm-rust/crates/cache/src/lib.rs | 13 +- litellm-rust/crates/cache/tests/caching.rs | 105 ++++++-- litellm-rust/crates/cache/tests/dual.rs | 215 +++++++++------ .../crates/python-bridge/src/cache/config.rs | 22 +- .../crates/python-bridge/src/cache/native.rs | 2 +- .../crates/python-bridge/src/cache/request.rs | 2 +- 21 files changed, 1136 insertions(+), 533 deletions(-) create mode 100644 litellm-rust/crates/cache-redis/src/topology.rs create mode 100644 litellm-rust/crates/cache/src/cache_type.rs diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 54d831378e4..85850c1d925 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -1,13 +1,14 @@ use std::{ cmp::Reverse, - collections::{BinaryHeap, HashMap}, + collections::{BinaryHeap, HashMap, HashSet}, + hash::Hash, sync::{Arc, Mutex}, time::{Duration, SystemTime, UNIX_EPOCH}, }; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, - Error, IncrementOperation, + BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache, + DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache, }; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; @@ -219,7 +220,7 @@ where key: &str, candidate: V, eligible: &[V], - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result { if self.max_size_in_memory == 0 { return Ok(candidate); @@ -239,14 +240,23 @@ where return Ok(existing.clone()); } let winner = existing.unwrap_or(candidate); - Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs)); + Self::set_expiration( + &mut state, + key, + now + self.get_ttl(&context).unwrap_or(self.default_ttl), + ); state.values.insert(key.into(), winner.clone()); Ok(winner) } } impl CounterCache for InMemoryCache { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { if self.max_size_in_memory == 0 { return Ok(amount); } @@ -255,7 +265,11 @@ impl CounterCache for InMemoryCache { Self::evict(&mut state, self.max_size_in_memory, now, key); let value = state.values.get(key).copied().unwrap_or_default() + amount; if !state.expirations.contains_key(key) { - Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs)); + Self::set_expiration( + &mut state, + key, + now + self.get_ttl(&context).unwrap_or(self.default_ttl), + ); } state.values.insert(key.into(), value); Ok(value) @@ -273,10 +287,7 @@ impl InMemoryCache { self.increment_cache( &operation.key, operation.amount, - CacheKwargs { - ttl: operation.ttl, - ..CacheKwargs::default() - }, + ExactCacheContext { ttl: operation.ttl }, ) }) .collect() @@ -285,28 +296,26 @@ impl InMemoryCache { impl BaseCache for InMemoryCache { type Value = V; + type Context = ExactCacheContext; - fn default_ttl(&self) -> Duration { - self.default_ttl + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl.or(Some(self.default_ttl)) } - fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { - let ttl = self.get_ttl(&kwargs); + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &ExactCacheContext, + ) -> Result<(), Error> { + let ttl = self.get_ttl(context).unwrap_or(self.default_ttl); self.set_cache(key, value, Some(ttl)).map(|_| ()) } - fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { self.get_cache(key) } - fn delete_cache(&self, key: &str) -> Result<(), Error> { - self.delete_cache(key) - } - - fn flush_cache(&self) -> Result<(), Error> { - self.flush_cache() - } - async fn disconnect(&self) -> Result<(), Error> { Ok(()) } @@ -320,6 +329,60 @@ impl BaseCache for InMemoryCache { } } +impl BatchCache for InMemoryCache {} + +impl DeleteCache for InMemoryCache { + fn delete_cache(&self, key: &str) -> Result<(), Error> { + InMemoryCache::delete_cache(self, key) + } +} + +impl FlushCache for InMemoryCache { + fn flush_cache(&self) -> Result<(), Error> { + InMemoryCache::flush_cache(self) + } +} + +impl TtlCache for InMemoryCache { + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + InMemoryCache::async_get_ttl(self, key).await + } +} + +impl SetCache for InMemoryCache> +where + T: Clone + Eq + Hash + Send + Sync + 'static, +{ + type SetValue = T; + type SetResult = Vec; + + async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result { + if self.max_size_in_memory == 0 { + return Ok(values); + } + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now, key); + let mut stored = state.values.get(key).cloned().unwrap_or_default(); + stored.extend(values.iter().cloned()); + if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value) + && measure(&stored)? > limit + { + return Ok(values); + } + if !state.expirations.contains_key(key) { + Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl)); + } + state.values.insert(key.into(), stored); + Ok(values) + } +} + #[cfg(test)] mod tests { use super::*; @@ -329,7 +392,7 @@ mod tests { let cache = InMemoryCache::::new(Some(4), None); for _ in 0..100 { cache - .increment_cache("counter", 1.0, CacheKwargs::default()) + .increment_cache("counter", 1.0, ExactCacheContext::default()) .unwrap(); } assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1); diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index c44590b63ca..0df0319b990 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashSet, sync::{ Arc, atomic::{AtomicU64, Ordering}, @@ -7,8 +8,8 @@ use std::{ }; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, - IncrementOperation, get_cache, set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error, + ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -152,39 +153,38 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() { let clock = clock(); let cache: CacheBackend> = Arc::new(cache(clock.clone(), 4)); let reader = Arc::clone(&cache); - let kwargs = CacheKwargs { + let context = ExactCacheContext { ttl: Some(Duration::from_secs(5)), - ..Default::default() }; - set_cache(cache.as_ref(), "sync", "first".into(), kwargs.clone()).unwrap(); + set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap(); assert_eq!( - get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), + get_cache(reader.as_ref(), "sync", &context).unwrap(), Some("first".into()) ); cache - .batch_cache_write("async", "second".into(), kwargs.clone()) + .batch_cache_write("async", "second".into(), context.clone()) .await .unwrap(); cache - .async_set_cache_pipeline(vec![("batch".into(), "third".into())], kwargs.clone()) + .async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone()) .await .unwrap(); drop(cache); for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] { assert_eq!( - reader.async_get_cache(key, &kwargs).await.unwrap(), + reader.async_get_cache(key, &context).await.unwrap(), Some(value.into()) ); } reader.async_delete_cache("async").await.unwrap(); assert_eq!( - reader.async_get_cache("async", &kwargs).await.unwrap(), + reader.async_get_cache("async", &context).await.unwrap(), None ); clock.store(106, Ordering::SeqCst); - assert_eq!(get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), None); + assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None); assert_eq!( - reader.async_get_cache("batch", &kwargs).await.unwrap(), + reader.async_get_cache("batch", &context).await.unwrap(), None ); } @@ -196,20 +196,19 @@ fn claims_are_atomic_and_refresh_eligible_winners() { let clock = clock.clone(); move || Duration::from_secs(clock.load(Ordering::SeqCst)) }); - let kwargs = CacheKwargs { + let context = ExactCacheContext { ttl: Some(Duration::from_secs(10)), - ..Default::default() }; assert_eq!( cache - .claim_cache("affinity", "first".to_string(), &[], kwargs.clone()) + .claim_cache("affinity", "first".to_string(), &[], context.clone()) .unwrap(), "first" ); clock.store(103, Ordering::SeqCst); assert_eq!( cache - .claim_cache("affinity", "second".to_string(), &[], kwargs.clone()) + .claim_cache("affinity", "second".to_string(), &[], context.clone()) .unwrap(), "first" ); @@ -224,7 +223,7 @@ fn claims_are_atomic_and_refresh_eligible_winners() { "affinity", "second".to_string(), &["first".to_string(), "second".to_string()], - kwargs, + context, ) .unwrap(), "first" @@ -239,11 +238,13 @@ fn claims_are_atomic_and_refresh_eligible_winners() { fn counters_increment_under_one_lock() { let cache = InMemoryCache::::default(); assert_eq!( - CounterCache::increment_cache(&cache, "counter", 1.5, CacheKwargs::default()).unwrap(), + CounterCache::increment_cache(&cache, "counter", 1.5, ExactCacheContext::default()) + .unwrap(), 1.5 ); assert_eq!( - CounterCache::increment_cache(&cache, "counter", 2.0, CacheKwargs::default()).unwrap(), + CounterCache::increment_cache(&cache, "counter", 2.0, ExactCacheContext::default()) + .unwrap(), 3.5 ); } @@ -263,7 +264,7 @@ fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc::new(Some(2), None); for key in ["a", "b", "a", "b"] { cache - .increment_cache(key, 1.0, CacheKwargs::default()) + .increment_cache(key, 1.0, ExactCacheContext::default()) .unwrap(); } assert_eq!(cache.get_cache("a").unwrap(), Some(2.0)); @@ -290,7 +291,7 @@ fn disabled_cache_does_not_retain_claims_or_counters() { let claims = InMemoryCache::::new(Some(0), None); assert_eq!( claims - .claim_cache("key", "first".into(), &[], CacheKwargs::default()) + .claim_cache("key", "first".into(), &[], ExactCacheContext::default()) .unwrap(), "first" ); @@ -299,7 +300,7 @@ fn disabled_cache_does_not_retain_claims_or_counters() { let counters = InMemoryCache::::new(Some(0), None); assert_eq!( counters - .increment_cache("key", 2.0, CacheKwargs::default()) + .increment_cache("key", 2.0, ExactCacheContext::default()) .unwrap(), 2.0 ); @@ -348,3 +349,20 @@ async fn increment_pipeline_preserves_operation_order() { ); assert_eq!(cache.get_cache("a").unwrap(), Some(3.0)); } + +#[tokio::test] +async fn set_capability_preserves_python_result_and_deduplicates_storage() { + let cache = InMemoryCache::>::new(None, None); + let inserted = vec!["a".into(), "a".into(), "b".into()]; + assert_eq!( + cache + .async_set_cache_sadd("members", inserted.clone(), None) + .await + .unwrap(), + inserted + ); + assert_eq!( + cache.get_cache("members").unwrap(), + Some(HashSet::from(["a".into(), "b".into()])) + ); +} diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 2249966dd79..a960c383bf4 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -4,14 +4,16 @@ use std::{ }; use litellm_cache::{ - BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, - ClaimCache, CounterCache, Error, + BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, + ClaimCache, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, }; use redis::Commands; mod operations; -pub use operations::{RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; +pub use operations::{ + RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, +}; const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); @@ -179,12 +181,7 @@ where } fn namespaced_key(&self, key: &str) -> String { - match &self.namespace { - Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { - format!("{namespace}:{key}") - } - _ => key.into(), - } + namespaced_key(self.namespace.as_deref(), key) } fn namespaced_pattern(&self) -> Result { @@ -260,20 +257,35 @@ where } } +fn namespaced_key(namespace: Option<&str>, key: &str) -> String { + match namespace { + Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { + format!("{namespace}:{key}") + } + _ => key.into(), + } +} + impl BaseCache for RedisCache where S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { type Value = S::Value; + type Context = ExactCacheContext; - fn default_ttl(&self) -> Duration { - self.default_ttl + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl.or(Some(self.default_ttl)) } - fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &ExactCacheContext, + ) -> Result<(), Error> { let payload = self.codec.encode(&value)?; - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(context).unwrap_or(self.default_ttl)); let key = self.namespaced_key(key); self.connections.execute(|connection| { connection @@ -282,7 +294,7 @@ where }) } - fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { let key = self.namespaced_key(key); let value = self.connections.execute(|connection| { connection @@ -292,48 +304,15 @@ where self.decode_response(value) } - fn get_cache_batch( - &self, - keys: &[String], - _: &CacheKwargs, - ) -> Result>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = self.connections.execute(|connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(connection) - .map_err(|_| Error::Unavailable) - })?; - values - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } - - fn delete_cache(&self, key: &str) -> Result<(), Error> { - let key = self.namespaced_key(key); - self.connections - .execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) - } - - fn flush_cache(&self) -> Result<(), Error> { - let pattern = self.namespaced_pattern()?; - self.connections - .execute(|connection| Self::flush_matching(connection, &pattern)) - } - async fn async_set_cache( &self, key: &str, value: Self::Value, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result<(), Error> { let payload = self.codec.encode(&value)?; let key = self.namespaced_key(key); - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) @@ -345,7 +324,7 @@ where async fn async_get_cache( &self, key: &str, - _: &CacheKwargs, + _: &ExactCacheContext, ) -> Result, Error> { let key = self.namespaced_key(key); let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { @@ -357,32 +336,10 @@ where self.decode_response(value) } - async fn async_get_cache_batch( - &self, - keys: Vec, - _: CacheKwargs, - ) -> Result>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - values - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } - async fn async_set_cache_pipeline( &self, cache_list: Vec<(String, Self::Value)>, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result<(), Error> { let entries = cache_list .into_iter() @@ -392,7 +349,7 @@ where .map(|payload| (self.namespaced_key(&key), payload)) }) .collect::, _>>()?; - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); Self::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); for (key, payload) in entries { @@ -410,22 +367,6 @@ where .await } - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { - connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) - }) - .await - } - - async fn async_flush_cache(&self) -> Result<(), Error> { - let pattern = self.namespaced_pattern()?; - Self::run_blocking(Arc::clone(&self.connections), move |connection| { - Self::flush_matching(connection, &pattern) - }) - .await - } - async fn disconnect(&self) -> Result<(), Error> { Ok(()) } @@ -457,26 +398,120 @@ where } } +impl BatchCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn batch_get_cache( + &self, + keys: &[String], + _: &ExactCacheContext, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = self.connections.execute(|connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + })?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } + + async fn async_batch_get_cache( + &self, + keys: Vec, + _: ExactCacheContext, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } +} + +impl DeleteCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn delete_cache(&self, key: &str) -> Result<(), Error> { + let key = self.namespaced_key(key); + self.connections + .execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + let key = self.namespaced_key(key); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) + }) + .await + } +} + +impl FlushCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn flush_cache(&self) -> Result<(), Error> { + let pattern = self.namespaced_pattern()?; + self.connections + .execute(|connection| Self::flush_matching(connection, &pattern)) + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + let pattern = self.namespaced_pattern()?; + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Self::flush_matching(connection, &pattern) + }) + .await + } +} + impl CounterCache for RedisCache where S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { let key = self.namespaced_key(key); - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); self.connections .execute(|connection| increment(connection, key, amount, ttl)) } - async fn async_increment_cache( + async fn async_increment( &self, key: &str, amount: f64, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result { let key = self.namespaced_key(key); - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); Self::run_blocking(Arc::clone(&self.connections), move |connection| { increment(connection, key, amount, ttl) }) @@ -570,10 +605,10 @@ where key: &str, candidate: S::Value, eligible: &[S::Value], - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result { let key = self.namespaced_key(key); - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); self.connections .execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl)) } @@ -583,10 +618,10 @@ where key: &str, candidate: S::Value, eligible: Vec, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result { let key = self.namespaced_key(key); - let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); let codec = self.codec.clone(); Self::run_blocking(Arc::clone(&self.connections), move |connection| { claim(connection, &codec, &key, candidate, &eligible, ttl) @@ -599,7 +634,9 @@ where mod tests { use std::time::Duration; - use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec}; + use litellm_cache::{ + BaseCache, CacheCodec, DeleteCache, ExactCacheContext, FlushCache, JsonCodec, + }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; @@ -648,10 +685,12 @@ mod tests { .with_namespace(Some("litellm-cache".into())); cache - .set_cache("key", value.clone(), CacheKwargs::default()) + .set_cache("key", value.clone(), &ExactCacheContext::default()) .unwrap(); assert_eq!( - cache.get_cache("key", &CacheKwargs::default()).unwrap(), + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), Some(value) ); cache.delete_cache("key").unwrap(); diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index f8c2bf4078c..d8d9ae24c4c 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -1,9 +1,12 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, Error, IncrementOperation}; +use litellm_cache::{ + CacheCodec, CacheScript, ClientInfoCache, Error, IncrementOperation, QueueCache, ScanCache, + ScriptCache, SetCache, TtlCache, +}; use redis::Commands; -use super::{ConnectionRef, RedisCache}; +use super::{ConnectionRef, Connections, RedisCache, namespaced_key}; const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", @@ -88,6 +91,46 @@ pub enum RedisLpopResult { Values(Vec>), } +pub struct RedisScript { + connections: Arc>, + namespace: Option, + source: String, +} + +impl CacheScript for RedisScript +where + C: redis::ConnectionLike + Send + 'static, +{ + type Argument = RedisArg; + type Output = redis::Value; + + async fn invoke( + &self, + keys: Vec, + arguments: Vec, + ) -> Result { + let keys = keys + .into_iter() + .map(|key| namespaced_key(self.namespace.as_deref(), &key)) + .collect::>(); + let connections = Arc::clone(&self.connections); + let source = self.source.clone(); + tokio::task::spawn_blocking(move || { + connections.execute(|connection| { + redis::cmd("EVAL") + .arg(source) + .arg(keys.len()) + .arg(keys) + .arg(arguments) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + }) + .await + .map_err(|_| Error::Unavailable)? + } +} + impl RedisCache where S: CacheCodec, @@ -498,3 +541,93 @@ fn increment_with_floor( .query(connection) .map_err(|_| Error::Unavailable) } + +impl TtlCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + RedisCache::async_get_ttl(self, key) + .await + .map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64))) + } +} + +impl ScanCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { + RedisCache::async_scan_iter(self, pattern, count).await + } +} + +impl ClientInfoCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type ClientList = String; + type Info = String; + + fn client_list(&self) -> Result { + RedisCache::client_list(self) + } + + fn info(&self) -> Result { + RedisCache::info(self) + } +} + +impl SetCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type SetValue = RedisArg; + type SetResult = usize; + + async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result { + RedisCache::async_set_cache_sadd(self, key, values, ttl).await + } +} + +impl QueueCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type QueueValue = RedisArg; + type PopResult = RedisLpopResult; + + async fn async_rpush(&self, key: &str, values: Vec) -> Result { + RedisCache::async_rpush(self, key, values).await + } + + async fn async_lpop(&self, key: &str, count: Option) -> Result { + RedisCache::async_lpop(self, key, count).await + } +} + +impl ScriptCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type Script = RedisScript; + + fn async_register_script(&self, source: String) -> Self::Script { + RedisScript { + connections: Arc::clone(&self.connections), + namespace: self.namespace.clone(), + source, + } + } +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 2548c7ac3c6..98f6bfd8ce5 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -1,3 +1,7 @@ mod cache; +mod topology; -pub use cache::{RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; +pub use cache::{ + RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, +}; +pub use topology::{RedisNode, RedisTopology}; diff --git a/litellm-rust/crates/cache-redis/src/topology.rs b/litellm-rust/crates/cache-redis/src/topology.rs new file mode 100644 index 00000000000..7f4ee48b222 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/topology.rs @@ -0,0 +1,14 @@ +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RedisNode { + pub host: String, + pub port: u16, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum RedisTopology { + #[default] + Standalone, + Cluster { + startup_nodes: Vec, + }, +} diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 3d755d0f44c..337f27984f8 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,8 +1,9 @@ use std::time::Duration; use litellm_cache::{ - BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache, - CounterCache, Error, IncrementOperation, JsonCodec, get_cache, set_cache, + BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheScript, ClaimCache, + CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec, + ScriptCache, get_cache, set_cache, }; use litellm_cache_redis::{ RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, @@ -48,12 +49,11 @@ fn generic_helpers_use_the_injected_codec_and_ttl() { ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); - let kwargs = CacheKwargs { + let context = ExactCacheContext { ttl: Some(Duration::from_millis(1500)), - ..Default::default() }; - set_cache(&cache, "counter", 7, kwargs.clone()).unwrap(); - assert_eq!(get_cache(&cache, "counter", &kwargs).unwrap(), Some(7)); + set_cache(&cache, "counter", 7, &context).unwrap(); + assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7)); } #[tokio::test] @@ -83,28 +83,27 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() { Some(Duration::from_secs(9)), TaggedByteCodec(42), ); - let kwargs = CacheKwargs::default(); + let context = ExactCacheContext::default(); cache - .batch_cache_write("counter", 7, kwargs.clone()) + .batch_cache_write("counter", 7, context.clone()) .await .unwrap(); assert_eq!( - cache.async_get_cache("counter", &kwargs).await.unwrap(), + cache.async_get_cache("counter", &context).await.unwrap(), Some(7) ); cache .async_set_cache_pipeline( vec![("batch".into(), 8)], - CacheKwargs { + ExactCacheContext { ttl: Some(Duration::from_millis(1500)), - ..Default::default() }, ) .await .unwrap(); cache.async_delete_cache("counter").await.unwrap(); assert_eq!( - cache.async_get_cache("counter", &kwargs).await.unwrap(), + cache.async_get_cache("counter", &context).await.unwrap(), None ); } @@ -117,30 +116,30 @@ async fn codec_errors_propagate_without_writing_partial_batches() { ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); - let kwargs = CacheKwargs::default(); + let context = ExactCacheContext::default(); assert_eq!( - cache.set_cache("invalid", 255, kwargs.clone()), + cache.set_cache("invalid", 255, &context), Err(Error::InvalidEntry) ); assert_eq!( - cache.async_set_cache("invalid", 255, kwargs.clone()).await, + cache.async_set_cache("invalid", 255, context.clone()).await, Err(Error::InvalidEntry) ); assert_eq!( cache .async_set_cache_pipeline( vec![("valid".into(), 7), ("invalid".into(), 255)], - kwargs.clone(), + context.clone(), ) .await, Err(Error::InvalidEntry) ); assert_eq!( - cache.get_cache("invalid", &kwargs), + cache.get_cache("invalid", &context), Err(Error::InvalidEntry) ); assert_eq!( - cache.async_get_cache("invalid", &kwargs).await, + cache.async_get_cache("invalid", &context).await, Err(Error::InvalidEntry) ); } @@ -155,12 +154,14 @@ fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() { let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) .with_namespace(Some("team".into())); assert_eq!( - cache.get_cache("key", &CacheKwargs::default()).unwrap(), + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), None ); assert_eq!( cache - .get_cache("team:key", &CacheKwargs::default()) + .get_cache("team:key", &ExactCacheContext::default()) .unwrap(), None ); @@ -221,9 +222,9 @@ async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() { assert_eq!( cache - .async_get_cache_batch( + .async_batch_get_cache( vec!["hit".into(), "miss".into(), "invalid".into()], - CacheKwargs::default(), + ExactCacheContext::default(), ) .await .unwrap(), @@ -327,6 +328,13 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { .arg("team:key"), Ok("team:key"), ), + MockCmd::new( + redis::cmd("EVAL") + .arg("return KEYS[1]") + .arg(1usize) + .arg("team:key"), + Ok("team:key"), + ), MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")), MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")), MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")), @@ -387,6 +395,14 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { .unwrap(), redis::Value::BulkString(b"team:key".to_vec()) ); + assert_eq!( + cache + .async_register_script("return KEYS[1]".into()) + .invoke(vec!["key".into()], Vec::new()) + .await + .unwrap(), + redis::Value::BulkString(b"team:key".to_vec()) + ); assert_eq!(cache.client_list().unwrap(), "id=1"); assert_eq!(cache.info().unwrap(), "redis_version:7"); cache.flushall().unwrap(); @@ -602,7 +618,7 @@ async fn claims_match_eligible_values_written_by_another_encoder() { "pin", candidate, vec![stored.clone()], - CacheKwargs::default() + ExactCacheContext::default() ) .await .unwrap(), @@ -630,7 +646,7 @@ fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() { "pin", candidate.clone(), &[serde_json::json!({"model_id": "a"})], - CacheKwargs::default() + ExactCacheContext::default() ) .unwrap(), candidate @@ -654,7 +670,7 @@ fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() { "pin", serde_json::json!({"model_id": "b"}), &[], - CacheKwargs::default() + ExactCacheContext::default() ) .unwrap(), serde_json::json!({"model_id": "a"}) @@ -679,7 +695,7 @@ async fn async_increment_runs_the_atomic_script() { assert_eq!( cache - .async_increment_cache("counter", 2.5, CacheKwargs::default()) + .async_increment("counter", 2.5, ExactCacheContext::default()) .await .unwrap(), 4.5 diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 1fd2bb809de..606c21410c7 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -1,6 +1,6 @@ use std::{sync::Mutex, time::Duration}; -use litellm_cache::{BaseCache, Error}; +use litellm_cache::{BaseCache, Error, ExactCacheContext}; use serde_json::Value; use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; @@ -18,7 +18,7 @@ impl WriteBuffer { } } - pub async fn async_store>( + pub async fn async_store>( &self, cache: &ResponseCache, request: &ResponseCacheRequest, diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index f18a48863c5..e50e68cdabb 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,6 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error}; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, +}; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; @@ -9,7 +11,7 @@ use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, - pub kwargs: CacheKwargs, + pub context: ExactCacheContext, pub max_age: Option, } @@ -24,17 +26,17 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - kwargs: CacheKwargs::default(), + context: ExactCacheContext::default(), max_age: None, } } } -pub struct ResponseCache> { +pub struct ResponseCache> { backend: Arc, } -impl> ResponseCache { +impl> ResponseCache { pub fn new(backend: Arc) -> Self { Self { backend } } @@ -43,11 +45,14 @@ impl> ResponseCache { &self.backend } - pub fn default_ttl(&self) -> Duration { - self.backend.default_ttl() + pub fn default_ttl(&self) -> Option { + self.backend.get_ttl(&ExactCacheContext::default()) } - pub async fn async_flush(&self) -> Result<(), Error> { + pub async fn async_flush(&self) -> Result<(), Error> + where + B: FlushCache, + { self.backend.async_flush_cache().await } @@ -65,7 +70,7 @@ impl> ResponseCache { } let entry = match self .backend - .get_cache(&cache_key(&request.key), &request.kwargs) + .get_cache(&cache_key(&request.key), &request.context) { Ok(entry) => entry, Err(Error::InvalidEntry) => None, @@ -84,7 +89,7 @@ impl> ResponseCache { } let entry = match self .backend - .async_get_cache(&cache_key(&request.key), &request.kwargs) + .async_get_cache(&cache_key(&request.key), &request.context) .await { Ok(entry) => entry, @@ -98,7 +103,10 @@ impl> ResponseCache { &self, requests: &[ResponseCacheRequest], now: Duration, - ) -> Result { + ) -> Result + where + B: BatchCache, + { let readable = requests .iter() .enumerate() @@ -109,7 +117,7 @@ impl> ResponseCache { .map(|(_, request)| cache_key(&request.key)) .collect::>(); let entries = if let Some((_, request)) = readable.first() { - self.backend.get_cache_batch(&keys, &request.kwargs)? + self.backend.batch_get_cache(&keys, &request.context)? } else { Vec::new() }; @@ -120,7 +128,10 @@ impl> ResponseCache { &self, requests: &[ResponseCacheRequest], now: Duration, - ) -> Result { + ) -> Result + where + B: BatchCache, + { let readable = requests .iter() .enumerate() @@ -132,7 +143,7 @@ impl> ResponseCache { .collect::>(); let entries = if let Some((_, request)) = readable.first() { self.backend - .async_get_cache_batch(keys, request.kwargs.clone()) + .async_batch_get_cache(keys, request.context.clone()) .await? } else { Vec::new() @@ -155,7 +166,7 @@ impl> ResponseCache { timestamp: Some(now.as_secs_f64()), response, }, - request.kwargs.clone(), + &request.context, ) } @@ -175,7 +186,7 @@ impl> ResponseCache { timestamp: Some(now.as_secs_f64()), response, }, - request.kwargs.clone(), + request.context.clone(), ) .await } @@ -210,26 +221,29 @@ impl> ResponseCache { timestamp: Some(now.as_secs_f64()), response, }, - request.kwargs, + request.context, ) }) .collect::>(); let Some((_, _, first_kwargs)) = writable.first() else { return Ok(()); }; - if writable.iter().all(|(_, _, kwargs)| kwargs == first_kwargs) { - let kwargs = first_kwargs.clone(); + if writable + .iter() + .all(|(_, _, context)| context == first_kwargs) + { + let context = first_kwargs.clone(); let cache_list = writable .into_iter() .map(|(key, entry, _)| (key, entry)) .collect(); return self .backend - .async_set_cache_pipeline(cache_list, kwargs) + .async_set_cache_pipeline(cache_list, context) .await; } - for (key, entry, kwargs) in writable { - self.backend.async_set_cache(&key, entry, kwargs).await?; + for (key, entry, context) in writable { + self.backend.async_set_cache(&key, entry, context).await?; } Ok(()) } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 94b25626f86..e4f78dae8b2 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -43,7 +43,7 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { )); let cache = ResponseCache::new(backend.clone()); let mut request = request(); - request.kwargs.ttl = Some(Duration::from_secs(10)); + request.context.ttl = Some(Duration::from_secs(10)); request.max_age = Some(Duration::from_secs(5)); cache .store( @@ -336,7 +336,7 @@ fn response_codec_preserves_values_without_timestamps() { assert_eq!(entry.response, raw); let backend = Arc::new(InMemoryCache::default()); - BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, Default::default()).unwrap(); + BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap(); let cache = ResponseCache::new(backend); assert_eq!( cache.lookup(&request(), Duration::from_secs(100)).unwrap(), diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index d6ef8052c4f..8bd69ba5ad6 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -1,7 +1,6 @@ use std::{future::Future, time::Duration}; use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; use crate::Error; @@ -12,10 +11,25 @@ pub enum BatchEntry { Invalid, } -#[derive(Clone, Debug, Default, PartialEq)] -pub struct CacheKwargs { +pub trait CacheContext: Clone + Send + Sync + 'static { + fn ttl(&self) -> Option; + + fn with_ttl(&self, ttl: Option) -> Self; +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ExactCacheContext { pub ttl: Option, - pub extras: Map, +} + +impl CacheContext for ExactCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { ttl } + } } #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] @@ -35,78 +49,44 @@ pub struct CacheConnectionResult { pub trait BaseCache: Send + Sync { type Value: Clone + Send + Sync + 'static; + type Context: CacheContext; - fn default_ttl(&self) -> Duration { - Duration::from_secs(60) - } + fn get_ttl(&self, context: &Self::Context) -> Option; - fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration { - kwargs.ttl.unwrap_or_else(|| self.default_ttl()) - } - - fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>; - - fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>; - - fn get_cache_batch( + fn set_cache( &self, - keys: &[String], - kwargs: &CacheKwargs, - ) -> Result>, Error> { - keys.iter() - .map(|key| match self.get_cache(key, kwargs) { - Ok(Some(value)) => Ok(BatchEntry::Hit(value)), - Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), - Err(error) => Err(error), - }) - .collect() - } + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error>; + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error>; fn async_set_cache( &self, key: &str, value: Self::Value, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + Send { - async move { self.set_cache(key, value, kwargs) } + async move { self.set_cache(key, value, &context) } } fn async_get_cache( &self, key: &str, - kwargs: &CacheKwargs, + context: &Self::Context, ) -> impl Future, Error>> + Send { - async move { self.get_cache(key, kwargs) } - } - - fn async_get_cache_batch( - &self, - keys: Vec, - kwargs: CacheKwargs, - ) -> impl Future>, Error>> + Send { - async move { - let mut entries = Vec::with_capacity(keys.len()); - for key in keys { - entries.push(match self.async_get_cache(&key, &kwargs).await { - Ok(Some(value)) => BatchEntry::Hit(value), - Ok(None) => BatchEntry::Miss, - Err(Error::InvalidEntry) => BatchEntry::Invalid, - Err(error) => return Err(error), - }); - } - Ok(entries) - } + async move { self.get_cache(key, context) } } fn async_set_cache_pipeline( &self, - cache_list: Vec<(String, Self::Value)>, - kwargs: CacheKwargs, + entries: Vec<(String, Self::Value)>, + context: Self::Context, ) -> impl Future> + Send { async move { - for (key, value) in cache_list { - self.async_set_cache(&key, value, kwargs.clone()).await?; + for (key, value) in entries { + self.async_set_cache(&key, value, context.clone()).await?; } Ok(()) } @@ -116,21 +96,9 @@ pub trait BaseCache: Send + Sync { &self, key: &str, value: Self::Value, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + Send { - self.async_set_cache(key, value, kwargs) - } - - fn delete_cache(&self, key: &str) -> Result<(), Error>; - - fn async_delete_cache(&self, key: &str) -> impl Future> + Send { - async move { self.delete_cache(key) } - } - - fn flush_cache(&self) -> Result<(), Error>; - - fn async_flush_cache(&self) -> impl Future> + Send { - async move { self.flush_cache() } + self.async_set_cache(key, value, context) } fn disconnect(&self) -> impl Future> + Send; diff --git a/litellm-rust/crates/cache/src/cache_type.rs b/litellm-rust/crates/cache/src/cache_type.rs new file mode 100644 index 00000000000..f0a97c04fd5 --- /dev/null +++ b/litellm-rust/crates/cache/src/cache_type.rs @@ -0,0 +1,85 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)] +pub enum CacheType { + #[serde(rename = "local")] + Local, + #[serde(rename = "redis")] + Redis, + #[serde(rename = "redis-semantic")] + RedisSemantic, + #[serde(rename = "valkey-semantic")] + ValkeySemantic, + #[serde(rename = "s3")] + S3, + #[serde(rename = "disk")] + Disk, + #[serde(rename = "qdrant-semantic")] + QdrantSemantic, + #[serde(rename = "azure-blob")] + AzureBlob, + #[serde(rename = "gcs")] + Gcs, +} + +impl CacheType { + pub const ALL: [Self; 9] = [ + Self::Local, + Self::Redis, + Self::RedisSemantic, + Self::ValkeySemantic, + Self::S3, + Self::Disk, + Self::QdrantSemantic, + Self::AzureBlob, + Self::Gcs, + ]; + + pub const fn as_python_name(self) -> &'static str { + match self { + Self::Local => "local", + Self::Redis => "redis", + Self::RedisSemantic => "redis-semantic", + Self::ValkeySemantic => "valkey-semantic", + Self::S3 => "s3", + Self::Disk => "disk", + Self::QdrantSemantic => "qdrant-semantic", + Self::AzureBlob => "azure-blob", + Self::Gcs => "gcs", + } + } + + pub fn from_python_name(value: &str) -> Option { + Self::ALL + .into_iter() + .find(|cache_type| cache_type.as_python_name() == value) + } +} + +#[cfg(test)] +mod tests { + use super::CacheType; + + #[test] + fn every_python_cache_type_has_one_round_trip_identity() { + let names = CacheType::ALL.map(CacheType::as_python_name); + assert_eq!( + names, + [ + "local", + "redis", + "redis-semantic", + "valkey-semantic", + "s3", + "disk", + "qdrant-semantic", + "azure-blob", + "gcs", + ] + ); + assert_eq!( + names.map(CacheType::from_python_name), + CacheType::ALL.map(Some) + ); + } +} diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 4ec94244ac5..fc7f46d943e 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -1,23 +1,23 @@ use std::sync::Arc; pub use crate::BaseCache as Cache; -use crate::{BaseCache, CacheKwargs, Error}; +use crate::{BaseCache, Error}; pub fn get_cache( cache: &B, key: &str, - kwargs: &CacheKwargs, + context: &B::Context, ) -> Result, Error> { - cache.get_cache(key, kwargs) + cache.get_cache(key, context) } pub fn set_cache( cache: &B, key: &str, value: B::Value, - kwargs: CacheKwargs, + context: &B::Context, ) -> Result<(), Error> { - cache.set_cache(key, value, kwargs) + cache.set_cache(key, value, context) } pub type CacheBackend = Arc; diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index 1613a5786f4..f7307e5c7bd 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -1,24 +1,77 @@ -use std::future::Future; +use std::{future::Future, time::Duration}; -use crate::{BaseCache, CacheKwargs, Error}; +use crate::{BaseCache, BatchEntry, Error}; #[derive(Clone, Debug, PartialEq)] pub struct IncrementOperation { pub key: String, pub amount: f64, - pub ttl: Option, + pub ttl: Option, +} + +pub trait BatchCache: BaseCache { + fn batch_get_cache( + &self, + keys: &[String], + context: &Self::Context, + ) -> Result>, Error> { + keys.iter() + .map(|key| match self.get_cache(key, context) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + }) + .collect() + } + + fn async_batch_get_cache( + &self, + keys: Vec, + context: Self::Context, + ) -> impl Future>, Error>> + Send { + async move { + let mut entries = Vec::with_capacity(keys.len()); + for key in keys { + entries.push(match self.async_get_cache(&key, &context).await { + Ok(Some(value)) => BatchEntry::Hit(value), + Ok(None) => BatchEntry::Miss, + Err(Error::InvalidEntry) => BatchEntry::Invalid, + Err(error) => return Err(error), + }); + } + Ok(entries) + } + } +} + +pub trait DeleteCache: BaseCache { + fn delete_cache(&self, key: &str) -> Result<(), Error>; + + fn async_delete_cache(&self, key: &str) -> impl Future> + Send { + async move { self.delete_cache(key) } + } +} + +pub trait FlushCache: BaseCache { + fn flush_cache(&self) -> Result<(), Error>; + + fn async_flush_cache(&self) -> impl Future> + Send { + async move { self.flush_cache() } + } } pub trait CounterCache: BaseCache { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; + fn increment_cache(&self, key: &str, amount: f64, context: Self::Context) + -> Result; - fn async_increment_cache( + fn async_increment( &self, key: &str, amount: f64, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + Send { - async move { self.increment_cache(key, amount, kwargs) } + async move { self.increment_cache(key, amount, context) } } } @@ -31,7 +84,7 @@ where key: &str, candidate: Self::Value, eligible: &[Self::Value], - kwargs: CacheKwargs, + context: Self::Context, ) -> Result; fn async_claim_cache( @@ -39,8 +92,78 @@ where key: &str, candidate: Self::Value, eligible: Vec, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + Send { - async move { self.claim_cache(key, candidate, &eligible, kwargs) } + async move { self.claim_cache(key, candidate, &eligible, context) } } } + +pub trait TtlCache: BaseCache { + fn async_get_ttl( + &self, + key: &str, + ) -> impl Future, Error>> + Send; +} + +pub trait SetCache: BaseCache { + type SetValue: Clone + Send + Sync + 'static; + type SetResult: Send + Sync + 'static; + + fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> impl Future> + Send; +} + +pub trait QueueCache: BaseCache { + type QueueValue: Clone + Send + Sync + 'static; + type PopResult: Send + Sync + 'static; + + fn async_rpush( + &self, + key: &str, + values: Vec, + ) -> impl Future> + Send; + + fn async_lpop( + &self, + key: &str, + count: Option, + ) -> impl Future> + Send; +} + +pub trait ScanCache: BaseCache { + fn async_scan_iter( + &self, + pattern: &str, + count: usize, + ) -> impl Future, Error>> + Send; +} + +pub trait ClientInfoCache: BaseCache { + type ClientList: Send + Sync + 'static; + type Info: Send + Sync + 'static; + + fn client_list(&self) -> Result; + + fn info(&self) -> Result; +} + +pub trait CacheScript: Send + Sync + 'static { + type Argument: Clone + Send + Sync + 'static; + type Output: Send + Sync + 'static; + + fn invoke( + &self, + keys: Vec, + arguments: Vec, + ) -> impl Future> + Send; +} + +pub trait ScriptCache: BaseCache { + type Script: CacheScript; + + fn async_register_script(&self, source: String) -> Self::Script; +} diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index 9858c5e2748..d68d4b2b69f 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -1,7 +1,8 @@ use std::{sync::Arc, time::Duration}; use crate::{ - BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ClaimCache, + CounterCache, DeleteCache, Error, FlushCache, }; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -94,19 +95,17 @@ impl DualCache { } } - fn promotion_kwargs(&self, kwargs: &CacheKwargs) -> CacheKwargs { - CacheKwargs { - ttl: self.promotion_ttl.or(kwargs.ttl), - extras: kwargs.extras.clone(), - } + fn promotion_context(&self, context: &C) -> C { + context.with_ttl(self.promotion_ttl.or(context.ttl())) } } -impl DualCache +impl DualCache where V: Clone + Send + Sync + 'static, - L1: BaseCache, - L2: BaseCache, + C: CacheContext, + L1: BaseCache, + L2: BaseCache, { fn missing(entries: &[BatchEntry]) -> Vec { entries @@ -119,7 +118,7 @@ where fn merge_batch( &self, keys: &[String], - kwargs: &CacheKwargs, + context: &C, mut entries: Vec>, missing: Vec, remote: Vec>, @@ -129,8 +128,9 @@ where } for (index, entry) in missing.into_iter().zip(remote) { if let BatchEntry::Hit(value) = &entry { + let promotion_context = self.promotion_context(context); self.l1 - .set_cache(&keys[index], value.clone(), self.promotion_kwargs(kwargs))?; + .set_cache(&keys[index], value.clone(), &promotion_context)?; } entries[index] = entry; } @@ -138,46 +138,105 @@ where } } -impl BaseCache for DualCache +impl BaseCache for DualCache where V: Clone + Send + Sync + 'static, - L1: BaseCache, - L2: BaseCache, + C: CacheContext, + L1: BaseCache, + L2: BaseCache, { type Value = V; + type Context = C; - fn default_ttl(&self) -> Duration { - self.l2.default_ttl() + fn get_ttl(&self, context: &Self::Context) -> Option { + self.l2.get_ttl(context) } - fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { + fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> { if self.writes_remote() { - self.remote(self.l2.set_cache(key, value.clone(), kwargs.clone()))?; + self.remote(self.l2.set_cache(key, value.clone(), context))?; } - self.l1.set_cache(key, value, kwargs) + self.l1.set_cache(key, value, context) } - fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { - if let Some(value) = self.l1.get_cache(key, kwargs)? { + fn get_cache(&self, key: &str, context: &C) -> Result, Error> { + if let Some(value) = self.l1.get_cache(key, context)? { return Ok(Some(value)); } if !self.reads_remote() { return Ok(None); } - let value = self.remote(self.l2.get_cache(key, kwargs))?.flatten(); + let value = self.remote(self.l2.get_cache(key, context))?.flatten(); if let Some(value) = &value { - self.l1 - .set_cache(key, value.clone(), self.promotion_kwargs(kwargs))?; + let promotion_context = self.promotion_context(context); + self.l1.set_cache(key, value.clone(), &promotion_context)?; } Ok(value) } - fn get_cache_batch( + async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> { + if self.writes_remote() { + self.remote( + self.l2 + .async_set_cache(key, value.clone(), context.clone()) + .await, + )?; + } + self.l1.async_set_cache(key, value, context).await + } + + async fn async_get_cache(&self, key: &str, context: &C) -> Result, Error> { + if let Some(value) = self.l1.async_get_cache(key, context).await? { + return Ok(Some(value)); + } + if !self.reads_remote() { + return Ok(None); + } + let value = self + .remote(self.l2.async_get_cache(key, context).await)? + .flatten(); + if let Some(value) = &value { + self.l1 + .async_set_cache(key, value.clone(), self.promotion_context(context)) + .await?; + } + Ok(value) + } + + async fn async_set_cache_pipeline( &self, - keys: &[String], - kwargs: &CacheKwargs, - ) -> Result>, Error> { - let entries = self.l1.get_cache_batch(keys, kwargs)?; + entries: Vec<(String, V)>, + context: C, + ) -> Result<(), Error> { + if self.writes_remote() { + self.remote( + self.l2 + .async_set_cache_pipeline(entries.clone(), context.clone()) + .await, + )?; + } + self.l1.async_set_cache_pipeline(entries, context).await + } + + async fn disconnect(&self) -> Result<(), Error> { + self.l2.disconnect().await?; + self.l1.disconnect().await + } + + async fn test_connection(&self) -> Result { + self.l2.test_connection().await + } +} + +impl BatchCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: BatchCache, + L2: BatchCache, +{ + fn batch_get_cache(&self, keys: &[String], context: &C) -> Result>, Error> { + let entries = self.l1.batch_get_cache(keys, context)?; let missing = Self::missing(&entries); if missing.is_empty() || !self.reads_remote() { return Ok(entries); @@ -186,49 +245,20 @@ where .iter() .map(|index| keys[*index].clone()) .collect::>(); - match self.remote(self.l2.get_cache_batch(&remote_keys, kwargs))? { - Some(remote) => self.merge_batch(keys, kwargs, entries, missing, remote), + match self.remote(self.l2.batch_get_cache(&remote_keys, context))? { + Some(remote) => self.merge_batch(keys, context, entries, missing, remote), None => Ok(entries), } } - async fn async_set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { - if self.writes_remote() { - self.remote( - self.l2 - .async_set_cache(key, value.clone(), kwargs.clone()) - .await, - )?; - } - self.l1.async_set_cache(key, value, kwargs).await - } - - async fn async_get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { - if let Some(value) = self.l1.async_get_cache(key, kwargs).await? { - return Ok(Some(value)); - } - if !self.reads_remote() { - return Ok(None); - } - let value = self - .remote(self.l2.async_get_cache(key, kwargs).await)? - .flatten(); - if let Some(value) = &value { - self.l1 - .async_set_cache(key, value.clone(), self.promotion_kwargs(kwargs)) - .await?; - } - Ok(value) - } - - async fn async_get_cache_batch( + async fn async_batch_get_cache( &self, keys: Vec, - kwargs: CacheKwargs, + context: C, ) -> Result>, Error> { let entries = self .l1 - .async_get_cache_batch(keys.clone(), kwargs.clone()) + .async_batch_get_cache(keys.clone(), context.clone()) .await?; let missing = Self::missing(&entries); if missing.is_empty() || !self.reads_remote() { @@ -237,29 +267,22 @@ where let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect(); match self.remote( self.l2 - .async_get_cache_batch(remote_keys, kwargs.clone()) + .async_batch_get_cache(remote_keys, context.clone()) .await, )? { - Some(remote) => self.merge_batch(&keys, &kwargs, entries, missing, remote), + Some(remote) => self.merge_batch(&keys, &context, entries, missing, remote), None => Ok(entries), } } +} - async fn async_set_cache_pipeline( - &self, - cache_list: Vec<(String, V)>, - kwargs: CacheKwargs, - ) -> Result<(), Error> { - if self.writes_remote() { - self.remote( - self.l2 - .async_set_cache_pipeline(cache_list.clone(), kwargs.clone()) - .await, - )?; - } - self.l1.async_set_cache_pipeline(cache_list, kwargs).await - } - +impl DeleteCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: DeleteCache, + L2: DeleteCache, +{ fn delete_cache(&self, key: &str) -> Result<(), Error> { if self.writes_remote() { self.remote(self.l2.delete_cache(key))?; @@ -273,7 +296,15 @@ where } self.l1.async_delete_cache(key).await } +} +impl FlushCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: FlushCache, + L2: FlushCache, +{ fn flush_cache(&self) -> Result<(), Error> { if self.writes_remote() { self.remote(self.l2.flush_cache())?; @@ -287,65 +318,47 @@ where } self.l1.async_flush_cache().await } - - async fn disconnect(&self) -> Result<(), Error> { - self.l2.disconnect().await?; - self.l1.disconnect().await - } - - async fn test_connection(&self) -> Result { - self.l2.test_connection().await - } } -impl CounterCache for DualCache +impl CounterCache for DualCache where - L1: BaseCache, - L2: CounterCache, + C: CacheContext, + L1: BaseCache, + L2: CounterCache, { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { - let value = self.l2.increment_cache(key, amount, kwargs.clone())?; - self.l1.set_cache(key, value, kwargs)?; + fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result { + let value = self.l2.increment_cache(key, amount, context.clone())?; + self.l1.set_cache(key, value, &context)?; Ok(value) } - async fn async_increment_cache( - &self, - key: &str, - amount: f64, - kwargs: CacheKwargs, - ) -> Result { + async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result { let value = self .l2 - .async_increment_cache(key, amount, kwargs.clone()) + .async_increment(key, amount, context.clone()) .await?; - self.l1.async_set_cache(key, value, kwargs).await?; + self.l1.async_set_cache(key, value, context).await?; Ok(value) } } -impl ClaimCache for DualCache +impl ClaimCache for DualCache where V: Clone + PartialEq + Send + Sync + 'static, - L1: ClaimCache, - L2: ClaimCache, + C: CacheContext, + L1: ClaimCache, + L2: ClaimCache, { - fn claim_cache( - &self, - key: &str, - candidate: V, - eligible: &[V], - kwargs: CacheKwargs, - ) -> Result { + fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result { match self.remote( self.l2 - .claim_cache(key, candidate.clone(), eligible, kwargs.clone()), + .claim_cache(key, candidate.clone(), eligible, context.clone()), )? { Some(winner) => { - self.l1.set_cache(key, winner.clone(), kwargs)?; + self.l1.set_cache(key, winner.clone(), &context)?; Ok(winner) } - None => self.l1.claim_cache(key, candidate, eligible, kwargs), + None => self.l1.claim_cache(key, candidate, eligible, context), } } @@ -354,20 +367,22 @@ where key: &str, candidate: V, eligible: Vec, - kwargs: CacheKwargs, + context: C, ) -> Result { match self.remote( self.l2 - .async_claim_cache(key, candidate.clone(), eligible.clone(), kwargs.clone()) + .async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone()) .await, )? { Some(winner) => { - self.l1.async_set_cache(key, winner.clone(), kwargs).await?; + self.l1 + .async_set_cache(key, winner.clone(), context) + .await?; Ok(winner) } None => { self.l1 - .async_claim_cache(key, candidate, eligible, kwargs) + .async_claim_cache(key, candidate, eligible, context) .await } } diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index cdd5589aa39..ce9f93b6dc4 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -1,14 +1,21 @@ mod base_cache; +mod cache_type; mod caching; mod capabilities; mod codec; -pub mod dual; +mod dual; mod error; pub use base_cache::{ - BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, + BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, + ExactCacheContext, }; +pub use cache_type::CacheType; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; -pub use capabilities::{ClaimCache, CounterCache, IncrementOperation}; +pub use capabilities::{ + BatchCache, CacheScript, ClaimCache, ClientInfoCache, CounterCache, DeleteCache, FlushCache, + IncrementOperation, QueueCache, ScanCache, ScriptCache, SetCache, TtlCache, +}; pub use codec::{CacheCodec, JsonCodec}; +pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 5c27aa9ff4e..9180ee9d0dc 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,20 +1,69 @@ use std::{sync::Mutex, time::Duration}; -use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache, +}; struct TestCache { default_ttl: Duration, - writes: Mutex>, + writes: Mutex>, +} + +#[derive(Clone)] +struct SemanticContext { + ttl: Option, + query: String, +} + +impl CacheContext for SemanticContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + ttl, + query: self.query.clone(), + } + } +} + +struct SemanticCache; + +impl BaseCache for SemanticCache { + type Value = String; + type Context = SemanticContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> { + Ok(()) + } + + fn get_cache(&self, _: &str, context: &Self::Context) -> Result, Error> { + Ok((context.query == "matching prompt").then(|| "semantic hit".into())) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } } impl BaseCache for TestCache { type Value = String; + type Context = ExactCacheContext; - fn default_ttl(&self) -> Duration { - self.default_ttl + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl.or(Some(self.default_ttl)) } - fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> { + fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> { Err(Error::Unavailable) } @@ -22,7 +71,7 @@ impl BaseCache for TestCache { &self, key: &str, value: Self::Value, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result<(), Error> { if key == "unavailable" { return Err(Error::Unavailable); @@ -30,22 +79,14 @@ impl BaseCache for TestCache { self.writes .lock() .unwrap() - .push((key.into(), value, kwargs)); + .push((key.into(), value, context)); Ok(()) } - fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { Ok(None) } - fn delete_cache(&self, _: &str) -> Result<(), Error> { - Ok(()) - } - - fn flush_cache(&self) -> Result<(), Error> { - Ok(()) - } - async fn disconnect(&self) -> Result<(), Error> { Ok(()) } @@ -62,15 +103,26 @@ fn ttl_uses_default_and_allows_per_call_override() { writes: Mutex::default(), }; assert_eq!( - cache.get_ttl(&CacheKwargs::default()), - Duration::from_secs(60) + cache.get_ttl(&ExactCacheContext::default()), + Some(Duration::from_secs(60)) ); assert_eq!( - cache.get_ttl(&CacheKwargs { + cache.get_ttl(&ExactCacheContext { ttl: Some(Duration::from_secs(5)), - ..Default::default() }), - Duration::from_secs(5) + Some(Duration::from_secs(5)) + ); +} + +#[test] +fn associated_context_preserves_backend_specific_lookup_inputs() { + let context = SemanticContext { + ttl: None, + query: "matching prompt".into(), + }; + assert_eq!( + get_cache(&SemanticCache, "shared-key", &context).unwrap(), + Some("semantic hit".into()) ); } @@ -81,12 +133,11 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { writes: Mutex::default(), }; let entry = String::from("cached"); - let kwargs = CacheKwargs { + let context = ExactCacheContext { ttl: Some(Duration::from_secs(5)), - ..Default::default() }; cache - .batch_cache_write("single", entry.clone(), kwargs.clone()) + .batch_cache_write("single", entry.clone(), context.clone()) .await .unwrap(); assert_eq!( @@ -97,7 +148,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { ("unavailable".into(), entry.clone()), ("skipped".into(), entry.clone()), ], - kwargs.clone(), + context.clone(), ) .await, Err(Error::Unavailable) @@ -105,8 +156,8 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { assert_eq!( *cache.writes.lock().unwrap(), vec![ - ("single".into(), entry.clone(), kwargs.clone()), - ("first".into(), entry, kwargs), + ("single".into(), entry.clone(), context.clone()), + ("first".into(), entry, context), ] ); } diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index 2e1e72c7119..e7e8927f8d0 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -4,8 +4,8 @@ use std::{ }; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, - dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy}, + BaseCache, BatchCache, CacheConnectionResult, ClaimCache, CounterCache, DeleteCache, DualCache, + Error, ExactCacheContext, FlushCache, ReadPolicy, RemoteFailurePolicy, WritePolicy, }; struct TestCache { @@ -27,26 +27,21 @@ where V: Clone + Send + Sync + 'static, { type Value = V; + type Context = ExactCacheContext; - fn set_cache(&self, _: &str, value: V, _: CacheKwargs) -> Result<(), Error> { + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl.or(Some(Duration::from_secs(60))) + } + + fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> { *self.value.lock().unwrap() = Some(value); Ok(()) } - fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { Ok(self.value.lock().unwrap().clone()) } - fn delete_cache(&self, _: &str) -> Result<(), Error> { - *self.value.lock().unwrap() = None; - Ok(()) - } - - fn flush_cache(&self) -> Result<(), Error> { - *self.value.lock().unwrap() = None; - Ok(()) - } - async fn disconnect(&self) -> Result<(), Error> { Ok(()) } @@ -56,8 +51,30 @@ where } } +impl BatchCache for TestCache where V: Clone + Send + Sync + 'static {} + +impl DeleteCache for TestCache +where + V: Clone + Send + Sync + 'static, +{ + fn delete_cache(&self, _: &str) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } +} + +impl FlushCache for TestCache +where + V: Clone + Send + Sync + 'static, +{ + fn flush_cache(&self) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } +} + impl CounterCache for TestCache { - fn increment_cache(&self, _: &str, amount: f64, _: CacheKwargs) -> Result { + fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result { if self.fail { return Err(Error::Unavailable); } @@ -77,7 +94,7 @@ where _: &str, candidate: V, eligible: &[V], - _: CacheKwargs, + _: ExactCacheContext, ) -> Result { if self.fail { return Err(Error::Unavailable); @@ -100,11 +117,12 @@ fn failed_l2_increment_leaves_l1_unchanged() { let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true))); assert_eq!( - cache.increment_cache("counter", 2.0, CacheKwargs::default()), + cache.increment_cache("counter", 2.0, ExactCacheContext::default()), Err(Error::Unavailable) ); assert_eq!( - l1.get_cache("counter", &CacheKwargs::default()).unwrap(), + l1.get_cache("counter", &ExactCacheContext::default()) + .unwrap(), Some(10.0) ); } @@ -121,9 +139,8 @@ fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { "affinity", "second".into(), &["first".into(), "second".into()], - CacheKwargs { + ExactCacheContext { ttl: Some(Duration::from_secs(60)), - ..Default::default() }, ) .unwrap(), @@ -135,12 +152,17 @@ struct SyncPanics(TestCache); impl BaseCache for SyncPanics { type Value = String; + type Context = ExactCacheContext; - fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> { + fn get_ttl(&self, context: &Self::Context) -> Option { + self.0.get_ttl(context) + } + + fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { panic!("sync L2 write on an async path") } - fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { panic!("sync L2 read on an async path") } @@ -148,54 +170,30 @@ impl BaseCache for SyncPanics { &self, key: &str, value: String, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result<(), Error> { - self.0.set_cache(key, value, kwargs) + self.0.set_cache(key, value, &context) } async fn async_get_cache( &self, key: &str, - kwargs: &CacheKwargs, + context: &ExactCacheContext, ) -> Result, Error> { - self.0.get_cache(key, kwargs) - } - - async fn async_get_cache_batch( - &self, - keys: Vec, - kwargs: CacheKwargs, - ) -> Result>, Error> { - assert_eq!(keys, ["missing"]); - Ok(vec![match self.0.get_cache("missing", &kwargs)? { - Some(value) => litellm_cache::BatchEntry::Hit(value), - None => litellm_cache::BatchEntry::Miss, - }]) + self.0.get_cache(key, context) } async fn async_set_cache_pipeline( &self, cache_list: Vec<(String, String)>, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result<(), Error> { for (key, value) in cache_list { - self.0.set_cache(&key, value, kwargs.clone())?; + self.0.set_cache(&key, value, &context)?; } Ok(()) } - fn delete_cache(&self, _: &str) -> Result<(), Error> { - panic!("sync L2 delete on an async path") - } - - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - self.0.delete_cache(key) - } - - fn flush_cache(&self) -> Result<(), Error> { - panic!("sync L2 flush on an async path") - } - async fn disconnect(&self) -> Result<(), Error> { Ok(()) } @@ -205,6 +203,36 @@ impl BaseCache for SyncPanics { } } +impl BatchCache for SyncPanics { + async fn async_batch_get_cache( + &self, + keys: Vec, + context: ExactCacheContext, + ) -> Result>, Error> { + assert_eq!(keys, ["missing"]); + Ok(vec![match self.0.get_cache("missing", &context)? { + Some(value) => litellm_cache::BatchEntry::Hit(value), + None => litellm_cache::BatchEntry::Miss, + }]) + } +} + +impl DeleteCache for SyncPanics { + fn delete_cache(&self, _: &str) -> Result<(), Error> { + panic!("sync L2 delete on an async path") + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + self.0.delete_cache(key) + } +} + +impl FlushCache for SyncPanics { + fn flush_cache(&self) -> Result<(), Error> { + panic!("sync L2 flush on an async path") + } +} + #[tokio::test] async fn async_operations_use_the_async_l2_methods() { let l1 = Arc::new(TestCache::new(None, false)); @@ -215,36 +243,36 @@ async fn async_operations_use_the_async_l2_methods() { false, ))), ); - let kwargs = CacheKwargs::default(); + let context = ExactCacheContext::default(); assert_eq!( - cache.async_get_cache("missing", &kwargs).await.unwrap(), + cache.async_get_cache("missing", &context).await.unwrap(), Some("remote".into()) ); assert_eq!( - l1.get_cache("missing", &kwargs).unwrap(), + l1.get_cache("missing", &context).unwrap(), Some("remote".into()) ); l1.delete_cache("missing").unwrap(); assert_eq!( cache - .async_get_cache_batch(vec!["missing".into()], kwargs.clone()) + .async_batch_get_cache(vec!["missing".into()], context.clone()) .await .unwrap(), [litellm_cache::BatchEntry::Hit("remote".to_string())] ); cache - .async_set_cache("missing", "written".into(), kwargs.clone()) + .async_set_cache("missing", "written".into(), context.clone()) .await .unwrap(); cache - .async_set_cache_pipeline(vec![("missing".into(), "piped".into())], kwargs.clone()) + .async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone()) .await .unwrap(); cache.async_delete_cache("missing").await.unwrap(); assert_eq!( - cache.async_get_cache("missing", &kwargs).await.unwrap(), + cache.async_get_cache("missing", &context).await.unwrap(), None ); } @@ -253,20 +281,17 @@ struct Unavailable; impl BaseCache for Unavailable { type Value = String; + type Context = ExactCacheContext; - fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> { + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { Err(Error::Unavailable) } - fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { - Err(Error::Unavailable) - } - - fn delete_cache(&self, _: &str) -> Result<(), Error> { - Err(Error::Unavailable) - } - - fn flush_cache(&self) -> Result<(), Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { Err(Error::Unavailable) } @@ -279,13 +304,27 @@ impl BaseCache for Unavailable { } } +impl BatchCache for Unavailable {} + +impl DeleteCache for Unavailable { + fn delete_cache(&self, _: &str) -> Result<(), Error> { + Err(Error::Unavailable) + } +} + +impl FlushCache for Unavailable { + fn flush_cache(&self) -> Result<(), Error> { + Err(Error::Unavailable) + } +} + impl ClaimCache for Unavailable { fn claim_cache( &self, _: &str, _: String, _: &[String], - _: CacheKwargs, + _: ExactCacheContext, ) -> Result { Err(Error::InvalidEntry) } @@ -293,24 +332,25 @@ impl ClaimCache for Unavailable { #[test] fn remote_failure_policy_selects_propagation_or_the_local_tier() { - let kwargs = CacheKwargs::default(); + let context = ExactCacheContext::default(); let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); assert_eq!( - strict.set_cache("key", "value".into(), kwargs.clone()), + strict.set_cache("key", "value".into(), &context), Err(Error::Unavailable) ); - assert_eq!(strict.get_cache("key", &kwargs), Err(Error::Unavailable)); + assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable)); let l1 = Arc::new(TestCache::new(None, false)); let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable)) .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); - assert_eq!(degraded.get_cache("key", &kwargs), Ok(None)); - degraded - .set_cache("key", "value".into(), kwargs.clone()) - .unwrap(); - assert_eq!(degraded.get_cache("key", &kwargs), Ok(Some("value".into()))); + assert_eq!(degraded.get_cache("key", &context), Ok(None)); + degraded.set_cache("key", "value".into(), &context).unwrap(); + assert_eq!( + degraded.get_cache("key", &context), + Ok(Some("value".into())) + ); degraded.delete_cache("key").unwrap(); - assert_eq!(l1.get_cache("key", &kwargs), Ok(None)); + assert_eq!(l1.get_cache("key", &context), Ok(None)); } #[test] @@ -321,7 +361,12 @@ fn claim_fallback_does_not_hide_non_availability_errors() { ) .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); assert_eq!( - cache.claim_cache("affinity", "second".into(), &[], CacheKwargs::default()), + cache.claim_cache( + "affinity", + "second".into(), + &[], + ExactCacheContext::default() + ), Err(Error::InvalidEntry) ); } @@ -332,11 +377,9 @@ fn local_only_policies_never_touch_l2() { let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone()) .with_read_policy(ReadPolicy::LocalOnly) .with_write_policy(WritePolicy::LocalOnly); - let kwargs = CacheKwargs::default(); + let context = ExactCacheContext::default(); - assert_eq!(cache.get_cache("key", &kwargs), Ok(None)); - cache - .set_cache("key", "local".into(), kwargs.clone()) - .unwrap(); - assert_eq!(l2.get_cache("key", &kwargs), Ok(Some("remote".into()))); + assert_eq!(cache.get_cache("key", &context), Ok(None)); + cache.set_cache("key", "local".into(), &context).unwrap(); + assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into()))); } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 6e694c20705..0e7d6aee11d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,5 +1,6 @@ use std::time::Duration; +use litellm_cache::CacheType; use pyo3::{ exceptions::{PyTypeError, PyValueError}, prelude::*, @@ -127,21 +128,30 @@ impl NativeCacheConfig { .extract::()?, }; let backend = facade.getattr("cache")?; - match backend_name.as_str() { - "local" => project_memory(&backend).map(|backend| { + match CacheType::from_python_name(&backend_name) { + Some(CacheType::Local) => project_memory(&backend).map(|backend| { CacheConfigProjection::Native(Box::new(Self { policy, backend: CacheBackendConfig::Memory(backend), })) }), - "redis" => match project_redis(&backend)? { + Some(CacheType::Redis) => match project_redis(&backend)? { Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { policy, backend: CacheBackendConfig::Redis(Box::new(backend)), }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, - _ => Ok(CacheConfigProjection::Unsupported( + Some( + CacheType::RedisSemantic + | CacheType::ValkeySemantic + | CacheType::S3 + | CacheType::Disk + | CacheType::QdrantSemantic + | CacheType::AzureBlob + | CacheType::Gcs, + ) + | None => Ok(CacheConfigProjection::Unsupported( UnsupportedCacheConfig::Backend, )), } @@ -149,10 +159,10 @@ impl NativeCacheConfig { pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { if service.default_ttl() - != match &self.backend { + != Some(match &self.backend { CacheBackendConfig::Memory(config) => config.default_ttl, CacheBackendConfig::Redis(config) => config.default_ttl, - } + }) { return Some("facade and native backend default TTLs 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 6a3835ac84d..a9475429e45 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -53,7 +53,7 @@ impl NativeResponseCache { } } - pub fn default_ttl(&self) -> Duration { + pub fn default_ttl(&self) -> Option { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index d8793abd115..0c5343a63d0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -24,7 +24,7 @@ fn request_input(input: RequestInput) -> PyResult { if let Some(controls) = input.controls { request.controls = controls; } - request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; + request.context.ttl = input.ttl_seconds.map(duration).transpose()?; request.max_age = input.max_age_seconds.map(duration).transpose()?; Ok(request) } From bd8df0ed8f8f1ebb0e5529148578a85b030c9393 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:54:23 -0700 Subject: [PATCH 52/56] test: keep the Qwen brand missing-key tests off the network and split the ChatGPT fields test back out --- .../llms/dashscope/test_qwen_brand_aliases.py | 16 +++++++++++++--- .../public_endpoints/test_public_endpoints.py | 3 +++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py index 6f264ea1424..69ae66d3818 100644 --- a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py +++ b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py @@ -95,6 +95,13 @@ BRAND_CASES = [ def clear_dashscope_family_env(monkeypatch): for env_var in DASHSCOPE_FAMILY_ENV_VARS: monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + + +@pytest.fixture +def no_provider_traffic(respx_mock, monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock class TestQwenBrandProviderResolution: @@ -267,22 +274,25 @@ class TestQwenBrandUserFacingNames: assert self.RETIRED_MAINLAND_NAME not in message @pytest.mark.parametrize("brand", BRAND_CASES) - def test_embedding_without_key_names_brand(self, brand): + def test_embedding_without_key_names_brand(self, brand, no_provider_traffic): with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: litellm.embedding(model=f"{brand['provider']}/text-embedding-v4", input=["hello"]) assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + assert no_provider_traffic.calls.call_count == 0 @pytest.mark.parametrize("brand", BRAND_CASES) - def test_rerank_without_key_names_brand(self, brand): + def test_rerank_without_key_names_brand(self, brand, no_provider_traffic): with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: litellm.rerank(model=f"{brand['provider']}/gte-rerank-v2", query="q", documents=["a", "b"]) assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + assert no_provider_traffic.calls.call_count == 0 @pytest.mark.parametrize("brand", BRAND_CASES) - def test_image_generation_without_key_names_brand(self, brand): + def test_image_generation_without_key_names_brand(self, brand, no_provider_traffic): with pytest.raises(litellm.APIConnectionError, match=re.escape(brand["display_name"])) as exc_info: litellm.image_generation(model=f"{brand['provider']}/qwen-image", prompt="a cup of coffee") assert self.RETIRED_MAINLAND_NAME not in str(exc_info.value) + assert no_provider_traffic.calls.call_count == 0 @pytest.mark.parametrize("brand", BRAND_CASES) @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index ba88fbd3951..0f1ff3b024d 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -357,6 +357,9 @@ def test_qwen_mainland_provider_fields_carry_the_qianwen_brand(): assert mainland_fields["api_key"]["label"] == "Qianwen AI Platform API Key" assert "Qianwen AI Platform" in mainland_fields["api_base"]["tooltip"] assert "Qwen AI Platform" not in json.dumps(mainland) + + +def test_chatgpt_provider_fields(): app_instance = FastAPI() app_instance.include_router(router) test_client = TestClient(app_instance) From a8745f24a2bdbfcd8257a141a5e69976639794d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:56:21 -0700 Subject: [PATCH 53/56] fix(ci): let the install smoke test boot its key-less proxy config The install smoke test starts the proxy on test_config_no_auth.yaml, which has no master key on purpose, and #42019's boot check now refuses that, so the three installing_litellm_on_python jobs have been red on main since 2026-09-20. Pass the documented local-dev override to the proxy child so the test keeps its no-auth config and the boot check stays as it is --- tests/local_testing/test_basic_python_version.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index ef500fdff42..f032486debd 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -3,6 +3,7 @@ import os import subprocess import time import traceback +from typing import Final import pytest @@ -253,6 +254,7 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None): raise filepath = os.path.dirname(os.path.abspath(__file__)) config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + proxy_env: Final = {**os.environ, "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"} server_process = subprocess.Popen( [ "uv", @@ -266,6 +268,7 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None): *extra_proxy_args, ], cwd=PROJECT_ROOT, + env=proxy_env, ) # Allow some time for the server to start (increased for CI environments) From b53f9ad65827eac2004b73d9f61eaeda5bb5981c Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:58:59 +0000 Subject: [PATCH 54/56] fix(fal_ai): validate returned image dimensions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/cost_calculator.py | 13 ++----- .../fal_ai/image_generation/transformation.py | 4 +- litellm/types/utils.py | 1 + litellm/utils.py | 1 + .../test_fal_ai_flux_dev_transformation.py | 38 +++++++++++++++++++ .../llms/fal_ai/test_cost_calculator.py | 27 +++++++++++++ 6 files changed, 73 insertions(+), 11 deletions(-) diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 114bbabe4e6..fd7d82d314f 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -23,7 +23,6 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( ) _OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) def _keyed_size(optional_params: Mapping[str, object]) -> str | None: @@ -51,7 +50,7 @@ def _image_dimensions(image: object) -> tuple[int, int] | None: provider_specific_fields: Final = _OBJECT_MAP.validate_python(raw_provider_specific_fields) width: Final = provider_specific_fields.get("width") height: Final = provider_specific_fields.get("height") - if not isinstance(width, int) or not isinstance(height, int): + if type(width) is not int or width <= 0 or type(height) is not int or height <= 0: return None return width, height @@ -132,13 +131,9 @@ def cost_calculator( ) if all(cost is not None for cost in keyed_costs): return sum(cost for cost in keyed_costs if cost is not None) - model_info: Final = next( - ( - entry - for key in (f"{litellm.LlmProviders.FAL_AI.value}/{normalized_model}", normalized_model) - if (entry := _entry(key)) is not None - ), - _EMPTY_ENTRY, + model_info: Final = litellm.get_model_info( + model=normalized_model, + custom_llm_provider=litellm.LlmProviders.FAL_AI.value, ) raw_output_cost_per_image: Final = model_info.get("output_cost_per_image") output_cost_per_image: Final = ( diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index a5afdc41e9a..fd8e280da1c 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -47,8 +47,8 @@ def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]: height: Final = image_map.get("height") content_type: Final = image_map.get("content_type") provider_specific_fields: Final[FalImageProviderSpecificFields] = { - **({"width": width} if isinstance(width, int) else {}), - **({"height": height} if isinstance(height, int) else {}), + **({"width": width} if isinstance(width, int) and type(width) is int and width > 0 else {}), + **({"height": height} if isinstance(height, int) and type(height) is int and height > 0 else {}), **({"content_type": content_type} if isinstance(content_type, str) else {}), } return ImageObject( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c661aa0839a..db9678eab43 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -315,6 +315,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x output output_cost_per_character_above_128k_tokens: float | None # only for vertex ai models output_cost_per_image: float | None + output_cost_per_pixel: ReadOnly[float | None] output_cost_per_image_token: float | None output_cost_per_video_token: float | None # for gemini omni models with video output output_vector_size: int | None diff --git a/litellm/utils.py b/litellm/utils.py index 9a80b115d4b..b54f700c67d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6087,6 +6087,7 @@ def _get_model_info_helper( output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), + output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), output_vector_size=_model_info.get("output_vector_size", None), diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py index 273b9151a0a..675d502240e 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py @@ -86,3 +86,41 @@ def test_flux_dev_response_omits_provider_specific_fields_when_fal_omits_metadat encoding=None, ) assert response.data[0].provider_specific_fields is None + + +@pytest.mark.parametrize( + "invalid_field, invalid_value, expected_fields", + ( + ("width", True, {"height": 768, "content_type": "image/png"}), + ("width", 0, {"height": 768, "content_type": "image/png"}), + ("width", -1, {"height": 768, "content_type": "image/png"}), + ("height", True, {"width": 1024, "content_type": "image/png"}), + ("height", 0, {"width": 1024, "content_type": "image/png"}), + ("height", -1, {"width": 1024, "content_type": "image/png"}), + ), +) +def test_flux_dev_response_drops_invalid_dimension_metadata(invalid_field, invalid_value, expected_fields): + metadata = {"width": 1024, "height": 768, "content_type": "image/png"} + metadata[invalid_field] = invalid_value + raw = httpx.Response( + 200, + json={ + "images": [ + { + "url": "https://fal.media/a.png", + **metadata, + } + ] + }, + ) + response = FalAIFluxDevConfig().transform_image_generation_response( + model="fal-ai/flux/dev", + raw_response=raw, + model_response=ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert response.data[0].provider_specific_fields == expected_fields diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py index e387586e417..6fb34d9f88e 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -123,6 +123,33 @@ def test_flux_dev_cost_uses_response_megapixels_per_image(): assert cost == pytest.approx(output_cost_per_pixel * 1_048_576 * (1 + 2 + 1)) +@pytest.mark.parametrize( + "dimensions", + ( + ((True, 1024),), + ((1024, 0),), + ((-1, 1024),), + ), +) +def test_flux_dev_invalid_response_dimensions_use_flat_price(dimensions): + model = "fal_ai/fal-ai/flux/dev" + cost = cost_calculator( + model=model, + image_response=_image_response_with_dimensions(dimensions), + optional_params={}, + ) + assert cost == litellm.model_cost[model]["output_cost_per_image"] * len(dimensions) + + +def test_unknown_fal_model_raises_when_flat_pricing_is_needed(): + with pytest.raises(Exception, match="isn't mapped yet"): + cost_calculator( + model="fal_ai/fal-ai/unknown-model", + image_response=_image_response(), + optional_params={}, + ) + + def test_image_edit_call_type_routes_to_fal_keyed_pricing(): model = "openai/gpt-image-2.5/flare/edit" cost = CostCalculatorUtils.route_image_generation_cost_calculator( From a2bb1f7e57c5e75f21f40ee7c6dec38457540a4d Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:01:21 +0000 Subject: [PATCH 55/56] chore(prices): sync OpenRouter prices: 1 model openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 97a38ac1657..5978e954c23 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -42971,21 +42971,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.19242e-07, + "input_cost_per_token": 9.15936e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.838484e-06, + "output_cost_per_token": 1.831872e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.66035e-08, + "cache_read_input_token_cost": 7.6328e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 97a38ac1657..5978e954c23 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -42971,21 +42971,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.19242e-07, + "input_cost_per_token": 9.15936e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.838484e-06, + "output_cost_per_token": 1.831872e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.66035e-08, + "cache_read_input_token_cost": 7.6328e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From 4a825a7259a60ccf9d4d89d900efd786c41631c0 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:01:35 +0000 Subject: [PATCH 56/56] chore(prices): sync AWS Bedrock prices: 1 model [enrichment failed: AWS Bedrock, 16 held] zai.glm-5: supports_vision, supports_audio_input, supports_response_schema --- litellm/model_prices_and_context_window_backup.json | 5 ++++- model_prices_and_context_window.json | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 97a38ac1657..69ee55d26ab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54468,7 +54468,10 @@ "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "zai.glm-4.7-flash": { "input_cost_per_token": 7e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 97a38ac1657..69ee55d26ab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54468,7 +54468,10 @@ "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_response_schema": true, + "supports_vision": false }, "zai.glm-4.7-flash": { "input_cost_per_token": 7e-08,