From 97031dc8eeac09b3576c6a40dffdce5d59084de5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 8 Oct 2025 17:41:04 -0700 Subject: [PATCH] Fix - (openrouter): move cache_control to content blocks for claude/gemini (#15345) * test_openrouter_transform_request_with_cache_control * fix CacheControlSupportedModels * test_openrouter_transform_request_with_cache_control_gemini --- .../llms/openrouter/chat/transformation.py | 66 +++++- .../test_openrouter_chat_transformation.py | 203 ++++++++++++++++++ 2 files changed, 268 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index bf57218c91..a47689a30b 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -6,6 +6,7 @@ Calls done in OpenAI/openai.py as OpenRouter is openai-compatible. Docs: https://openrouter.ai/docs/parameters """ +from enum import Enum from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union import httpx @@ -20,6 +21,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig from ..common_utils import OpenRouterException +class CacheControlSupportedModels(str, Enum): + """Models that support cache_control in content blocks.""" + CLAUDE = "claude" + GEMINI = "gemini" + + class OpenrouterConfig(OpenAIGPTConfig): def map_openai_params( self, @@ -48,19 +55,73 @@ class OpenrouterConfig(OpenAIGPTConfig): ) return mapped_openai_params + def _supports_cache_control_in_content(self, model: str) -> bool: + """ + Check if the model supports cache_control in content blocks. + + Returns: + bool: True if model supports cache_control (Claude or Gemini models) + """ + model_lower = model.lower() + return any( + supported_model.value in model_lower + for supported_model in CacheControlSupportedModels + ) + def remove_cache_control_flag_from_messages_and_tools( self, model: str, messages: List[AllMessageValues], tools: Optional[List["ChatCompletionToolParam"]] = None, ) -> Tuple[List[AllMessageValues], Optional[List["ChatCompletionToolParam"]]]: - if "claude" in model.lower(): # don't remove 'cache_control' flag + if self._supports_cache_control_in_content(model): return messages, tools else: return super().remove_cache_control_flag_from_messages_and_tools( model, messages, tools ) + def _move_cache_control_to_content( + self, messages: List[AllMessageValues] + ) -> List[AllMessageValues]: + """ + Move cache_control from message level to content blocks. + OpenRouter requires cache_control to be inside content blocks, not at message level. + + When cache_control is at message level, it's added to ALL content blocks + to cache the entire message content. + """ + transformed_messages = [] + for message in messages: + message_copy = dict(message) + cache_control = message_copy.pop("cache_control", None) + + if cache_control is not None: + content = message_copy.get("content") + + if isinstance(content, list): + # Content is already a list, add cache_control to all blocks + if len(content) > 0: + content_copy = [] + for block in content: + block_copy = dict(block) + block_copy["cache_control"] = cache_control + content_copy.append(block_copy) + message_copy["content"] = content_copy + else: + # Content is a string, convert to structured format + message_copy["content"] = [ + { + "type": "text", + "text": content, + "cache_control": cache_control, + } + ] + + transformed_messages.append(message_copy) + + return transformed_messages + def transform_request( self, model: str, @@ -75,6 +136,9 @@ class OpenrouterConfig(OpenAIGPTConfig): Returns: dict: The transformed request. Sent as the body of the API call. """ + if self._supports_cache_control_in_content(model): + messages = self._move_cache_control_to_content(messages) + extra_body = optional_params.pop("extra_body", {}) response = super().transform_request( model, messages, optional_params, litellm_params, headers diff --git a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py index 02eb872aba..ff81ba84c0 100644 --- a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -114,3 +114,206 @@ def test_openrouter_cache_control_flag_removal(): headers={}, ) assert transformed_request["messages"][0].get("cache_control") is None + + + +def test_openrouter_transform_request_with_cache_control(): + """ + Test transform_request moves cache_control from message level to content blocks (string content). + + Input: + { + "role": "user", + "content": "what are the key terms...", + "cache_control": {"type": "ephemeral"} + } + + Expected Output: + { + "role": "user", + "content": [ + { + "type": "text", + "text": "what are the key terms...", + "cache_control": {"type": "ephemeral"} + } + ] + } + """ + import json + config = OpenrouterConfig() + + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are an AI assistant tasked with analyzing legal documents." + }, + { + "type": "text", + "text": "Here is the full text of a complex legal agreement" + } + ] + }, + { + "role": "user", + "content": "what are the key terms and conditions in this agreement?", + "cache_control": {"type": "ephemeral"} + } + ] + + transformed_request = config.transform_request( + model="openrouter/anthropic/claude-3-5-sonnet-20240620", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + print("\n=== Transformed Request ===") + print(json.dumps(transformed_request, indent=4, default=str)) + + assert "messages" in transformed_request + assert len(transformed_request["messages"]) == 2 + + user_message = transformed_request["messages"][1] + assert user_message["role"] == "user" + assert isinstance(user_message["content"], list) + assert user_message["content"][0]["type"] == "text" + assert user_message["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_openrouter_transform_request_with_cache_control_list_content(): + """ + Test transform_request moves cache_control to all content blocks when content is already a list. + + Input: + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a historian..."}, + {"type": "text", "text": "HUGE TEXT BODY"} + ], + "cache_control": {"type": "ephemeral"} + } + + Expected Output: + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a historian...", + "cache_control": {"type": "ephemeral"} + }, + { + "type": "text", + "text": "HUGE TEXT BODY", + "cache_control": {"type": "ephemeral"} + } + ] + } + """ + import json + config = OpenrouterConfig() + + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a historian studying the fall of the Roman Empire." + }, + { + "type": "text", + "text": "HUGE TEXT BODY" + } + ], + "cache_control": {"type": "ephemeral"} + }, + { + "role": "user", + "content": "What triggered the collapse?" + } + ] + + transformed_request = config.transform_request( + model="openrouter/anthropic/claude-3-5-sonnet-20240620", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + print("\n=== Transformed Request (List Content) ===") + print(json.dumps(transformed_request, indent=4, default=str)) + + assert "messages" in transformed_request + assert len(transformed_request["messages"]) == 2 + + system_message = transformed_request["messages"][0] + assert system_message["role"] == "system" + assert isinstance(system_message["content"], list) + assert len(system_message["content"]) == 2 + assert system_message["content"][0]["cache_control"] == {"type": "ephemeral"} + assert system_message["content"][1]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in system_message + + +def test_openrouter_transform_request_with_cache_control_gemini(): + """ + Test transform_request moves cache_control to content blocks for Gemini models. + + Input: + { + "role": "user", + "content": "Analyze this data", + "cache_control": {"type": "ephemeral"} + } + + Expected Output: + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this data", + "cache_control": {"type": "ephemeral"} + } + ] + } + """ + import json + config = OpenrouterConfig() + + messages = [ + { + "role": "user", + "content": "Analyze this data", + "cache_control": {"type": "ephemeral"} + } + ] + + transformed_request = config.transform_request( + model="openrouter/google/gemini-2.0-flash-exp:free", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + print("\n=== Transformed Request (Gemini) ===") + print(json.dumps(transformed_request, indent=4, default=str)) + + assert "messages" in transformed_request + assert len(transformed_request["messages"]) == 1 + + user_message = transformed_request["messages"][0] + assert user_message["role"] == "user" + assert isinstance(user_message["content"], list) + assert user_message["content"][0]["type"] == "text" + assert user_message["content"][0]["cache_control"] == {"type": "ephemeral"} + \ No newline at end of file