mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 08:18:15 +00:00
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
This commit is contained in:
Ishaan Jaff
@@ -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
|
||||
|
||||
@@ -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"}
|
||||
|
||||
Reference in New Issue
Block a user