From 76bf0cd5797a47b21c6d96ebd8ce645a197b519c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 23 Jul 2026 03:43:01 +0000 Subject: [PATCH 001/149] test(e2e): replace custom endpoints_client with provider SDK clients The llm_translation suite drove /v1/responses, /v1/messages, /embeddings, /v1/images/generations, /v1/moderations and /v1/audio/* through a bespoke endpoints_client wrapper that no customer runs. Tests now call the proxy the way customers do: the OpenAI SDK for the OpenAI-compatible surface and the Anthropic SDK for /v1/messages, wired through a session-scoped sdk fixture (sdk_clients.py) that points both SDKs at the proxy with a virtual key. Endpoints no official SDK covers keep the shared typed transport: rerank moves onto ProxyClient (RerankBody/RerankResponse in models.py) and the passthrough header test parses with the shared AnthropicMessagesResponse model. Files that only used endpoints_client for model registration now use the proxy fixture directly. endpoints_client.py is deleted; anthropic joins the e2e-dev dependency group so the lint env resolves the SDK imports. Resolves LIT-4577 --- pyproject.toml | 1 + tests/e2e/CLAUDE.md | 2 + tests/e2e/CONTRIBUTING.md | 2 + tests/e2e/llm_translation/conftest.py | 10 +- tests/e2e/llm_translation/endpoints_client.py | 385 ------------------ tests/e2e/llm_translation/sdk_clients.py | 53 +++ .../llm_translation/test_audio_speech_e2e.py | 99 ++--- .../test_audio_transcriptions_e2e.py | 25 +- .../e2e/llm_translation/test_cache_control.py | 5 +- .../test_credential_messages_e2e.py | 31 +- .../test_custom_pricing_e2e.py | 33 +- .../test_embeddings_endpoint_e2e.py | 92 ++--- .../test_image_generation_e2e.py | 73 ++-- .../test_messages_azure_foundry_e2e.py | 195 ++++----- .../e2e/llm_translation/test_messages_e2e.py | 176 ++++---- ...st_messages_mid_conversation_system_e2e.py | 124 +++--- ...onversation_system_native_providers_e2e.py | 136 ++++--- .../llm_translation/test_moderations_e2e.py | 61 +-- .../e2e/llm_translation/test_ocr_rust_e2e.py | 10 +- .../test_passthrough_headers_e2e.py | 8 +- tests/e2e/llm_translation/test_rerank_e2e.py | 48 +-- .../e2e/llm_translation/test_responses_e2e.py | 355 +++++++--------- .../test_responses_metadata_e2e.py | 69 ++-- tests/e2e/models.py | 19 + tests/e2e/proxy_client.py | 12 + uv.lock | 4 +- 26 files changed, 828 insertions(+), 1200 deletions(-) delete mode 100644 tests/e2e/llm_translation/endpoints_client.py create mode 100644 tests/e2e/llm_translation/sdk_clients.py diff --git a/pyproject.toml b/pyproject.toml index 62bd37c3db6..f152a6a73cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -198,6 +198,7 @@ e2e-dev = [ "playwright==1.61.0", "websockets>=15.0.1,<16.0", "locust==2.45.0", + "anthropic==0.84.0", ] proxy-dev = [ "prisma==0.11.0", diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0e39664e358..fa73cbacd96 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -58,6 +58,8 @@ That snippet only conveys intent. What you actually write uses the real harness: Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). The SDKs raise their own typed exceptions on failure, which is exactly the customer-observable contract; management routes (model/key CRUD, spend read-back) and endpoints no official SDK covers (e.g. `/v1/rerank`, `/v1/ocr`, custom passthrough paths) stay on the shared transport. Raw HTTP client imports remain banned either way + The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index fc43769aca8..fa0174f686c 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -122,6 +122,8 @@ That snippet only conveys intent. What you actually write uses the real harness: Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). Management routes and endpoints no official SDK covers stay on the shared transport, and raw HTTP client imports remain banned either way + The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py index f35ecf0760d..9fd45799773 100644 --- a/tests/e2e/llm_translation/conftest.py +++ b/tests/e2e/llm_translation/conftest.py @@ -2,14 +2,16 @@ The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared -ProxyClient, so the `resources` fixture cleans up keys this suite creates. +ProxyClient, so the `resources` fixture cleans up keys this suite creates. The +`sdk` fixture hands tests real provider SDK clients (OpenAI, Anthropic) pointed +at the proxy, the way customers actually call it. """ import pytest -from endpoints_client import EndpointsClient, build_endpoints_client from passthrough_client import PassthroughClient, build_client from proxy_client import ProxyClient +from sdk_clients import SdkClients, build_sdk_clients def pytest_configure(config: pytest.Config) -> None: @@ -25,5 +27,5 @@ def client(proxy: ProxyClient) -> PassthroughClient: @pytest.fixture(scope="session") -def endpoints_client(proxy: ProxyClient) -> EndpointsClient: - return build_endpoints_client(proxy) +def sdk() -> SdkClients: + return build_sdk_clients() diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py deleted file mode 100644 index ace621d03b3..00000000000 --- a/tests/e2e/llm_translation/endpoints_client.py +++ /dev/null @@ -1,385 +0,0 @@ -"""Client for the non-chat inference endpoints (responses, messages, rerank, -embeddings, audio speech, image generation). - -Each test registers the deployment it needs through /model/new (deleted on -teardown), so nothing is hardcoded into the gateway config, then drives the -endpoint with `send` and parses the provider-native body with a suite-local model -so the assertion is on real content, not just a 200. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -from pydantic import BaseModel - -from proxy_client import ProxyClient -from e2e_http import BinaryStream, Result, StreamingResponse -from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock - -__all__ = [ - "CacheControl", - "RichMessage", - "TextBlock", -] - - -class FunctionParameterProperty(BaseModel): - type: str - description: str | None = None - - -class FunctionParameters(BaseModel): - type: Literal["object"] = "object" - properties: dict[str, FunctionParameterProperty] - required: list[str] = [] - - -class ResponsesFunctionTool(BaseModel): - type: Literal["function"] = "function" - name: str - description: str | None = None - parameters: FunctionParameters - - -class ResponsesInputTextPart(BaseModel): - type: Literal["input_text"] = "input_text" - text: str - - -class ResponsesInputImagePart(BaseModel): - type: Literal["input_image"] = "input_image" - image_url: str - - -ResponsesInputContentPart = ResponsesInputTextPart | ResponsesInputImagePart - - -class ResponsesInputMessage(BaseModel): - role: Literal["user", "assistant", "system"] = "user" - content: list[ResponsesInputContentPart] - - -ResponsesInput = str | list[ResponsesInputMessage] - - -class ResponsesRequest(BaseModel): - model: str - input: ResponsesInput - instructions: str | None = None - stream: bool = False - tools: list[ResponsesFunctionTool] | None = None - - -class MessagesRequest(BaseModel): - model: str - max_tokens: int - messages: list[ChatMessage] - - -class RichMessagesRequest(BaseModel): - model: str - max_tokens: int = 64 - system: list[TextBlock] - messages: list[RichMessage] - - -class EmbeddingsRequest(BaseModel): - model: str - input: str - - -class RerankRequest(BaseModel): - model: str - query: str - documents: list[str] - top_n: int - - -class SpeechRequest(BaseModel): - model: str - input: str - voice: str - - -class ImageRequest(BaseModel): - model: str - prompt: str - n: int = 1 - size: str = "1024x1024" - - -class TranscriptionForm(BaseModel): - model: str - response_format: str = "json" - - -class ModerationRequest(BaseModel): - model: str - input: str - - -class ResponsesOutputContent(BaseModel): - type: str | None = None - text: str | None = None - - -class ResponsesOutputItem(BaseModel): - type: str | None = None - content: list[ResponsesOutputContent] = [] - name: str | None = None - arguments: str | None = None - call_id: str | None = None - - -class ResponsesResult(BaseModel): - id: str | None = None - status: str | None = None - model: str | None = None - output: list[ResponsesOutputItem] = [] - - @property - def text(self) -> str: - return "".join( - content.text or "" for item in self.output for content in item.content - ) - - @property - def function_calls(self) -> tuple[ResponsesOutputItem, ...]: - return tuple( - item - for item in self.output - if item.type == "function_call" - and item.name is not None - and item.arguments is not None - ) - - -class ResponsesStreamEvent(BaseModel): - event_id: str | None = None - - -class ResponsesStreamEventType(BaseModel): - type: str - - -class ResponsesOutputTextDeltaEvent(ResponsesStreamEvent): - type: Literal["response.output_text.delta"] - delta: str - - -class AnthropicContentBlock(BaseModel): - type: str | None = None - text: str | None = None - - -class MessagesUsage(BaseModel): - input_tokens: int = 0 - output_tokens: int = 0 - cache_creation_input_tokens: int = 0 - cache_read_input_tokens: int = 0 - - -class MessagesResult(BaseModel): - id: str | None = None - role: str | None = None - model: str | None = None - content: list[AnthropicContentBlock] = [] - usage: MessagesUsage = MessagesUsage() - - @property - def text(self) -> str: - return "".join(block.text or "" for block in self.content) - - -class EmbeddingItem(BaseModel): - embedding: list[float] = [] - - -class EmbeddingsResult(BaseModel): - data: list[EmbeddingItem] = [] - - @property - def first_vector(self) -> tuple[float, ...]: - return tuple(self.data[0].embedding) if self.data else () - - -class RerankItem(BaseModel): - index: int | None = None - relevance_score: float | None = None - - -class RerankResult(BaseModel): - results: list[RerankItem] = [] - - -class ImageItem(BaseModel): - url: str | None = None - b64_json: str | None = None - - -class ImagesResult(BaseModel): - data: list[ImageItem] = [] - - -class TranscriptionResult(BaseModel): - text: str = "" - - -class ModerationResultItem(BaseModel): - flagged: bool - categories: dict[str, bool] = {} - - @property - def flagged_categories(self) -> tuple[str, ...]: - return tuple(name for name, hit in self.categories.items() if hit) - - -class ModerationResult(BaseModel): - results: list[ModerationResultItem] = [] - - @property - def first(self) -> ModerationResultItem | None: - return self.results[0] if self.results else None - - -@dataclass(frozen=True, slots=True) -class EndpointsClient: - proxy: ProxyClient - - def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: - return self.proxy.create_model(model_name, litellm_params) - - def delete_model(self, model_id: str) -> None: - self.proxy.delete_model(model_id) - - def _send( - self, path: str, key: str, body: BaseModel, *, stream: bool = False - ) -> StreamingResponse: - return self.proxy.transport.send( - path, - headers=self.proxy.transport.bearer(key), - json=body, - stream=stream, - ) - - def responses( - self, key: str, model: str, text: str, *, stream: bool = False - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=text, - instructions="You are a helpful assistant", - stream=stream, - ), - stream=stream, - ) - - def responses_vision( - self, key: str, model: str, text: str, image_url: str - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=[ - ResponsesInputMessage( - content=[ - ResponsesInputTextPart(text=text), - ResponsesInputImagePart(image_url=image_url), - ] - ) - ], - instructions="You are a helpful assistant", - ), - ) - - def responses_with_tools( - self, key: str, model: str, text: str, tools: list[ResponsesFunctionTool] - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=text, - instructions="You are a helpful assistant", - tools=tools, - ), - ) - - def messages( - self, key: str, model: str, text: str, *, max_tokens: int = 64 - ) -> StreamingResponse: - return self._send( - "/v1/messages", - key, - MessagesRequest( - model=model, - max_tokens=max_tokens, - messages=[ChatMessage(role="user", content=text)], - ), - ) - - def embeddings(self, key: str, model: str, text: str) -> StreamingResponse: - return self._send("/embeddings", key, EmbeddingsRequest(model=model, input=text)) - - def rerank( - self, key: str, model: str, query: str, documents: list[str], top_n: int - ) -> StreamingResponse: - return self._send( - "/v1/rerank", - key, - RerankRequest(model=model, query=query, documents=documents, top_n=top_n), - ) - - def audio_speech( - self, key: str, model: str, text: str, *, voice: str = "alloy" - ) -> StreamingResponse: - return self._send( - "/v1/audio/speech", key, SpeechRequest(model=model, input=text, voice=voice) - ) - - def audio_speech_stream( - self, key: str, model: str, text: str, *, voice: str = "alloy" - ) -> BinaryStream: - return self.proxy.transport.stream_binary( - "/v1/audio/speech", - headers=self.proxy.transport.bearer(key), - json=SpeechRequest(model=model, input=text, voice=voice), - ) - - def transcribe( - self, key: str, model: str, *, filename: str, content: bytes - ) -> Result[TranscriptionResult]: - return self.proxy.transport.upload( - "/v1/audio/transcriptions", - headers=self.proxy.transport.bearer(key), - form=TranscriptionForm(model=model), - filename=filename, - content=content, - file_content_type="audio/wav", - response_type=TranscriptionResult, - ) - - def moderations(self, key: str, model: str, text: str) -> Result[ModerationResult]: - return self.proxy.transport.post( - "/v1/moderations", - headers=self.proxy.transport.bearer(key), - json=ModerationRequest(model=model, input=text), - response_type=ModerationResult, - ) - - def images(self, key: str, model: str, prompt: str) -> StreamingResponse: - return self._send( - "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) - ) - - -def build_endpoints_client(proxy: ProxyClient) -> EndpointsClient: - return EndpointsClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/sdk_clients.py b/tests/e2e/llm_translation/sdk_clients.py new file mode 100644 index 00000000000..8a6b187f979 --- /dev/null +++ b/tests/e2e/llm_translation/sdk_clients.py @@ -0,0 +1,53 @@ +"""Real provider SDK clients pointed at the proxy, connected the way customers +connect (LIT-4577). + +The OpenAI SDK drives the OpenAI-compatible surface (/responses, /embeddings, +/images/generations, /moderations, /audio/*) and the Anthropic SDK drives +/v1/messages, each authenticated with a litellm virtual key. Errors surface as +the SDK's own exceptions, exactly what an end user sees. Retries are disabled +so a proxy fault fails the test instead of being papered over, and the timeout +matches the shared transport's request budget. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +from anthropic import Anthropic +from openai import OpenAI + +from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT + + +def response_header(headers: Mapping[str, str], name: str) -> str | None: + """Typed read of an SDK response header: httpx.Headers.get returns Any and + httpx itself is a banned import in suite code, so tests read headers through + the Mapping[str, str] interface Headers fulfils.""" + return headers[name] if name in headers else None + + +@dataclass(frozen=True, slots=True) +class SdkClients: + base_url: str + request_timeout: float + + def openai(self, key: str) -> OpenAI: + return OpenAI( + base_url=self.base_url, + api_key=key, + timeout=self.request_timeout, + max_retries=0, + ) + + def anthropic(self, key: str) -> Anthropic: + return Anthropic( + base_url=self.base_url, + api_key=key, + timeout=self.request_timeout, + max_retries=0, + ) + + +def build_sdk_clients() -> SdkClients: + return SdkClients(base_url=PROXY_BASE_URL, request_timeout=REQUEST_TIMEOUT) diff --git a/tests/e2e/llm_translation/test_audio_speech_e2e.py b/tests/e2e/llm_translation/test_audio_speech_e2e.py index b95cef8db4d..859292ecc7b 100644 --- a/tests/e2e/llm_translation/test_audio_speech_e2e.py +++ b/tests/e2e/llm_translation/test_audio_speech_e2e.py @@ -1,9 +1,10 @@ """Live e2e: POST /v1/audio/speech returns audio, non-streamed and streamed. -The non-streamed call asserts an audio (not JSON) body. The streamed call consumes -the response the way a player would and asserts customer-observable streaming: -chunked transfer encoding (a buffered body would carry a content-length) with -non-zero audio bytes. +Both calls go through the real OpenAI SDK (LIT-4577). The non-streamed call +asserts an audio (not JSON) body. The streamed call consumes the response the +way a player would and asserts customer-observable streaming: chunked transfer +encoding (a buffered body would carry a content-length) with non-zero audio +bytes. """ from __future__ import annotations @@ -11,68 +12,70 @@ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients, response_header pytestmark = pytest.mark.e2e +def _register(proxy: ProxyClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model( + model, + LiteLLMParamsBody(model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY"), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + class TestAudioSpeech: @pytest.mark.covers("llm.audio_speech.openai.basic.nonstream.works") def test_audio_speech_returns_audio( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-speech-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, "e2e-speech") + client = sdk.openai(resources.key()) - result = endpoints_client.audio_speech(key, model, "Hello!") - require_successful_call(result) - assert "audio" in (result.content_type or ""), ( - f"/audio/speech content-type is not audio: {result.content_type!r}" + response = client.audio.speech.with_raw_response.create( + model=model, voice="alloy", input="Hello!" ) - assert result.body, "/audio/speech returned an empty body" + content_type = response_header(response.headers, "content-type") + assert "audio" in (content_type or ""), ( + f"/audio/speech content-type is not audio: {content_type!r}" + ) + assert response.content, "/audio/speech returned an empty body" @pytest.mark.covers("llm.audio_speech.openai.basic.stream.works") def test_audio_speech_streams_audio_chunks( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-speech-stream-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, "e2e-speech-stream") + client = sdk.openai(resources.key()) - result = endpoints_client.audio_speech_stream( - key, - model, - "Streaming speech should arrive in several audio chunks so a client can " - "begin playback well before the whole clip has finished generating.", + with client.audio.speech.with_streaming_response.create( + model=model, + voice="alloy", + input=( + "Streaming speech should arrive in several audio chunks so a client can " + "begin playback well before the whole clip has finished generating." + ), + ) as response: + content_type = response_header(response.headers, "content-type") + transfer_encoding = response_header(response.headers, "transfer-encoding") + content_length = response_header(response.headers, "content-length") + total_bytes = sum(len(chunk) for chunk in response.iter_bytes(chunk_size=8192)) + + assert "audio" in (content_type or ""), ( + f"/audio/speech content-type is not audio: {content_type!r}" ) - assert result.ok, ( - f"/audio/speech stream failed (status {result.status_code}); body={result.error_body}" + assert "chunked" in (transfer_encoding or ""), ( + f"/audio/speech did not stream: transfer-encoding={transfer_encoding!r}, " + f"content-length={content_length!r} (a buffered body is not a stream)" ) - assert "audio" in (result.content_type or ""), ( - f"/audio/speech content-type is not audio: {result.content_type!r}" - ) - assert result.chunked, ( - f"/audio/speech did not stream: transfer-encoding={result.transfer_encoding!r}, " - f"content-length={result.content_length!r} (a buffered body is not a stream)" - ) - assert result.content_length is None, ( - f"/audio/speech advertised content-length={result.content_length!r} on a " + assert content_length is None, ( + f"/audio/speech advertised content-length={content_length!r} on a " f"streamed response (a buffered body is not a stream)" ) - assert result.total_bytes > 0, "/audio/speech stream returned no audio bytes" + assert total_bytes > 0, "/audio/speech stream returned no audio bytes" diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py index af6123dc46a..617192f88c7 100644 --- a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -1,8 +1,9 @@ """Live e2e: POST /v1/audio/transcriptions turns speech into text. Registers an OpenAI speech-to-text deployment at runtime and uploads a spoken -weather question (the realtime suite's 24kHz WAV fixture) as multipart, asserting -the returned transcript is non-empty and mentions the word it was asked about. +weather question (the realtime suite's 24kHz WAV fixture) through the real +OpenAI SDK (LIT-4577), asserting the returned transcript is non-empty and +mentions the word it was asked about. """ from __future__ import annotations @@ -12,10 +13,10 @@ from pathlib import Path import pytest from e2e_config import unique_marker -from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -27,24 +28,22 @@ WEATHER_WAV = ( class TestAudioTranscriptions: @pytest.mark.covers("llm.audio_transcriptions.openai.basic.nonstream.works") def test_audio_transcriptions_returns_text( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: model = f"e2e-transcribe-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="openai/gpt-4o-mini-transcribe", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) - result = unwrap( - endpoints_client.transcribe( - key, model, filename=WEATHER_WAV.name, content=WEATHER_WAV.read_bytes() - ) + transcription = client.audio.transcriptions.create( + model=model, file=(WEATHER_WAV.name, WEATHER_WAV.read_bytes(), "audio/wav") ) - text = result.text.strip() + text = transcription.text.strip() assert text, "/audio/transcriptions returned an empty transcript" assert "weather" in text.lower(), ( f"transcript of a spoken weather question does not mention weather: {text!r}" diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index a2c17b0fb66..888cbb422eb 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -15,7 +15,7 @@ service_tier lives in test_provider_features_e2e.py. The provider-native cache_control request shape is not expressible with the shared ``ChatBody`` (whose content is a plain string), so the cacheable body is -built from the typed content blocks shared in ``endpoints_client.py``. +built from the typed content blocks shared in ``models.py``. """ from __future__ import annotations @@ -27,9 +27,8 @@ from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, unwrap -from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager -from models import ChatResponse, LiteLLMParamsBody, Usage +from models import CacheControl, ChatResponse, LiteLLMParamsBody, RichMessage, TextBlock, Usage from passthrough_client import PassthroughClient import os diff --git a/tests/e2e/llm_translation/test_credential_messages_e2e.py b/tests/e2e/llm_translation/test_credential_messages_e2e.py index 49ea748430e..d4bf4566cc0 100644 --- a/tests/e2e/llm_translation/test_credential_messages_e2e.py +++ b/tests/e2e/llm_translation/test_credential_messages_e2e.py @@ -7,43 +7,48 @@ import os import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager from models import CredentialCreateBody, LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e class TestCredentialBackedMessages: @pytest.mark.covers("mgmt.credential.new.serves_request") - def test_credential_backed_messages(self, endpoints_client: EndpointsClient, resources: ResourceManager) -> None: + def test_credential_backed_messages( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: marker = unique_marker() credential_name = f"e2e-cred-{marker}" model = f"e2e-cred-messages-{marker}" anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") assert anthropic_api_key, "ANTHROPIC_API_KEY must be set for this live e2e test" - endpoints_client.proxy.create_credential( + proxy.create_credential( CredentialCreateBody( credential_name=credential_name, credential_values={"api_key": anthropic_api_key}, ) ) - resources.defer(lambda: endpoints_client.proxy.delete_credential(credential_name)) + resources.defer(lambda: proxy.delete_credential(credential_name)) - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="anthropic/claude-haiku-4-5", litellm_credential_name=credential_name, ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) - key = resources.key() - result = endpoints_client.messages(key, model, "reply with one word") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" - assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + client = sdk.anthropic(resources.key()) + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "reply with one word"}], + ) + assert message.role == "assistant", f"unexpected role: {message.role!r}" + text = "".join(block.text for block in message.content if block.type == "text") + assert text.strip(), f"/v1/messages returned no text: {message.content!r}" diff --git a/tests/e2e/llm_translation/test_custom_pricing_e2e.py b/tests/e2e/llm_translation/test_custom_pricing_e2e.py index b4ff631a56b..1cebf90fa21 100644 --- a/tests/e2e/llm_translation/test_custom_pricing_e2e.py +++ b/tests/e2e/llm_translation/test_custom_pricing_e2e.py @@ -25,7 +25,6 @@ from pydantic import BaseModel, RootModel from e2e_config import unique_marker from proxy_client import ProxyClient from e2e_http import Success, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import ( ChatBody, @@ -71,7 +70,7 @@ def _approx_equal(actual: float, expected: float) -> bool: def _provision( - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, prefix: str, *, @@ -84,7 +83,7 @@ def _provision( marker keeps the name unique so concurrent runs on the shared proxy never collide.""" model_name = f"{prefix}-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model_name, LiteLLMParamsBody( model=BACKEND_MODEL, @@ -93,15 +92,15 @@ def _provision( output_cost_per_token=output_cost_per_token, ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model_name def _provision_custom_priced( - endpoints_client: EndpointsClient, resources: ResourceManager + proxy: ProxyClient, resources: ResourceManager ) -> str: return _provision( - endpoints_client, + proxy, resources, "custom-priced-flash", input_cost_per_token=CUSTOM_INPUT_RATE, @@ -151,14 +150,14 @@ def _poll_breakdown_row(proxy: ProxyClient, key: str, response_id: str | None) - class TestCustomPricing: def test_custom_pricing_is_billed_at_configured_rate( self, - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, ) -> None: - model = _provision_custom_priced(endpoints_client, resources) + model = _provision_custom_priced(proxy, resources) chat = unwrap( - endpoints_client.proxy.chat( + proxy.chat( scoped_key, ChatBody( model=model, @@ -172,7 +171,7 @@ class TestCustomPricing: ) ) - row = _poll_breakdown_row(endpoints_client.proxy, scoped_key, chat.id) + row = _poll_breakdown_row(proxy, scoped_key, chat.id) assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll breakdown = row.metadata.cost_breakdown @@ -195,10 +194,10 @@ class TestCustomPricing: ) def test_model_info_reports_custom_pricing( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model = _provision_custom_priced(endpoints_client, resources) - entry = _model_info_entry(endpoints_client.proxy.model_info(), model) + model = _provision_custom_priced(proxy, resources) + entry = _model_info_entry(proxy.model_info(), model) assert entry.litellm_params.input_cost_per_token == CUSTOM_INPUT_RATE, ( f"/model/info litellm_params input rate " @@ -210,20 +209,20 @@ class TestCustomPricing: ) def test_custom_pricing_is_isolated_from_sibling_deployment( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: # Register the override first so its rate is in the backend cost map before # the sibling resolves; a leak (LIT-3897) would then poison the sibling. - custom = _provision_custom_priced(endpoints_client, resources) + custom = _provision_custom_priced(proxy, resources) sibling = _provision( - endpoints_client, + proxy, resources, "base-flash", input_cost_per_token=None, output_cost_per_token=None, ) - entries = {entry.model_name: entry for entry in endpoints_client.proxy.model_info()} + entries = {entry.model_name: entry for entry in proxy.model_info()} custom_entry = entries.get(custom) sibling_entry = entries.get(sibling) assert custom_entry is not None, f"{custom} absent from /model/info" diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 157caedd561..dac69591341 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,9 +1,10 @@ """Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex. -Each test registers the deployment it needs at runtime (deleted on teardown) and -asserts a non-empty, non-zero vector came back. The LIT-3167 guard in -tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking is -covered by tests/e2e/quota_management/spend_tracking/. +Each test registers the deployment it needs at runtime (deleted on teardown), +drives the endpoint with the real OpenAI SDK (LIT-4577), and asserts a +non-empty, non-zero vector came back. The LIT-3167 guard in +tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking +is covered by tests/e2e/quota_management/spend_tracking/. """ from __future__ import annotations @@ -11,79 +12,74 @@ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EmbeddingsResult, EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e +def _assert_embedding_vector( + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + params: LiteLLMParamsBody, +) -> None: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) + + embeddings = client.embeddings.create(model=model, input="Say this is a test!") + assert embeddings.data, f"/embeddings returned no data: {embeddings!r}" + vector = embeddings.data[0].embedding + assert vector, f"/embeddings returned no vector: {embeddings!r}" + assert any(component != 0.0 for component in vector), "embedding vector is all zeros" + + class TestEmbeddingsEndpoint: @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings", LiteLLMParamsBody( model="openai/text-embedding-3-small", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.covers("llm.embeddings.bedrock.basic.nonstream.works") def test_bedrock_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-bedrock-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-bedrock", LiteLLMParamsBody( model="bedrock/amazon.titan-embed-text-v2:0", aws_region_name="us-west-2" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") def test_vertex_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-vertex-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-vertex", LiteLLMParamsBody( model="vertex_ai/gemini-embedding-2", vertex_project="os.environ/VERTEXAI_PROJECT", vertex_location="us-central1", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 1ba78a7e083..c5204162501 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -1,8 +1,7 @@ """Live e2e: POST /v1/images/generations returns an image. -Registers an OpenAI image deployment at runtime and asserts the response carries a -generated image (url or base64). Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an image deployment at runtime, drives it through the real OpenAI SDK +(LIT-4577), and asserts the response carries a generated image (url or base64). """ from __future__ import annotations @@ -10,50 +9,58 @@ from __future__ import annotations import pytest from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, ImagesResult from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e -def _assert_image_returned(body: str) -> None: - parsed = ImagesResult.model_validate_json(body) - assert parsed.data, f"/images/generations returned no data: {body[:300]}" - first = parsed.data[0] - assert first.b64_json or first.url, ( - f"generated image has neither b64_json nor url: {body[:300]}" - ) +def _assert_image_returned( + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + params: LiteLLMParamsBody, +) -> None: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) + + images = client.images.generate(model=model, prompt="Draw a cute cat", n=1, size="1024x1024") + data = images.data or [] + assert data, f"/images/generations returned no data: {images!r}" + first = data[0] + assert first.b64_json or first.url, "generated image has neither b64_json nor url" class TestImageGeneration: @pytest.mark.covers("llm.images_generations.openai.basic.nonstream.works") def test_image_generation_returns_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-image-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="openai/gpt-image-1-mini", api_key="os.environ/OPENAI_API_KEY" - ), + _assert_image_returned( + proxy, + resources, + sdk, + "e2e-image", + LiteLLMParamsBody(model="openai/gpt-image-1-mini", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.images(key, model, "Draw a cute cat") - require_successful_call(result) - _assert_image_returned(result.body) - - @pytest.mark.covers("llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"]) + @pytest.mark.covers( + "llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"] + ) def test_bedrock_image_generation_returns_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") - model = f"e2e-bedrock-image-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_image_returned( + proxy, + resources, + sdk, + "e2e-bedrock-image", LiteLLMParamsBody( model="bedrock/amazon.titan-image-generator-v2:0", aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", @@ -61,9 +68,3 @@ class TestImageGeneration: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.images(key, model, "Draw a cute cat") - require_successful_call(result) - _assert_image_returned(result.body) diff --git a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py index d8d44820e80..90568d8c951 100644 --- a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py +++ b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py @@ -1,9 +1,9 @@ """Live e2e: POST /v1/messages routed to Azure AI Foundry Anthropic deployments. Registers `azure_ai/` deployments at runtime and drives the Messages -endpoint through the gateway across the behaviors an Anthropic client relies on: -a basic completion, a streamed completion, and tool use (non-streaming and -streaming). Auth is the Azure API key (`x-api-key`); the deployment reads +endpoint through the gateway with the real Anthropic SDK (LIT-4577) across the +behaviors an Anthropic client relies on: a basic completion, a streamed +completion, and tool use (non-streaming and streaming). The deployment reads `AZURE_AI_API_BASE` / `AZURE_AI_API_KEY` from the proxy env, so no secret is sent in the request. """ @@ -11,60 +11,50 @@ sent in the request. from __future__ import annotations import pytest +from anthropic.types import RawMessageStreamEvent, ToolParam from e2e_config import EXPECT_RUST, unique_marker -from e2e_http import StreamingResponse, require_successful_call, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager -from models import ( - AnthropicCustomTool, - AnthropicMessagesBody, - ChatMessage, - JsonSchemaProperty, - LiteLLMParamsBody, - ToolInputSchema, -) +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e AZURE_FOUNDRY_MODEL = "azure_ai/claude-haiku-4-5" -WEATHER_TOOL = AnthropicCustomTool( - name="get_weather", - description="Get the current weather for a city.", - input_schema=ToolInputSchema( - properties={"city": JsonSchemaProperty(type="string")}, - required=["city"], - ), -) +WEATHER_TOOL: ToolParam = { + "name": "get_weather", + "description": "Get the current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, +} -def _assert_streamed_ok(result: StreamingResponse) -> None: - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("content_block_delta" in event for event in result.stream_events), ( - "stream carried no content deltas" +def _assert_rust_served(headers: dict[str, str]) -> None: + if not EXPECT_RUST: + return + assert headers.get("x-litellm-rust") == "true", ( + "E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the " + "Rust path, but the response carried no x-litellm-rust marker. The request " + "still succeeded, which is exactly the failure mode: a gateway whose native " + f"extension is unavailable falls back to Python silently. headers={headers}" ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" - ) - if EXPECT_RUST: - assert result.headers.get("x-litellm-rust") == "true", ( - "E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the " - "Rust path, but the response carried no x-litellm-rust marker. The request " - "still succeeded, which is exactly the failure mode: a gateway whose native " - f"extension is unavailable falls back to Python silently. headers={result.headers}" - ) + + +def _assert_streamed_ok(event_types: list[str]) -> None: + assert event_types, "stream produced no SSE events" + assert "content_block_delta" in event_types, "stream carried no content deltas" + assert "message_stop" in event_types, "stream never reached message_stop" class TestAzureFoundryMessages: - def _register( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> tuple[str, str]: + def _register(self, proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-azure-foundry-messages-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model=AZURE_FOUNDRY_MODEL, @@ -72,91 +62,78 @@ class TestAzureFoundryMessages: api_key="os.environ/AZURE_AI_API_KEY", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key(models=[model]) + resources.defer(lambda: proxy.delete_model(model_id)) + return model @pytest.mark.covers("llm.messages.azure_foundry.basic.nonstream.works") def test_basic_nonstream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - messages=[ChatMessage(role="user", content="Reply with one word.")], - ), - ) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "Reply with one word."}], ) - assert response.content, f"no content blocks in response: {response}" - text = "".join(block.text or "" for block in response.content if block.type == "text") - assert text.strip(), f"/v1/messages returned no text: {response}" + assert message.content, f"no content blocks in response: {message!r}" + text = "".join(block.text for block in message.content if block.type == "text") + assert text.strip(), f"/v1/messages returned no text: {message.content!r}" @pytest.mark.covers("llm.messages.azure_foundry.basic.stream.works") def test_basic_stream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - stream=True, - messages=[ChatMessage(role="user", content="Count from one to three.")], - ), + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + raw = client.messages.with_raw_response.create( + model=model, + max_tokens=64, + stream=True, + messages=[{"role": "user", "content": "Count from one to three."}], ) - _assert_streamed_ok(result) + _assert_rust_served({name.lower(): value for name, value in raw.headers.items()}) + _assert_streamed_ok([event.type for event in raw.parse()]) @pytest.mark.covers("llm.messages.azure_foundry.tool_use.nonstream.works") def test_tool_use_nonstream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + message = client.messages.create( + model=model, + max_tokens=256, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], ) - assert response.content, f"no content blocks in response: {response}" - assert any(block.type == "tool_use" for block in response.content), ( - f"model did not call the tool: {response}" + assert message.content, f"no content blocks in response: {message!r}" + assert any(block.type == "tool_use" for block in message.content), ( + f"model did not call the tool: {message.content!r}" ) @pytest.mark.covers("llm.messages.azure_foundry.tool_use.stream.works") def test_tool_use_stream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - stream=True, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("tool_use" in event for event in result.stream_events), ( - "stream carried no tool_use block" - ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + stream = client.messages.create( + model=model, + max_tokens=256, + stream=True, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], ) + events: list[RawMessageStreamEvent] = list(stream) + event_types = [event.type for event in events] + assert event_types, "stream produced no SSE events" + assert any( + event.type == "content_block_start" and event.content_block.type == "tool_use" + for event in events + ), "stream carried no tool_use block" + assert "message_stop" in event_types, "stream never reached message_stop" diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44376218c6b..0c5734bc299 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -1,41 +1,35 @@ """Live e2e: POST /v1/messages (Anthropic Messages API) returns a real completion. -Registers an Anthropic deployment at runtime, drives the Messages endpoint through -the gateway, and asserts an assistant message with text came back, both -non-streaming and streamed. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an Anthropic deployment at runtime and drives the Messages endpoint +through the gateway with the real Anthropic SDK, the client customers actually +use (LIT-4577), asserting an assistant message with text came back, both +non-streaming and streamed. """ from __future__ import annotations import pytest +from anthropic.types import Message, ToolParam from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call, unwrap -from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager -from models import ( - AnthropicCustomTool, - AnthropicMessagesBody, - ChatMessage, - JsonSchemaProperty, - LiteLLMParamsBody, - SpendLogRow, - ToolInputSchema, -) +from models import LiteLLMParamsBody, SpendLogRow +from proxy_client import ProxyClient +from sdk_clients import SdkClients, response_header pytestmark = pytest.mark.e2e ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5" -WEATHER_TOOL = AnthropicCustomTool( - name="get_weather", - description="Get the current weather for a city.", - input_schema=ToolInputSchema( - properties={"city": JsonSchemaProperty(type="string")}, - required=["city"], - ), -) +WEATHER_TOOL: ToolParam = { + "name": "get_weather", + "description": "Get the current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, +} def _approx_equal(actual: float, expected: float) -> bool: @@ -43,60 +37,66 @@ def _approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") + + class TestAnthropicMessages: def _register( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> tuple[str, str]: - model = f"e2e-messages-{unique_marker()}" - model_id = endpoints_client.create_model( + self, proxy: ProxyClient, resources: ResourceManager, prefix: str = "e2e-messages" + ) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model( model, - LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" - ), + LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + return model @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.works") def test_messages_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key()) - result = endpoints_client.messages(key, model, "reply with one word") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" - assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "reply with one word"}], + ) + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert _text(message).strip(), f"/v1/messages returned no text: {message.content!r}" @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.cost_logged") def test_messages_logs_cost_matching_the_response_header( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("ANTHROPIC_API_KEY") - model = f"e2e-messages-cost-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + model = self._register(proxy, resources, prefix="e2e-messages-cost") key = resources.key() + client = sdk.anthropic(key) - result = endpoints_client.messages(key, model, f"reply with one word {unique_marker()}") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant" and parsed.text.strip(), ( - f"/v1/messages returned no assistant text: {result.body[:300]}" + raw = client.messages.with_raw_response.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": f"reply with one word {unique_marker()}"}], + ) + message = raw.parse() + assert message.role == "assistant" and _text(message).strip(), ( + f"/v1/messages returned no assistant text: {message.content!r}" ) # The customer reads per-request cost off the response header (LIT-4076), so # it must be present and positive on /v1/messages, not only /chat/completions. - header_cost = result.response_cost - assert header_cost is not None and header_cost > 0, ( - "x-litellm-response-cost header missing or non-positive on /v1/messages; " - f"headers={result.headers}" + raw_header_cost = response_header(raw.headers, "x-litellm-response-cost") + assert raw_header_cost is not None, ( + "x-litellm-response-cost header missing on /v1/messages; " + f"headers={dict(raw.headers)}" + ) + header_cost = float(raw_header_cost) + assert header_cost > 0, ( + f"x-litellm-response-cost header non-positive on /v1/messages: {header_cost}" ) # Correlate the spend row by the unique scoped key, not the Anthropic response @@ -107,7 +107,7 @@ class TestAnthropicMessages: def _priced(rows: list[SpendLogRow]) -> bool: return any(r.spend is not None and r.spend > 0 for r in rows) - rows = endpoints_client.proxy.poll_logs_for_key(key, predicate=_priced) + rows = proxy.poll_logs_for_key(key, predicate=_priced) priced = [r for r in rows if r.spend is not None and r.spend > 0] assert priced, ( f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}" @@ -123,50 +123,36 @@ class TestAnthropicMessages: @pytest.mark.covers("llm.messages.anthropic.basic.stream.works") def test_messages_streams_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key()) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - stream=True, - messages=[ChatMessage(role="user", content="Count from one to three.")], - ), - ) - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("content_block_delta" in event for event in result.stream_events), ( - "stream carried no content deltas" - ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" + stream = client.messages.create( + model=model, + max_tokens=64, + stream=True, + messages=[{"role": "user", "content": "Count from one to three."}], ) + event_types = [event.type for event in stream] + assert event_types, "stream produced no SSE events" + assert "content_block_delta" in event_types, "stream carried no content deltas" + assert "message_stop" in event_types, "stream never reached message_stop" @pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works") def test_messages_tool_use( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key()) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) + message = client.messages.create( + model=model, + max_tokens=256, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], ) - assert response.content, f"no content blocks in response: {response}" - assert any(block.type == "tool_use" for block in response.content), ( - f"model did not call the tool: {response}" + assert message.content, f"no content blocks in response: {message!r}" + assert any(block.type == "tool_use" for block in message.content), ( + f"model did not call the tool: {message.content!r}" ) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index 4b3191e60bb..b79f3027692 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -17,27 +17,28 @@ entry whose prefix spans ``system`` plus message turns is invalidated when the reminder is hoisted (the ``system`` field mutates and a turn disappears from ``messages``), while an entry ending at the system block itself would survive the hoist and mask the regression. + +Calls go through the real Anthropic SDK (LIT-4577). The SDK's ``MessageParam`` +type only admits user/assistant roles, so the system reminder turn is cast to +it; the SDK serializes the dict verbatim, which is exactly the wire shape under +test. """ from __future__ import annotations import time +from typing import cast import pytest +from anthropic import Anthropic +from anthropic.types import Message, MessageParam, TextBlockParam from pydantic import BaseModel from e2e_config import unique_marker -from e2e_http import Result, unwrap -from endpoints_client import ( - CacheControl, - EndpointsClient, - MessagesResult, - RichMessage, - RichMessagesRequest, - TextBlock, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -48,50 +49,51 @@ CACHE_PRIMING_DEADLINE_SECONDS = 60.0 CACHE_PRIMING_INTERVAL_SECONDS = 3.0 -def _cacheable_system_block(marker: str) -> TextBlock: +def _cacheable_system_block(marker: str) -> TextBlockParam: """A system prompt comfortably above Sonnet's 1024-token minimum cacheable size, unique per run so no other run's cache entry can satisfy the read.""" text = " ".join( f"Reference paragraph {index} for run {marker}." for index in range(300) ) - return TextBlock(text=text, cache_control=CacheControl()) + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} -def _user_turn(text: str, *, cached: bool = False) -> RichMessage: - block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return RichMessage(role="user", content=[block]) +def _user_turn(text: str, *, cached: bool = False) -> MessageParam: + block: TextBlockParam = ( + {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + if cached + else {"type": "text", "text": text} + ) + return {"role": "user", "content": [block]} -def _system_reminder_turn() -> RichMessage: - return RichMessage( - role="system", - content=[ - TextBlock( - text="Answer with exactly one word." - ) - ], +def _system_reminder_turn() -> MessageParam: + return cast( + "MessageParam", + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Answer with exactly one word.", + } + ], + }, ) -def _post_messages( - client: EndpointsClient, key: str, body: RichMessagesRequest -) -> Result[MessagesResult]: - return client.proxy.transport.post( - "/v1/messages", - headers=client.proxy.transport.bearer(key), - json=body, - response_type=MessagesResult, - ) +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") def _register_invoke_deployment( - client: EndpointsClient, resources: ResourceManager, bedrock_model: str + proxy: ProxyClient, resources: ResourceManager, bedrock_model: str ) -> str: model = f"e2e-midsys-{unique_marker()}" - model_id = client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model=bedrock_model, aws_region_name=AWS_REGION) ) - resources.defer(lambda: client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model @@ -114,7 +116,7 @@ class PrimedCache(BaseModel): def _prime_prompt_cache( - client: EndpointsClient, key: str, model: str, system_block: TextBlock + client: Anthropic, model: str, system_block: TextBlockParam ) -> PrimedCache: """Send first-turn calls (fresh cache-marked user turn each attempt, identical system prefix) until one both reads the system prefix back from @@ -125,17 +127,19 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = RichMessagesRequest( + usage = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[_user_turn(user_text, cached=True)], - ) - usage = unwrap(_post_messages(client, key, body)).usage - if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + ).usage + read_tokens = usage.cache_read_input_tokens or 0 + creation_tokens = usage.cache_creation_input_tokens or 0 + if read_tokens > 0 and creation_tokens > 0: return PrimedCache( first_user_text=user_text, - prefix_read_tokens=usage.cache_read_input_tokens, - first_turn_creation_tokens=usage.cache_creation_input_tokens, + prefix_read_tokens=read_tokens, + first_turn_creation_tokens=creation_tokens, ) if time.monotonic() >= deadline: pytest.fail( @@ -151,32 +155,30 @@ class TestBedrockInvokeMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_invoke_deployment( - endpoints_client, resources, FLAGGED_INVOKE_MODEL - ) - key = resources.key(models=[model]) + model = _register_invoke_deployment(proxy, resources, FLAGGED_INVOKE_MODEL) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( + second = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[ _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + {"role": "assistant", "content": [{"type": "text", "text": "OK."}]}, _user_turn("Reply with one word again.", cached=True), ], ) - second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) - assert second.text.strip(), ( + assert _text(second).strip(), ( f"{model}: reminder turn returned no completion text" ) - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: turn with a mid-conversation system reminder read " f"{second.usage.cache_read_input_tokens} cached tokens, expected at " f"least the {primed.full_prefix_tokens} cached on turn one " @@ -191,29 +193,27 @@ class TestBedrockInvokeMidConversationSystem: exercised_on=[], ) def test_unflagged_model_hoists_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_invoke_deployment( - endpoints_client, resources, UNFLAGGED_INVOKE_MODEL - ) - key = resources.key(models=[model]) + model = _register_invoke_deployment(proxy, resources, UNFLAGGED_INVOKE_MODEL) + client = sdk.anthropic(resources.key(models=[model])) - body = RichMessagesRequest( + completion = client.messages.create( model=model, - system=[TextBlock(text="You are terse.")], + max_tokens=64, + system=[{"type": "text", "text": "You are terse."}], messages=[ _user_turn(f"Say hi. Run {unique_marker()}."), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), + {"role": "assistant", "content": [{"type": "text", "text": "Hi."}]}, _user_turn("Say bye."), ], ) - completion = unwrap(_post_messages(endpoints_client, key, body)) assert completion.role == "assistant", ( f"{model}: unexpected role {completion.role!r}" ) - assert completion.text.strip(), ( + assert _text(completion).strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " f"that rejects role 'system' inside messages instead of being hoisted" diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 97d24e0564b..9a0af34fb5d 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -24,27 +24,28 @@ entry whose prefix spans ``system`` plus message turns is invalidated when the reminder is hoisted (the ``system`` field mutates and a turn disappears from ``messages``), while an entry ending at the system block itself would survive the hoist and mask the regression. + +Calls go through the real Anthropic SDK (LIT-4577). The SDK's ``MessageParam`` +type only admits user/assistant roles, so the system reminder turn is cast to +it; the SDK serializes the dict verbatim, which is exactly the wire shape under +test. """ from __future__ import annotations import time +from typing import cast import pytest +from anthropic import Anthropic +from anthropic.types import Message, MessageParam, TextBlockParam from pydantic import BaseModel from e2e_config import unique_marker -from e2e_http import Result, unwrap -from endpoints_client import ( - CacheControl, - EndpointsClient, - MessagesResult, - RichMessage, - RichMessagesRequest, - TextBlock, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -68,40 +69,47 @@ def _vertex_params(model: str) -> LiteLLMParamsBody: ) -def _cacheable_system_block(marker: str) -> TextBlock: +def _cacheable_system_block(marker: str) -> TextBlockParam: """A system prompt comfortably above the 1024-token minimum cacheable size, unique per run so no other run's cache entry can satisfy the read.""" text = " ".join(f"Reference paragraph {index} for run {marker}." for index in range(300)) - return TextBlock(text=text, cache_control=CacheControl()) + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} -def _user_turn(text: str, *, cached: bool = False) -> RichMessage: - block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return RichMessage(role="user", content=[block]) +def _user_turn(text: str, *, cached: bool = False) -> MessageParam: + block: TextBlockParam = ( + {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + if cached + else {"type": "text", "text": text} + ) + return {"role": "user", "content": [block]} -def _system_reminder_turn() -> RichMessage: - return RichMessage( - role="system", - content=[TextBlock(text="Answer with exactly one word.")], +def _system_reminder_turn() -> MessageParam: + return cast( + "MessageParam", + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Answer with exactly one word.", + } + ], + }, ) -def _post_messages(client: EndpointsClient, key: str, body: RichMessagesRequest) -> Result[MessagesResult]: - return client.proxy.transport.post( - "/v1/messages", - headers=client.proxy.transport.bearer(key), - json=body, - response_type=MessagesResult, - ) +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") def _register_deployment( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody ) -> str: model = f"e2e-midsys-{unique_marker()}" - model_id = client.create_model(model, params) - resources.defer(lambda: client.delete_model(model_id)) + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) return model @@ -124,7 +132,7 @@ class PrimedCache(BaseModel): def _prime_prompt_cache( - client: EndpointsClient, key: str, model: str, system_block: TextBlock + client: Anthropic, model: str, system_block: TextBlockParam ) -> PrimedCache: """Send first-turn calls (fresh cache-marked user turn each attempt, identical system prefix) until one both reads the system prefix back from @@ -135,17 +143,19 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = RichMessagesRequest( + usage = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[_user_turn(user_text, cached=True)], - ) - usage = unwrap(_post_messages(client, key, body)).usage - if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + ).usage + read_tokens = usage.cache_read_input_tokens or 0 + creation_tokens = usage.cache_creation_input_tokens or 0 + if read_tokens > 0 and creation_tokens > 0: return PrimedCache( first_user_text=user_text, - prefix_read_tokens=usage.cache_read_input_tokens, - first_turn_creation_tokens=usage.cache_creation_input_tokens, + prefix_read_tokens=read_tokens, + first_turn_creation_tokens=creation_tokens, ) if time.monotonic() >= deadline: pytest.fail( @@ -156,28 +166,31 @@ def _prime_prompt_cache( def _assert_flagged_model_keeps_cache( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + params: LiteLLMParamsBody, ) -> None: - model = _register_deployment(client, resources, params) - key = resources.key(models=[model]) + model = _register_deployment(proxy, resources, params) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( + second = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[ _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + {"role": "assistant", "content": [{"type": "text", "text": "OK."}]}, _user_turn("Reply with one word again.", cached=True), ], ) - second = unwrap(_post_messages(client, key, reminder_turn_body)) - assert second.text.strip(), f"{model}: reminder turn returned no completion text" - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert _text(second).strip(), f"{model}: reminder turn returned no completion text" + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: turn with a mid-conversation system reminder read " f"{second.usage.cache_read_input_tokens} cached tokens, expected at " f"least the {primed.full_prefix_tokens} cached on turn one " @@ -189,25 +202,28 @@ def _assert_flagged_model_keeps_cache( def _assert_unflagged_model_hoists_and_succeeds( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + params: LiteLLMParamsBody, ) -> None: - model = _register_deployment(client, resources, params) - key = resources.key(models=[model]) + model = _register_deployment(proxy, resources, params) + client = sdk.anthropic(resources.key(models=[model])) - body = RichMessagesRequest( + completion = client.messages.create( model=model, - system=[TextBlock(text="You are terse.")], + max_tokens=64, + system=[{"type": "text", "text": "You are terse."}], messages=[ _user_turn(f"Say hi. Run {unique_marker()}."), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), + {"role": "assistant", "content": [{"type": "text", "text": "Hi."}]}, _user_turn("Say bye."), ], ) - completion = unwrap(_post_messages(client, key, body)) assert completion.role == "assistant", f"{model}: unexpected role {completion.role!r}" - assert completion.text.strip(), ( + assert _text(completion).strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " f"role 'system' inside messages instead of being hoisted" @@ -223,19 +239,19 @@ class TestAzureFoundryMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - _assert_flagged_model_keeps_cache(endpoints_client, resources, _azure_params(self.FLAGGED_MODEL)) + _assert_flagged_model_keeps_cache(proxy, resources, sdk, _azure_params(self.FLAGGED_MODEL)) @pytest.mark.covers( "llm.messages.azure_foundry.mid_conversation_system.nonstream.works", exercised_on=[], ) def test_unflagged_model_hoists_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: _assert_unflagged_model_hoists_and_succeeds( - endpoints_client, resources, _azure_params(self.UNFLAGGED_MODEL) + proxy, resources, sdk, _azure_params(self.UNFLAGGED_MODEL) ) @@ -248,17 +264,17 @@ class TestVertexMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - _assert_flagged_model_keeps_cache(endpoints_client, resources, _vertex_params(self.FLAGGED_MODEL)) + _assert_flagged_model_keeps_cache(proxy, resources, sdk, _vertex_params(self.FLAGGED_MODEL)) @pytest.mark.covers( "llm.messages.vertex.mid_conversation_system.nonstream.works", exercised_on=[], ) def test_unflagged_model_hoists_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: _assert_unflagged_model_hoists_and_succeeds( - endpoints_client, resources, _vertex_params(self.UNFLAGGED_MODEL) + proxy, resources, sdk, _vertex_params(self.UNFLAGGED_MODEL) ) diff --git a/tests/e2e/llm_translation/test_moderations_e2e.py b/tests/e2e/llm_translation/test_moderations_e2e.py index 69cf4414a48..bc26d79dcab 100644 --- a/tests/e2e/llm_translation/test_moderations_e2e.py +++ b/tests/e2e/llm_translation/test_moderations_e2e.py @@ -1,19 +1,22 @@ """Live e2e: POST /v1/moderations classifies content against the provider policy. -Registers OpenAI's omni moderation model at runtime and asserts the product -promise on both sides of the decision: clearly violent text comes back flagged -with at least one policy category tripped, and benign text comes back not flagged. +Registers OpenAI's omni moderation model at runtime, drives it through the real +OpenAI SDK (LIT-4577), and asserts the product promise on both sides of the +decision: clearly violent text comes back flagged with at least one policy +category tripped, and benign text comes back not flagged. """ from __future__ import annotations import pytest +from openai.types import Moderation +from pydantic import TypeAdapter from e2e_config import unique_marker -from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -21,45 +24,49 @@ VIOLENT_TEXT = "I am going to find you and kill you, and I will hurt everyone yo BENIGN_TEXT = "I enjoyed the sunny afternoon and a relaxing walk in the park today." -def _register_moderation_model( - endpoints_client: EndpointsClient, resources: ResourceManager -) -> str: +def _register_moderation_model(proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-moderation-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="openai/omni-moderation-latest", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model +_CATEGORY_FLAGS = TypeAdapter(dict[str, bool | None]) + + +def _flagged_categories(item: Moderation) -> tuple[str, ...]: + flags = _CATEGORY_FLAGS.validate_python(item.categories.model_dump()) + return tuple(name for name, hit in flags.items() if hit) + + class TestModerations: @pytest.mark.covers("llm.moderations.openai.basic.nonstream.works") def test_moderations_flags_violent_content( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_moderation_model(endpoints_client, resources) - key = resources.key() + model = _register_moderation_model(proxy, resources) + client = sdk.openai(resources.key()) - result = unwrap(endpoints_client.moderations(key, model, VIOLENT_TEXT)) - item = result.first - assert item is not None, f"/moderations returned no results: {result}" - assert item.flagged, f"violent text was not flagged: {item}" - assert item.flagged_categories, ( - f"flagged result reported no true category: {item}" - ) + moderation = client.moderations.create(model=model, input=VIOLENT_TEXT) + assert moderation.results, f"/moderations returned no results: {moderation!r}" + item = moderation.results[0] + assert item.flagged, f"violent text was not flagged: {item!r}" + assert _flagged_categories(item), f"flagged result reported no true category: {item!r}" def test_moderations_passes_benign_content( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_moderation_model(endpoints_client, resources) - key = resources.key() + model = _register_moderation_model(proxy, resources) + client = sdk.openai(resources.key()) - result = unwrap(endpoints_client.moderations(key, model, BENIGN_TEXT)) - item = result.first - assert item is not None, f"/moderations returned no results: {result}" + moderation = client.moderations.create(model=model, input=BENIGN_TEXT) + assert moderation.results, f"/moderations returned no results: {moderation!r}" + item = moderation.results[0] assert not item.flagged, ( - f"benign text was flagged as {item.flagged_categories}: {item}" + f"benign text was flagged as {_flagged_categories(item)}: {item!r}" ) diff --git a/tests/e2e/llm_translation/test_ocr_rust_e2e.py b/tests/e2e/llm_translation/test_ocr_rust_e2e.py index cdbf1883314..986ac675e98 100644 --- a/tests/e2e/llm_translation/test_ocr_rust_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_rust_e2e.py @@ -22,9 +22,9 @@ import pytest from e2e_config import unique_marker from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody, OcrBody, OcrDocument, OcrResponse +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -143,14 +143,14 @@ def _assert_ocr_document(response: OcrResponse) -> None: class TestRustOcrGateway: @pytest.mark.parametrize("case", RUST_OCR_CASES, ids=_CASE_IDS) def test_rust_ocr_response( - self, endpoints_client: EndpointsClient, resources: ResourceManager, case: _OcrCase + self, proxy: ProxyClient, resources: ResourceManager, case: _OcrCase ) -> None: model = f"rust-ocr-{case.suffix}-{unique_marker()}" - model_id = endpoints_client.create_model(model, case.provider.litellm_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + model_id = proxy.create_model(model, case.provider.litellm_params()) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - response = unwrap(endpoints_client.proxy.ocr(key, OcrBody(model=model, document=case.document))) + response = unwrap(proxy.ocr(key, OcrBody(model=model, document=case.document))) _assert_ocr_document(response) diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py index 045988334d5..00ca810db15 100644 --- a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -18,9 +18,8 @@ from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import AuthHeaders, NoBody, require_successful_call, unwrap -from endpoints_client import MessagesResult from lifecycle import ResourceManager -from models import ChatMessage, KeyGenerateBody +from models import AnthropicMessagesResponse, ChatMessage, KeyGenerateBody from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e @@ -128,8 +127,9 @@ class TestPassthroughHeaders: json=_messages_body(), ) require_successful_call(result) - completion = MessagesResult.model_validate_json(result.body) - assert completion.text.strip(), ( + completion = AnthropicMessagesResponse.model_validate_json(result.body) + text = "".join(block.text or "" for block in (completion.content or [])) + assert text.strip(), ( f"static x-api-key must reach Anthropic for the call to succeed at all; got {result.body[:300]}" ) diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index 0857ff65a52..c0e4d7def42 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -1,8 +1,9 @@ """Live e2e: POST /v1/rerank ranks documents by relevance. Registers a Cohere rerank deployment at runtime and asserts the endpoint returns -scored results within the requested top_n. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +scored results within the requested top_n. No official OpenAI/Anthropic SDK +covers /v1/rerank, so the call rides the shared typed transport via +ProxyClient.rerank. """ from __future__ import annotations @@ -10,10 +11,10 @@ from __future__ import annotations import pytest from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, RerankResult +from e2e_http import unwrap from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import LiteLLMParamsBody, RerankBody, RerankResponse +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -26,39 +27,39 @@ DOCUMENTS = [ QUERY = "What is the capital of the United States?" -def _assert_top_n_scored(body: str) -> None: - parsed = RerankResult.model_validate_json(body) - assert parsed.results, f"/rerank returned no results: {body[:300]}" - assert len(parsed.results) <= 3, f"top_n=3 not honored: {body[:300]}" - assert parsed.results[0].relevance_score is not None, ( - f"top rerank result has no relevance_score: {body[:300]}" +def _assert_top_n_scored(response: RerankResponse) -> None: + assert response.results, f"/rerank returned no results: {response!r}" + assert len(response.results) <= 3, f"top_n=3 not honored: {response!r}" + assert response.results[0].relevance_score is not None, ( + f"top rerank result has no relevance_score: {response!r}" ) class TestRerank: @pytest.mark.covers("llm.rerank.cohere.basic.nonstream.works") def test_rerank_scores_top_n( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: model = f"e2e-rerank-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model="cohere/rerank-v3.5", api_key="os.environ/COHERE_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) - require_successful_call(result) - _assert_top_n_scored(result.body) + response = unwrap( + proxy.rerank(key, RerankBody(model=model, query=QUERY, documents=DOCUMENTS, top_n=3)) + ) + _assert_top_n_scored(response) @pytest.mark.covers("llm.rerank.bedrock.basic.nonstream.works", exercised_on=["rerank"]) def test_bedrock_rerank_scores_top_n( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"e2e-bedrock-rerank-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="bedrock/amazon.rerank-v1:0", @@ -67,9 +68,10 @@ class TestRerank: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) - require_successful_call(result) - _assert_top_n_scored(result.body) + response = unwrap( + proxy.rerank(key, RerankBody(model=model, query=QUERY, documents=DOCUMENTS, top_n=3)) + ) + _assert_top_n_scored(response) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index d24d2b53b71..93f844505d0 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -1,8 +1,8 @@ """Live e2e: POST /v1/responses returns a real completion. -Registers an OpenAI deployment at runtime, drives the Responses API through the -gateway, and asserts output text came back. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an OpenAI deployment at runtime and drives the Responses API through +the gateway with the real OpenAI SDK, the client customers actually use +(LIT-4577), asserting output text came back. """ from __future__ import annotations @@ -11,34 +11,45 @@ import json from typing import cast import pytest -from pydantic import BaseModel, ValidationError +from openai.types.responses import ( + FunctionToolParam, + Response, + ResponseFunctionToolCall, + ResponseInputParam, +) +from pydantic import BaseModel from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import ( - EndpointsClient, - FunctionParameterProperty, - FunctionParameters, - ResponsesFunctionTool, - ResponsesOutputTextDeltaEvent, - ResponsesResult, - ResponsesStreamEventType, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +INSTRUCTIONS = "You are a helpful assistant" +CAT_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg" -WEATHER_TOOL = ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), -) +WEATHER_TOOL: FunctionToolParam = { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + "strict": False, +} + + +def _openai_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY") + + +def _anthropic_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY") def _bedrock_params() -> LiteLLMParamsBody: @@ -50,6 +61,25 @@ def _bedrock_params() -> LiteLLMParamsBody: ) +def _register(proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody) -> str: + model = f"e2e-responses-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + +def _function_calls(response: Response) -> tuple[ResponseFunctionToolCall, ...]: + return tuple(item for item in response.output if isinstance(item, ResponseFunctionToolCall)) + + +def _assert_weather_call(response: Response) -> None: + function_call = next((call for call in _function_calls(response) if call.name == "get_weather"), None) + assert function_call is not None, f"no get_weather function call: {response.output!r}" + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + + class WeatherArguments(BaseModel): location: str @@ -57,242 +87,157 @@ class WeatherArguments(BaseModel): class TestResponses: @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") def test_responses_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS + ) + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.openai.basic.stream.works") def test_responses_streaming_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word", stream=True) - require_successful_call(result) - delta_events = tuple( - parsed - for event in result.stream_events - if (parsed := _parse_stream_event(event)) is not None + stream = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, stream=True + ) + events = list(stream) + assert events, "responses stream returned no events" + deltas = [event.delta for event in events if event.type == "response.output_text.delta"] + assert any(delta for delta in deltas), "responses stream returned no text deltas" + assert events[-1].type == "response.completed", ( + f"responses stream did not terminate with response.completed: {events[-1].type}" ) - - assert any(event.delta for event in delta_events), "responses stream returned no text deltas" - assert result.stream_events, "responses stream returned no events" - assert ( - ResponsesStreamEventType.model_validate_json(result.stream_events[-1]).type - == "response.completed" - ), "responses stream did not terminate with response.completed" @pytest.mark.covers("llm.responses.openai.basic.nonstream.cost_logged") def test_responses_logs_cost( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) + + raw = client.responses.with_raw_response.create( + model=model, input=f"reply with one word {unique_marker()}", instructions=INSTRUCTIONS + ) + response = raw.parse() + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" + assert raw.headers.get("x-litellm-call-id") and response.id, ( + f"missing response identifiers: id={response.id!r}, headers={dict(raw.headers)}" ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.responses(key, model, f"reply with one word {unique_marker()}") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" - assert result.call_id and parsed.id, f"missing response identifiers: {result.body[:300]}" - - rows = endpoints_client.proxy.poll_logs_for_request_id( - parsed.id, + rows = proxy.poll_logs_for_request_id( + response.id, predicate=lambda logged_rows: any((row.spend or 0) > 0 for row in logged_rows), ) row = next((logged_row for logged_row in rows if (logged_row.spend or 0) > 0), None) - assert row is not None, f"no costed spend row for response id {parsed.id}" + assert row is not None, f"no costed spend row for response id {response.id}" assert "gpt-4o-mini" in (row.model or ""), f"unexpected spend row model: {row.model}" @pytest.mark.covers("llm.responses.openai.tool_use.nonstream.works") def test_responses_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, - model, - "What is the weather in San Francisco? Use the get_weather tool.", - [ - ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), - ) - ], + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next( - (call for call in parsed.function_calls if call.name == "get_weather"), - None, - ) - assert function_call is not None, f"no get_weather function call: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.covers("llm.responses.openai.vision.nonstream.works") def test_responses_vision_describes_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + model = _register( + proxy, + resources, LiteLLMParamsBody(model="openai/gpt-4o", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + client = sdk.openai(resources.key()) - result = endpoints_client.responses_vision( - key, - model, - "What animal is shown in this image? Answer in one word", - "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg", + vision_input: ResponseInputParam = [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What animal is shown in this image? Answer in one word", + }, + {"type": "input_image", "image_url": CAT_IMAGE_URL, "detail": "auto"}, + ], + } + ] + response = client.responses.create(model=model, input=vision_input, instructions=INSTRUCTIONS) + text = response.output_text.strip().lower() + assert text, f"/responses vision returned no output text: {response.output!r}" + assert any(keyword in text for keyword in ("cat", "feline")), ( + f"vision response did not describe the image: {text[:300]}" ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - text = parsed.text.strip().lower() - assert text, f"/responses vision returned no output text: {result.body[:300]}" - assert any( - keyword in text - for keyword in ("cat", "feline") - ), f"vision response did not describe the image: {parsed.text[:300]}" @pytest.mark.covers("llm.responses.anthropic.basic.nonstream.works") def test_responses_anthropic_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _anthropic_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS + ) + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.anthropic.tool_use.nonstream.works") def test_responses_anthropic_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _anthropic_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, - model, - "What is the weather in San Francisco? Use the get_weather tool.", - [ - ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), - ) - ], + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next( - (call for call in parsed.function_calls if call.name == "get_weather"), - None, - ) - assert function_call is not None, f"no get_weather function call: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.covers("llm.responses.bedrock_converse.basic.nonstream.works") def test_responses_bedrock_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model(model, _bedrock_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _bedrock_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses over bedrock returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS + ) + assert response.output_text.strip(), ( + f"/responses over bedrock returned no output text: {response.output!r}" + ) @pytest.mark.covers("llm.responses.bedrock_converse.tool_use.nonstream.works") def test_responses_bedrock_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model(model, _bedrock_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _bedrock_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, model, "What is the weather in San Francisco? Use the get_weather tool.", [WEATHER_TOOL] + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next((call for call in parsed.function_calls if call.name == "get_weather"), None) - assert function_call is not None, f"no get_weather function call over bedrock: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" - - -def _parse_stream_event( - event: str, -) -> ResponsesOutputTextDeltaEvent | None: - try: - return ResponsesOutputTextDeltaEvent.model_validate_json(event) - except ValidationError: - return None + _assert_weather_call(response) diff --git a/tests/e2e/llm_translation/test_responses_metadata_e2e.py b/tests/e2e/llm_translation/test_responses_metadata_e2e.py index 6cf24348095..87727bde4d8 100644 --- a/tests/e2e/llm_translation/test_responses_metadata_e2e.py +++ b/tests/e2e/llm_translation/test_responses_metadata_e2e.py @@ -1,8 +1,8 @@ """Live e2e: /v1/responses with store + metadata (LIT-1201 customer path). -Customers attach metadata and store=true, then continue with previous_response_id. -Both turns must succeed, and any Redis keys written for the session must carry a -positive TTL (not unbounded). +Customers attach metadata and store=true through the OpenAI SDK, then continue +with previous_response_id. Both turns must succeed, and any Redis keys written +for the session must carry a positive TTL (not unbounded). """ from __future__ import annotations @@ -15,21 +15,14 @@ import pytest from pydantic import BaseModel, ConfigDict from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, ResponsesResult from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e - -class ResponsesMetadataBody(BaseModel): - model: str - input: str - store: bool = True - metadata: dict[str, str] - previous_response_id: str | None = None - instructions: str | None = "You are a helpful assistant." +INSTRUCTIONS = "You are a helpful assistant." class RedisKeyInfo(BaseModel): @@ -67,50 +60,42 @@ class TestResponsesMetadata: exercised_on=["responses"], ) def test_store_metadata_continues_and_redis_keys_have_ttl( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: # Anthropic avoids OpenAI/Gemini quota flakes; Responses translation still # exercises store + metadata + previous_response_id on the proxy. marker = unique_marker() model = f"e2e-resp-meta-{marker}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="anthropic/claude-haiku-4-5-20251001", api_key="os.environ/ANTHROPIC_API_KEY", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) - first = endpoints_client.proxy.transport.send( - "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), - json=ResponsesMetadataBody( - model=model, - input=f"Remember marker {marker}. Reply with one word.", - metadata={"session_id": marker, "customer": "e2e"}, - ), + first = client.responses.create( + model=model, + input=f"Remember marker {marker}. Reply with one word.", + store=True, + metadata={"session_id": marker, "customer": "e2e"}, + instructions=INSTRUCTIONS, ) - require_successful_call(first) - parsed = ResponsesResult.model_validate_json(first.body) - assert parsed.id, f"responses must return an id: {first.body[:300]}" - assert parsed.text.strip(), f"responses returned empty text: {first.body[:300]}" + assert first.id, f"responses must return an id: {first!r}" + assert first.output_text.strip(), f"responses returned empty text: {first.output!r}" - second = endpoints_client.proxy.transport.send( - "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), - json=ResponsesMetadataBody( - model=model, - input="Reply with the single word ok.", - previous_response_id=parsed.id, - metadata={"session_id": marker, "turn": "2"}, - ), + second = client.responses.create( + model=model, + input="Reply with the single word ok.", + store=True, + previous_response_id=first.id, + metadata={"session_id": marker, "turn": "2"}, + instructions=INSTRUCTIONS, ) - require_successful_call(second) - second_parsed = ResponsesResult.model_validate_json(second.body) - assert second_parsed.text.strip(), ( - f"previous_response_id follow-up returned empty text: {second.body[:300]}" + assert second.output_text.strip(), ( + f"previous_response_id follow-up returned empty text: {second.output!r}" ) time.sleep(1.0) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b3ea9346180..722722f4e19 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -364,6 +364,25 @@ class EmbedResponse(BaseModel): model: str | None = None +# ---------- rerank ---------- + + +class RerankBody(BaseModel): + model: str + query: str + documents: list[str] + top_n: int + + +class RerankItem(BaseModel): + index: int | None = None + relevance_score: float | None = None + + +class RerankResponse(BaseModel): + results: list[RerankItem] = [] + + # ---------- ocr ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..234c37d2f77 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -57,6 +57,8 @@ from models import ( ModelUpdateBody, OcrBody, OcrResponse, + RerankBody, + RerankResponse, SpendLogRow, SpendLogs, SpendLogsPage, @@ -291,6 +293,16 @@ class ProxyClient: response_type=OcrResponse, ) + def rerank(self, key: str, body: RerankBody) -> Result[RerankResponse]: + """POST /v1/rerank (Cohere-format). No official OpenAI/Anthropic SDK + covers this route, so it stays on the shared typed transport.""" + return self.transport.post( + "/v1/rerank", + headers=self.transport.bearer(key), + json=body, + response_type=RerankResponse, + ) + def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]: """POST /v1/messages/count_tokens (Anthropic-native). Sends the anthropic-version header so the native path accepts it; harmless on the diff --git a/uv.lock b/uv.lock index b3d6fccff26..caf976b9202 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-19T00:00:06.091071Z" +exclude-newer = "2026-07-20T03:20:37.107777782Z" exclude-newer-span = "P3D" [manifest] @@ -4300,6 +4300,7 @@ dev = [ { name = "vcrpy" }, ] e2e-dev = [ + { name = "anthropic" }, { name = "locust" }, { name = "playwright" }, { name = "websockets" }, @@ -4477,6 +4478,7 @@ dev = [ { name = "vcrpy", specifier = "==8.2.1" }, ] e2e-dev = [ + { name = "anthropic", specifier = "==0.84.0" }, { name = "locust", specifier = "==2.45.0" }, { name = "playwright", specifier = "==1.61.0" }, { name = "websockets", specifier = ">=15.0.1,<16.0" }, From 0e34c18d7ed9a846258fe6456117b8143134bfc6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 23 Jul 2026 06:32:29 +0000 Subject: [PATCH 002/149] docs(e2e): name the e2e-dev group in the documented suite run commands The llm_translation suite needs the e2e-dev dependency group at collection time (websockets for the realtime folder, now also the anthropic SDK for the sdk fixture). make bootstrap installs the group, but the documented pytest command did not name it, so a default dev-group environment failed collection. Naming the group on uv run makes the command work from any environment state --- tests/e2e/CONTRIBUTING.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index fa0174f686c..e542c123988 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -35,10 +35,10 @@ The suites run against a live proxy, so bring one up first by running the litell curl -fs http://localhost:4000/health/liveliness ``` -4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): +4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`). The suites' client dependencies (the provider SDKs, websockets) live in the `e2e-dev` dependency group; `make bootstrap` installs it, and naming the group on the run keeps the command working from any environment state: ```bash - uv run pytest tests/e2e/llm_translation/ -v + uv run --group e2e-dev pytest tests/e2e/llm_translation/ -v ``` The browser tests in the `management/` suite drive the dashboard the proxy serves at `/ui` through playwright, an optional dependency behind `importorskip` (the suite's API tests run without it). It lives in the `e2e-dev` dependency group; install it along with its browser: @@ -148,7 +148,7 @@ Before you push ```bash litellm --config .yml --port 4000 - uv run pytest tests/e2e// -v + uv run --group e2e-dev pytest tests/e2e// -v ``` 4. Capture screenshots of the test run and attach them to the PR as proof 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 003/149] 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 004/149] 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 005/149] 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 006/149] 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 a558a0b6a983a78fd24a7e1f3760482c079b6255 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 23:10:03 +0000 Subject: [PATCH 007/149] fix(proxy): renew budget reservation counter TTL while the request is in flight A reservation lives inside spend counter keys that expire on the Redis idle TTL (60s). A stream that outlives the TTL dropped its reservation, so a concurrent request on any worker was admitted against the DB floor until the stream reconciled. Renew the counter TTL with EXPIRE every ttl/2 while the reservation is open and stop once it is finalized, so an idle counter still expires on its own if the worker dies. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 15 +++++ litellm/proxy/proxy_server.py | 6 ++ .../spend_tracking/budget_reservation.py | 44 +++++++++++++- .../proxy/test_budget_reservation.py | 60 ++++++++++++++++++- 4 files changed, 123 insertions(+), 2 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 106c1580110..b1531e299d0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -76,6 +76,8 @@ class _AsyncRedisCommands(Protocol): def ttl(self, name: str) -> Awaitable[int]: ... + def expire(self, name: str, time: int) -> Awaitable[bool]: ... + def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... @@ -1795,6 +1797,19 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) return None + @_redis_circuit_breaker_guard + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + """EXPIRE an existing key without touching its value. False when the key is absent or Redis failed.""" + _used_ttl: Final = self.get_ttl(ttl=ttl) + if _used_ttl is None: + return False + try: + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) + except Exception as e: + verbose_logger.debug("Redis EXPIRE Error: %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) + return False + @_redis_circuit_breaker_guard async def async_rpush( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a617aec9f5c..33f099540b4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3222,6 +3222,12 @@ async def increment_spend_counter(counter_key: str, increment: float): return await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def refresh_spend_counter_ttl(counter_key: str) -> bool: + if spend_counter_cache.redis_cache is None: + return False + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + + async def _increment_spend_counter_cache(counter_key: str, increment: float): if spend_counter_cache.redis_cache is not None: try: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index d9f6e33c43c..bd9eb154839 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -99,6 +100,42 @@ def get_reserved_counter_keys(budget_reservation: dict | None) -> set: } +_lease_renewals: Final[set[asyncio.Task[None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks + + +def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], counter_keys: frozenset[str]) -> None: + """A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL + while the request is in flight so a request longer than the TTL does not drop its + reservation and admit concurrent requests against the DB floor on any worker.""" + from litellm.proxy.proxy_server import spend_counter_cache + + if spend_counter_cache.redis_cache is None or not counter_keys: + return + task: Final = asyncio.create_task( + _renew_reservation_lease( + budget_reservation=budget_reservation, + counter_keys=counter_keys, + interval=spend_counter_cache.redis_cache.default_ttl / 2, + ) + ) + _lease_renewals.add(task) + task.add_done_callback(_lease_renewals.discard) + + +async def _renew_reservation_lease( + budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float +) -> None: + from litellm.proxy.proxy_server import refresh_spend_counter_ttl + + deadline: Final = time.monotonic() + litellm.request_timeout + while time.monotonic() < deadline: + await asyncio.sleep(interval) + if budget_reservation.get("finalized") is True: + return + for counter_key in counter_keys: + await refresh_spend_counter_ttl(counter_key=counter_key) + + def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: UserAPIKeyAuth | None) -> bool: """ Whether an over-budget key's own ``max_budget`` reservation should be @@ -294,12 +331,17 @@ async def reserve_budget_for_request( llm_router=llm_router, input_token_counts=input_token_counts, ) - return { + budget_reservation: Final = { "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), } + _start_reservation_lease_renewal( + budget_reservation=budget_reservation, + counter_keys=frozenset(get_reserved_counter_keys(budget_reservation=budget_reservation)), + ) + return budget_reservation async def reconcile_budget_reservation( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..2812be2c53e 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1,5 +1,6 @@ import asyncio import threading +import time from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace @@ -2199,26 +2200,83 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: - def __init__(self) -> None: + """In-memory stand-in for RedisCache with real wall-clock key expiry.""" + + def __init__(self, default_ttl: float = 60.0) -> None: + self.default_ttl = default_ttl self.store: dict[str, float] = {} + self.expires_at: dict[str, float] = {} + self.refresh_count = 0 + + def _evict_expired(self, key: str) -> None: + if self.expires_at.get(key, float("inf")) <= time.monotonic(): + self.store.pop(key, None) + self.expires_at.pop(key, None) async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: + self._evict_expired(key) return self.store.get(key) async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = self.store.get(key, 0.0) + float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = max(self.store.get(key, float("-inf")), float(value)) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: self.store[key] = float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return True async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + self.expires_at.pop(key, None) + + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self._evict_expired(key) + if key not in self.store: + return False + self.refresh_count += 1 + self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) + return True + + +@pytest.mark.asyncio +async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( + spend_counter_state, +): + """A request that runs longer than the counter TTL must keep its reservation in Redis + (so a concurrent request on any worker still sees it), and renewal must stop once the + reservation is reconciled so an idle counter still expires on its own.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease" + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert await redis_cache.async_get_cache(key=counter_key) == pytest.approx(0.6) + concurrent = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert concurrent is not None + assert concurrent["reserved_cost"] == pytest.approx(0.4) + + await release_budget_reservation(reservation) + await release_budget_reservation(concurrent) + await asyncio.sleep(0.15) + refreshes_after_release = redis_cache.refresh_count + await asyncio.sleep(0.35) + assert redis_cache.refresh_count == refreshes_after_release + assert await redis_cache.async_get_cache(key=counter_key) is None @pytest.mark.asyncio From 3cd4a768aef300c62f66792242904c9bc801c00d Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 23:24:21 +0000 Subject: [PATCH 008/149] fix(proxy): keep the reservation lease alive across a failed Redis EXPIRE Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 9 ++----- litellm/proxy/proxy_server.py | 6 ++++- .../proxy/test_budget_reservation.py | 27 ++++++++++++++++++- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index b1531e299d0..5ed739b5d10 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1799,16 +1799,11 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: - """EXPIRE an existing key without touching its value. False when the key is absent or Redis failed.""" + """EXPIRE an existing key without touching its value. False when the key is absent.""" _used_ttl: Final = self.get_ttl(ttl=ttl) if _used_ttl is None: return False - try: - return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) - except Exception as e: - verbose_logger.debug("Redis EXPIRE Error: %s", e) - _record_swallowed_redis_failure(self._circuit_breaker, e) - return False + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) @_redis_circuit_breaker_guard async def async_rpush( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33f099540b4..476eda3d081 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3225,7 +3225,11 @@ async def increment_spend_counter(counter_key: str, increment: float): async def refresh_spend_counter_ttl(counter_key: str) -> bool: if spend_counter_cache.redis_cache is None: return False - return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + try: + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + except Exception as e: + verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e) + return False async def _increment_spend_counter_cache(counter_key: str, increment: float): diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 2812be2c53e..3c33ca11c91 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2202,11 +2202,13 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: """In-memory stand-in for RedisCache with real wall-clock key expiry.""" - def __init__(self, default_ttl: float = 60.0) -> None: + def __init__(self, default_ttl: float = 60.0, fail_first_refresh: bool = False) -> None: self.default_ttl = default_ttl self.store: dict[str, float] = {} self.expires_at: dict[str, float] = {} + self.refresh_attempts = 0 self.refresh_count = 0 + self.fail_first_refresh = fail_first_refresh def _evict_expired(self, key: str) -> None: if self.expires_at.get(key, float("inf")) <= time.monotonic(): @@ -2239,6 +2241,9 @@ class _ExpiringRedisCache: self.expires_at.pop(key, None) async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self.refresh_attempts += 1 + if self.fail_first_refresh and self.refresh_attempts == 1: + raise ConnectionError("Redis circuit breaker is open") self._evict_expired(key) if key not in self.store: return False @@ -2279,6 +2284,26 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( assert await redis_cache.async_get_cache(key=counter_key) is None +@pytest.mark.asyncio +async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( + spend_counter_state, +): + """One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the + rest of the request.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-blip", spend=0.0, max_budget=1.0) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert redis_cache.refresh_attempts >= 3 + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, From 668e0142664ac6fa6ab6a6c53925c79cf20b1b6e Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 9 Sep 2026 00:11:05 +0000 Subject: [PATCH 009/149] fix(proxy): stop reservation lease renewal once the request task is gone A streaming /v1/messages client disconnect skips reconciliation, so the lease kept the orphaned reservation alive until request_timeout instead of the plain 60s counter TTL the base branch had. Stop renewing when the request task that took the reservation is done. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/budget_reservation.py | 10 ++++++-- .../proxy/test_budget_reservation.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index bd9eb154839..c0af6685600 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -116,6 +116,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c budget_reservation=budget_reservation, counter_keys=counter_keys, interval=spend_counter_cache.redis_cache.default_ttl / 2, + request_task=asyncio.current_task(), ) ) _lease_renewals.add(task) @@ -123,14 +124,19 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c async def _renew_reservation_lease( - budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float + budget_reservation: Mapping[str, object], + counter_keys: frozenset[str], + interval: float, + request_task: asyncio.Task[object] | None, ) -> None: + """Stops on finalization or once the request task that took the reservation is gone, so a + disconnect path that skipped reconciliation falls back to the plain counter TTL.""" from litellm.proxy.proxy_server import refresh_spend_counter_ttl deadline: Final = time.monotonic() + litellm.request_timeout while time.monotonic() < deadline: await asyncio.sleep(interval) - if budget_reservation.get("finalized") is True: + if budget_reservation.get("finalized") is True or (request_task is not None and request_task.done()): return for counter_key in counter_keys: await refresh_spend_counter_ttl(counter_key=counter_key) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 3c33ca11c91..626c82ea397 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2304,6 +2304,29 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_reservation_lease_stops_when_request_task_ends_without_reconciling( + spend_counter_state, +): + """A request whose task ends without reconciling (client disconnect path that skips the + cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of + pinning the reservation until the request timeout.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-orphan", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease-orphan" + + reservation = await asyncio.create_task(_reserve(valid_token, 0.6, key_cache, proxy_logging_obj)) + assert reservation is not None + assert reservation["finalized"] is False + + await asyncio.sleep(0.5) + assert redis_cache.refresh_count == 0 + assert await redis_cache.async_get_cache(key=counter_key) is None + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, From 1a358318f8821878ea4623bfd7cb4af2dee3da3f Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 10:38:45 -0700 Subject: [PATCH 010/149] feat(ui): add upgrade banner with latest release changelog stats Adds GET /get/latest_release_info, which fetches the latest stable GitHub release once per hour per worker and buckets its PR bullets by conventional commit prefix into new features, fixes, and other updates. The Admin UI compares the running proxy version to that release and shows a dismissible top banner with the stat line when the proxy is behind. Dismissal is stored per release version in localStorage so the banner returns for the next one. Co-Authored-By: Claude Code --- litellm/proxy/proxy_server.py | 4 + .../latest_release_endpoints.py | 138 +++++++++++++ .../test_latest_release_endpoints.py | 191 ++++++++++++++++++ .../latestRelease/useLatestReleaseInfo.ts | 16 ++ .../src/app/(dashboard)/layout.tsx | 3 + .../src/components/UpgradeBanner.test.tsx | 108 ++++++++++ .../src/components/UpgradeBanner.tsx | 77 +++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 54 +++++ .../src/utils/versionUtils.test.ts | 53 +++++ .../src/utils/versionUtils.ts | 23 +++ 10 files changed, 667 insertions(+) create mode 100644 litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py create mode 100644 tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts create mode 100644 ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx create mode 100644 ui/litellm-dashboard/src/components/UpgradeBanner.tsx create mode 100644 ui/litellm-dashboard/src/utils/versionUtils.test.ts create mode 100644 ui/litellm-dashboard/src/utils/versionUtils.ts diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7219b373dc3..65146e8744f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -647,6 +647,9 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.types_utils.utils import get_instance_fn +from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + router as latest_release_endpoints_router, +) from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, ) @@ -18368,6 +18371,7 @@ app.include_router(debugging_endpoints_router) app.include_router(rust_control_plane_router) app.include_router(ui_crud_endpoints_router) app.include_router(user_banner_endpoints_router) +app.include_router(latest_release_endpoints_router) app.include_router(team_callback_router) app.include_router(budget_management_router) app.include_router(model_management_router) diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py new file mode 100644 index 00000000000..c0306e65aee --- /dev/null +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -0,0 +1,138 @@ +import re +from collections import Counter +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, Protocol + +import httpx +from fastapi import APIRouter, Depends +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router: Final = APIRouter() + +LATEST_RELEASE_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/latest" +LATEST_RELEASE_FETCH_TIMEOUT_SECONDS: Final = 5 +LATEST_RELEASE_CACHE_TTL_SECONDS: Final = 60 * 60 +LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS: Final = 5 * 60 +LATEST_RELEASE_CACHE_KEY: Final = "latest_release_info" + +_RELEASE_BULLET_PATTERN: Final = re.compile(r"^\*\s+([A-Za-z]+)(\([^)]*\))?!?:\s") + +_Bucket = Literal["new_features", "bug_fixes", "other_updates"] +_PREFIX_BUCKETS: Final[Mapping[str, _Bucket]] = MappingProxyType({"feat": "new_features", "fix": "bug_fixes"}) + + +class LatestReleaseInfo(BaseModel): + version: str + new_features: int + bug_fixes: int + other_updates: int + release_url: str + + +@dataclass(frozen=True, slots=True) +class LatestReleaseUnavailable: + reason: str + + +class _GitHubRelease(BaseModel): + tag_name: str + html_url: str + body: str + + +class _AsyncGetClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... + + +_latest_release_cache: Final = InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) + + +def _default_client() -> _AsyncGetClient: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + return get_async_httpx_client(llm_provider=httpxSpecialProvider.UI) + + +def _default_cache() -> InMemoryCache: + return _latest_release_cache + + +def count_release_bullets(body: str) -> Counter[_Bucket]: + """ + Bucket a release body's ``* type(scope): title by @user in `` bullets by conventional-commit type. + Lines without that shape (headings, "New Contributors" entries) are skipped, not counted as other. + """ + return Counter( + _PREFIX_BUCKETS.get(match.group(1).lower(), "other_updates") + for line in body.splitlines() + if (match := _RELEASE_BULLET_PATTERN.match(line)) is not None + ) + + +def parse_latest_release(response: httpx.Response) -> LatestReleaseInfo | LatestReleaseUnavailable: + if response.status_code != 200: + return LatestReleaseUnavailable(reason=f"GitHub responded with status {response.status_code}") + try: + release: Final = _GitHubRelease.model_validate_json(response.content) + except ValidationError as e: + return LatestReleaseUnavailable(reason=f"GitHub release payload was not the expected shape: {e}") + counts: Final = count_release_bullets(release.body) + return LatestReleaseInfo( + version=release.tag_name.removeprefix("v"), + new_features=counts["new_features"], + bug_fixes=counts["bug_fixes"], + other_updates=counts["other_updates"], + release_url=release.html_url, + ) + + +async def fetch_latest_release(client: _AsyncGetClient) -> LatestReleaseInfo | LatestReleaseUnavailable: + try: + response: Final = await client.get(LATEST_RELEASE_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + except httpx.HTTPError as e: + return LatestReleaseUnavailable(reason=f"{type(e).__name__}: {e}") + return parse_latest_release(response) + + +async def get_latest_release_info( + client: _AsyncGetClient, cache: InMemoryCache +) -> LatestReleaseInfo | LatestReleaseUnavailable: + cached: Final = cache.get_cache(LATEST_RELEASE_CACHE_KEY) + if isinstance(cached, (LatestReleaseInfo, LatestReleaseUnavailable)): + return cached + result: Final = await fetch_latest_release(client) + ttl: Final = ( + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS + if isinstance(result, LatestReleaseUnavailable) + else LATEST_RELEASE_CACHE_TTL_SECONDS + ) + cache.set_cache(LATEST_RELEASE_CACHE_KEY, result, ttl=ttl) + return result + + +@router.get( + "/get/latest_release_info", + tags=["UI Settings"], # mutable-ok: FastAPI's route decorator only accepts a list + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator only accepts a list + response_model=LatestReleaseInfo | None, +) +async def latest_release_info( + client: _AsyncGetClient = Depends(_default_client), + cache: InMemoryCache = Depends(_default_cache), +) -> LatestReleaseInfo | None: + """ + Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + """ + result: Final = await get_latest_release_info(client=client, cache=cache) + if isinstance(result, LatestReleaseUnavailable): + verbose_proxy_logger.warning("LiteLLM: latest release info unavailable: %s", result.reason) + return None + return result diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py new file mode 100644 index 00000000000..0bbcdfe317d --- /dev/null +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py @@ -0,0 +1,191 @@ +import json +import time +from typing import Final + +import httpx +import pytest +from fastapi.testclient import TestClient + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + LATEST_RELEASE_CACHE_KEY, + LATEST_RELEASE_CACHE_TTL_SECONDS, + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS, + LATEST_RELEASE_URL, + LatestReleaseInfo, + LatestReleaseUnavailable, + _default_cache, + _default_client, + count_release_bullets, + get_latest_release_info, +) + +SAMPLE_BODY: Final = """## What's Changed +* feat(proxy): add upgrade banner by @kerry in https://github.com/BerriAI/litellm/pull/1 +* fix(azure): retry on 429 by @a in https://github.com/BerriAI/litellm/pull/2 +* fix: handle empty body by @b in https://github.com/BerriAI/litellm/pull/3 +* Feat(ui)!: drop legacy theme by @c in https://github.com/BerriAI/litellm/pull/4 +* chore(deps): bump httpx by @d in https://github.com/BerriAI/litellm/pull/5 +* docs: fix typo by @e in https://github.com/BerriAI/litellm/pull/6 +* Litellm dev 09 08 2026 by @f in https://github.com/BerriAI/litellm/pull/7 +* refactor(router) : spaced colon does not match by @g in https://github.com/BerriAI/litellm/pull/8 + +## New Contributors +* @kerry made their first contribution in https://github.com/BerriAI/litellm/pull/1 + +**Full Changelog**: https://github.com/BerriAI/litellm/compare/v1.101.0...v1.102.0 +""" + +SAMPLE_RELEASE: Final = { + "tag_name": "v1.102.0", + "html_url": "https://github.com/BerriAI/litellm/releases/tag/v1.102.0", + "body": SAMPLE_BODY, +} +EXPECTED_INFO: Final = { + "version": "1.102.0", + "new_features": 2, + "bug_fixes": 2, + "other_updates": 2, + "release_url": SAMPLE_RELEASE["html_url"], +} + + +class _RecordingClient: + def __init__(self, outcomes: list[httpx.Response | Exception]) -> None: + self._outcomes = outcomes + self.calls: list[tuple[str, float | None]] = [] + + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.calls.append((url, timeout)) + outcome = self._outcomes[min(len(self.calls) - 1, len(self._outcomes) - 1)] + if isinstance(outcome, Exception): + raise outcome + return outcome + + +def _github_response(status: int = 200, payload: object = SAMPLE_RELEASE) -> httpx.Response: + return httpx.Response(status, content=json.dumps(payload).encode()) + + +def _fresh_cache() -> InMemoryCache: + return InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) + + +def _override_dependencies(client: _RecordingClient, cache: InMemoryCache, role: LitellmUserRoles) -> None: + async def auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_id="test-user", user_role=role) + + app.dependency_overrides[user_api_key_auth] = auth + app.dependency_overrides[_default_client] = lambda: client + app.dependency_overrides[_default_cache] = lambda: cache + + +@pytest.fixture +def http_client(): + yield TestClient(app) + app.dependency_overrides.pop(user_api_key_auth, None) + app.dependency_overrides.pop(_default_client, None) + app.dependency_overrides.pop(_default_cache, None) + + +class TestCountReleaseBullets: + def test_buckets_by_conventional_commit_type(self): + counts = count_release_bullets(SAMPLE_BODY) + assert counts["new_features"] == 2 + assert counts["bug_fixes"] == 2 + assert counts["other_updates"] == 2 + + def test_ignores_non_bullet_lines_and_contributor_entries(self): + assert ( + sum(count_release_bullets("## What's Changed\n\n* @x made their first contribution in url\n").values()) == 0 + ) + + def test_empty_body_yields_zero_counts(self): + counts = count_release_bullets("") + assert (counts["new_features"], counts["bug_fixes"], counts["other_updates"]) == (0, 0, 0) + + +class TestGetLatestReleaseInfo: + @pytest.mark.asyncio + async def test_fetches_and_parses_github_release(self): + client = _RecordingClient([_github_response()]) + result = await get_latest_release_info(client=client, cache=_fresh_cache()) + assert isinstance(result, LatestReleaseInfo) + assert result.model_dump() == EXPECTED_INFO + assert client.calls == [(LATEST_RELEASE_URL, 5)] + + @pytest.mark.asyncio + async def test_second_call_within_ttl_does_not_refetch(self): + client = _RecordingClient([_github_response()]) + cache = _fresh_cache() + first = await get_latest_release_info(client=client, cache=cache) + second = await get_latest_release_info(client=client, cache=cache) + assert first == second + assert len(client.calls) == 1 + + @pytest.mark.asyncio + async def test_success_is_cached_for_the_full_ttl(self): + cache = _fresh_cache() + await get_latest_release_info(client=_RecordingClient([_github_response()]), cache=cache) + remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() + assert LATEST_RELEASE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_CACHE_TTL_SECONDS + + @pytest.mark.asyncio + async def test_failure_is_cached_briefly_so_github_is_not_hammered(self): + client = _RecordingClient([httpx.ConnectError("boom")]) + cache = _fresh_cache() + first = await get_latest_release_info(client=client, cache=cache) + second = await get_latest_release_info(client=client, cache=cache) + assert isinstance(first, LatestReleaseUnavailable) + assert first == second + assert len(client.calls) == 1 + remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() + assert ( + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "response", + [ + _github_response(status=403, payload={"message": "rate limited"}), + _github_response(status=500, payload={}), + _github_response(payload={"tag_name": "v1.0.0"}), + httpx.Response(200, content=b"not json"), + ], + ids=["rate_limited", "server_error", "missing_fields", "not_json"], + ) + async def test_bad_github_responses_are_unavailable(self, response: httpx.Response): + result = await get_latest_release_info(client=_RecordingClient([response]), cache=_fresh_cache()) + assert isinstance(result, LatestReleaseUnavailable) + + +class TestLatestReleaseInfoEndpoint: + def test_returns_release_stats_for_authenticated_user(self, http_client): + _override_dependencies(_RecordingClient([_github_response()]), _fresh_cache(), LitellmUserRoles.INTERNAL_USER) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() == EXPECTED_INFO + + def test_returns_null_when_github_is_unreachable(self, http_client): + _override_dependencies( + _RecordingClient([httpx.ConnectError("boom")]), _fresh_cache(), LitellmUserRoles.PROXY_ADMIN + ) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() is None + + def test_repeated_requests_reuse_cache(self, http_client): + client = _RecordingClient([_github_response()]) + _override_dependencies(client, _fresh_cache(), LitellmUserRoles.PROXY_ADMIN) + assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO + assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO + assert len(client.calls) == 1 + + def test_rejects_unauthenticated_requests(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-1234") + response = TestClient(app).get("/get/latest_release_info") + assert response.status_code in (401, 403) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts new file mode 100644 index 00000000000..5186baa605c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts @@ -0,0 +1,16 @@ +import { $api } from "@/lib/http/api"; +import type { components } from "@/lib/http/schema"; + +export type LatestReleaseInfo = components["schemas"]["LatestReleaseInfo"]; + +export const useLatestReleaseInfo = (accessToken: string | null | undefined) => + $api.useQuery( + "get", + "/get/latest_release_info", + {}, + { + enabled: Boolean(accessToken), + staleTime: 60 * 60 * 1000, + retry: false, + }, + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 98f2a36d6f3..928fc5f2410 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -12,6 +12,7 @@ import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner"; import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; +import { UpgradeBanner } from "@/components/UpgradeBanner"; import { uiHref } from "@/utils/uiHref"; import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; import { createApiClient } from "@/lib/http/client"; @@ -115,6 +116,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
@@ -134,6 +136,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
{children}
diff --git a/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx new file mode 100644 index 00000000000..be031ea8ad5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx @@ -0,0 +1,108 @@ +import React from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { describeRelease, UpgradeBanner, UpgradeBannerView } from "./UpgradeBanner"; +import type { LatestReleaseInfo } from "@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo"; + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({ + useHealthReadinessDetails: vi.fn(), +})); +vi.mock("@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo", () => ({ + useLatestReleaseInfo: vi.fn(), +})); + +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useLatestReleaseInfo } from "@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo"; + +const RELEASE: LatestReleaseInfo = { + version: "1.103.0", + new_features: 12, + bug_fixes: 30, + other_updates: 8, + release_url: "https://github.com/BerriAI/litellm/releases/tag/v1.103.0", +}; + +describe("describeRelease", () => { + it("lists features, fixes, and other updates in the agreed order", () => { + expect(describeRelease(RELEASE)).toBe("12 new features, 30 fixes, and 8 other updates"); + }); + + it("singularises counts of one", () => { + expect(describeRelease({ ...RELEASE, new_features: 1, bug_fixes: 1, other_updates: 1 })).toBe( + "1 new feature, 1 fix, and 1 other update", + ); + }); +}); + +describe("UpgradeBannerView", () => { + beforeEach(() => { + localStorage.clear(); + }); + + afterEach(() => { + localStorage.clear(); + }); + + it("renders nothing while either version is unknown", () => { + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + const { container: noRelease } = render(); + expect(noRelease).toBeEmptyDOMElement(); + }); + + it("renders nothing when the running version is up to date or ahead", () => { + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + const { container: ahead } = render(); + expect(ahead).toBeEmptyDOMElement(); + }); + + it("shows the latest version, the stat line, and the current version when behind", () => { + render(); + const alert = screen.getByRole("alert"); + expect(alert).toHaveTextContent("The latest version is v1.103.0: 12 new features, 30 fixes, and 8 other updates"); + expect(alert).toHaveTextContent("Your current version is v1.102.0"); + expect(screen.getByRole("link", { name: "v1.103.0" })).toHaveAttribute("href", RELEASE.release_url); + }); + + it("dismissing hides the banner and keeps it hidden on remount for the same release", () => { + const { unmount } = render(); + fireEvent.click(screen.getByRole("button", { name: "Close" })); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + unmount(); + + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + }); + + it("reappears once a newer release ships after a dismissal", () => { + const { unmount } = render(); + fireEvent.click(screen.getByRole("button", { name: "Close" })); + unmount(); + + render(); + expect(screen.getByRole("alert")).toHaveTextContent("The latest version is v1.104.0"); + }); +}); + +describe("UpgradeBanner", () => { + beforeEach(() => { + localStorage.clear(); + }); + + it("feeds both hooks the access token and renders from their data", () => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data: { litellm_version: "1.102.0" } } as any); + vi.mocked(useLatestReleaseInfo).mockReturnValue({ data: RELEASE } as any); + render(); + expect(useHealthReadinessDetails).toHaveBeenCalledWith("token"); + expect(useLatestReleaseInfo).toHaveBeenCalledWith("token"); + expect(screen.getByRole("alert")).toHaveTextContent("The latest version is v1.103.0"); + }); + + it("renders nothing when the release endpoint returns null", () => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data: { litellm_version: "1.102.0" } } as any); + vi.mocked(useLatestReleaseInfo).mockReturnValue({ data: null } as any); + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UpgradeBanner.tsx b/ui/litellm-dashboard/src/components/UpgradeBanner.tsx new file mode 100644 index 00000000000..cc57223eaff --- /dev/null +++ b/ui/litellm-dashboard/src/components/UpgradeBanner.tsx @@ -0,0 +1,77 @@ +"use client"; + +import React, { useState } from "react"; +import { ArrowUpCircle, X } from "lucide-react"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { + type LatestReleaseInfo, + useLatestReleaseInfo, +} from "@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo"; +import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; +import { isNewerVersion } from "@/utils/versionUtils"; + +const DISMISS_KEY_PREFIX = "litellm:upgradeBannerDismissed:"; + +interface UpgradeBannerProps { + accessToken: string | null; +} + +interface UpgradeBannerViewProps { + currentVersion: string | null | undefined; + latestRelease: LatestReleaseInfo | null | undefined; +} + +const plural = (count: number, singular: string, pluralForm: string): string => + `${count} ${count === 1 ? singular : pluralForm}`; + +export const describeRelease = ({ new_features, bug_fixes, other_updates }: LatestReleaseInfo): string => + [ + plural(new_features, "new feature", "new features"), + plural(bug_fixes, "fix", "fixes"), + `and ${plural(other_updates, "other update", "other updates")}`, + ].join(", "); + +export const UpgradeBannerView: React.FC = ({ currentVersion, latestRelease }) => { + const [locallyDismissed, setLocallyDismissed] = useState(false); + + if (!currentVersion || !latestRelease || !isNewerVersion(currentVersion, latestRelease.version)) { + return null; + } + + const dismissKey = `${DISMISS_KEY_PREFIX}${latestRelease.version}`; + if (locallyDismissed || getLocalStorageItem(dismissKey) === "true") { + return null; + } + + const handleClose = () => { + setLocalStorageItem(dismissKey, "true"); + setLocallyDismissed(true); + }; + + return ( + + + + The latest version is{" "} + + v{latestRelease.version} + + : {describeRelease(latestRelease)} + + Your current version is v{currentVersion} + + + + + ); +}; + +export const UpgradeBanner: React.FC = ({ accessToken }) => { + const { data: healthData } = useHealthReadinessDetails(accessToken); + const { data: latestRelease } = useLatestReleaseInfo(accessToken); + return ; +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..cd12fd10528 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -5109,6 +5109,27 @@ export interface paths { patch?: never; trace?: never; }; + "/get/latest_release_info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Latest Release Info + * @description Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + * Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + */ + get: operations["latest_release_info_get_latest_release_info_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/get/mcp_semantic_filter_settings": { parameters: { query?: never; @@ -28352,6 +28373,19 @@ export interface components { } & { [key: string]: unknown; }; + /** LatestReleaseInfo */ + LatestReleaseInfo: { + /** Bug Fixes */ + bug_fixes: number; + /** New Features */ + new_features: number; + /** Other Updates */ + other_updates: number; + /** Release Url */ + release_url: string; + /** Version */ + version: string; + }; /** ListAccessGroupsResponse */ ListAccessGroupsResponse: { /** Access Groups */ @@ -47428,6 +47462,26 @@ export interface operations { }; }; }; + latest_release_info_get_latest_release_info_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LatestReleaseInfo"] | null; + }; + }; + }; + }; get_mcp_semantic_filter_settings_get_mcp_semantic_filter_settings_get: { parameters: { query?: never; diff --git a/ui/litellm-dashboard/src/utils/versionUtils.test.ts b/ui/litellm-dashboard/src/utils/versionUtils.test.ts new file mode 100644 index 00000000000..714c7165d6d --- /dev/null +++ b/ui/litellm-dashboard/src/utils/versionUtils.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from "vitest"; +import { isNewerVersion, parseReleaseVersion } from "./versionUtils"; + +describe("parseReleaseVersion", () => { + it("reads the numeric components with or without a leading v", () => { + expect(parseReleaseVersion("1.102.0")).toEqual([1, 102, 0]); + expect(parseReleaseVersion("v1.102.0")).toEqual([1, 102, 0]); + }); + + it("ignores a prerelease suffix", () => { + expect(parseReleaseVersion("1.102.0-dev.3")).toEqual([1, 102, 0]); + expect(parseReleaseVersion("1.102.0.rc1")).toEqual([1, 102, 0]); + }); + + it("returns null for strings that are not a release version", () => { + expect(parseReleaseVersion("")).toBeNull(); + expect(parseReleaseVersion("latest")).toBeNull(); + expect(parseReleaseVersion("1.102")).toBeNull(); + }); +}); + +describe("isNewerVersion", () => { + it("is false when the versions are equal", () => { + expect(isNewerVersion("1.102.0", "1.102.0")).toBe(false); + expect(isNewerVersion("1.102.0", "v1.102.0")).toBe(false); + }); + + it("is true when the latest version is ahead on any component", () => { + expect(isNewerVersion("1.102.0", "1.102.1")).toBe(true); + expect(isNewerVersion("1.102.5", "1.103.0")).toBe(true); + expect(isNewerVersion("1.999.9", "2.0.0")).toBe(true); + }); + + it("is false when the running version is already ahead", () => { + expect(isNewerVersion("1.103.0", "1.102.9")).toBe(false); + expect(isNewerVersion("2.0.0", "1.999.9")).toBe(false); + }); + + it("compares components numerically rather than as strings", () => { + expect(isNewerVersion("1.9.0", "1.10.0")).toBe(true); + expect(isNewerVersion("1.10.0", "1.9.0")).toBe(false); + }); + + it("treats a prerelease of the latest version as not behind", () => { + expect(isNewerVersion("1.102.0-dev.1", "1.102.0")).toBe(false); + expect(isNewerVersion("1.101.0-dev.1", "1.102.0")).toBe(true); + }); + + it("is false when either version cannot be parsed", () => { + expect(isNewerVersion("unknown", "1.102.0")).toBe(false); + expect(isNewerVersion("1.102.0", "")).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/versionUtils.ts b/ui/litellm-dashboard/src/utils/versionUtils.ts new file mode 100644 index 00000000000..62fcb97c0ef --- /dev/null +++ b/ui/litellm-dashboard/src/utils/versionUtils.ts @@ -0,0 +1,23 @@ +const RELEASE_VERSION_PATTERN = /^v?(\d+)\.(\d+)\.(\d+)/; + +export const parseReleaseVersion = (version: string): readonly [number, number, number] | null => { + const match = RELEASE_VERSION_PATTERN.exec(version.trim()); + if (!match) { + return null; + } + return [Number(match[1]), Number(match[2]), Number(match[3])]; +}; + +export const isNewerVersion = (current: string, latest: string): boolean => { + const currentParts = parseReleaseVersion(current); + const latestParts = parseReleaseVersion(latest); + if (!currentParts || !latestParts) { + return false; + } + for (let i = 0; i < 3; i += 1) { + if (latestParts[i] !== currentParts[i]) { + return latestParts[i] > currentParts[i]; + } + } + return false; +}; From 7edc79583139309ef7571e3bb5b77828d9cb535e Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:21:39 -0700 Subject: [PATCH 011/149] fix(ui): fix CI failures on upgrade banner PR Use Annotated[..., Depends(...)] instead of a call in the default value to clear the B008 ruff budget. Extract an inline test object over the no-large-inline-object-arg eslint budget. Mock UpgradeBanner in the layout test, matching the other top banners, since it now renders through react-query hooks that need a QueryClientProvider the test doesn't set up. Co-Authored-By: Claude Code --- litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py | 6 +++--- ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx | 4 ++++ ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx | 5 ++--- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py index c0306e65aee..4a01728dfe1 100644 --- a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -3,7 +3,7 @@ from collections import Counter from collections.abc import Awaitable, Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Literal, Protocol +from typing import Annotated, Final, Literal, Protocol import httpx from fastapi import APIRouter, Depends @@ -124,8 +124,8 @@ async def get_latest_release_info( response_model=LatestReleaseInfo | None, ) async def latest_release_info( - client: _AsyncGetClient = Depends(_default_client), - cache: InMemoryCache = Depends(_default_cache), + client: Annotated[_AsyncGetClient, Depends(_default_client)], + cache: Annotated[InMemoryCache, Depends(_default_cache)], ) -> LatestReleaseInfo | None: """ Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7f1f4cc4bd5..e1fb6e21596 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -37,6 +37,10 @@ vi.mock("@/components/UserBanner", () => ({ UserBanner: () => null, })); +vi.mock("@/components/UpgradeBanner", () => ({ + UpgradeBanner: () => null, +})); + vi.mock("@/contexts/ThemeContext", () => ({ ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, })); diff --git a/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx index be031ea8ad5..6aa384600c8 100644 --- a/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx @@ -28,9 +28,8 @@ describe("describeRelease", () => { }); it("singularises counts of one", () => { - expect(describeRelease({ ...RELEASE, new_features: 1, bug_fixes: 1, other_updates: 1 })).toBe( - "1 new feature, 1 fix, and 1 other update", - ); + const singularCounts = { ...RELEASE, new_features: 1, bug_fixes: 1, other_updates: 1 }; + expect(describeRelease(singularCounts)).toBe("1 new feature, 1 fix, and 1 other update"); }); }); From daa665e578fcdde3e9a7d1a780581d0e3558b682 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 21:46:25 -0700 Subject: [PATCH 012/149] build(deps): re-suppress GHSA-h7x2-h6g9-p789 in osv-scan, mlflow still has no fixed release (#41036) Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- osv-scanner.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/osv-scanner.toml b/osv-scanner.toml index 3e070fc8cf7..482254d4da6 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -5,5 +5,5 @@ reason = "diskcache has no fixed release published; remove this entry once one e [[IgnoredVulns]] id = "GHSA-h7x2-h6g9-p789" -ignoreUntil = 2026-09-14 -reason = "mlflow has no fixed release published; remove this entry once one exists" +ignoreUntil = 2026-10-14 +reason = "mlflow has no fixed release published (3.16.0, 2026-09-04, and master still store gateway secret api_base unvalidated); remove this entry once one exists" From bf8df3ab022b81d54b9d8dcc3d97f8002d1bc3a5 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 16:18:46 +0200 Subject: [PATCH 013/149] hibp support in password policy --- litellm/proxy/_types.py | 11 ++ litellm/proxy/auth/password_policy.py | 68 +++++++++ .../internal_user_endpoints.py | 9 +- litellm/proxy/proxy_server.py | 3 +- litellm/types/llms/custom_http.py | 1 + .../proxy/auth/test_onboarding.py | 140 +++++++++++++++++- .../proxy/auth/test_password_policy.py | 134 +++++++++++++++++ .../test_internal_user_endpoints.py | 34 +++++ tests/test_litellm/proxy/test__types.py | 16 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 10 files changed, 409 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b315a2beac9..9b38b12ab07 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1808,6 +1808,17 @@ class NewUserRequest(GenerateRequestBase): send_invite_email: bool | None = None sso_user_id: str | None = None organizations: list[str] | None = None + password: str | None = None + + @field_validator("password") + @classmethod + def password_not_supported(cls, value: str | None) -> str | None: + if value is not None: + raise ValueError( + "password cannot be set via /user/new. Users set their own password through an " + "invitation link (POST /invitation/new)." + ) + return value class NewUserResponse(GenerateKeyResponse): diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index ab7a565894a..80c3ebeedca 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -4,13 +4,26 @@ Applied at every path that persists a new or changed password for a DB-backed user (``/user/update``, ``/user/bulk_update``, and the invitation onboarding claim flow), so the strength bar is configured in one place instead of per-endpoint. + +Also screens new passwords against known data breaches via the +haveibeenpwned.com (HIBP) k-anonymity range API: only the first 5 characters +of the password's SHA-1 hash ever leave the proxy, and the check fails open +(allows the password) when HIBP is unreachable. """ +import hashlib from collections.abc import Mapping from dataclasses import dataclass from typing import Final +from litellm._logging import verbose_proxy_logger +from litellm._version import version +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.types.llms.custom_http import httpxSpecialProvider + +HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" +HIBP_TIMEOUT_SECONDS: Final = 5.0 DEFAULT_MIN_LENGTH: Final = 12 MIN_ALLOWED_LENGTH: Final = 8 @@ -90,3 +103,58 @@ def validate_password_policy(password: str, general_settings: Mapping[str, objec param="password", code=400, ) + + +def _hibp_client() -> AsyncHTTPHandler: + return get_async_httpx_client( + llm_provider=httpxSpecialProvider.PasswordBreachCheck, + params={"timeout": HIBP_TIMEOUT_SECONDS}, + ) + + +def _is_suffix_in_range_response(response_body: str, hash_suffix: str) -> bool: + for line in response_body.upper().splitlines(): + entry_suffix, _, count = line.strip().partition(":") + if entry_suffix == hash_suffix: + return int(count.strip() or "0") > 0 + return False + + +async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool: + # usedforsecurity=False: SHA-1 is only a lookup key into the HIBP dataset, so no security property rests on it + sha1_hex: Final = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + try: + response: Final = await client.get( + f"{HIBP_RANGE_API_BASE}/{sha1_hex[:5]}", + headers={"Add-Padding": "true", "User-Agent": f"litellm-proxy/{version}"}, + ) + response.raise_for_status() + breached: Final = _is_suffix_in_range_response(response.text, sha1_hex[5:]) + except Exception as e: + verbose_proxy_logger.warning("Breached-password check skipped, HIBP lookup failed: %s", e) + return False + return breached + + +async def validate_password_not_breached( + password: str, + general_settings: Mapping[str, object], + client: AsyncHTTPHandler | None = None, +) -> None: + """Raise ``ProxyException`` (400) if ``password`` appears in a known data breach. + + Fails open: an unreachable or misbehaving HIBP allows the password.""" + check_enabled: Final = general_settings.get("password_policy_check_breached_passwords", True) is not False + if not check_enabled: + return + if not await _is_password_breached(password, client if client is not None else _hibp_client()): + return + raise ProxyException( + message=( + "This password appears in known data breaches and cannot be used. " + "Please choose a different password." + ), + type=ProxyErrorTypes.validation_error, + param="password", + code=400, + ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e3efda507f6..f22bad70817 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -29,7 +29,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object -from litellm.proxy.auth.password_policy import validate_password_policy +from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import ( @@ -162,10 +162,11 @@ def _team_membership_table( return team_membership_table -def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: +async def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: """Validate and hash password field in-place if present.""" if "password" in data and data["password"] is not None: validate_password_policy(data["password"], general_settings) + await validate_password_not_breached(data["password"], general_settings) data["password"] = hash_password(data["password"]) @@ -561,7 +562,7 @@ async def new_user( # generate_key_helper_fn only forwards object_permission_id, so without this the entitlement # the caller sent would be dropped on the floor. data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client) - _hash_password_in_dict(data_json, general_settings) + data_json.pop("password", None) # always None: NewUserRequest.password_not_supported rejects any other value teams = data.teams if teams is None: teams = check_if_default_team_set() @@ -1449,7 +1450,7 @@ async def _update_single_user_helper( data_json: Final[dict] = user_request.model_dump(exclude_unset=True) non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) - _hash_password_in_dict(non_default_values, general_settings) + await _hash_password_in_dict(non_default_values, general_settings) existing_user_row: BaseModel | None = None if user_request.user_id: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..d2b486b4410 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -331,7 +331,7 @@ from litellm.proxy.auth.model_checks import ( get_mcp_server_ids, get_team_models, ) -from litellm.proxy.auth.password_policy import validate_password_policy +from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( _fetch_global_spend_with_event_coordination, user_api_key_auth, @@ -16369,6 +16369,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): ) validate_password_policy(data.password, general_settings) + await validate_password_not_breached(data.password, general_settings) hashed_pw: Final = hash_password(data.password) current_time = litellm.utils.get_utc_datetime() async with prisma_client.db.tx() as tx: diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index d80d7410aae..793893451df 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -31,6 +31,7 @@ class httpxSpecialProvider(str, Enum): UI = "ui" Sandbox = "sandbox" ModelCostMap = "model_cost_map" + PasswordBreachCheck = "password_breach_check" VerifyTypes = str | bool | ssl.SSLContext diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index 524b655b465..8939baedd50 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -8,15 +8,20 @@ Covers the security behavior of: session key only after the password is written """ +import hashlib from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch +import httpx import jwt import pytest +import respx from fastapi import HTTPException import litellm -from litellm.proxy._types import InvitationClaim +from litellm.proxy._types import InvitationClaim, ProxyException + +_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False} # --------------------------------------------------------------------------- # Helpers @@ -386,7 +391,9 @@ async def test_claim_token_rejects_concurrent_reuse_before_password_write(): with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.master_key", "sk-test"), - patch("litellm.proxy.proxy_server.general_settings", {}), + patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + "litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK + ), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", new_callable=AsyncMock, @@ -426,7 +433,9 @@ async def test_claim_token_sets_accepted_at_after_password_written(): with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.master_key", "sk-test"), - patch("litellm.proxy.proxy_server.general_settings", {}), + patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + "litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK + ), patch("litellm.proxy.proxy_server.premium_user", False), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", @@ -483,7 +492,9 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.master_key", "sk-test"), - patch("litellm.proxy.proxy_server.general_settings", {}), + patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + "litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK + ), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", new_callable=AsyncMock, @@ -505,3 +516,124 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): } assert rollback_kwargs["data"]["accepted_at"] is None assert rollback_kwargs["data"]["is_accepted"] is False + + +# --------------------------------------------------------------------------- +# POST /onboarding/claim_token - password policy +# --------------------------------------------------------------------------- + + +def _hibp_url_for(password: str) -> str: + sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + return f"https://api.pwnedpasswords.com/range/{sha1[:5]}" + + +def _hibp_suffix_for(password: str) -> str: + return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()[5:] + + +@pytest.mark.asyncio +async def test_claim_token_rejects_short_password_before_consuming_invite(): + """Default policy requires 12 characters; the invite must stay claimable.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite, _make_user()) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="Sh0rt!pw", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above + ): + with pytest.raises(ProxyException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.code == "400" + assert "at least 12 characters" in exc_info.value.message + prisma.db.litellm_invitationlink.update_many.assert_not_called() + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_claim_token_rejects_breached_password_before_consuming_invite(): + """A password found in the HIBP corpus must be rejected and never stored.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + password = "P@ssword123456" + respx.get(_hibp_url_for(password)).mock( + return_value=httpx.Response(200, text=f"{_hibp_suffix_for(password)}:1387") + ) + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite, _make_user()) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password=password, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above + ): + with pytest.raises(ProxyException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.code == "400" + assert "data breaches" in exc_info.value.message + prisma.db.litellm_invitationlink.update_many.assert_not_called() + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_claim_token_fails_open_when_hibp_unreachable(): + """An HIBP outage must never block onboarding: the claim proceeds.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + password = "NewP@ssw0rd-2026" + respx.get(_hibp_url_for(password)).mock(side_effect=httpx.ConnectError("no route to host")) + + invite = _make_invite(is_accepted=False) + user = _make_user() + prisma = _make_prisma(invite, user) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password=password, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: same as above + patch( # test-quality-ok: same as above + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "sk-generated-key", "user_id": "user-123"}, + ), + patch( # test-quality-ok: same as above + "litellm.proxy.proxy_server.get_custom_url", + return_value="http://localhost:4000/", + ), + patch( # test-quality-ok: same as above + "litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation", + return_value=False, + ), + patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), # test-quality-ok: same as above + ): + result = await claim_onboarding_link(data=data, request=request) + + assert "token" in result + prisma.db.litellm_usertable.update.assert_called_once() diff --git a/tests/test_litellm/proxy/auth/test_password_policy.py b/tests/test_litellm/proxy/auth/test_password_policy.py index f6e7d443907..edc88d21218 100644 --- a/tests/test_litellm/proxy/auth/test_password_policy.py +++ b/tests/test_litellm/proxy/auth/test_password_policy.py @@ -2,22 +2,54 @@ Tests for the configurable password-strength policy in `litellm.proxy.auth.password_policy`, enforced on every path that persists a new or changed password for a locally-managed user. + +The breach-check (HIBP) tests inject a real AsyncHTTPHandler wrapping an +httpx.MockTransport, so no network is touched and nothing is monkeypatched. """ +import hashlib + +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.password_policy import ( DEFAULT_MIN_LENGTH, MIN_ALLOWED_LENGTH, PasswordPolicy, get_password_policy, + validate_password_not_breached, validate_password_policy, ) STRONG_PASSWORD = "Str0ng!Passw0rd" +def _sha1_upper(password: str) -> str: + return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + + +def _client_with_transport(handler) -> AsyncHTTPHandler: + http_handler = AsyncHTTPHandler() + http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + return http_handler + + +def _client_never_called() -> AsyncHTTPHandler: + def handler(request: httpx.Request) -> httpx.Response: + raise AssertionError(f"unexpected HTTP call to {request.url}") + + return _client_with_transport(handler) + + +def _client_returning(body: str, status_code: int = 200) -> AsyncHTTPHandler: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(status_code, text=body) + + return _client_with_transport(handler) + + def test_get_password_policy_defaults_to_pif_baseline(): policy = get_password_policy({}) assert policy == PasswordPolicy( @@ -134,3 +166,105 @@ def test_validate_password_policy_rejects_unicode_letter_as_special_character(): def test_validate_password_policy_accepts_real_special_character_with_unicode_letters(): """Same base password as the rejection test above, plus an actual symbol.""" assert validate_password_policy("Passwörd1234!", {}) is None + + +@pytest.mark.asyncio +async def test_breach_check_skipped_when_disabled(): + result = await validate_password_not_breached( + password="password12345", # breached in reality, but the check is off + general_settings={"password_policy_check_breached_passwords": False}, + client=_client_never_called(), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_rejects_breached_password(): + password = "correct horse battery staple" + sha1 = _sha1_upper(password) + body = f"AAAA000000000000000000000000000000A:0\r\n{sha1[5:]}:42\r\nBBBB000000000000000000000000000000B:7" + + with pytest.raises(ProxyException) as exc_info: + await validate_password_not_breached(password=password, general_settings={}, client=_client_returning(body)) + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert "data breaches" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_only_sha1_prefix_leaves_the_proxy(): + password = "a very secret password" + sha1 = _sha1_upper(password) + captured_requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + result = await validate_password_not_breached( + password=password, general_settings={}, client=_client_with_transport(handler) + ) + assert result is None + + (request,) = captured_requests + assert request.url.path == f"/range/{sha1[:5]}" + assert sha1[5:] not in str(request.url) + assert request.headers["Add-Padding"] == "true" + assert "litellm" in request.headers["User-Agent"] + + +@pytest.mark.asyncio +async def test_ignores_padding_entries_with_zero_count(): + """HIBP padding entries (requested via Add-Padding) carry count 0 and must + not be treated as breaches when they collide with the password's suffix.""" + password = "a padded-away password" + sha1 = _sha1_upper(password) + + result = await validate_password_not_breached( + password=password, general_settings={}, client=_client_returning(f"{sha1[5:]}:0") + ) + assert result is None + + +@pytest.mark.asyncio +async def test_accepts_password_absent_from_breach_corpus(): + result = await validate_password_not_breached( + password="a genuinely novel password", + general_settings={}, + client=_client_returning("0018A45C4D1DEF81644B54AB7F969B88D65:1\r\n00D4F6E8FA6EECAD2A3AA415EEC418D38EC:2"), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_breach_check_fails_open_on_network_error(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("no route to host") + + result = await validate_password_not_breached( + password="password12345", # breached, but HIBP is unreachable + general_settings={}, + client=_client_with_transport(handler), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_breach_check_fails_open_on_http_error_status(): + result = await validate_password_not_breached( + password="password12345", + general_settings={}, + client=_client_returning("service unavailable", status_code=503), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_breach_check_fails_open_on_malformed_response_body(): + result = await validate_password_not_breached( + password="password12345", + general_settings={}, + client=_client_returning(f"{_sha1_upper('password12345')[5:]}:not-a-number"), + ) + assert result is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0d8b19345f1..e7b5172e0fb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1,9 +1,12 @@ +import hashlib import json from datetime import datetime, timezone from types import SimpleNamespace from typing import Final +import httpx import pytest +import respx from fastapi.testclient import TestClient from fastapi import HTTPException from pytest_mock import MockerFixture @@ -4502,6 +4505,11 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo _update_single_user_helper, ) + mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses + "litellm.proxy.proxy_server.general_settings", + {"password_policy_check_breached_passwords": False}, + ) + mock_prisma_client = _admin_prisma existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user"} @@ -4519,3 +4527,29 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo written_data = mock_prisma_client.update_data.call_args.kwargs["data"] assert written_data.get("password") is not None assert written_data["password"] != strong_password + + +@pytest.mark.asyncio +@respx.mock +async def test_user_update_rejects_breached_password(_admin_prisma): + """A strength-passing password found in the HIBP corpus must be rejected + before it ever reaches the DB write.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + password = "Str0ng!Passw0rd" + sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + respx.get(f"https://api.pwnedpasswords.com/range/{sha1[:5]}").mock( + return_value=httpx.Response(200, text=f"{sha1[5:]}:1387") + ) + + user_request = UpdateUserRequest(user_id="target-user", password=password) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) + + assert exc_info.value.code == "400" + assert "data breaches" in exc_info.value.message + _admin_prisma.db.litellm_usertable.find_first.assert_not_called() diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 26bb1533da4..8f9c44d7a38 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -10,6 +10,7 @@ from litellm.proxy._types import ( LiteLLM_AuditLogs, LiteLLM_TeamMembership, LitellmUserRoles, + NewUserRequest, OrganizationMemberUpdateRequest, ResetSpendRequest, UpdateKeyRequest, @@ -277,3 +278,18 @@ def test_team_membership_budget_table_present_still_works(): } result = LiteLLM_TeamMembership.model_validate(data) assert result.litellm_budget_table is None + + +def test_new_user_request_loudly_rejects_a_password(): + """ + /user/new has never persisted a password (the field used to be silently + dropped). Sending one must now fail visibly so the dead path cannot be + revived without going through the password policy. + """ + with pytest.raises(ValidationError, match="invitation link"): + NewUserRequest(user_email="alice@example.com", password="hunter2hunter2") + + +def test_new_user_request_without_password_still_works(): + request = NewUserRequest(user_email="alice@example.com") + assert request.password is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..0cc49382498 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32869,6 +32869,8 @@ export interface components { object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null; /** Organizations */ organizations?: string[] | null; + /** Password */ + password?: string | null; /** * Permissions * @default {} From e2ea7e97a58300c2d885c1f1c04854799fbbfcf6 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 16:54:26 +0200 Subject: [PATCH 014/149] fix(auth): document /user/new password rejection and format password_policy --- litellm/proxy/auth/password_policy.py | 3 +-- litellm/proxy/management_endpoints/internal_user_endpoints.py | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 80c3ebeedca..84680fa9020 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -151,8 +151,7 @@ async def validate_password_not_breached( return raise ProxyException( message=( - "This password appears in known data breaches and cannot be used. " - "Please choose a different password." + "This password appears in known data breaches and cannot be used. Please choose a different password." ), type=ProxyErrorTypes.validation_error, param="password", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f22bad70817..409877bdc6f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -495,6 +495,7 @@ async def new_user( - prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts. - organizations: List[str] - List of organization id's the user is a member of - budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. + - password: Optional[str] - Not supported; any value is rejected with a 422. Users set their own password through an invitation link (POST /invitation/new). Returns: - key: (str) The generated api key for the user - expires: (datetime) Datetime object for when key expires. From 1f0ab3d176af4955a1c75b1b44371a310529e169 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 16:57:28 +0200 Subject: [PATCH 015/149] fix(auth): drop general_settings import left unused in new_user --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 409877bdc6f..ddef2127945 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -515,7 +515,7 @@ async def new_user( ``` """ try: - from litellm.proxy.proxy_server import _license_check, general_settings, prisma_client + from litellm.proxy.proxy_server import _license_check, prisma_client if prisma_client is None: raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) From fcf7cb6e0c7fc5d39c7f4017630c48becb086ca9 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 17:12:27 +0200 Subject: [PATCH 016/149] fix(ui): regenerate schema.d.ts for the new_user password docstring --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0cc49382498..c5ee810069d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16798,6 +16798,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts. * - organizations: List[str] - List of organization id's the user is a member of * - budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. + * - password: Optional[str] - Not supported; any value is rejected with a 422. Users set their own password through an invitation link (POST /invitation/new). * Returns: * - key: (str) The generated api key for the user * - expires: (datetime) Datetime object for when key expires. From a9a0bcb9f84909cb41e3670b881292d2e415c1b8 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Fri, 4 Sep 2026 21:13:14 +0200 Subject: [PATCH 017/149] move hibp url to constants --- litellm/constants.py | 3 +++ litellm/proxy/auth/password_policy.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..59a7c3d1e40 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,3 +2039,6 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) + +# API endpoint for breached password k-anonymity search +HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" \ No newline at end of file diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 84680fa9020..958547d6fc7 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -18,11 +18,11 @@ from typing import Final from litellm._logging import verbose_proxy_logger from litellm._version import version +from litellm.constants import HIBP_RANGE_API_BASE from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.types.llms.custom_http import httpxSpecialProvider -HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" HIBP_TIMEOUT_SECONDS: Final = 5.0 DEFAULT_MIN_LENGTH: Final = 12 From 0bb0218d0b60dddb1a05ad2e5ef498e417e210e9 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 16:27:06 +0200 Subject: [PATCH 018/149] fix(auth): screen bulk-update passwords concurrently before any db write /user/bulk_update awaited a separate HIBP lookup for each user in the batch, so a degraded-slow HIBP (5s timeout per lookup) could stretch a 500-user batch to ~2500s and time out the request after some updates had already persisted. validate_passwords_bulk dedupes the batch's passwords, strength-checks first, then fires every needed HIBP lookup concurrently, bounding the worst case at one 5s timeout window. bulk_update_processed_users now screens the whole batch before the serial update loop, so a rejected password fails only its own entry and validation failures precede any persistence. --- litellm/constants.py | 2 +- litellm/proxy/auth/password_policy.py | 79 ++- .../internal_user_endpoints.py | 43 +- .../proxy/auth/test_password_policy.py | 75 +++ .../test_internal_user_endpoints.py | 575 +++++++----------- 5 files changed, 390 insertions(+), 384 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 59a7c3d1e40..5f8fa203b37 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2041,4 +2041,4 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" EMPTY_MAPPING: Final = MappingProxyType({}) # API endpoint for breached password k-anonymity search -HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" \ No newline at end of file +HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 958547d6fc7..c27f276252a 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -11,9 +11,11 @@ of the password's SHA-1 hash ever leave the proxy, and the check fails open (allows the password) when HIBP is unreachable. """ +import asyncio import hashlib -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Final from litellm._logging import verbose_proxy_logger @@ -136,6 +138,33 @@ async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool return breached +def is_breach_check_enabled(general_settings: Mapping[str, object]) -> bool: + return general_settings.get("password_policy_check_breached_passwords", True) is not False + + +async def is_password_breached( + password: str, + general_settings: Mapping[str, object], + client: AsyncHTTPHandler | None = None, +) -> bool: + """False when the check is disabled, the password is absent from the HIBP + corpus, or HIBP is unreachable (fail open).""" + if not is_breach_check_enabled(general_settings): + return False + return await _is_password_breached(password, client if client is not None else _hibp_client()) + + +def breached_password_error() -> ProxyException: + return ProxyException( + message=( + "This password appears in known data breaches and cannot be used. Please choose a different password." + ), + type=ProxyErrorTypes.validation_error, + param="password", + code=400, + ) + + async def validate_password_not_breached( password: str, general_settings: Mapping[str, object], @@ -144,16 +173,42 @@ async def validate_password_not_breached( """Raise ``ProxyException`` (400) if ``password`` appears in a known data breach. Fails open: an unreachable or misbehaving HIBP allows the password.""" - check_enabled: Final = general_settings.get("password_policy_check_breached_passwords", True) is not False - if not check_enabled: + if not await is_password_breached(password, general_settings, client): return - if not await _is_password_breached(password, client if client is not None else _hibp_client()): - return - raise ProxyException( - message=( - "This password appears in known data breaches and cannot be used. Please choose a different password." - ), - type=ProxyErrorTypes.validation_error, - param="password", - code=400, + raise breached_password_error() + + +def _strength_verdict(password: str, general_settings: Mapping[str, object]) -> ProxyException | None: + try: + validate_password_policy(password, general_settings) + except ProxyException as e: + return e + return None + + +async def validate_passwords_bulk( + passwords: Sequence[str], + general_settings: Mapping[str, object], + client: AsyncHTTPHandler | None = None, +) -> Mapping[str, ProxyException | None]: + """Per-unique-password policy verdicts for a batch: the ProxyException to + surface, or None when the password is acceptable. + + Deduplicates first, then issues every needed HIBP lookup concurrently, so a + batch caller pays one HIBP timeout window in the worst case instead of one + per password (each lookup still fails open independently).""" + unique_passwords: Final = tuple(dict.fromkeys(passwords)) + strength_verdicts: Final[Mapping[str, ProxyException | None]] = MappingProxyType( + {password: _strength_verdict(password, general_settings) for password in unique_passwords} + ) + to_screen: Final = tuple(password for password in unique_passwords if strength_verdicts[password] is None) + breached_flags: Final = await asyncio.gather( + *(is_password_breached(password, general_settings, client) for password in to_screen) + ) + breached_passwords: Final = frozenset(password for password, breached in zip(to_screen, breached_flags) if breached) + return MappingProxyType( + { + password: breached_password_error() if password in breached_passwords else strength_verdicts[password] + for password in unique_passwords + } ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ddef2127945..58a8ba0b09d 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -27,9 +27,14 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object -from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy +from litellm.proxy.auth.password_policy import ( + validate_password_not_breached, + validate_password_policy, + validate_passwords_bulk, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import ( @@ -162,11 +167,17 @@ def _team_membership_table( return team_membership_table -async def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: - """Validate and hash password field in-place if present.""" +async def _hash_password_in_dict( + data: dict, general_settings: Mapping[str, object], password_prevalidated: bool = False +) -> None: + """Validate and hash password field in-place if present. + + ``password_prevalidated`` skips the policy checks for callers that already + validated the password (the bulk path screens its whole batch upfront).""" if "password" in data and data["password"] is not None: - validate_password_policy(data["password"], general_settings) - await validate_password_not_breached(data["password"], general_settings) + if not password_prevalidated: + validate_password_policy(data["password"], general_settings) + await validate_password_not_breached(data["password"], general_settings) data["password"] = hash_password(data["password"]) @@ -1429,6 +1440,7 @@ async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, + password_prevalidated: bool = False, ) -> dict[str, Any]: """ Helper function to update a single user. @@ -1451,7 +1463,7 @@ async def _update_single_user_helper( data_json: Final[dict] = user_request.model_dump(exclude_unset=True) non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) - await _hash_password_in_dict(non_default_values, general_settings) + await _hash_password_in_dict(non_default_values, general_settings, password_prevalidated=password_prevalidated) existing_user_row: BaseModel | None = None if user_request.user_id: @@ -1700,19 +1712,38 @@ async def bulk_update_processed_users( users_to_update: list[UpdateUserRequest], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, + hibp_client: AsyncHTTPHandler | None = None, ) -> BulkUpdateUserResponse: + from litellm.proxy.proxy_server import general_settings + results: Final[list[UserUpdateResult]] = [] successful_updates = 0 failed_updates = 0 + # Screen the batch's passwords upfront and concurrently: done per-user + # inside the loop below, each HIBP lookup would be awaited serially and a + # degraded-slow HIBP could stretch a full batch to minutes, timing out the + # request after some updates already persisted. + password_verdicts: Final = await validate_passwords_bulk( + tuple(u.password for u in users_to_update if u.password is not None), + general_settings, + client=hibp_client, + ) + # Process each user update independently try: for user_request in users_to_update: try: + if ( + user_request.password is not None + and (password_error := password_verdicts.get(user_request.password)) is not None + ): + raise password_error response = await _update_single_user_helper( user_request=user_request, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, + password_prevalidated=True, ) # Record success results.append( diff --git a/tests/test_litellm/proxy/auth/test_password_policy.py b/tests/test_litellm/proxy/auth/test_password_policy.py index edc88d21218..f9f5025b57f 100644 --- a/tests/test_litellm/proxy/auth/test_password_policy.py +++ b/tests/test_litellm/proxy/auth/test_password_policy.py @@ -7,6 +7,7 @@ The breach-check (HIBP) tests inject a real AsyncHTTPHandler wrapping an httpx.MockTransport, so no network is touched and nothing is monkeypatched. """ +import asyncio import hashlib import httpx @@ -21,6 +22,7 @@ from litellm.proxy.auth.password_policy import ( get_password_policy, validate_password_not_breached, validate_password_policy, + validate_passwords_bulk, ) STRONG_PASSWORD = "Str0ng!Passw0rd" @@ -268,3 +270,76 @@ async def test_breach_check_fails_open_on_malformed_response_body(): client=_client_returning(f"{_sha1_upper('password12345')[5:]}:not-a-number"), ) assert result is None + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_screens_concurrently(): + """All HIBP lookups for a batch must be in flight at once: each handler + call stalls until every expected request has arrived, and a handler that + gives up waiting reports the password as breached. Serial awaiting (the + old per-user behavior) leaves each earlier request waiting forever for the + later ones, so every verdict comes back as a breach and the test fails.""" + passwords = ("Uniqu3!Passw0rd-a", "Uniqu3!Passw0rd-b", "Uniqu3!Passw0rd-c") + suffix_by_prefix = {_sha1_upper(p)[:5]: _sha1_upper(p)[5:] for p in passwords} + all_arrived = asyncio.Event() + arrivals: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + arrivals.append(request.url.path) + if len(arrivals) == len(passwords): + all_arrived.set() + try: + await asyncio.wait_for(all_arrived.wait(), timeout=5) + except TimeoutError: + return httpx.Response(200, text=f"{suffix_by_prefix[request.url.path.rsplit('/', 1)[-1]]}:1") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + verdicts = await validate_passwords_bulk(passwords, {}, client=_client_with_transport(handler)) + assert set(arrivals) == {f"/range/{prefix}" for prefix in suffix_by_prefix} + assert all(verdicts[p] is None for p in passwords) + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_deduplicates_lookups(): + """500 users sharing one password must cost exactly one HIBP lookup.""" + password = "Sh@red-Passw0rd!" + request_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal request_count + request_count += 1 + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + verdicts = await validate_passwords_bulk((password,) * 500, {}, client=_client_with_transport(handler)) + assert request_count == 1 + assert verdicts == {password: None} + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_mixed_verdicts(): + """Weak passwords are rejected without an HIBP lookup; breached ones get + the breach error; acceptable ones map to None.""" + breached = "Br3ached!Passw0rd" + clean = "Cl3an!!Passw0rd42" + weak = "short1!" + breached_sha1 = _sha1_upper(breached) + looked_up_prefixes: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + looked_up_prefixes.append(request.url.path.rsplit("/", 1)[-1]) + if request.url.path == f"/range/{breached_sha1[:5]}": + return httpx.Response(200, text=f"{breached_sha1[5:]}:99") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + verdicts = await validate_passwords_bulk((breached, clean, weak), {}, client=_client_with_transport(handler)) + assert _sha1_upper(weak)[:5] not in looked_up_prefixes + assert verdicts[clean] is None + assert "data breaches" in verdicts[breached].message + assert verdicts[breached].code == "400" + assert "12 characters" in verdicts[weak].message + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_empty_batch_makes_no_lookups(): + verdicts = await validate_passwords_bulk((), {}, client=_client_never_called()) + assert verdicts == {} diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index e7b5172e0fb..024d7b8300f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -11,6 +11,7 @@ from fastapi.testclient import TestClient from fastapi import HTTPException from pytest_mock import MockerFixture +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( LiteLLM_UserTableFiltered, @@ -67,9 +68,7 @@ async def test_ui_view_users_with_null_email(mocker, caplog): # Proxy admin: no org filter, no get_user_object call response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN), user_id="test_user", user_email=None, team_id=None, @@ -77,9 +76,7 @@ async def test_ui_view_users_with_null_email(mocker, caplog): page_size=50, ) - assert response == [ - LiteLLM_UserTableFiltered(user_id="test-user-null-email", user_email=None) - ] + assert response == [LiteLLM_UserTableFiltered(user_id="test-user-null-email", user_email=None)] @pytest.mark.asyncio @@ -103,9 +100,7 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), user_id=None, user_email="foo", team_id=None, @@ -128,9 +123,7 @@ async def test_ui_view_users_org_admin_filtered_by_org(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -268,9 +261,7 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -401,9 +392,7 @@ async def test_ui_view_users_flag_on_team_admin_org_member_no_team_id(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -462,9 +451,7 @@ async def test_ui_view_users_flag_on_team_admin_not_in_org_resolves_via_key_team async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -507,9 +494,7 @@ async def test_ui_view_users_flag_on_team_admin_not_in_org_resolves_via_key_team # No team_id query param, but team_id on the API key response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="team-admin-no-org", user_role=None, team_id=tid - ), + user_api_key_dict=UserAPIKeyAuth(user_id="team-admin-no-org", user_role=None, team_id=tid), user_id=None, user_email="u", team_id=None, @@ -538,13 +523,9 @@ def test_user_daily_activity_types(): # Assert all fields in SpendMetrics are reported in DailySpendMetadata as "total_" for field in spend_metrics.__dict__: if field.startswith("total_"): - assert hasattr( - daily_spend_metadata, field - ), f"Field {field} is not reported in DailySpendMetadata" + assert hasattr(daily_spend_metadata, field), f"Field {field} is not reported in DailySpendMetadata" else: - assert not hasattr( - daily_spend_metadata, field - ), f"Field {field} is reported in DailySpendMetadata" + assert not hasattr(daily_spend_metadata, field), f"Field {field} is reported in DailySpendMetadata" @pytest.mark.asyncio @@ -591,9 +572,7 @@ async def test_get_users_includes_timestamps(mocker): # Call get_users function directly with proxy admin auth admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users( - page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None - ) + response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) print("user /list response: ", response) @@ -654,14 +633,10 @@ async def test_get_users_redacts_scim_enterprise_metadata(mocker): ) admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users( - page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None - ) + response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) listed = response["users"][0] - assert listed.metadata == { - "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} - } + assert listed.metadata == {"scim_metadata": {"givenName": "Jane", "familyName": "Doe"}} assert "scim_enterprise" not in (listed.metadata or {}) @@ -853,9 +828,7 @@ async def test_new_user_license_over_limit(mocker): mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) # Create test request data - user_request = NewUserRequest( - user_email="test@example.com", user_role="internal_user" - ) + user_request = NewUserRequest(user_email="test@example.com", user_role="internal_user") # Mock user_api_key_dict mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin") @@ -916,9 +889,7 @@ async def test_new_user_license_gate_counts_only_billable_users(mocker): request = NewUserRequest(user_role="internal_user") # 2 active + 3 deactivated -> billable 2, not over max_users 2: gate passes - mocker.patch( - "litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3) - ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3)) with pytest.raises(ProxyException) as passed: await new_user(data=request, user_api_key_dict=admin) assert key_gen.call_count == 1 @@ -926,9 +897,7 @@ async def test_new_user_license_gate_counts_only_billable_users(mocker): # 3 active, 0 deactivated -> billable 3, over max_users 2: gate blocks key_gen.reset_mock() - mocker.patch( - "litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0) - ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0)) with pytest.raises(ProxyException) as blocked: await new_user(data=request, user_api_key_dict=admin) assert blocked.value.code == 403 or blocked.value.code == "403" @@ -978,14 +947,10 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) # Test Case 1: INTERNAL_USER trying to create PROXY_ADMIN - user_request = NewUserRequest( - user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_request = NewUserRequest(user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN) # Mock user_api_key_dict with non-admin role - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER) # Call new_user function and expect ProxyException with pytest.raises(ProxyException) as exc_info: @@ -993,9 +958,7 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): # Verify the exception details assert exc_info.value.code == 403 or exc_info.value.code == "403" - assert "Only proxy admins can create administrative users" in str( - exc_info.value.message - ) + assert "Only proxy admins can create administrative users" in str(exc_info.value.message) assert "proxy_admin" in str(exc_info.value.message) assert "proxy_admin_viewer" in str(exc_info.value.message) assert str(LitellmUserRoles.PROXY_ADMIN) in str(exc_info.value.message) @@ -1008,15 +971,11 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): ) with pytest.raises(ProxyException) as exc_info2: - await new_user( - data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict - ) + await new_user(data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict) # Verify the exception details assert exc_info2.value.code == 403 or exc_info2.value.code == "403" - assert "Only proxy admins can create administrative users" in str( - exc_info2.value.message - ) + assert "Only proxy admins can create administrative users" in str(exc_info2.value.message) assert str(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) in str(exc_info2.value.message) @@ -1055,9 +1014,7 @@ async def test_new_user_non_admin_permissions_non_empty_rejected(mocker): user_role=LitellmUserRoles.INTERNAL_USER, permissions={"get_spend_routes": True}, ) - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(ProxyException) as exc_info: await new_user(data=data, user_api_key_dict=caller) @@ -1101,9 +1058,7 @@ async def test_new_user_non_admin_permissions_explicit_empty_rejected(mocker): permissions={}, ) assert "permissions" in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(ProxyException) as exc_info: await new_user(data=data, user_api_key_dict=caller) @@ -1156,9 +1111,7 @@ async def test_new_user_non_admin_omits_permissions_succeeds(mocker): user_role=LitellmUserRoles.INTERNAL_USER, ) assert "permissions" not in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) result = await new_user(data=data, user_api_key_dict=caller) assert result is not None @@ -1232,14 +1185,10 @@ async def test_update_single_user_non_admin_permissions_rejected(mocker): user_id="alice", permissions={"get_spend_routes": True}, ) - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc_info: - await _update_single_user_helper( - user_request=data, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=data, user_api_key_dict=caller) assert exc_info.value.status_code == 403 assert "permissions" in str(exc_info.value.detail) @@ -1261,14 +1210,10 @@ async def test_update_single_user_non_admin_permissions_explicit_empty_rejected( data = UpdateUserRequest(user_id="alice", permissions={}) assert "permissions" in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc_info: - await _update_single_user_helper( - user_request=data, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=data, user_api_key_dict=caller) assert exc_info.value.status_code == 403 assert "permissions" in str(exc_info.value.detail) @@ -1324,15 +1269,11 @@ async def test_user_info_url_encoding_plus_character(mocker): mock_request.url.query = "user_id=machine-user+alp-air-admin-b58-b@tempus.com" # Mock user_api_key_dict - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_admin", user_role="proxy_admin" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin", user_role="proxy_admin") # Call user_info function with the URL-decoded user_id (as FastAPI would pass it) # FastAPI would normally convert + to space, but our fix should handle this - decoded_user_id = ( - "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us - ) + decoded_user_id = "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us expected_user_id = "machine-user+alp-air-admin-b58-b@tempus.com" response = await user_info( @@ -1383,9 +1324,7 @@ async def test_user_info_nonexistent_user(mocker): mock_request = mocker.MagicMock(spec=Request) # Mock user_api_key_dict - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_admin", user_role="proxy_admin" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin", user_role="proxy_admin") # Call user_info function with a non-existent user_id nonexistent_user_id = "nonexistent-user@example.com" @@ -1423,14 +1362,10 @@ async def test_user_info_no_user_id_view_only_admin_gets_proxy_admin_payload(moc mock_get_user_info_for_proxy_admin, ) - viewer = UserAPIKeyAuth( - user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ) + viewer = UserAPIKeyAuth(user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) mock_request = mocker.MagicMock(spec=Request) - response = await user_info( - user_id=None, user_api_key_dict=viewer, request=mock_request - ) + response = await user_info(user_id=None, user_api_key_dict=viewer, request=mock_request) mock_get_user_info_for_proxy_admin.assert_awaited_once_with(user_api_key_dict=viewer) assert response is admin_payload @@ -1457,9 +1392,7 @@ async def test_new_user_default_teams_flow(mocker): mock_prisma_client.db.litellm_usertable.count = mock_count persisted_user_row = mocker.MagicMock() persisted_user_row.teams = ["96fed65b-0182-4ff4-8429-2721cd7d42af"] - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - return_value=persisted_user_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=persisted_user_row) # Mock duplicate checks to pass async def mock_check_duplicate_user_email(*args, **kwargs): @@ -1527,26 +1460,20 @@ async def test_new_user_default_teams_flow(mocker): ) # Create test request data WITHOUT teams (teams should come from defaults) - user_request = NewUserRequest( - user_email="test@example.com", user_role="internal_user" - ) + user_request = NewUserRequest(user_email="test@example.com", user_role="internal_user") # Mock user_api_key_dict mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin") # Call new_user function - response = await new_user( - data=user_request, user_api_key_dict=mock_user_api_key_dict - ) + response = await new_user(data=user_request, user_api_key_dict=mock_user_api_key_dict) # Verify generate_key_helper_fn was called WITHOUT teams mock_generate_key_helper_fn.assert_called_once() call_kwargs = mock_generate_key_helper_fn.call_args.kwargs # Teams should be removed from the data passed to generate_key_helper_fn - assert ( - "teams" not in call_kwargs - ), "Teams should not be passed to generate_key_helper_fn" + assert "teams" not in call_kwargs, "Teams should not be passed to generate_key_helper_fn" assert call_kwargs["request_type"] == "user" assert call_kwargs["user_email"] == "test@example.com" assert call_kwargs["user_role"] == "internal_user" @@ -1591,24 +1518,16 @@ def test_update_internal_new_user_params_proxy_admin_role(): try: # Create test data with PROXY_ADMIN role - data = NewUserRequest( - user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + data = NewUserRequest(user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN.value) data_json = data.model_dump(exclude_unset=True) # Call the function result = _update_internal_new_user_params(data_json=data_json, data=data) # Assertions - default params should NOT be applied for PROXY_ADMIN - assert ( - "max_budget" not in result - ), "Default max_budget should NOT be applied to PROXY_ADMIN" - assert ( - "models" not in result - ), "Default models should NOT be applied to PROXY_ADMIN" - assert ( - "tpm_limit" not in result - ), "Default tpm_limit should NOT be applied to PROXY_ADMIN" + assert "max_budget" not in result, "Default max_budget should NOT be applied to PROXY_ADMIN" + assert "models" not in result, "Default models should NOT be applied to PROXY_ADMIN" + assert "tpm_limit" not in result, "Default tpm_limit should NOT be applied to PROXY_ADMIN" # These should still work assert result["user_email"] == "admin@example.com" @@ -1722,15 +1641,9 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): user_email_clause = where_clause.get("user_email", {}) # Check that the query structure is correct for case insensitive search - assert ( - "equals" in user_email_clause - ), "Query should use 'equals' for case insensitive search" - assert ( - user_email_clause.get("mode") == "insensitive" - ), "Query should use 'insensitive' mode" - assert ( - user_email_clause.get("equals") == "user@example.com" - ), "Query should search for the provided email" + assert "equals" in user_email_clause, "Query should use 'equals' for case insensitive search" + assert user_email_clause.get("mode") == "insensitive", "Query should use 'insensitive' mode" + assert user_email_clause.get("equals") == "user@example.com", "Query should search for the provided email" return mock_existing_user # Return existing user to simulate duplicate @@ -1741,9 +1654,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): await _check_duplicate_user_email("user@example.com", mock_prisma_client) assert exc_info.value.status_code == 409 - assert "User with email User@Example.com already exists" in str( - exc_info.value.detail - ) + assert "User with email User@Example.com already exists" in str(exc_info.value.detail) # Test Case 2: No duplicate found async def mock_find_first_no_duplicate(*args, **kwargs): @@ -1768,9 +1679,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): pytest.fail(f"Should not raise exception when no duplicate found, but got: {e}") # Test Case 3: None email should not cause issues - await _check_duplicate_user_email( - None, mock_prisma_client - ) # Should not raise exception + await _check_duplicate_user_email(None, mock_prisma_client) # Should not raise exception @pytest.mark.asyncio @@ -1880,9 +1789,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): # Verify dashboard key is not in results result_team_ids = [key.get("team_id") for key in result] - assert ( - UI_SESSION_TOKEN_TEAM_ID not in result_team_ids - ), "Dashboard key should be filtered out" + assert UI_SESSION_TOKEN_TEAM_ID not in result_team_ids, "Dashboard key should be filtered out" # Verify regular keys are included assert "regular-team" in result_team_ids, "Regular team key should be included" @@ -1892,9 +1799,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): result_tokens = [key.get("token") for key in result] assert "sk-regular-token" in result_tokens, "Regular key should be included" assert "sk-no-team-token" in result_tokens, "No-team key should be included" - assert ( - "sk-dashboard-token" not in result_tokens - ), "Dashboard key should not be included" + assert "sk-dashboard-token" not in result_tokens, "Dashboard key should not be included" def test_process_keys_for_user_info_handles_none_keys(monkeypatch): @@ -2331,9 +2236,7 @@ async def test_get_user_daily_activity_non_admin_cannot_view_other_users(monkeyp ) assert exc_info.value.status_code == 403 - assert "Non-admin users can only view their own spend data" in str( - exc_info.value.detail - ) + assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) # Case 2: Non-admin omits user_id — should default to their own user_id mock_response = MagicMock() @@ -2620,39 +2523,23 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): async def mock_find_unique(*args, **kwargs): return mock_user_row - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock find_many for teams (no teams) - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) # Mock all delete_many calls - mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock( - return_value=1 - ) - mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1) + mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Call delete_user data = DeleteUserRequest(user_ids=["admin-creator"]) - user_api_key_dict = UserAPIKeyAuth( - user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_api_key_dict = UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN) await delete_user(data=data, user_api_key_dict=user_api_key_dict) @@ -2661,9 +2548,7 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): call_kwargs = mock_prisma_client.db.litellm_invitationlink.delete_many.call_args where_clause = call_kwargs.kwargs.get("where") or call_kwargs[1].get("where") - assert ( - "OR" in where_clause - ), "Should use OR to match user_id, created_by, and updated_by" + assert "OR" in where_clause, "Should use OR to match user_id, created_by, and updated_by" or_conditions = where_clause["OR"] assert len(or_conditions) == 3, "Should have 3 OR conditions" @@ -2706,9 +2591,7 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): async def mock_find_unique(*args, **kwargs): return mock_target_user - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Caller (org_admin_user) administers org-A. caller_membership = mocker.MagicMock() @@ -2734,16 +2617,12 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): return [caller_membership] return [] - mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock( - side_effect=mock_find_memberships - ) + mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock(side_effect=mock_find_memberships) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) data = DeleteUserRequest(user_ids=["victim"]) - user_api_key_dict = UserAPIKeyAuth( - user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN - ) + user_api_key_dict = UserAPIKeyAuth(user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc: await delete_user(data=data, user_api_key_dict=user_api_key_dict) @@ -2751,11 +2630,8 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): # Critical: no delete_many calls should have executed. assert ( - not hasattr( - mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls" - ) - or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) - == 0 + not hasattr(mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls") + or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) == 0 ) @@ -2774,9 +2650,7 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): mock_prisma_client = mocker.MagicMock() # user_email lookup yields None → would silently create pre-fix. - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=None) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest( @@ -2790,9 +2664,7 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=org_admin - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=org_admin) assert exc.value.status_code == 404 @@ -2836,17 +2708,13 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -2900,17 +2768,13 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -2919,9 +2783,7 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): ) assert isinstance(response, UserInfoV2Response) - assert response.metadata == { - "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} - } + assert response.metadata == {"scim_metadata": {"givenName": "Jane", "familyName": "Doe"}} assert "scim_enterprise" not in (response.metadata or {}) @@ -2990,17 +2852,13 @@ async def test_user_info_v2_internal_user_can_query_self(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER) response = await user_info_v2( request=mock_request, @@ -3035,17 +2893,13 @@ async def test_user_info_v2_internal_user_cannot_query_other(mocker): return mock_caller_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3092,17 +2946,13 @@ async def test_user_info_v2_no_user_id_defaults_to_self(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER) # Call without user_id response = await user_info_v2( @@ -3130,17 +2980,13 @@ async def test_user_info_v2_nonexistent_user_returns_404(mocker): async def mock_find_unique(*args, **kwargs): return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3188,17 +3034,13 @@ async def test_user_info_v2_response_shape(mocker): async def mock_find_unique(*args, **kwargs): return mock_user_row - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -3233,9 +3075,7 @@ async def test_user_info_v2_response_shape(mocker): # The dashboard's user edit form hydrates its per-model budget rows from # these two, so dropping them makes a save replace the user's budgets. - assert response_dict["model_max_budget"] == { - "gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"} - } + assert response_dict["model_max_budget"] == {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}} assert response_dict["model_max_budget_usage"] == { "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} } @@ -3294,9 +3134,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): return mock_target return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock team with caller as admin mock_team = mocker.MagicMock() @@ -3313,17 +3151,13 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): async def mock_find_many_teams(*args, **kwargs): return [mock_team] - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - side_effect=mock_find_many_teams - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(side_effect=mock_find_many_teams) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - team_admin_key = UserAPIKeyAuth( - user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + team_admin_key = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER) response = await user_info_v2( request=mock_request, @@ -3363,9 +3197,7 @@ async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker): return mock_target return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock team where caller is admin mock_team = mocker.MagicMock() @@ -3381,17 +3213,13 @@ async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker): async def mock_find_many_teams(*args, **kwargs): return [mock_team] - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - side_effect=mock_find_many_teams - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(side_effect=mock_find_many_teams) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - team_admin_key = UserAPIKeyAuth( - user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + team_admin_key = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3441,18 +3269,14 @@ async def test_user_info_v2_url_encoding_plus_character(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) mock_request.url.query = f"user_id={expected_user_id}" - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) # Simulate FastAPI converting + to space decoded_user_id = "machine-user admin@example.com" @@ -3549,9 +3373,7 @@ def test_enforce_user_info_access_admin_bypass(): _enforce_user_info_access, ) - admin = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value) # Should not raise even when querying a different user _enforce_user_info_access(user_id="someone_else", user_api_key_dict=admin) @@ -3590,9 +3412,7 @@ def test_enforce_user_info_access_owner_allowed(): _enforce_user_info_access, ) - user = UserAPIKeyAuth( - user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value - ) + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) _enforce_user_info_access(user_id="alice", user_api_key_dict=user) @@ -3604,9 +3424,7 @@ def test_enforce_user_info_access_no_user_id_allowed(): _enforce_user_info_access, ) - user = UserAPIKeyAuth( - user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value - ) + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) _enforce_user_info_access(user_id=None, user_api_key_dict=user) @@ -3663,9 +3481,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budge "max_budget": 100, } existing_user.user_id = "user-1" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest.model_validate({"user_id": "user-1", budget_field: budget_value}) @@ -3675,9 +3491,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budge ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=caller) assert exc.value.status_code == 403 assert budget_field in str(exc.value.detail) mock_prisma_client.update_data.assert_not_called() @@ -3699,9 +3513,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker): "spend": 50.0, } existing_user.user_id = "user-1" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest( @@ -3714,9 +3526,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker): ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=caller) assert exc.value.status_code == 403 assert "spend" in str(exc.value.detail) @@ -3735,12 +3545,8 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): "max_budget": 100, } existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user", "max_budget": 500} - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user", "max_budget": 500}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3754,9 +3560,7 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): user_role=LitellmUserRoles.PROXY_ADMIN, ) - result = await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + result = await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) assert result is not None @@ -3772,12 +3576,8 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user", "spend": -25.0} - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user", "spend": -25.0}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3792,13 +3592,9 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): # without raising the recurring budget ceiling. Future changes should # continue allowing negative spend counters. user_request = UpdateUserRequest(user_id="target-user", spend=-25) - admin_caller = UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") @@ -3815,9 +3611,7 @@ async def test_user_update_rejects_non_finite_spend(mocker): existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mock_prisma_client.update_data = mocker.AsyncMock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3827,14 +3621,10 @@ async def test_user_update_rejects_non_finite_spend(mocker): ) user_request = UpdateUserRequest(user_id="target-user", spend=float("nan")) - admin_caller = UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) assert exc.value.status_code == 400 mock_prisma_client.update_data.assert_not_called() mock_invalidate.assert_not_awaited() @@ -3854,9 +3644,7 @@ async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): mock_prisma_client = mocker.MagicMock() find_many = mocker.AsyncMock( return_value=[ - SimpleNamespace( - user_id="u1", user_email="alice@example.com", user_alias="Alice" - ), + SimpleNamespace(user_id="u1", user_email="alice@example.com", user_alias="Alice"), SimpleNamespace(user_id="u2", user_email=None, user_alias="bob-alias"), ] ) @@ -4055,19 +3843,13 @@ def _object_permission_mocks(mocker, existing_object_permission_id=None): } existing_user.user_id = "target-user" existing_user.object_permission_id = existing_object_permission_id - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock(return_value=None) mock_prisma_client.db.litellm_objectpermissiontable.upsert = mocker.AsyncMock( return_value=SimpleNamespace(object_permission_id="perm-new") ) mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user"} - ) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -4103,9 +3885,7 @@ async def test_user_update_persists_mcp_entitlement_and_links_it(mocker): "mcp_tool_permissions": {"github": ["list_issues"]}, }, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) upsert_kwargs = mock_prisma_client.db.litellm_objectpermissiontable.upsert.call_args.kwargs @@ -4139,9 +3919,7 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): user_id="target-user", object_permission={"mcp_tool_permissions": {"github": []}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} @@ -4174,9 +3952,7 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): await _update_single_user_helper( user_request=UpdateUserRequest(user_id="target-user", object_permission={}), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) written = mock_prisma_client.update_data.call_args.kwargs["data"] @@ -4213,9 +3989,7 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock user_id="target-user", object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} @@ -4246,9 +4020,7 @@ async def test_non_admin_cannot_clear_their_own_mcp_entitlement(mocker): with pytest.raises(HTTPException) as exc: await _update_single_user_helper( user_request=UpdateUserRequest(user_id="target-user", object_permission={}), - user_api_key_dict=UserAPIKeyAuth( - user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER - ), + user_api_key_dict=UserAPIKeyAuth(user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER), ) assert exc.value.status_code == 403 @@ -4276,9 +4048,7 @@ async def test_non_admin_cannot_rewrite_their_own_mcp_entitlement(mocker): user_id="target-user", object_permission={"mcp_servers": [], "mcp_tool_permissions": {}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER - ), + user_api_key_dict=UserAPIKeyAuth(user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER), ) assert exc.value.status_code == 403 @@ -4295,9 +4065,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): return_value=SimpleNamespace(object_permission_id="perm-created") ) mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.count = mocker.AsyncMock(return_value=0) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -4306,9 +4074,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): ) mock_generate = mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", - new=mocker.AsyncMock( - return_value={"user_id": "new-human", "token": "sk-x", "expires": None} - ), + new=mocker.AsyncMock(return_value={"user_id": "new-human", "token": "sk-x", "expires": None}), ) mocker.patch( "litellm.proxy.hooks.user_management_event_hooks.UserManagementEventHooks.async_user_created_hook", @@ -4320,9 +4086,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): user_id="new-human", object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) created = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] @@ -4364,16 +4128,12 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): response = await user_info_v2( request=SimpleNamespace(query_params={}), user_id="human-1", - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) assert response.object_permission is not None assert response.object_permission.mcp_servers == ["github"] - assert response.object_permission.mcp_tool_permissions == { - "github": ["list_issues"] - } + assert response.object_permission.mcp_tool_permissions == {"github": ["list_issues"]} @pytest.mark.asyncio @@ -4389,9 +4149,7 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): ], ids=["supplied", "omitted", "empty"], ) -async def test_user_new_persists_model_max_budget( - monkeypatch, model_max_budget, expected_written -): +async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, expected_written): """ /user/new used to echo model_max_budget back while writing {} to the user row, so a per-model budget looked configured and was read by nothing. @@ -4553,3 +4311,90 @@ async def test_user_update_rejects_breached_password(_admin_prisma): assert exc_info.value.code == "400" assert "data breaches" in exc_info.value.message _admin_prisma.db.litellm_usertable.find_first.assert_not_called() + + +def _hibp_client_with_handler(handler) -> AsyncHTTPHandler: + """A real AsyncHTTPHandler over httpx.MockTransport (the DI seam used + throughout test_password_policy.py), so no network is touched.""" + http_handler = AsyncHTTPHandler() + http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + return http_handler + + +@pytest.mark.asyncio +async def test_bulk_update_breached_password_fails_only_that_user(_admin_prisma, mocker): + """In a bulk batch, a breached password fails only its own entry, before + any DB write for it; sibling entries with acceptable passwords persist.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + bulk_update_processed_users, + ) + + breached = "Br3ached!Passw0rd" + clean = "NewP@ssw0rd123" + breached_sha1 = hashlib.sha1(breached.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == f"/range/{breached_sha1[:5]}": + return httpx.Response(200, text=f"{breached_sha1[5:]}:1387") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "user-clean"} + existing_user.user_id = "user-clean" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-clean"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await bulk_update_processed_users( + users_to_update=[ + UpdateUserRequest(user_id="user-breached", password=breached), + UpdateUserRequest(user_id="user-clean", password=clean), + ], + user_api_key_dict=admin_caller, + hibp_client=_hibp_client_with_handler(handler), + ) + + assert response.successful_updates == 1 + assert response.failed_updates == 1 + by_user = {r.user_id: r for r in response.results} + assert by_user["user-breached"].success is False + assert "data breaches" in by_user["user-breached"].error + assert by_user["user-clean"].success is True + (write_call,) = mock_prisma_client.update_data.call_args_list + assert write_call.kwargs["user_id"] == "user-clean" + + +@pytest.mark.asyncio +async def test_bulk_update_screens_shared_password_with_single_lookup(_admin_prisma, mocker): + """A batch where every user gets the same password costs one HIBP lookup, + not one per user (the serial per-user checks this regresses against).""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + bulk_update_processed_users, + ) + + lookup_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal lookup_count + lookup_count += 1 + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "user-0"} + existing_user.user_id = "user-0" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-0"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await bulk_update_processed_users( + users_to_update=[UpdateUserRequest(user_id=f"user-{i}", password="NewP@ssw0rd123") for i in range(5)], + user_api_key_dict=admin_caller, + hibp_client=_hibp_client_with_handler(handler), + ) + + assert response.successful_updates == 5 + assert lookup_count == 1 From 1d18d11fcf69388787b824b4d519e9855eede23e Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 16:41:36 +0200 Subject: [PATCH 019/149] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 58a8ba0b09d..4f037ffd43d 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -574,7 +574,7 @@ async def new_user( # generate_key_helper_fn only forwards object_permission_id, so without this the entitlement # the caller sent would be dropped on the floor. data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client) - data_json.pop("password", None) # always None: NewUserRequest.password_not_supported rejects any other value + data_json.pop("password", None) teams = data.teams if teams is None: teams = check_if_default_team_set() From 5bb2c9e76f565bb4c4fb82ec527841a6340862e2 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Wed, 9 Sep 2026 13:46:29 +0200 Subject: [PATCH 020/149] fix(auth): annotate the strict-rule suppressions the merged gates now count The staging merge brought BLE001 into the strict ruff set and lowered the LIT002 ceiling, so the HIBP fail-open except and the params/headers dicts in password_policy.py now need their noqa and mutable-ok reasons. The headers dict moves to an annotated Final so the suppression fits the line limit. --- litellm/proxy/auth/password_policy.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index c27f276252a..7f06a0993d3 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -110,7 +110,7 @@ def validate_password_policy(password: str, general_settings: Mapping[str, objec def _hibp_client() -> AsyncHTTPHandler: return get_async_httpx_client( llm_provider=httpxSpecialProvider.PasswordBreachCheck, - params={"timeout": HIBP_TIMEOUT_SECONDS}, + params={"timeout": HIBP_TIMEOUT_SECONDS}, # mutable-ok: callee takes a bare dict (PEP 589) ) @@ -125,14 +125,18 @@ def _is_suffix_in_range_response(response_body: str, hash_suffix: str) -> bool: async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool: # usedforsecurity=False: SHA-1 is only a lookup key into the HIBP dataset, so no security property rests on it sha1_hex: Final = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + headers: Final = { # mutable-ok: callee takes a bare dict (PEP 589) + "Add-Padding": "true", + "User-Agent": f"litellm-proxy/{version}", + } try: response: Final = await client.get( f"{HIBP_RANGE_API_BASE}/{sha1_hex[:5]}", - headers={"Add-Padding": "true", "User-Agent": f"litellm-proxy/{version}"}, + headers=headers, ) response.raise_for_status() breached: Final = _is_suffix_in_range_response(response.text, sha1_hex[5:]) - except Exception as e: + except Exception as e: # noqa: BLE001 # fail-open: any HIBP failure skips the check, never breaks the caller verbose_proxy_logger.warning("Breached-password check skipped, HIBP lookup failed: %s", e) return False return breached From d79a893e3776923ac33f802975e01e2557e130a5 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Fri, 4 Sep 2026 09:02:34 +0200 Subject: [PATCH 021/149] feat(auth): add self-service change-password endpoint Admin password sets on /user/update and per-user /user/bulk_update stay supported and policy-enforced. The request model hides the password from repr so management alerts never format the plaintext, and the all_users bulk path rejects passwords instead of writing one plaintext value to every row. --- litellm/proxy/_types.py | 14 +- litellm/proxy/auth/route_checks.py | 17 +- .../internal_user_endpoints.py | 12 +- .../password_endpoints.py | 117 ++++++++ litellm/proxy/proxy_server.py | 4 + .../proxy/auth/test_route_checks.py | 213 ++++++------- .../test_internal_user_endpoints.py | 29 ++ .../test_password_endpoints.py | 283 ++++++++++++++++++ tests/test_litellm/proxy/test__types.py | 26 ++ .../ChangePasswordForm.integration.test.tsx | 68 +++++ .../change-password/ChangePasswordForm.tsx | 100 +++++++ .../app/(dashboard)/change-password/page.tsx | 7 + .../app/(dashboard)/hooks/useAuthorized.ts | 1 + .../Navbar/UserDropdown/UserDropdown.test.tsx | 52 +++- .../Navbar/UserDropdown/UserDropdown.tsx | 23 +- .../SidebarAccountMenu.test.tsx | 43 +++ .../SidebarAccountMenu/SidebarAccountMenu.tsx | 24 +- .../src/components/networking.tsx | 14 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 78 ++++- 19 files changed, 1004 insertions(+), 121 deletions(-) create mode 100644 litellm/proxy/management_endpoints/password_endpoints.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b38b12ab07..7731ca736d5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -864,6 +864,7 @@ class LiteLLMRoutes(enum.Enum): "/organization/daily/activity", "/user/available_roles", # read-only role metadata; any authenticated user may read "/user/list", # org admins checked in endpoint; non-admins get 403 + "/user/password/change", # endpoint only ever writes the caller's own row "/model/{model_id}/update", "/prompt/list", "/prompt/info", @@ -1841,7 +1842,8 @@ class NewUserResponse(GenerateKeyResponse): class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest - password: str | None = None + # repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model + password: str | None = Field(default=None, repr=False) spend: float | None = None metadata: dict | None = None user_alias: str | None = None @@ -1871,6 +1873,16 @@ class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): return values +class ChangePasswordRequest(LiteLLMPydanticObjectBase): + current_password: str = Field(repr=False) + new_password: str = Field(repr=False) + + +class ChangePasswordResponse(LiteLLMPydanticObjectBase): + user_id: str + message: str + + class DeleteUserRequest(LiteLLMPydanticObjectBase): user_ids: list[str] # required diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 953e3cf3e88..e4a36e73373 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -796,7 +796,8 @@ class RouteChecks: in the codebase is automatically readable by Admin Viewer without needing to remember to add it to an allowlist. 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): - - Allow `/user/update` only when restricted to user_email/password. + - Allow `/user/update` only when restricted to user_email. + - Allow `/user/password/change` (endpoint only writes the caller's own row). - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. - Otherwise allow only if the route is in admin_viewer_routes / global_spend_tracking_routes (legacy explicit-allow set). @@ -816,10 +817,10 @@ class RouteChecks: if request_data is not None and isinstance(request_data, dict): _params_updated: Final = request_data.keys() for param in _params_updated: - if param not in ["user_email", "password"]: + if param != "user_email": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated", + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated", ) elif route in _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES or ( route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) @@ -838,21 +839,25 @@ class RouteChecks: return # ── Unsafe HTTP method: explicit checks ────────────────────────── - # Allow `/user/update` for self-service email / password change. + # Allow `/user/update` for self-service email change. if route == "/user/update": if request_data is not None and isinstance(request_data, dict): for param in request_data: - if param not in ["user_email", "password"]: + if param != "user_email": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=( f"user not allowed to access this route, role= {_user_role}. " f"Trying to access: {route} and updating invalid param: {param}. " - "only user_email and password can be updated" + "only user_email can be updated" ), ) return + # Self-service password change; the endpoint only writes the caller's own row. + if route == "/user/password/change": + return + # Hard-block known write routes regardless of HTTP method (defensive # — these are POSTs in practice, but pinning them here protects # against future GET-shaped writes). diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 4f037ffd43d..368d5ed0c10 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1644,7 +1644,7 @@ async def user_update( Parameters: - user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated. - user_email: Optional[str] - Specify a user email. - - password: Optional[str] - Specify a user password. + - password: Optional[str] - Set the user's password (admin only). Must satisfy the configured password policy. Users change their own password with POST /user/password/change. - user_alias: Optional[str] - A descriptive name for you to know who this user id refers to. - teams: Optional[list] - specify a list of team id's a user belongs to. - send_invite_email: Optional[bool] - Specify if an invite email should be sent. @@ -1881,6 +1881,16 @@ async def bulk_user_update( status_code=403, detail="Only proxy admins can update all users at once.", ) + if data.user_updates.password is not None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Setting one password for all users is not supported. " + "Use per-user updates via the 'users' list instead." + ) + }, + ) # Optimized path for updating all users directly in database all_users_in_db: Final = await _user_table(prisma_client).find_many(order={"created_at": "desc"}) diff --git a/litellm/proxy/management_endpoints/password_endpoints.py b/litellm/proxy/management_endpoints/password_endpoints.py new file mode 100644 index 00000000000..78e53d7a777 --- /dev/null +++ b/litellm/proxy/management_endpoints/password_endpoints.py @@ -0,0 +1,117 @@ +""" +Self-service password management. + +/user/password/change + +Deliberately NOT wrapped in `management_endpoint_wrapper`: the wrapper emits +request kwargs to OTEL spans, which would log plaintext passwords. The audit +signal is emitted by hand below, with field names only, never values. +""" + +from typing import TYPE_CHECKING, Final + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + ChangePasswordRequest, + ChangePasswordResponse, + CommonProxyErrors, + LitellmTableNames, + UserAPIKeyAuth, +) +from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.utils import hash_password, verify_password +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.user_repository import UserRepository + +if TYPE_CHECKING: + from prisma import models as prisma_models + + from litellm.proxy.utils import PrismaClient + +router: Final = APIRouter() + +_PASSWORD_CHANGED_AUDIT_VALUES: Final = '{"fields_changed": ["password"]}' + + +def _user_table( + prisma_client: "PrismaClient | None", +) -> "TableActions[prisma_models.LiteLLM_UserTable]": + user_table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table + return user_table + + +@router.post( + "/user/password/change", + tags=["Internal User management"], + dependencies=(Depends(user_api_key_auth),), +) +async def change_password( + data: ChangePasswordRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> ChangePasswordResponse: + """ + Change the calling user's own password. + + Requires the current password. The new password must satisfy the + configured password policy (`general_settings.password_policy_*`: minimum + length, character classes, and, when enabled, breached-password screening + via haveibeenpwned.com). + + Parameters: + - current_password: str - The user's current password. + - new_password: str - The password to change to. + """ + from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + user_id: Final = user_api_key_dict.user_id + if user_id is None: + raise HTTPException( + status_code=400, + detail={"error": "No user is associated with this session, so there is no password to change."}, + ) + + user_row: Final = await _user_table(prisma_client).find_first(where={"user_id": user_id}) + stored_password: Final = user_row.password if user_row is not None else None + if stored_password is None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + "This account has no password set, so there is no password to change. " + "Passwords are set through an invitation link (POST /invitation/new)." + ) + }, + ) + + if not verify_password(data.current_password, stored_password): + raise HTTPException(status_code=400, detail={"error": "Current password is incorrect."}) + + validate_password_policy(data.new_password, general_settings) + await validate_password_not_breached(data.new_password, general_settings) + + await _user_table(prisma_client).update( + where={"user_id": user_id}, + data={"password": hash_password(data.new_password)}, + ) + + verbose_proxy_logger.info("Password changed via /user/password/change for user_id=%s", user_id) + await create_object_audit_log( + object_id=user_id, + action="updated", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.USER_TABLE_NAME, + after_value=_PASSWORD_CHANGED_AUDIT_VALUES, + ) + return ChangePasswordResponse(user_id=user_id, message="Password updated successfully.") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d2b486b4410..d803e8fd4be 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -549,6 +549,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.password_endpoints import ( + router as password_management_router, +) from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, ) @@ -18765,6 +18768,7 @@ app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) app.include_router(internal_user_router) +app.include_router(password_management_router) app.include_router(team_router) app.include_router(ui_sso_router) app.include_router(organization_router) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c8b3d789665..4f58e3c86ff 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2,7 +2,6 @@ import os from datetime import datetime from unittest.mock import MagicMock, patch - import pytest from fastapi import HTTPException, Request @@ -38,7 +37,7 @@ def test_non_admin_config_update_route_rejected(): request.query_params = {} # Test that calling /config/update route raises HTTPException with 403 status - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -49,9 +48,8 @@ def test_non_admin_config_update_route_rejected(): ) # Verify the exception is raised with the correct message - assert ( - "Only proxy admin can be used to generate, delete, update info for new keys/users/teams" - in str(exc_info.value) + assert "Only proxy admin can be used to generate, delete, update info for new keys/users/teams" in str( + exc_info.value ) assert "Route=/config/update" in str(exc_info.value) assert "Your role=internal_user" in str(exc_info.value) @@ -130,7 +128,7 @@ def test_user_banner_update_rejected_for_non_admin(): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -703,9 +701,7 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route): valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"]) - result = RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token - ) + result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) assert result is True @@ -730,9 +726,7 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route): valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"]) - result = RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token - ) + result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) assert result is True @@ -802,18 +796,14 @@ def test_google_routes_with_dynamic_model_names_accessible_to_internal_users(): ) # If no exception is raised, the test passes except Exception as e: - pytest.fail( - f"Internal user should be able to access Google generateContent route. Got error: {str(e)}" - ) + pytest.fail(f"Internal user should be able to access Google generateContent route. Got error: {e!s}") def test_virtual_key_allowed_routes_with_multiple_litellm_routes_member_names(): """Test that virtual key works with multiple LiteLLMRoutes member names in allowed_routes""" # Create a UserAPIKeyAuth with multiple LiteLLMRoutes member names - valid_token = UserAPIKeyAuth( - user_id="test_user", allowed_routes=["openai_routes", "info_routes"] - ) + valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["openai_routes", "info_routes"]) # Test that routes from both groups are allowed result1 = RouteChecks.is_virtual_key_allowed_to_call_route( @@ -867,13 +857,9 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit(): ) # Test that explicit routes are allowed - result1 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/chat/completions", valid_token=valid_token - ) + result1 = RouteChecks.is_virtual_key_allowed_to_call_route(route="/chat/completions", valid_token=valid_token) - result2 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/custom/route", valid_token=valid_token - ) + result2 = RouteChecks.is_virtual_key_allowed_to_call_route(route="/custom/route", valid_token=valid_token) assert result1 is True assert result2 is True @@ -1241,9 +1227,7 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): ) assert exc_info.value.status_code == 403 - assert "Virtual key is not allowed to call this route" in str( - exc_info.value.detail - ) + assert "Virtual key is not allowed to call this route" in str(exc_info.value.detail) def test_check_passthrough_route_access_key_metadata_exact_match(): @@ -1702,9 +1686,7 @@ def test_videos_route_accessible_to_internal_users(): ) # If no exception is raised, the test passes except Exception as e: - pytest.fail( - f"Internal user should be able to access /v1/videos route. Got error: {str(e)}" - ) + pytest.fail(f"Internal user should be able to access /v1/videos route. Got error: {e!s}") def test_videos_route_with_virtual_key_llm_api_routes(): @@ -1726,12 +1708,8 @@ def test_videos_route_with_virtual_key_llm_api_routes(): ] for route in test_routes: - result = RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token - ) - assert ( - result is True - ), f"Virtual key with llm_api_routes should be able to access {route}" + result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) + assert result is True, f"Virtual key with llm_api_routes should be able to access {route}" def test_non_proxy_admin_wildcard_allowed_routes(): @@ -1802,9 +1780,7 @@ def test_proxy_admin_viewer_can_access_global_spend_tags(): ) # If no exception is raised, the test passes except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {e!s}") # Routes returning proxy-wide spend across every team / customer / api_key. @@ -1832,7 +1808,7 @@ def test_internal_user_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -1861,7 +1837,7 @@ def test_internal_user_view_only_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, @@ -1963,9 +1939,7 @@ def test_proxy_admin_viewer_can_access_audit_logs(route): request_data={}, ) except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access {route} route. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access {route} route. Got error: {e!s}") # ── Admin Viewer parity: Logs page endpoints ────────────────────────────────── @@ -2028,9 +2002,7 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route): request_data={}, ) except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}") @pytest.mark.parametrize( @@ -2140,7 +2112,7 @@ def test_internal_user_blocked_from_admin_viewer_logs_routes(route): if route not in INTERNAL_USER_BLOCKED_SUBSET: return - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2216,9 +2188,7 @@ def test_proxy_admin_viewer_can_access_settings_read_endpoints(route): request_data={}, ) except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}") # ── Admin Viewer parity: default-allow GET semantics ───────────────────────── @@ -2417,9 +2387,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: ) local_file = os.path.abspath(local_file) - spec = importlib.util.spec_from_file_location( - "local_enterprise_route_checks", local_file - ) + spec = importlib.util.spec_from_file_location("local_enterprise_route_checks", local_file) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod.EnterpriseRouteChecks @@ -2430,9 +2398,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2448,9 +2414,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2466,9 +2430,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2479,9 +2441,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks.should_call_route("/v1/chat/completions") assert exc_info.value.status_code == 403 - assert "LLM API routes are disabled for this instance." in str( - exc_info.value.detail - ) + assert "LLM API routes are disabled for this instance." in str(exc_info.value.detail) @patch("litellm.proxy.proxy_server.premium_user", True) def test_should_embeddings_still_blocked_when_llm_api_disabled(self): @@ -2489,9 +2449,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2509,9 +2467,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2530,9 +2486,7 @@ def test_route_in_additional_public_routes_wildcard_match(): from litellm.proxy.auth.auth_utils import route_in_additonal_public_routes with ( - patch( - "litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]} - ), + patch("litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]}), patch("litellm.proxy.proxy_server.premium_user", True), ): # Wildcard should match subpaths @@ -2624,7 +2578,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re ) # /config/update is still blocked - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2712,8 +2666,6 @@ def test_available_roles_accessible_to_non_admin_users(user_role): # ── _user_is_org_admin tests ────────────────────────────────────────────────── - - def _make_org_admin_user(org_id: str) -> LiteLLM_UserTable: membership = LiteLLM_OrganizationMembershipTable( user_id="org-admin-user", @@ -2836,9 +2788,7 @@ async def test_add_team_org_context_noop_when_org_id_already_present(): raise AssertionError("must not resolve when organization_id is present") body = {"team_id": "team-1", "organization_id": "org-explicit"} - out = await add_team_org_context_to_request_body( - route="/team/update", request_body=body, fetch_team_org_id=fetch - ) + out = await add_team_org_context_to_request_body(route="/team/update", request_body=body, fetch_team_org_id=fetch) assert out == body @@ -2850,9 +2800,7 @@ async def test_add_team_org_context_noop_for_other_routes(): raise AssertionError("must not resolve for a non-opted-in route") body = {"team_id": "team-1"} - out = await add_team_org_context_to_request_body( - route="/team/delete", request_body=body, fetch_team_org_id=fetch - ) + out = await add_team_org_context_to_request_body(route="/team/delete", request_body=body, fetch_team_org_id=fetch) assert out == body @@ -2865,9 +2813,7 @@ async def test_add_team_org_context_noop_when_team_has_no_org(): return None body = {"team_id": "team-1"} - out = await add_team_org_context_to_request_body( - route="/team/update", request_body=body, fetch_team_org_id=fetch - ) + out = await add_team_org_context_to_request_body(route="/team/update", request_body=body, fetch_team_org_id=fetch) assert out == body @@ -3151,9 +3097,7 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): # Removing the endpoint should clean up openai_routes # remove_endpoint_routes takes endpoint_id (UUID portion of # the route key "{id}:exact:{path}:{methods}") - registered = ( - InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() - ) + registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() endpoint_ids = {k.split(":")[0] for k in registered} for eid in endpoint_ids: InitPassThroughEndpointHelpers.remove_endpoint_routes(eid) @@ -3163,9 +3107,7 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): LiteLLMRoutes.openai_routes.value[:] = original_routes # Clean up any routes registered during this test to avoid # polluting the module-level _registered_pass_through_routes - registered = ( - InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() - ) + registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() for k in registered: InitPassThroughEndpointHelpers.remove_endpoint_routes(k.split(":")[0]) @@ -3196,8 +3138,7 @@ def test_provider_name_substring_not_classified_as_llm_route(route): from litellm.proxy.auth.route_checks import RouteChecks assert RouteChecks.is_llm_api_route(route=route) is False, ( - f"{route!r} should NOT be classified as an LLM API route — " - "provider-name substring match bypass" + f"{route!r} should NOT be classified as an LLM API route — provider-name substring match bypass" ) @@ -3219,9 +3160,7 @@ def test_legitimate_passthrough_routes_still_classified_as_llm_route(route): """Legitimate passthrough routes must still pass is_llm_api_route.""" from litellm.proxy.auth.route_checks import RouteChecks - assert ( - RouteChecks.is_llm_api_route(route=route) is True - ), f"{route!r} should be classified as an LLM API route" + assert RouteChecks.is_llm_api_route(route=route) is True, f"{route!r} should be classified as an LLM API route" @pytest.mark.parametrize( @@ -3279,7 +3218,7 @@ def test_internal_user_blocked_from_search_tool_writes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -3655,12 +3594,7 @@ def test_agent_inference_routes_stay_llm_api(route): def test_agent_routes_union_still_covers_both_halves(route): """Keys configured with allowed_routes=["agent_routes"] must keep both halves.""" - assert ( - RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.agent_routes.value - ) - is True - ) + assert RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_routes.value) is True @pytest.mark.parametrize("route", AGENT_MANAGEMENT_ROUTES) @@ -3714,6 +3648,75 @@ def test_agent_registry_route_gate_open_to_non_admin_roles(user_role, method, ro valid_token=valid_token, request_data={}, ) + + +def test_proxy_admin_viewer_user_update_password_param_rejected(): + """The self-service /user/update password carve-out is closed: non-admins + change their own password through /user/password/change, which verifies + the current password. Admin password sets don't pass through this check.""" + with pytest.raises(HTTPException) as exc_info: + RouteChecks._check_proxy_admin_viewer_access( + route="/user/update", + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + request_data={"password": "hunter2hunter2"}, + ) + assert exc_info.value.status_code == 403 + assert "password" in str(exc_info.value.detail) + + +def test_proxy_admin_viewer_user_update_user_email_still_allowed(): + request = MagicMock(spec=Request) + request.method = "POST" + + allowed = RouteChecks._check_proxy_admin_viewer_access( + route="/user/update", + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + request_data={"user_email": "viewer@example.com"}, + request=request, + ) + + assert allowed is None + + +def test_proxy_admin_viewer_can_change_own_password(): + request = MagicMock(spec=Request) + request.method = "POST" + + allowed = RouteChecks._check_proxy_admin_viewer_access( + route="/user/password/change", + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + request_data={"current_password": "a", "new_password": "b"}, + request=request, + ) + + assert allowed is None + + +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_non_admin_roles_can_change_own_password(user_role): + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + + allowed = RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="test_user", user_role=user_role), + _user_role=user_role, + route="/user/password/change", + request=request, + valid_token=valid_token, + request_data={"current_password": "a", "new_password": "b"}, + ) + + assert allowed is None + + TEAM_CALLBACK_ROUTES = ( "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback", "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback/langfuse", diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 024d7b8300f..a443235fe15 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4313,6 +4313,35 @@ async def test_user_update_rejects_breached_password(_admin_prisma): _admin_prisma.db.litellm_usertable.find_first.assert_not_called() +@pytest.mark.asyncio +async def test_bulk_update_all_users_rejects_a_password(_admin_prisma): + """The all_users fast path writes user_updates straight to update_many, + bypassing _update_single_user_helper. A password riding along would be + stored as unvalidated plaintext on every row, so it must be rejected + before any DB access.""" + from fastapi import HTTPException + + from litellm.proxy._types import UpdateUserRequestNoUserIDorEmail + from litellm.proxy.management_endpoints.internal_user_endpoints import bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkUpdateUserRequest, + ) + + data = BulkUpdateUserRequest( + all_users=True, + user_updates=UpdateUserRequestNoUserIDorEmail(password="Str0ng!Passw0rd"), + ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(HTTPException) as exc_info: + await bulk_user_update(data=data, user_api_key_dict=admin_caller) + + assert exc_info.value.status_code == 400 + assert "not supported" in str(exc_info.value.detail) + _admin_prisma.db.litellm_usertable.find_many.assert_not_called() + _admin_prisma.db.litellm_usertable.update_many.assert_not_called() + + def _hibp_client_with_handler(handler) -> AsyncHTTPHandler: """A real AsyncHTTPHandler over httpx.MockTransport (the DI seam used throughout test_password_policy.py), so no network is touched.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py new file mode 100644 index 00000000000..c7e0a385ae9 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py @@ -0,0 +1,283 @@ +""" +Tests for POST /user/password/change (litellm/proxy/management_endpoints/password_endpoints.py). + +HIBP traffic is intercepted with respx; no test here touches the network. +""" + +import hashlib +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +import respx +from fastapi import HTTPException + +from litellm.proxy._types import LitellmTableNames, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.management_endpoints.password_endpoints import change_password +from litellm.proxy.utils import hash_password, verify_password + +CURRENT_PASSWORD = "OldP@ssw0rd-2026" +NEW_PASSWORD = "NewP@ssw0rd-2026" + +_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False} + + +def _make_user_row(password: str | None) -> MagicMock: + user = MagicMock() + user.user_id = "user-123" + user.password = password + return user + + +def _make_prisma(user: MagicMock | None) -> MagicMock: + prisma = MagicMock() + prisma.db.litellm_usertable.find_first = AsyncMock(return_value=user) + prisma.db.litellm_usertable.update = AsyncMock(return_value=user) + return prisma + + +def _caller(user_id: str | None = "user-123") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id=user_id) + + +def _hibp_url_for(password: str) -> str: + sha1 = hashlib.sha1(password.encode(), usedforsecurity=False).hexdigest().upper() + return f"https://api.pwnedpasswords.com/range/{sha1[:5]}" + + +def _hibp_suffix_for(password: str) -> str: + return hashlib.sha1(password.encode(), usedforsecurity=False).hexdigest().upper()[5:] + + +@pytest.mark.asyncio +async def test_change_password_success_writes_new_scrypt_hash(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + response = await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert response.user_id == "user-123" + update_kwargs = prisma.db.litellm_usertable.update.call_args.kwargs + assert update_kwargs["where"] == {"user_id": "user-123"} + stored = update_kwargs["data"]["password"] + assert stored != NEW_PASSWORD + assert verify_password(NEW_PASSWORD, stored) + + +@pytest.mark.asyncio +async def test_change_password_rejects_wrong_current_password(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 400 + assert "Current password is incorrect" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_rejects_session_without_user(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(user=None) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(user_id=None), + ) + + assert exc_info.value.status_code == 400 + prisma.db.litellm_usertable.find_first.assert_not_called() + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_rejects_account_without_password(): + """SSO users and the env-credential admin have no DB password row to change.""" + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(password=None)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 400 + assert "no password set" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_enforces_min_length(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(ProxyException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password="Short1!"), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert "at least 12 characters" in exc_info.value.message + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_change_password_rejects_breached_password(): + """With the default policy, the new password is screened against HIBP.""" + from litellm.proxy._types import ChangePasswordRequest + + breached_password = "Password123!" + respx.get(_hibp_url_for(breached_password)).mock( + return_value=httpx.Response(200, text=f"{_hibp_suffix_for(breached_password)}:1") + ) + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(ProxyException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=breached_password), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert "data breaches" in exc_info.value.message + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_change_password_verifies_current_password_before_hibp_lookup(): + """A caller who fails current-password verification must not trigger any + HIBP traffic. The HIBP check fails open on errors, so an unmocked lookup + could not prove ordering; instead the route is registered and asserted + uncalled.""" + from litellm.proxy._types import ChangePasswordRequest + + hibp_route = respx.get(_hibp_url_for(NEW_PASSWORD)).mock(return_value=httpx.Response(200, text="")) + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 400 + assert "Current password is incorrect" in exc_info.value.detail["error"] + assert not hibp_route.called + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_success_emits_redacted_audit_log(): + """A successful change must land in the audit trail as field names only; + the plaintext passwords must never reach the audit call.""" + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + audit_mock = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.management_endpoints.password_endpoints.create_object_audit_log", audit_mock), + ): + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + audit_mock.assert_awaited_once() + audit_kwargs = audit_mock.await_args.kwargs + assert audit_kwargs["object_id"] == "user-123" + assert audit_kwargs["action"] == "updated" + assert audit_kwargs["table_name"] == LitellmTableNames.USER_TABLE_NAME + assert audit_kwargs["after_value"] == '{"fields_changed": ["password"]}' + assert CURRENT_PASSWORD not in str(audit_kwargs) + assert NEW_PASSWORD not in str(audit_kwargs) + + +@pytest.mark.asyncio +async def test_change_password_failure_emits_no_audit_log(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + audit_mock = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.management_endpoints.password_endpoints.create_object_audit_log", audit_mock), + ): + with pytest.raises(HTTPException): + await change_password( + data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + audit_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_change_password_requires_db(): + from litellm.proxy._types import ChangePasswordRequest + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 500 diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 8f9c44d7a38..4ea30c7bf87 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -5,6 +5,7 @@ from pydantic import ValidationError from litellm.proxy._types import ( ROLES_WITHIN_ORG, + ChangePasswordRequest, GenerateKeyRequest, KeyRequest, LiteLLM_AuditLogs, @@ -293,3 +294,28 @@ def test_new_user_request_loudly_rejects_a_password(): def test_new_user_request_without_password_still_works(): request = NewUserRequest(user_email="alice@example.com") assert request.password is None + + +def test_update_user_request_accepts_a_password(): + """Admins set user passwords through /user/update; the value must survive + model validation so the endpoint can policy-check and hash it.""" + request = UpdateUserRequest(user_id="user-123", password="hunter2hunter2") + assert request.password == "hunter2hunter2" + + +def test_update_user_request_password_hidden_from_repr(): + """management_endpoint_wrapper string-formats endpoint kwargs into Slack + alerts, so the model's repr/str must never contain the plaintext password.""" + request = UpdateUserRequest(user_id="user-123", password="hunter2hunter2") + assert "hunter2hunter2" not in repr(request) + assert "hunter2hunter2" not in str(request) + + +def test_change_password_request_passwords_hidden_from_repr(): + """Any accidental str()/repr() of the request model (debug logs, exception + handlers, a future management_endpoint_wrapper) must never contain either + plaintext password.""" + request = ChangePasswordRequest(current_password="hunter2hunter2", new_password="NewP@ssw0rd-2026") + for rendered in (repr(request), str(request)): + assert "hunter2hunter2" not in rendered + assert "NewP@ssw0rd-2026" not in rendered diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx new file mode 100644 index 00000000000..e78f170cf0c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx @@ -0,0 +1,68 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import ChangePasswordForm from "./ChangePasswordForm"; + +const mockChangePasswordCall = vi.fn(); +const mockToastSuccess = vi.fn(); + +vi.mock("@/components/networking", () => ({ + changePasswordCall: (...args: unknown[]) => mockChangePasswordCall(...args), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-session-token" }), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { + success: (...args: unknown[]) => mockToastSuccess(...args), + fromError: vi.fn(), + }, +})); + +const fillForm = (values: { current: string; next: string; confirm: string }) => { + fireEvent.change(screen.getByLabelText("Current Password"), { target: { value: values.current } }); + fireEvent.change(screen.getByLabelText("New Password"), { target: { value: values.next } }); + fireEvent.change(screen.getByLabelText("Confirm New Password"), { target: { value: values.confirm } }); +}; + +const submit = () => fireEvent.click(screen.getByRole("button", { name: "Change Password" })); + +describe("ChangePasswordForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("sends the current and new password to the change endpoint and resets on success", async () => { + mockChangePasswordCall.mockResolvedValue({ user_id: "user-123", message: "Password updated successfully." }); + render(); + + fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" }); + submit(); + + expect(await screen.findByLabelText("Current Password")).toHaveValue(""); + expect(mockChangePasswordCall).toHaveBeenCalledWith("sk-session-token", "OldP@ssw0rd-2026", "NewP@ssw0rd-2026"); + expect(mockToastSuccess).toHaveBeenCalled(); + }); + + it("blocks submission when the confirmation does not match", async () => { + render(); + + fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "Different-2026" }); + submit(); + + expect(await screen.findByText("New passwords do not match")).toBeInTheDocument(); + expect(mockChangePasswordCall).not.toHaveBeenCalled(); + }); + + it("shows the proxy's rejection message unwrapped", async () => { + mockChangePasswordCall.mockRejectedValue(new Error("{'error': 'Current password is incorrect.'}")); + render(); + + fillForm({ current: "wrong-password", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" }); + submit(); + + expect(await screen.findByText("Current password is incorrect.")).toBeInTheDocument(); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx new file mode 100644 index 00000000000..4c51b7f3d15 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx @@ -0,0 +1,100 @@ +"use client"; + +import React, { useState } from "react"; +import { CircleAlert } from "lucide-react"; +import { z } from "zod/v4"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent } from "@/components/ui/card"; +import { FieldGroup } from "@/components/ui/field"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { changePasswordCall } from "@/components/networking"; +import { extractProxyErrorMessage } from "@/lib/http/client"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { toast } from "@/lib/toast"; + +const changePasswordSchema = z + .object({ + currentPassword: z.string().min(1, "Current password is required"), + newPassword: z.string().min(1, "New password is required"), + confirmNewPassword: z.string().min(1, "Confirm your new password"), + }) + .refine((values) => values.newPassword === values.confirmNewPassword, { + message: "New passwords do not match", + path: ["confirmNewPassword"], + }); + +type ChangePasswordValues = z.infer; + +export function ChangePasswordForm() { + const { accessToken } = useAuthorized(); + const form = useZodForm(changePasswordSchema, { + defaultValues: { currentPassword: "", newPassword: "", confirmNewPassword: "" }, + }); + const [isPending, setIsPending] = useState(false); + const [submitError, setSubmitError] = useState(null); + + const handleSubmit = async (values: ChangePasswordValues) => { + if (!accessToken) return; + setSubmitError(null); + setIsPending(true); + try { + await changePasswordCall(accessToken, values.currentPassword, values.newPassword); + toast.success("Password updated"); + form.reset(); + } catch (error) { + setSubmitError(extractProxyErrorMessage(error)); + } finally { + setIsPending(false); + } + }; + + return ( +
+ + +

Change Password

+

+ Enter your current password and choose a new one. The new password must meet this proxy's password + policy. +

+ +
+ + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {submitError && ( + + + {submitError} + + )} + +
+ +
+
+ + +
+ ); +} + +export default ChangePasswordForm; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx new file mode 100644 index 00000000000..0a6ae926ceb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx @@ -0,0 +1,7 @@ +"use client"; + +import ChangePasswordForm from "./ChangePasswordForm"; + +export default function ChangePasswordPage() { + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 40d1ec09d1f..089153cec76 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -50,6 +50,7 @@ const useAuthorized = () => { isViewOnly: isViewOnlySessionRole(decoded?.user_role), premiumUser: decoded?.premium_user ?? null, disabledPersonalKeyCreation: decoded?.disabled_non_admin_personal_key_creation ?? null, + loginMethod: decoded?.login_method ?? null, showSSOBanner: decoded?.login_method === "username_password", }; }; diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx index cad5ced340e..4bdf0da3407 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx @@ -3,13 +3,25 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../../tests/test-utils"; import UserDropdown from "./UserDropdown"; -let mockUseAuthorizedImpl = () => ({ +let mockUseAuthorizedImpl: () => { + userId: string | null; + userEmail: string | null; + userRoleLabel: string; + premiumUser: boolean; + loginMethod?: string | null; +} = () => ({ userId: "test-user-id", userEmail: "test@example.com", userRoleLabel: "Admin", premiumUser: false, }); +const mockRouterPush = vi.fn(); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ push: mockRouterPush }), +})); + let mockUseDisableShowPromptsImpl = () => false; let mockGetLocalStorageItemImpl = (key: string): string | null => { @@ -143,6 +155,44 @@ describe("UserDropdown", () => { expect(mockOnLogout).toHaveBeenCalledTimes(1); }); + it("should navigate to the change-password page for username/password sessions", async () => { + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRoleLabel: "Admin", + premiumUser: false, + loginMethod: "username_password", + }); + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(getAccountTrigger()); + + await user.click(await screen.findByText("Change Password")); + + expect(mockRouterPush).toHaveBeenCalledWith(expect.stringContaining("change-password")); + }); + + it("should hide the change-password entry for SSO sessions", async () => { + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRoleLabel: "Admin", + premiumUser: false, + loginMethod: "sso", + }); + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(getAccountTrigger()); + + await waitFor(() => { + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); + }); + + expect(screen.queryByText("Change Password")).not.toBeInTheDocument(); + }); + it("should toggle hide new feature indicators switch", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx index 95f76dbb2cc..9c02defc778 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx @@ -9,7 +9,9 @@ import { setLocalStorageItem, } from "@/utils/localStorageUtils"; import { navAccountDisplayName } from "@/components/Navbar/navDisplayName"; -import { ChevronDown, ChevronsUpDown, Crown, LogOut, Mail, ShieldCheck, User } from "lucide-react"; +import { uiHref } from "@/utils/uiHref"; +import { ChevronDown, ChevronsUpDown, Crown, KeyRound, LogOut, Mail, ShieldCheck, User } from "lucide-react"; +import { useRouter } from "next/navigation"; import { Avatar, AvatarFallback } from "@/components/ui/avatar"; import { Badge } from "@/components/ui/badge"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; @@ -63,7 +65,9 @@ interface UserDropdownProps { } const UserDropdown: React.FC = ({ onLogout, variant = "navbar", collapsed = false }) => { - const { userId, userEmail, userRoleLabel: userRole, premiumUser } = useAuthorized(); + const { userId, userEmail, userRoleLabel: userRole, premiumUser, loginMethod } = useAuthorized(); + const router = useRouter(); + const [open, setOpen] = useState(false); const disableShowPrompts = useDisableShowPrompts(); const disableBlogPosts = useDisableBlogPosts(); const disableBouncingIcon = useDisableBouncingIcon(); @@ -197,7 +201,7 @@ const UserDropdown: React.FC = ({ onLogout, variant = "navbar const displayName = navAccountDisplayName(userEmail, userId); return ( - + {variant === "sidebar" ? ( = ({ onLogout, variant = "navbar > {renderUserInfoSection()} + {loginMethod === "username_password" && ( + + )} + )} + + + )} + {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 072/149] 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 073/149] 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 074/149] 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 075/149] 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 076/149] 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 077/149] 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 078/149] 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 079/149] 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 080/149] 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 b415f2263a6b5cb3b26071c977c55bcc91604a8b Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 08:59:46 +0000 Subject: [PATCH 081/149] test(proxy): restore pipeline methods on expiring Redis fake after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/test_budget_reservation.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6905585fd5e..86bf896188f 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1,7 +1,7 @@ import asyncio import threading import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -11,6 +11,7 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( AgenticAnthropicStreamingIterator, @@ -2400,6 +2401,14 @@ class _ExpiringRedisCache: self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) return True + async def async_increment_pipeline( + self, increment_list: Sequence[RedisPipelineIncrementOperation], **kwargs: object + ) -> list[float]: + return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list] + + def get_ttl(self, **kwargs: object) -> int | None: + return int(self.default_ttl) + @pytest.mark.asyncio async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( @@ -2475,15 +2484,6 @@ async def test_reservation_lease_stops_when_request_task_ends_without_reconcilin assert redis_cache.refresh_count == 0 assert await redis_cache.async_get_cache(key=counter_key) is None - async def async_increment_pipeline(self, increment_list, **kwargs): - results = [] - for op in increment_list: - results.append(await self.async_increment(op["key"], op["increment_value"])) - return results - - def get_ttl(self, **kwargs) -> None: - return None - class _TeamMembershipFloorDb: """Stands in for `prisma_client.db`: only the team-membership row exists and its spend is the DB floor.""" From ef55eb6bdf0ccb69533d9541266c5108e0176609 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 09:07:08 +0000 Subject: [PATCH 082/149] 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 083/149] 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 ea04267912f5066e8de7b857626d9588bdac39cd Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 09:30:20 +0000 Subject: [PATCH 084/149] fix(proxy): return an immutable bucket mapping from count_release_bullets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../ui_crud_endpoints/latest_release_endpoints.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py index 61124e9f27e..ad5cc8efc31 100644 --- a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -4,7 +4,7 @@ from collections import Counter from collections.abc import Awaitable, Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Annotated, Final, Literal, Protocol +from typing import Annotated, Final, Literal, Protocol, TypeAlias import httpx from fastapi import APIRouter, Depends @@ -25,7 +25,7 @@ LATEST_RELEASE_CACHE_KEY: Final = "latest_release_info" _RELEASE_BULLET_PATTERN: Final = re.compile(r"^\*\s+(?:([A-Za-z]+)(?:\([^)]*\))?!?:\s)?\S") _NEW_CONTRIBUTOR_PATTERN: Final = re.compile(r"^\*\s+@\S+ made their first contribution\b") -_Bucket = Literal["new_features", "bug_fixes", "other_updates"] +_Bucket: TypeAlias = Literal["new_features", "bug_fixes", "other_updates"] _PREFIX_BUCKETS: Final[Mapping[str, _Bucket]] = MappingProxyType({"feat": "new_features", "fix": "bug_fixes"}) @@ -81,9 +81,9 @@ def _bucket_for(line: str) -> _Bucket | None: return "other_updates" if prefix is None else _PREFIX_BUCKETS.get(prefix.lower(), "other_updates") -def count_release_bullets(body: str) -> Counter[_Bucket]: +def count_release_bullets(body: str) -> Mapping[_Bucket, int]: """Bucket release-note bullets by conventional-commit type or ``other_updates``.""" - return Counter(bucket for line in body.splitlines() if (bucket := _bucket_for(line)) is not None) + return MappingProxyType(Counter(bucket for line in body.splitlines() if (bucket := _bucket_for(line)) is not None)) def parse_latest_release(response: httpx.Response) -> LatestReleaseInfo | LatestReleaseUnavailable: @@ -96,9 +96,9 @@ def parse_latest_release(response: httpx.Response) -> LatestReleaseInfo | Latest counts: Final = count_release_bullets(release.body) return LatestReleaseInfo( version=release.tag_name.removeprefix("v"), - new_features=counts["new_features"], - bug_fixes=counts["bug_fixes"], - other_updates=counts["other_updates"], + new_features=counts.get("new_features", 0), + bug_fixes=counts.get("bug_fixes", 0), + other_updates=counts.get("other_updates", 0), release_url=release.html_url, ) From 3777b0b0d9393b8739d126893203a9080b8038bf Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 09:54:29 +0000 Subject: [PATCH 085/149] test(proxy): document why the release info route check asserts by not raising Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/auth/test_route_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index ff9995df1ca..87bf4595af5 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -127,7 +127,7 @@ def test_user_banner_read_open_to_non_admin_roles(role): LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, ], ) -def test_latest_release_info_read_open_to_non_admin_roles(role): +def test_latest_release_info_read_open_to_non_admin_roles(role): # test-quality-ok: allowed path returns None, not raising is the observable user_obj = LiteLLM_UserTable( user_id="test_user", user_email="test@example.com", From 95cf7066d16186e94a8f27828ac38db55ef45cf7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 20:38:13 -0700 Subject: [PATCH 086/149] 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 087/149] 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 088/149] 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 089/149] 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 090/149] 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 091/149] 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 092/149] 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 093/149] 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 094/149] 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 095/149] 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 096/149] 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 097/149] 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 1f1b61173d79ae86b2fdafd48111030bea7871f0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 18:10:41 +0000 Subject: [PATCH 098/149] fix(otel v2): map OCR page markdown onto the generation output Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/model/payloads.py | 19 ++++++++- .../otel/test_otel_v2_sources_of_truth.py | 39 +++++++++++++++++++ .../otel/test_otel_v2_vendor_mappers.py | 22 +++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index c23b3291365..3164e0977b7 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -427,7 +427,7 @@ class LLMCallSpanData: # plain ``.get`` — no repeated ``isinstance`` guards. raw_response: Final = payload.get("response") response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) - choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) + choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) or _ocr_choices(response) # ``finish_reasons`` is metadata, not content, so derive it from # ``choices_out`` before gating. The raw message/choice bodies are only # retained when content capture is enabled (see ``capture_span_content``); @@ -752,6 +752,23 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: return (choice,) +def _ocr_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + """An ``OCRResponse`` ``pages`` list folded into one chat-shaped assistant choice.""" + markdowns: Final = tuple( + text for page in _dicts(response.get("pages")) if (text := as_str(page.get("markdown"))) is not None + ) + if not markdowns: + return () + message: Final[_AssistantMessage] = { + "role": "assistant", + "content": "\n\n".join(markdowns), + "refusal": None, + "tool_calls": None, + } + choice: Final[_Choice] = {"message": message, "finish_reason": None} + return (choice,) + + def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None: texts: Final = tuple( text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 17de3cf1e8a..01f2a13d252 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -872,6 +872,45 @@ def test_chat_choices_win_over_a_responses_output_list(): assert data.finish_reasons == ("stop",) +def _ocr_payload(pages: list[object]): + return _sample_payload( + call_type="aocr", + custom_llm_provider="mistral", + model="mistral-ocr-latest", + messages=None, + response={"object": "ocr", "model": "mistral-ocr-latest", "pages": pages, "usage_info": {"pages_processed": 2}}, + ) + + +def test_ocr_pages_become_one_assistant_choice_joined_in_page_order(): + data = LLMCallSpanData.from_standard_logging_payload( + _ocr_payload([{"index": 0, "markdown": "# Invoice"}, {"index": 1, "markdown": "Total: 42"}]), + capture_content=True, + ) + + assert data.choices_out == ( + { + "message": {"role": "assistant", "content": "# Invoice\n\nTotal: 42", "refusal": None, "tool_calls": None}, + "finish_reason": None, + }, + ) + assert data.finish_reasons == () + + +def test_ocr_output_follows_the_content_capture_gate(): + data = LLMCallSpanData.from_standard_logging_payload(_ocr_payload([{"index": 0, "markdown": "# Invoice"}])) + + assert data.choices_out == () + + +def test_ocr_pages_without_markdown_stay_empty(): + data = LLMCallSpanData.from_standard_logging_payload( + _ocr_payload([{"index": 0, "images": []}, "not-a-page"]), capture_content=True + ) + + assert data.choices_out == () + + def test_request_identity_prefers_canonical_team_keys(): from litellm.integrations.otel.model.payloads import RequestIdentity diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index 4e375de0494..9fa198c4ec5 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -227,6 +227,28 @@ def test_langfuse_mapper_renders_a_responses_api_call_from_the_standard_logging_ assert attrs["langfuse.observation.type"] == "generation" +def test_langfuse_mapper_renders_an_ocr_call_with_the_page_markdown_as_output(): + payload = { + "call_type": "aocr", + "custom_llm_provider": "mistral", + "model": "mistral-ocr-latest", + "messages": None, + "response": { + "object": "ocr", + "model": "mistral-ocr-latest", + "pages": [{"index": 0, "markdown": "# Invoice"}, {"index": 1, "markdown": "Total: 42"}], + "usage_info": {"pages_processed": 2}, + }, + } + data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True) + attrs = LangfuseMapper().map(data) + + assert json.loads(attrs["langfuse.observation.output"]) == [ + {"role": "assistant", "content": "# Invoice\n\nTotal: 42", "refusal": None, "tool_calls": None} + ] + assert attrs["langfuse.observation.type"] == "generation" + + # --------------------------------------------------------------------------- # # Weave # --------------------------------------------------------------------------- # From 1b6b704ddd639898895b1add55cdb9579a278c61 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:15:44 -0700 Subject: [PATCH 099/149] 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 100/149] 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 101/149] 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 102/149] 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 103/149] 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 104/149] 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 105/149] 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 106/149] 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 107/149] 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 108/149] 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 109/149] 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 110/149] 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 111/149] 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 0e1605010052a54b5a2b9e01fee89499e10858d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:21:50 -0700 Subject: [PATCH 112/149] fix(anthropic): forward Claude Code safeguards and dangerous-tool-use beta to Bedrock Invoke and Vertex on /v1/messages Claude Code's server-side auto-mode classifier sends a `safeguards` body field together with the `dangerous-tool-use-2026-09-03` beta. PR #42152 made the first-party anthropic route pass them through, but the beta header mapping left the other two Claude platforms at null, so Bedrock Invoke dropped both (classifier silently disabled) and Vertex forwarded the body field without the beta, which the platform rejects with "safeguards: Extra inputs are not permitted" (a 400 Claude Code hides by retrying without them). Map the beta for bedrock and vertex_ai in the beta headers config and add `safeguards` to the Bedrock Invoke request allowlist so the pair reaches both platforms unchanged. Nothing is injected: a client that sends `safeguards` without the beta still gets the platform's 400, exactly as api.anthropic.com answers it. --- litellm/anthropic_beta_headers_config.json | 6 ++ litellm/types/llms/bedrock.py | 1 + ...erimental_pass_through_messages_handler.py | 95 +++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 43 +++++++++ .../test_anthropic_beta_headers_filtering.py | 14 +++ 5 files changed, 159 insertions(+) diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index eb31cc17a15..c4c60ae715b 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -11,6 +11,7 @@ "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", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", @@ -44,6 +45,7 @@ "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", + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": "files-api-2025-04-14", @@ -76,6 +78,7 @@ "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": null, + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -109,6 +112,7 @@ "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", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -142,6 +146,7 @@ "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", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": null, "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -175,6 +180,7 @@ "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", + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 10082cf2373..f7518fefae4 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1236,6 +1236,7 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): thinking: dict metadata: dict output_config: dict + safeguards: list # `context_management` is allowed for Bedrock InvokeModel only when it # carries `compact_20260112` edits paired with the `compact-2026-01-12` 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 e8bfcb86bf6..8d940e5efad 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 @@ -1544,3 +1544,98 @@ async def test_anthropic_messages_streaming_forwards_safeguards_and_keeps_safegu assert captured["body"]["safeguards"] == safeguards assert events[0]["message"]["safeguard_results"] == safeguard_results assert [e for e in events if e["type"] == "message_delta"][0]["delta"]["safeguard_results"] == safeguard_results + + +def _claude_code_auto_mode_request() -> tuple[list[dict[str, object]], list[dict[str, object]]]: + """Shapes are what Claude Code 2.1.278 sends and Bedrock Invoke / Vertex rawPredict return, captured 2026-09-21.""" + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] + return safeguards, safeguard_results + + +def _upstream_answering_with(safeguard_results: list[dict[str, object]], captured: dict[str, object]) -> AsyncHTTPHandler: + def upstream_records_the_request(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + captured["anthropic-beta"] = request.headers.get("anthropic-beta") + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + "safeguard_results": safeguard_results, + }, + request=request, + ) + + upstream = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream_records_the_request)) + return upstream + + +@pytest.mark.asyncio +async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_bedrock_invoke( + local_beta_headers_config, +): + """Bedrock Invoke takes betas in the body's `anthropic_beta` and 400s on `safeguards` without the beta.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + safeguards, safeguard_results = _claude_code_auto_mode_request() + captured: dict[str, object] = {} + + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="bedrock/us.anthropic.claude-sonnet-5", + custom_llm_provider="bedrock", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=_upstream_answering_with(safeguard_results, captured), + safeguards=safeguards, + extra_headers={"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, + ) + + assert captured["body"]["safeguards"] == safeguards + assert captured["body"]["anthropic_beta"] == ["dangerous-tool-use-2026-09-03"] + assert response["safeguard_results"] == safeguard_results + + +@pytest.mark.asyncio +async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( + local_beta_headers_config, +): + """Vertex rawPredict takes the beta as the `anthropic-beta` header and 400s on `safeguards` without it.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + safeguards, safeguard_results = _claude_code_auto_mode_request() + captured: dict[str, object] = {} + + with patch.object(VertexBase, "_ensure_access_token", return_value=("test-token", "test-project")): + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="vertex_ai/claude-sonnet-5", + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="global", + vertex_credentials="{}", + client=_upstream_answering_with(safeguard_results, captured), + safeguards=safeguards, + extra_headers={"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, + ) + + assert captured["body"]["safeguards"] == safeguards + assert "anthropic_beta" not in captured["body"] + assert set(captured["anthropic-beta"].split(",")) == { + "dangerous-tool-use-2026-09-03", + "interleaved-thinking-2025-05-14", + } + assert response["safeguard_results"] == safeguard_results diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 7be005c0efe..e3c4a84cc60 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1651,6 +1651,49 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): assert set(result).issubset(cfg.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS) +def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta(local_beta_headers_config): + """ + Claude Code's server-side auto-mode classifier sends `safeguards` alongside the + dangerous-tool-use-2026-09-03 beta. Bedrock Invoke accepts the pair, answers + "safeguards: Extra inputs are not permitted" for the field alone, and returns + `safeguard_results: []` for the beta alone, so both must reach it unchanged. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + anthropic_messages_optional_request_params={"max_tokens": 64, "safeguards": safeguards}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, + ) + + assert result["safeguards"] == safeguards + assert result["anthropic_beta"] == ["dangerous-tool-use-2026-09-03"] + + +def test_bedrock_messages_stream_decoder_keeps_safeguard_results(): + """Bedrock streams the classifier verdicts on message_start and on the final message_delta, exactly as api.anthropic.com does.""" + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="us.anthropic.claude-sonnet-5") + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] + + message_delta = decoder._chunk_parser( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 1}, + "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3, "outputTokenCount": 1}, + } + ) + + assert isinstance(message_delta, dict) + assert message_delta["delta"]["safeguard_results"] == safeguard_results + + def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): """ In proxy deployments the client (e.g. Claude Code) doesn't know the backend diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index 3c967283abf..d600d2b734b 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -442,6 +442,20 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["thinking-binding-controls-2026-08-01"] + @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "vertex_ai"]) + def test_dangerous_tool_use_forwarded(self, provider): + """Claude Code's server-side auto-mode classifier sends `safeguards` together with + dangerous-tool-use-2026-09-03. Bedrock Invoke and Vertex rawPredict both answer + "safeguards: Extra inputs are not permitted" when the body field arrives without + the beta (probed 2026-09-21), so dropping the header turned every auto-mode turn + into a 400 on Vertex and silently disabled the classifier on Bedrock.""" + filtered = filter_and_transform_beta_headers( + beta_headers=["dangerous-tool-use-2026-09-03"], + provider=provider, + ) + + assert filtered == ["dangerous-tool-use-2026-09-03"] + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ From 34293483059d74af7138b62aa55865454237c165 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:21:53 +0000 Subject: [PATCH 113/149] 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 114/149] 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 10d343c3eee7472341c04e7ad5a1438f8eeb1a9d Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 19:06:13 +0000 Subject: [PATCH 115/149] fix(proxy): detach stored credential when model editor selects None (LIT-7597) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_management_endpoints.py | 45 +- .../proxy/auth/test_model_checks.py | 86 ++++ .../test_model_management_endpoints.py | 399 ++++++++++++++++++ .../src/components/ModelInfoEditForm.tsx | 16 +- .../src/components/model_info_view.test.tsx | 113 ++++- .../src/components/model_info_view.tsx | 12 +- 6 files changed, 651 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ea124776d0b..a8242904636 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -28,6 +28,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.ptu_pricing import ( CUSTOM_PRICING_FIELDS, PTU_EMPTIED_PRICING_FIELDS, @@ -145,7 +146,7 @@ if TYPE_CHECKING: from prisma import types as prisma_types router: Final = APIRouter() -CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points"}) +CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points", "litellm_credential_name"}) NULL_CLEARABLE_LITELLM_PARAMS: Final = frozenset((*SPECIAL_MODEL_INFO_PARAMS, *CLEARABLE_LITELLM_PARAMS)) @@ -332,6 +333,28 @@ def _raise_on_strategy_router_write_violation( ) +def _raise_on_invalid_credential_name(litellm_params: updateLiteLLMParams | None) -> None: + if litellm_params is None or "litellm_credential_name" not in litellm_params.model_fields_set: + return + credential_name: Final = litellm_params.litellm_credential_name + if credential_name is None: + return + if credential_name == "": + raise ProxyException( + message="litellm_credential_name cannot be an empty string. Send null to detach the stored credential or omit the field to leave it unchanged.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + if CredentialAccessor.find_credential(credential_name) is None: + raise ProxyException( + message=f"Credential '{credential_name}' not found. Create it via /credentials before attaching it to a model.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + + AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" _STORED_LITELLM_PARAMS_SQL: Final = ( @@ -1111,6 +1134,7 @@ async def patch_model( user_api_key_dict=user_api_key_dict, existing_litellm_params=db_model.litellm_params, ) + _raise_on_invalid_credential_name(patch_data.litellm_params) ModelManagementAuthChecks.can_user_set_aws_session_tags( litellm_params=patch_data.litellm_params, @@ -1921,21 +1945,28 @@ class ModelManagementAuthChecks: user_api_key_dict: UserAPIKeyAuth, existing_litellm_params: GenericLiteLLMParams | None = None, ) -> Literal[True]: - if litellm_params is None or litellm_params.litellm_credential_name is None: + if litellm_params is None: return True - if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None: - existing_credential_name: Final = decrypt_value_helper( + if "litellm_credential_name" not in litellm_params.model_fields_set: + return True + existing_credential_name: Final = ( + decrypt_value_helper( value=existing_litellm_params.litellm_credential_name, key="litellm_credential_name", exception_type="debug", return_original_value=True, ) - if litellm_params.litellm_credential_name == existing_credential_name: - return True + if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None + else None + ) + requested_credential_name: Final = litellm_params.litellm_credential_name + if requested_credential_name == existing_credential_name: + return True if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return True + action: Final = "detach" if requested_credential_name is None else "attach" raise ProxyException( - message=f"Only a proxy admin can attach a stored credential (litellm_credential_name) to a model. Your role={user_api_key_dict.user_role}.", + message=f"Only a proxy admin can {action} a stored credential (litellm_credential_name) on a model. Your role={user_api_key_dict.user_role}.", type=ProxyErrorTypes.auth_error.value, code=status.HTTP_403_FORBIDDEN, param="litellm_credential_name", diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f10622e954b..13171a42cda 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -523,6 +523,92 @@ def test_wildcard_credential_hydration_preserves_missing_credential_name( } +def test_hydrate_credential_name_none_leaves_params_untouched(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ) + ], + ) + params = LiteLLM_Params(model="openai/gpt-4o", litellm_credential_name=None) + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key is None + assert result.litellm_credential_name is None + + +def test_hydrate_replaced_credential_uses_new_credential_values(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ), + CredentialItem( + credential_name="other-credential", + credential_info={}, + credential_values={"api_key": "sk-other"}, + ), + ], + ) + params = LiteLLM_Params(model="openai/gpt-4o", litellm_credential_name="other-credential") + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key == "sk-other" + assert result.litellm_credential_name is None + + +def test_hydrate_inline_api_key_wins_over_stored_credential(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ) + ], + ) + params = LiteLLM_Params( + model="openai/gpt-4o", + api_key="sk-inline", + litellm_credential_name="shared-credential", + ) + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key == "sk-inline" + + @pytest.mark.asyncio async def test_get_available_models_for_user_expands_query_team_wildcard( monkeypatch, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 376309d8a7e..72db04dfbaa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -308,6 +308,46 @@ class TestModelManagementAuthChecks: ) assert result is True + def test_can_user_attach_credential_non_admin_explicit_null_clear_fails(self): + from litellm.proxy._types import ProxyException + from litellm.types.router import updateLiteLLMParams as litellm_params + + with pytest.raises(ProxyException) as exc_info: + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + ) + + assert exc_info.value.code == "403" + assert exc_info.value.param == "litellm_credential_name" + + def test_can_user_attach_credential_admin_explicit_null_clear_succeeds(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + ) + + assert result is True + + def test_can_user_attach_credential_null_without_existing_allows_any_role(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params(model="test_model"), + ) + + assert result is True + def test_can_user_attach_credential_unchanged_encrypted_existing_allows_any_role(self, monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") encrypted_name = encrypt_value_helper(value="shared-credential") @@ -4000,6 +4040,365 @@ class TestUpdateDBModelClearCacheControlInjectionPoints: assert params["tpm"] == 10 +class TestUpdateDBModelClearCredentialName: + def test_explicit_null_removes_stored_credential_name(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=None) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + info: Final = json.loads(result["model_info"]) + assert "litellm_credential_name" not in params + assert params["model"] == "openai/gpt-4o" + assert params["api_base"] == "https://api.openai.com/v1" + assert params["api_key"] == "sk-real" + assert params["tpm"] == 100 + assert info["team_id"] == "team-keep" + assert info["access_groups"] == ["prod"] + + def test_omitted_credential_name_keeps_stored_association(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(tpm=10)) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert params["litellm_credential_name"] == "shared-credential" + assert params["tpm"] == 10 + + def test_null_clear_on_model_without_credential_is_noop(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_base="https://api.openai.com/v1"), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=None) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert "litellm_credential_name" not in params + assert params["model"] == "openai/gpt-4o" + assert params["api_base"] == "https://api.openai.com/v1" + + def test_null_credential_clear_alongside_pricing_clear(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + input_cost_per_token=0.000001, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", input_cost_per_token=0.000001), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams( + litellm_credential_name=None, + input_cost_per_token=None, + ) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + info: Final = json.loads(result["model_info"]) + assert "litellm_credential_name" not in params + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + + def test_replace_credential_name_keeps_other_params(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name="other-credential") + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert params["litellm_credential_name"] == "other-credential" + assert params["api_base"] == "https://api.openai.com/v1" + assert params["api_key"] == "sk-real" + assert params["tpm"] == 100 + + +class TestPatchModelCredentialName: + @staticmethod + async def _patch_model( + monkeypatch, + db_model: Deployment, + user_api_key_dict: UserAPIKeyAuth, + credential_name: str | None, + ) -> list[dict[str, object]]: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model, update_db_model + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ), + CredentialItem( + credential_name="other-credential", + credential_info={}, + credential_values={"api_key": "sk-other"}, + ), + ], + ) + persisted: Final[list[dict[str, object]]] = [] + + async def persist_model(**kwargs): + row: Final = update_db_model(db_model=kwargs["db_model"], updated_patch=kwargs["patch_data"]) + persisted.append(row) + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + return updated_row + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=db_model), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=persist_model), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.raise_if_reload_degraded_serving" + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", + new=AsyncMock(), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.live_model_ids_snapshot", + return_value=frozenset(), + ), + ): + await patch_model( + model_id="dep-cred-1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=credential_name) + ), + user_api_key_dict=user_api_key_dict, + ) + + return persisted + + @pytest.mark.asyncio + async def test_patch_model_rejects_empty_string_credential_name(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "", + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "litellm_credential_name" + assert "empty" in exc_info.value.message.lower() + + @staticmethod + def _admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + @staticmethod + def _team_admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="team-admin", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-keep") + + @pytest.mark.asyncio + async def test_patch_model_rejects_unknown_credential_name(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "ghost-credential", + ) + + assert exc_info.value.code == "400" + assert "not found" in exc_info.value.message.lower() + + @pytest.mark.asyncio + async def test_patch_model_replaces_credential_name_and_preserves_other_params(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), "other-credential") + params: Final = json.loads(persisted[0]["litellm_params"]) + assert params["litellm_credential_name"] == "other-credential" + assert params["api_base"] == "https://api.openai.com/v1" + + @pytest.mark.asyncio + async def test_patch_model_admin_null_clear_persists_without_credential(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), None) + params: Final = json.loads(persisted[0]["litellm_params"]) + assert "litellm_credential_name" not in params + assert params["api_base"] == "https://api.openai.com/v1" + + @pytest.mark.asyncio + async def test_patch_model_rejects_non_admin_explicit_null_clear(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model(monkeypatch, db_model, self._team_admin_user(), None) + + assert exc_info.value.code == "403" + assert exc_info.value.param == "litellm_credential_name" + + @pytest.mark.asyncio + async def test_patch_model_clear_then_reattach_round_trip(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + cleared: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), None) + cleared_model: Final = Deployment.model_validate( + { + "model_name": db_model.model_name, + "litellm_params": json.loads(cleared[0]["litellm_params"]), + "model_info": json.loads(cleared[0]["model_info"]), + } + ) + reattached: Final = await self._patch_model( + monkeypatch, + cleared_model, + self._admin_user(), + "shared-credential", + ) + params: Final = json.loads(reattached[0]["litellm_params"]) + assert params["litellm_credential_name"] == "shared-credential" + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index d07e49e4712..b4fefe1d2c3 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -102,7 +102,7 @@ export interface ModelEditFormValues { vector_store_ids?: string[]; tags?: string[]; health_check_model?: string | null; - litellm_credential_name?: string; + litellm_credential_name?: string | null; litellm_extra_params?: string; model_info?: string; team_id?: string; @@ -139,7 +139,7 @@ const modelEditShape = { vector_store_ids: z.array(z.string()).optional(), tags: z.array(z.string()).optional(), health_check_model: z.string().nullish(), - litellm_credential_name: textish, + litellm_credential_name: z.string().nullish(), litellm_extra_params: textish, model_info: textish, team_id: textish, @@ -254,7 +254,7 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [], // antd never mounted this field for a non-wildcard model, so the key must be absent, not null. ...(isWildcardModel ? { health_check_model: localModelData.model_info?.health_check_model } : {}), - litellm_credential_name: localModelData.litellm_params?.litellm_credential_name || "", + litellm_credential_name: localModelData.litellm_params?.litellm_credential_name ?? null, litellm_extra_params: JSON.stringify( Object.fromEntries( Object.entries(localModelData.litellm_params || {}).filter( @@ -635,8 +635,8 @@ const ModelInfoEditForm: React.FC = ({ {isEditing ? ( {({ id, value, onChange, onBlur }) => { - const items = [ - { value: "", label: "None" }, + const items: { value: string | null; label: string }[] = [ + { value: null, label: "None" }, ...credentialsList.map((credential) => ({ value: credential.credential_name, label: credential.credential_name, @@ -645,15 +645,15 @@ const ModelInfoEditForm: React.FC = ({ return (