diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index f683ad0e377..ec3a0a1ac24 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -10,6 +10,7 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import Final, Literal +from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import CachedContentRequestBody from litellm.utils import is_cached_message @@ -112,6 +113,14 @@ def separate_cached_messages( """ Returns separated cached and non-cached messages. + Vertex reads cachedContent as a prefix of contents, so context caching is only + used when the cached block is a true prefix of the conversation. Leading system + messages are folded into the cached block (they are preserved as the cache's + system_instruction downstream). If any other message sits before the cached + block, or a system message sits after it, no caching is applied: earlier + messages would be re-sent after the cached block, reordering the conversation, + and a system message outside the cache is dropped alongside cachedContent. + Args: messages: List of messages to be separated. @@ -136,8 +145,27 @@ def separate_cached_messages( first_cached_idx: Final = filtered_messages[0][0] last_cached_idx: Final = filtered_messages[last_continuous_block_idx][0] - cached_messages = messages[first_cached_idx : last_cached_idx + 1] - non_cached_messages = messages[:first_cached_idx] + messages[last_cached_idx + 1 :] + prefix_messages: Final = messages[:first_cached_idx] + suffix_messages: Final = messages[last_cached_idx + 1 :] + + # Only "system" counts as foldable: the Vertex/Gemini mapping does not + # fold "developer" messages into system_instruction, so a developer + # message around the block intentionally disables caching. + non_system_prefix: Final = any(m["role"] != "system" for m in prefix_messages) + trailing_system: Final = any(m["role"] == "system" for m in suffix_messages) + if non_system_prefix or trailing_system: + # Not a true prefix: send the conversation uncached so nothing is + # reordered and no system message is silently dropped. + verbose_logger.debug( + "Vertex context caching skipped: cached block is not a true prefix (%s)", + "non-system message before the cached block" + if non_system_prefix + else "system message after the cached block", + ) + return [], list(messages) + + cached_messages = list(prefix_messages) + messages[first_cached_idx : last_cached_idx + 1] + non_cached_messages = list(suffix_messages) else: non_cached_messages = messages diff --git a/tests/unit/llms/vertex_ai/context_caching/test_separate_cached_messages_prefix.py b/tests/unit/llms/vertex_ai/context_caching/test_separate_cached_messages_prefix.py new file mode 100644 index 00000000000..ef1e2438e47 --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_separate_cached_messages_prefix.py @@ -0,0 +1,124 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/47046 + +Vertex reads cachedContent as a prefix of contents. When the cache_control +block is not a prefix of the conversation, the conversation was reordered +(earlier messages re-sent after the cached block) and a system message +outside the cached block was silently dropped under modify_params. +""" + +from litellm.llms.vertex_ai.context_caching.transformation import ( + separate_cached_messages, + transform_openai_messages_to_gemini_context_caching, +) + +CACHED = {"cache_control": {"type": "ephemeral"}} + + +def _system(text="you are helpful"): + return {"role": "system", "content": text} + + +def _user(text, cached=False): + msg = {"role": "user", "content": text} + if cached: + msg["cache_control"] = CACHED["cache_control"] + return msg + + +def _assistant(text, cached=False): + msg = {"role": "assistant", "content": text} + if cached: + msg["cache_control"] = CACHED["cache_control"] + return msg + + +def test_non_prefix_block_sends_conversation_uncached(): + # Exact repro from the issue: cache_control on a message that is not the + # start of the conversation. + messages = [ + _system(), + _user("What's the weather in Tokyo today?"), + _assistant("Sunny."), + _user("And Nagoya tomorrow?", cached=True), + ] + cached, non_cached = separate_cached_messages(messages=messages) + assert cached == [] + assert non_cached == messages + + +def test_non_system_prefix_sends_conversation_uncached(): + messages = [ + _user("Q1"), + _assistant("A1"), + _user("Q2", cached=True), + ] + cached, non_cached = separate_cached_messages(messages=messages) + assert cached == [] + assert non_cached == messages + + +def test_leading_system_message_folded_into_cached_block(): + # A system message before the block can be preserved inside the cache's + # system_instruction, so caching still applies with nothing reordered. + messages = [ + _system(), + _user("Q1", cached=True), + _assistant("A1", cached=True), + _user("Q2"), + ] + cached, non_cached = separate_cached_messages(messages=messages) + assert cached == messages[:3] + assert non_cached == [messages[3]] + + +def test_true_prefix_block_still_caches(): + messages = [ + _user("Q1", cached=True), + _assistant("A1", cached=True), + _user("Q2"), + ] + cached, non_cached = separate_cached_messages(messages=messages) + assert cached == messages[:2] + assert non_cached == [messages[2]] + + +def test_trailing_system_message_sends_conversation_uncached(): + # A system message after the cached block would be extracted and dropped + # alongside cachedContent, so caching must be skipped. + messages = [ + _user("Q1", cached=True), + _system("late instruction"), + _user("Q2"), + ] + cached, non_cached = separate_cached_messages(messages=messages) + assert cached == [] + assert non_cached == messages + + +def test_no_cache_control_unchanged(): + messages = [_system(), _user("Q1"), _assistant("A1")] + cached, non_cached = separate_cached_messages(messages=messages) + assert cached == [] + assert non_cached == messages + + +def test_cache_body_preserves_folded_system_message(): + # The folded system message must land in the cache's system_instruction, + # not in contents and not dropped. + cached, _ = separate_cached_messages( + messages=[_system(), _user("Q1", cached=True)] + ) + body = transform_openai_messages_to_gemini_context_caching( + model="gemini-2.0-flash", + messages=cached, + custom_llm_provider="gemini", + cache_key="test-key", + vertex_project=None, + vertex_location=None, + ) + # Preserved as system_instruction when the model supports it, otherwise + # merged into contents - but never dropped. + import json as _json + + assert "you are helpful" in _json.dumps(body)