diff --git a/litellm/litellm_core_utils/prompt_templates/server_tools.py b/litellm/litellm_core_utils/prompt_templates/server_tools.py index ce8e8114eb2..392b036cf87 100644 --- a/litellm/litellm_core_utils/prompt_templates/server_tools.py +++ b/litellm/litellm_core_utils/prompt_templates/server_tools.py @@ -151,6 +151,22 @@ def trailing_system_messages(data: Mapping[str, object], route: ServerToolRoute) ) +def uncached_system_directive(message: Mapping[str, object]) -> Mapping[str, object]: + return { # mutable-ok: Provider wire format requires native JSON containers. + **{key: value for key, value in message.items() if key != "cache_control"}, + **( + { # mutable-ok: Provider wire format requires native JSON containers. + "content": [ # mutable-ok: Provider wire format requires native JSON containers. + {key: value for key, value in _OBJECT.validate_python(block).items() if key != "cache_control"} + for block in _items(message.get("content")) + ], + } + if isinstance(message.get("content"), (list, tuple)) + else {} # mutable-ok: Provider wire format requires native JSON containers. + ), + } + + def append_server_reference(data: Mapping[str, object], route: ServerToolRoute, reference: str) -> Mapping[str, object]: field: Final = "input" if route == "aresponses" else "messages" messages: Final = _items(data.get(field)) diff --git a/litellm/proxy/memory/continuation.py b/litellm/proxy/memory/continuation.py index 5d7ba0fdf1d..c916b76c532 100644 --- a/litellm/proxy/memory/continuation.py +++ b/litellm/proxy/memory/continuation.py @@ -177,7 +177,20 @@ class MemoryContinuations: if patch.replaces > index + 1 or patch.replaces < 1: raise HTTPException(status_code=409, detail="Invalid memory continuation") prefix: Final = result[: -(patch.replaces - 1)] if patch.replaces > 1 else result - return _append_items(prefix, patch.replacement, self.route) + # Clients move cache breakpoints between turns. Reuse their current + # directives rather than restoring an obsolete cached copy. + directives: Final = MappingProxyType( + { + prefix_hashes((item,), self.route)[0]: item + for item in items[index + 1 - patch.replaces : index + 1] + if item.get("role") == "system" + } + ) + replacement: Final = tuple( + directives.get(prefix_hashes((item,), self.route)[0], item) if item.get("role") == "system" else item + for item in patch.replacement + ) + return _append_items(prefix, replacement, self.route) return reduce(apply, range(len(items)), ()) diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index 2fc7433e296..29ae52d0118 100644 --- a/litellm/proxy/memory/gateway.py +++ b/litellm/proxy/memory/gateway.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import ( continue_server_tools, inject_server_tools, trailing_system_messages, + uncached_system_directive, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.sse_keepalive import wrap_passthrough_sse_bytes_with_keepalive_pings @@ -147,7 +148,10 @@ class GatewayMemoryLoop: **self.data, **( { # mutable-ok: Native provider JSON containers. - "messages": [*messages, *directives], # mutable-ok: Provider request JSON. + "messages": [ # mutable-ok: Provider request JSON. + *messages, + *(uncached_system_directive(directive) for directive in directives), + ], } if directives and messages[-len(directives) :] != directives else {} # mutable-ok: Native provider JSON containers. diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index 54d4dfbffcf..650c71939e5 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -524,6 +524,7 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_ body = await incoming.json() observed.append(body) messages = body["messages"] + assert json.dumps(body).count('"cache_control"') == 4 assert all( message["role"] != "system" or messages[index + 1]["role"] == "assistant" for index, message in enumerate(messages[:-1]) @@ -540,8 +541,16 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_ "role": "user", "content": [{"type": "text", "text": "Read README.md", "cache_control": {"type": "ephemeral"}}], } - directive = {"role": "system", "content": "Use concise answers"} + directive = { + "role": "system", + "content": [{"type": "text", "text": "Use concise answers", "cache_control": {"type": "ephemeral"}}], + } + earlier_directive = {"role": "system", "content": [{"type": "text", "text": "Use concise answers"}]} original = { + "system": [ + {"type": "text", "text": "Cached prefix " + str(i), "cache_control": {"type": "ephemeral"}} + for i in range(2) + ], "messages": [prefix, directive], "tools": [{"name": "Read", "input_schema": {"type": "object"}}], "tool_choice": {"type": "tool", "name": "Read"}, @@ -558,7 +567,7 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_ "tool_choice": {"type": "none"}, "messages": [ prefix, - directive, + earlier_directive, {"role": "assistant", "content": [client_call]}, { "role": "user", @@ -572,7 +581,8 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_ pass assert len(observed) == 2 assert observed[0]["messages"][0] == observed[1]["messages"][0] == prefix - assert observed[1]["messages"].count(directive) == 2 + assert observed[1]["messages"].count(directive) == 1 + assert observed[1]["messages"].count(earlier_directive) == 1 assert observed[1]["messages"].count({"role": "assistant", "content": [client_call]}) == 1 assert observed[1]["messages"][-1] == directive assert observed[1]["messages"][-3]["content"][0]["tool_use_id"] == "client_read" @@ -580,16 +590,32 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_ @pytest.mark.asyncio -async def test_claude_output_directives_reach_search_answer_and_reflection_rounds(prisma_edge: MagicMock) -> None: +@pytest.mark.parametrize("cached_directive", [False, True]) +async def test_claude_output_directives_reach_search_answer_and_reflection_rounds( + prisma_edge: MagicMock, cached_directive: bool +) -> None: provider = FastAPI() observed = [] - directive = {"role": "system", "content": [], "output_config": {"effort": "low"}} + directive = { + "role": "system", + "content": [{"type": "text", "text": "Reply briefly", "cache_control": {"type": "ephemeral"}}] + if cached_directive + else [], + "output_config": {"effort": "low"}, + } @provider.post("/v1/messages") async def model(incoming: Request): body = await incoming.json() observed.append(body) - assert body["messages"][-1] == directive + assert json.dumps(body).count('"cache_control"') == 3 + int(cached_directive) + last = body["messages"][-1] + assert last["role"] == "system" and last["output_config"] == {"effort": "low"} + assert last["content"] == ( + [{"type": "text", "text": "Reply briefly"}] + if cached_directive and len(observed) > 1 + else directive["content"] + ) if len(observed) == 1: content = [ {"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {"query": "demo"}} @@ -606,7 +632,19 @@ async def test_claude_output_directives_reach_search_answer_and_reflection_round "content": content, } - original = {"messages": [{"role": "user", "content": "My demo port?"}, directive]} + original = { + "system": [ + {"type": "text", "text": "Cached prefix " + str(i), "cache_control": {"type": "ephemeral"}} + for i in range(2) + ], + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "My demo port?", "cache_control": {"type": "ephemeral"}}], + }, + directive, + ], + } loop = GatewayMemoryLoop(provider, request(), original, "anthropic_messages", store(prisma_edge)) async for _ in loop.run(): pass