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..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,6 +1795,58 @@ class TestContextCachingEndpoints: self.mock_async_client.post.assert_not_called() +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", "expected_cached", "expected_non_cached"), + [ + pytest.param( + [_SYSTEM, _Q1, _A1, _MARKED_Q2], + [_SYSTEM, _Q1, _A1, _MARKED_Q2], + [], + id="marked-latest-turn", + ), + pytest.param( + [_Q1, _A1, _MARKED_Q2, _A2, _Q3], + [_Q1, _A1, _MARKED_Q2], + [_A2, _Q3], + id="marked-middle-turn", + ), + pytest.param( + [_MARKED_SYSTEM, _Q1], + [_MARKED_SYSTEM], + [_Q1], + id="marked-system", + ), + ], +) +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 == expected_cached + assert non_cached_messages == expected_non_cached + + def test_cached_messages_end_on_supported_turn(): from litellm.llms.vertex_ai.context_caching.transformation import ( cached_messages_end_on_supported_turn, @@ -2545,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."}]}