mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(sap): preserve cache_control on messages and tools for prompt caching
This commit is contained in:
Devin AI
parent
daf22ec871
commit
53308754bb
3 changed files with 257 additions and 21 deletions
|
|
@ -1,32 +1,24 @@
|
|||
from typing import Union, Literal, Optional
|
||||
from collections.abc import Sequence
|
||||
from enum import Enum
|
||||
import warnings
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
|
||||
def validate_different_content(v: Union[str, dict, list]) -> str:
|
||||
if v in ((), {}, []):
|
||||
return ""
|
||||
elif isinstance(v, dict) and "text" in v:
|
||||
return v["text"]
|
||||
elif isinstance(v, list):
|
||||
new_v = []
|
||||
for item in v:
|
||||
if isinstance(item, dict) and "text" in item:
|
||||
if item["text"]:
|
||||
new_v.append(item["text"])
|
||||
elif isinstance(item, str):
|
||||
new_v.append(item)
|
||||
return "\n".join(new_v)
|
||||
elif isinstance(v, str):
|
||||
return v
|
||||
raise ValueError("Content must be a string")
|
||||
class CacheControl(BaseModel):
|
||||
"""
|
||||
Prompt caching breakpoint, as accepted by SAP orchestration for Anthropic Claude and Amazon Nova models.
|
||||
"""
|
||||
|
||||
type_: Literal["ephemeral"] = Field(default="ephemeral", alias="type")
|
||||
ttl: Literal["5m", "1h"] | None = None
|
||||
|
||||
|
||||
class TextContent(BaseModel):
|
||||
type_: Literal["text"] = Field(default="text", alias="type")
|
||||
text: str
|
||||
cache_control: CacheControl | None = None
|
||||
|
||||
|
||||
class ImageURLContent(BaseModel):
|
||||
|
|
@ -37,6 +29,45 @@ class ImageURLContent(BaseModel):
|
|||
class ImageContent(BaseModel):
|
||||
type_: Literal["image_url"] = Field(default="image_url", alias="type")
|
||||
image_url: ImageURLContent
|
||||
cache_control: CacheControl | None = None
|
||||
|
||||
|
||||
def _to_text_content(item: object) -> TextContent | None:
|
||||
if isinstance(item, str):
|
||||
return TextContent(type="text", text=item)
|
||||
if isinstance(item, dict) and item.get("text"):
|
||||
return TextContent.model_validate({**item, "type": "text"})
|
||||
return None
|
||||
|
||||
|
||||
def _to_content(items: tuple[object, ...]) -> str | Sequence[TextContent]:
|
||||
blocks = tuple(block for block in map(_to_text_content, items) if block is not None)
|
||||
if any(block.cache_control is not None for block in blocks):
|
||||
return blocks
|
||||
return "\n".join(block.text for block in blocks)
|
||||
|
||||
|
||||
def validate_different_content(v: object) -> str | Sequence[TextContent]:
|
||||
"""
|
||||
Flatten content blocks into a plain string, unless a block carries a prompt caching breakpoint,
|
||||
in which case the blocks are kept so that `cache_control` reaches the model.
|
||||
"""
|
||||
if v in ((), {}, []):
|
||||
return ""
|
||||
if isinstance(v, str):
|
||||
return v
|
||||
if isinstance(v, dict) and "text" in v:
|
||||
return _to_content((v,))
|
||||
if isinstance(v, (list, tuple)):
|
||||
return _to_content(tuple(v))
|
||||
raise ValueError("Content must be a string")
|
||||
|
||||
|
||||
def flatten_content_to_text(v: object) -> str:
|
||||
content = validate_different_content(v)
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
return "\n".join(block.text for block in content)
|
||||
|
||||
|
||||
class FunctionObj(BaseModel):
|
||||
|
|
@ -70,10 +101,14 @@ class FunctionTool(BaseModel):
|
|||
class ChatCompletionTool(BaseModel):
|
||||
type_: Literal["function"] = Field(default="function", alias="type")
|
||||
function: FunctionTool
|
||||
cache_control: CacheControl | None = None
|
||||
|
||||
def model_dump(self, **kwargs) -> dict:
|
||||
kwargs["exclude_unset"] = False
|
||||
return super().model_dump(**kwargs)
|
||||
dumped = super().model_dump(**kwargs)
|
||||
if self.cache_control is None:
|
||||
return {key: value for key, value in dumped.items() if key != "cache_control"}
|
||||
return {**dumped, "cache_control": self.cache_control.model_dump(by_alias=True, exclude_none=True)}
|
||||
|
||||
|
||||
class MessageToolCall(BaseModel):
|
||||
|
|
@ -88,7 +123,7 @@ class SAPMessage(BaseModel):
|
|||
"""
|
||||
|
||||
role: Literal["system", "developer"] = "system"
|
||||
content: str
|
||||
content: str | Sequence[TextContent]
|
||||
|
||||
_content_validator = field_validator("content", mode="before")(validate_different_content)
|
||||
|
||||
|
|
@ -104,13 +139,13 @@ class SAPAssistantMessage(BaseModel):
|
|||
refusal: str = ""
|
||||
tool_calls: list[MessageToolCall] = []
|
||||
|
||||
_content_validator = field_validator("content", mode="before")(validate_different_content)
|
||||
_content_validator = field_validator("content", mode="before")(flatten_content_to_text)
|
||||
|
||||
|
||||
class SAPToolChatMessage(BaseModel):
|
||||
role: Literal["tool"] = "tool"
|
||||
tool_call_id: str
|
||||
content: str
|
||||
content: str | Sequence[TextContent]
|
||||
|
||||
_content_validator = field_validator("content", mode="before")(validate_different_content)
|
||||
|
||||
|
|
|
|||
|
|
@ -204,3 +204,67 @@ async def test_sap_chat_required_headers(
|
|||
f"Header '{header_name}' has incorrect value. "
|
||||
f"Expected: '{expected_value}', Got: '{request.headers[header_name]}'"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sap_chat_forwards_cache_control(
|
||||
respx_mock,
|
||||
sap_api_response,
|
||||
fake_token_creator,
|
||||
fake_deployment_url,
|
||||
):
|
||||
"""`cache_control` set by the caller must reach SAP orchestration (issue #34797)."""
|
||||
import json
|
||||
|
||||
import litellm
|
||||
|
||||
litellm.disable_aiohttp_transport = True
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.sap.chat.transformation.GenAIHubOrchestrationConfig.deployment_url",
|
||||
new_callable=PropertyMock,
|
||||
return_value=fake_deployment_url,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.sap.chat.transformation.get_token_creator",
|
||||
return_value=fake_token_creator,
|
||||
),
|
||||
):
|
||||
route = respx_mock.post(f"{fake_deployment_url}/v2/completion")
|
||||
route.respond(json=sap_api_response)
|
||||
|
||||
await litellm.acompletion(
|
||||
model="sap/anthropic--claude-4.5-sonnet",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "large static prompt",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Hello"},
|
||||
],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "cached_tool", "parameters": {}},
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert route.called
|
||||
body = json.loads(route.calls[0].request.content)
|
||||
prompt = body["config"]["modules"]["prompt_templating"]["prompt"]
|
||||
assert prompt["template"][0]["content"] == [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "large static prompt",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
assert prompt["tools"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
|
|
|
|||
|
|
@ -639,3 +639,140 @@ class TestSAPTransformationIntegration:
|
|||
config["config"]["modules"][1]["translation"]["input"]["type"]
|
||||
== "sap_document_translation"
|
||||
)
|
||||
|
||||
|
||||
class TestSAPPromptCaching:
|
||||
"""`cache_control` breakpoints must survive transformation (issue #34797)."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_config(self):
|
||||
from litellm.llms.sap.chat.transformation import GenAIHubOrchestrationConfig
|
||||
|
||||
config = GenAIHubOrchestrationConfig()
|
||||
config.token_creator = lambda: "Bearer TEST_TOKEN"
|
||||
config._base_url = "https://api.test-sap.com"
|
||||
config._resource_group = "test-group"
|
||||
|
||||
return config
|
||||
|
||||
def test_cache_control_preserved_on_system_user_and_tool_messages(self, mock_config):
|
||||
config = mock_config.transform_request(
|
||||
model="anthropic--claude-4.5-sonnet",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "large static prompt",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "question",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "tool result",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
litellm_params={},
|
||||
optional_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
template = config["config"]["modules"]["prompt_templating"]["prompt"]["template"]
|
||||
assert list(template[0]["content"]) == [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "large static prompt",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
assert list(template[1]["content"]) == [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "question",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
]
|
||||
assert list(template[2]["content"]) == [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "tool result",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
|
||||
def test_cache_control_preserved_on_tool_definitions(self, mock_config):
|
||||
config = mock_config.transform_request(
|
||||
model="anthropic--claude-4.5-sonnet",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
optional_params={
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "cached_tool", "parameters": {}},
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "uncached_tool", "parameters": {}},
|
||||
},
|
||||
]
|
||||
},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
tools = config["config"]["modules"]["prompt_templating"]["prompt"]["tools"]
|
||||
assert tools[0]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in tools[1]
|
||||
|
||||
def test_content_blocks_without_cache_control_stay_flattened(self, mock_config):
|
||||
config = mock_config.transform_request(
|
||||
model="anthropic--claude-4.5-sonnet",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "part one"},
|
||||
{"type": "text", "text": "part two"},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "answer",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "tool result"},
|
||||
],
|
||||
litellm_params={},
|
||||
optional_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
template = config["config"]["modules"]["prompt_templating"]["prompt"]["template"]
|
||||
assert template[0]["content"] == "part one\npart two"
|
||||
assert template[1]["content"] == "answer"
|
||||
assert template[2]["content"] == "tool result"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue