mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(responses): propagate cache_control through responses->chat completion bridge
User/assistant messages, tool call outputs, and function calls built from Responses API input items were dropping cache_control, breaking prompt caching for Bedrock/Anthropic when routed through the litellm_completion transform. Also propagate cache_control onto the last content element of user/assistant messages, since Anthropic/Bedrock read it from the content block rather than (or in addition to) the message level. Adds the missing cache_control field to ChatCompletionToolMessage and ChatCompletionResponseMessage TypedDicts to keep the assignments type-safe.
This commit is contained in:
parent
56a63b4b29
commit
2fd4bd2e09
3 changed files with 154 additions and 8 deletions
|
|
@ -1421,14 +1421,21 @@ class LiteLLMCompletionResponsesConfig:
|
|||
# Since guardrails skip None content anyway, we return empty list to exclude it from structured messages
|
||||
if content is None:
|
||||
return []
|
||||
return [
|
||||
GenericChatCompletionMessage(
|
||||
role=_input_item_role(input_item),
|
||||
content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(
|
||||
content
|
||||
),
|
||||
)
|
||||
]
|
||||
msg = GenericChatCompletionMessage(
|
||||
role=input_item.get("role") or "user",
|
||||
content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(
|
||||
content
|
||||
),
|
||||
)
|
||||
if input_item.get("cache_control"):
|
||||
msg["cache_control"] = input_item.get("cache_control")
|
||||
# Also propagate to the last content element for providers
|
||||
# that read cache_control from content blocks (Anthropic, Bedrock)
|
||||
if isinstance(msg["content"], list) and len(msg["content"]) > 0:
|
||||
last_element = msg["content"][-1]
|
||||
if isinstance(last_element, dict):
|
||||
last_element["cache_control"] = input_item["cache_control"]
|
||||
return [msg]
|
||||
|
||||
@staticmethod
|
||||
def _reasoning_text_from_content(input_item: Mapping[str, object]) -> str | None:
|
||||
|
|
@ -1646,6 +1653,9 @@ class LiteLLMCompletionResponsesConfig:
|
|||
content=_normalize_function_call_output_to_tool_content(tool_call_output.get("output")),
|
||||
tool_call_id=str(call_id),
|
||||
)
|
||||
output_cache_control = tool_call_output.get("cache_control")
|
||||
if output_cache_control:
|
||||
tool_output_message["cache_control"] = output_cache_control
|
||||
|
||||
_tool_use_definition: Final = TOOL_CALLS_CACHE.get_cache(
|
||||
key=tool_call_output.get("call_id") or "",
|
||||
|
|
@ -1742,6 +1752,9 @@ class LiteLLMCompletionResponsesConfig:
|
|||
role="assistant",
|
||||
content=None, # Function calls don't have content
|
||||
)
|
||||
function_call_cache_control = function_call.get("cache_control")
|
||||
if function_call_cache_control:
|
||||
chat_completion_response_message["cache_control"] = function_call_cache_control
|
||||
|
||||
return [chat_completion_response_message]
|
||||
|
||||
|
|
|
|||
|
|
@ -878,6 +878,7 @@ class ChatCompletionToolMessage(TypedDict):
|
|||
role: Literal["tool"]
|
||||
content: str | Iterable[ToolMessageContentPart] # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
tool_call_id: str
|
||||
cache_control: NotRequired[ChatCompletionCachedContent]
|
||||
|
||||
|
||||
class ChatCompletionFunctionMessage(TypedDict):
|
||||
|
|
@ -910,6 +911,7 @@ class ChatCompletionDeveloperMessage(OpenAIChatCompletionDeveloperMessage, total
|
|||
class GenericChatCompletionMessage(TypedDict, total=False):
|
||||
role: Required[str]
|
||||
content: Required[str | list]
|
||||
cache_control: ChatCompletionCachedContent
|
||||
|
||||
|
||||
ValidUserMessageContentTypes = [
|
||||
|
|
@ -1100,6 +1102,7 @@ class ChatCompletionResponseMessage(TypedDict, total=False):
|
|||
provider_specific_fields: dict | None
|
||||
reasoning_content: str | None
|
||||
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None
|
||||
cache_control: ChatCompletionCachedContent
|
||||
|
||||
|
||||
class ChatCompletionUsageBlock(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -5172,3 +5172,133 @@ async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool
|
|||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "tool_choice={'type': 'file_search'}" in str(exc_info.value)
|
||||
class TestCacheControlPreservationOnMessages:
|
||||
"""Tests that cache_control is preserved when transforming Responses API input items
|
||||
to Chat Completion messages. This is critical for providers like Bedrock/Anthropic
|
||||
that use cache_control for prompt caching."""
|
||||
|
||||
def test_cache_control_on_user_message(self):
|
||||
"""cache_control on a user message should be passed through to GenericChatCompletionMessage"""
|
||||
input_item = {
|
||||
"role": "user",
|
||||
"content": "Hello, world!",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
|
||||
result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
|
||||
input_item=input_item
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
msg = result[0]
|
||||
assert msg["role"] == "user"
|
||||
assert msg.get("cache_control") == {"type": "ephemeral"}
|
||||
|
||||
def test_cache_control_on_assistant_message(self):
|
||||
"""cache_control on an assistant message should be passed through"""
|
||||
input_item = {
|
||||
"role": "assistant",
|
||||
"content": "I can help with that.",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
|
||||
result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
|
||||
input_item=input_item
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
msg = result[0]
|
||||
assert msg["role"] == "assistant"
|
||||
assert msg.get("cache_control") == {"type": "ephemeral"}
|
||||
|
||||
def test_no_cache_control_when_absent(self):
|
||||
"""Messages without cache_control should not have the field set"""
|
||||
input_item = {
|
||||
"role": "user",
|
||||
"content": "Hello",
|
||||
}
|
||||
|
||||
result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
|
||||
input_item=input_item
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
msg = result[0]
|
||||
assert "cache_control" not in msg
|
||||
|
||||
def test_cache_control_on_function_call_output(self):
|
||||
"""cache_control on a function_call_output should be passed through to tool message"""
|
||||
TOOL_CALLS_CACHE.set_cache(
|
||||
key="call_123",
|
||||
value={
|
||||
"id": "call_123",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "NYC"}'},
|
||||
},
|
||||
)
|
||||
|
||||
input_item = {
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_123",
|
||||
"output": '{"temp": 72}',
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
|
||||
result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
|
||||
input_item=input_item
|
||||
)
|
||||
|
||||
# Should have assistant message + tool message
|
||||
assert len(result) == 2
|
||||
tool_msg = result[1]
|
||||
assert tool_msg["role"] == "tool"
|
||||
assert tool_msg.get("cache_control") == {"type": "ephemeral"}
|
||||
|
||||
def test_cache_control_on_function_call(self):
|
||||
"""cache_control on a function_call input should be passed through to assistant message"""
|
||||
input_item = {
|
||||
"type": "function_call",
|
||||
"call_id": "call_456",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "NYC"}',
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
|
||||
result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
|
||||
input_item=input_item
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
msg = result[0]
|
||||
assert msg["role"] == "assistant"
|
||||
assert msg.get("cache_control") == {"type": "ephemeral"}
|
||||
|
||||
def test_cache_control_preserved_in_full_input_transform(self):
|
||||
"""cache_control should survive the full transform_responses_api_input_to_messages path"""
|
||||
input_items = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the weather?",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Let me check.",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Thanks!",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
]
|
||||
|
||||
result = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=input_items,
|
||||
responses_api_request={},
|
||||
)
|
||||
|
||||
# 3 messages, no system
|
||||
assert len(result) == 3
|
||||
assert "cache_control" not in result[0]
|
||||
assert result[1].get("cache_control") == {"type": "ephemeral"}
|
||||
assert result[2].get("cache_control") == {"type": "ephemeral"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue