diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index cff1bebceb..289963e917 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -304,7 +304,7 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools + messages=cached_messages, tools=tools, model=model ) google_cache_name = self.check_cache( cache_key=generated_cache_key, @@ -433,7 +433,7 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools + messages=cached_messages, tools=tools, model=model ) google_cache_name = await self.async_check_cache( cache_key=generated_cache_key, diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 88d1b59c5b..e9d14d4e18 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -187,9 +187,9 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache == "existing_cache_name" - # Verify cache key was generated with tools + # Verify cache key was generated with tools and model mock_cache_obj.get_cache_key.assert_called_once_with( - messages=cached_messages, tools=self.sample_tools + messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro" ) @pytest.mark.parametrize( @@ -460,9 +460,9 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache == "existing_cache_name" - # Verify cache key was generated with tools + # Verify cache key was generated with tools and model mock_cache_obj.get_cache_key.assert_called_once_with( - messages=cached_messages, tools=self.sample_tools + messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro" ) @pytest.mark.asyncio