From f70f01b8fa671a1a5b5f40f540dae7bf1ff4247c Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:28:54 +0000 Subject: [PATCH 1/2] fix(vertex_ai): cache the whole conversation prefix up to the cache_control block so Gemini sees turns in order Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../context_caching/transformation.py | 9 +++-- .../test_vertex_ai_context_caching.py | 40 +++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index a4601af944b..6d0b6d45436 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -110,16 +110,17 @@ def separate_cached_messages( messages: list[AllMessageValues], ) -> tuple[list[AllMessageValues], list[AllMessageValues]]: """ - Returns the first continuous block of cached messages, plus every system message once - any message is cached, since generateContent rejects a system_instruction beside a - cachedContent, and the remaining messages. + Returns the conversation prefix ending at the first continuous block of cached messages, + plus every system message once any message is cached, since generateContent rejects a + system_instruction beside a cachedContent, and the remaining messages. Gemini reads + cachedContent as a prefix of contents, so the cache must hold everything before the block. """ filtered_messages: Final = [(idx, message) for idx, message in enumerate(messages) if is_cached_message(message)] if not filtered_messages: return [], messages last_continuous_block_idx: Final = get_first_continuous_block_idx(filtered_messages) - cached_block: Final = range(filtered_messages[0][0], filtered_messages[last_continuous_block_idx][0] + 1) + cached_block: Final = range(filtered_messages[last_continuous_block_idx][0] + 1) goes_in_cache: Final = tuple( idx in cached_block or message.get("role") == "system" for idx, message in enumerate(messages) ) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index bc8421c80f8..aaa22a06766 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1790,6 +1790,46 @@ class TestContextCachingEndpoints: self.mock_async_client.post.assert_not_called() +def _marked(role: Literal["system", "user"], text: str) -> dict: + return {"role": role, "content": [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]} + + +@pytest.mark.parametrize( + "messages", + [ + pytest.param( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "A1"}, + _marked("user", "Q2"), + ], + id="marked-latest-turn", + ), + pytest.param( + [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "A1"}, + _marked("user", "Q2"), + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "Q3"}, + ], + id="marked-middle-turn", + ), + pytest.param( + [_marked("system", "sys"), {"role": "user", "content": "Q1"}], + id="marked-system", + ), + ], +) +def test_separate_cached_messages_keeps_conversation_order(messages): + from litellm.llms.vertex_ai.context_caching.transformation import separate_cached_messages + + cached_messages, non_cached_messages = separate_cached_messages(messages) + + assert cached_messages + non_cached_messages == messages + + def test_cached_messages_end_on_supported_turn(): from litellm.llms.vertex_ai.context_caching.transformation import ( cached_messages_end_on_supported_turn, From e4cff7345b8935c87121e62c5137e09a1a14a9a2 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:36:20 +0000 Subject: [PATCH 2/2] test(vertex_ai): assert the exact cached prefix and remaining turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_context_caching.py | 61 ++++++++++++------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index aaa22a06766..f1dd2798f6f 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -6,7 +6,6 @@ import httpx import pytest import litellm - from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import VertexAIError @@ -14,6 +13,12 @@ from litellm.llms.vertex_ai.context_caching.known_cached_contents import KnownCa from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( ContextCachingEndpoints, ) +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAssistantMessage, + ChatCompletionSystemMessage, + ChatCompletionUserMessage, +) @pytest.fixture @@ -1790,44 +1795,56 @@ class TestContextCachingEndpoints: self.mock_async_client.post.assert_not_called() -def _marked(role: Literal["system", "user"], text: str) -> dict: - return {"role": role, "content": [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]} +def _marked_user(text: str) -> ChatCompletionUserMessage: + return {"role": "user", "content": [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]} + + +_SYSTEM: Final[ChatCompletionSystemMessage] = {"role": "system", "content": "sys"} +_Q1: Final[ChatCompletionUserMessage] = {"role": "user", "content": "Q1"} +_A1: Final[ChatCompletionAssistantMessage] = {"role": "assistant", "content": "A1"} +_A2: Final[ChatCompletionAssistantMessage] = {"role": "assistant", "content": "A2"} +_Q3: Final[ChatCompletionUserMessage] = {"role": "user", "content": "Q3"} +_MARKED_Q2: Final = _marked_user("Q2") +_MARKED_SYSTEM: Final[ChatCompletionSystemMessage] = { + "role": "system", + "content": [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}], +} @pytest.mark.parametrize( - "messages", + ("messages", "expected_cached", "expected_non_cached"), [ pytest.param( - [ - {"role": "system", "content": "sys"}, - {"role": "user", "content": "Q1"}, - {"role": "assistant", "content": "A1"}, - _marked("user", "Q2"), - ], + [_SYSTEM, _Q1, _A1, _MARKED_Q2], + [_SYSTEM, _Q1, _A1, _MARKED_Q2], + [], id="marked-latest-turn", ), pytest.param( - [ - {"role": "user", "content": "Q1"}, - {"role": "assistant", "content": "A1"}, - _marked("user", "Q2"), - {"role": "assistant", "content": "A2"}, - {"role": "user", "content": "Q3"}, - ], + [_Q1, _A1, _MARKED_Q2, _A2, _Q3], + [_Q1, _A1, _MARKED_Q2], + [_A2, _Q3], id="marked-middle-turn", ), pytest.param( - [_marked("system", "sys"), {"role": "user", "content": "Q1"}], + [_MARKED_SYSTEM, _Q1], + [_MARKED_SYSTEM], + [_Q1], id="marked-system", ), ], ) -def test_separate_cached_messages_keeps_conversation_order(messages): +def test_separate_cached_messages_caches_prefix_through_marked_block( + messages: list[AllMessageValues], + expected_cached: list[AllMessageValues], + expected_non_cached: list[AllMessageValues], +) -> None: from litellm.llms.vertex_ai.context_caching.transformation import separate_cached_messages cached_messages, non_cached_messages = separate_cached_messages(messages) - assert cached_messages + non_cached_messages == messages + assert cached_messages == expected_cached + assert non_cached_messages == expected_non_cached def test_cached_messages_end_on_supported_turn(): @@ -2585,7 +2602,9 @@ def test_a_plain_system_prompt_goes_into_the_gemini_cache_instead_of_beside_it(l ) created: Final = next( - json.loads(call.content) for call in calls if call.method == "POST" and call.url.path.endswith("/cachedContents") + json.loads(call.content) + for call in calls + if call.method == "POST" and call.url.path.endswith("/cachedContents") ) generated: Final = next(json.loads(call.content) for call in calls if call.url.path.endswith(":generateContent")) assert created["system_instruction"] == {"parts": [{"text": "You are terse."}]}