Merge da7f912d04 into e16d79d635

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
 

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)