mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 2095cf41bb into 6085538b33
This commit is contained in:
commit
00a394dfbd
6 changed files with 272 additions and 9 deletions
|
|
@ -743,7 +743,22 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
|
||||
@classmethod
|
||||
def target_consumes_cache_control(cls, model: str) -> bool:
|
||||
return cls.is_anthropic_claude_model(model) or cls.is_bedrock_arn_model(model) or "gemini" in model.lower()
|
||||
if cls.is_anthropic_claude_model(model) or cls.is_bedrock_arn_model(model):
|
||||
return True
|
||||
if "gemini" in model.lower():
|
||||
from litellm.utils import is_explicitly_disabled_factory
|
||||
|
||||
if "/" not in model:
|
||||
if (
|
||||
is_explicitly_disabled_factory(f"gemini/{model}", None, "supports_prompt_caching")
|
||||
or is_explicitly_disabled_factory(f"vertex_ai/{model}", None, "supports_prompt_caching")
|
||||
):
|
||||
return False
|
||||
|
||||
if is_explicitly_disabled_factory(model, None, "supports_prompt_caching"):
|
||||
return False
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def translate_thinking_for_model(
|
||||
|
|
|
|||
|
|
@ -130,6 +130,14 @@ def _context_caching_skip_reason(
|
|||
cache_model: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
|
||||
) -> str | None:
|
||||
from litellm.utils import is_explicitly_disabled_factory
|
||||
|
||||
if is_explicitly_disabled_factory(deployment_model, custom_llm_provider, "supports_prompt_caching") or (
|
||||
cache_model != deployment_model
|
||||
and is_explicitly_disabled_factory(cache_model, custom_llm_provider, "supports_prompt_caching")
|
||||
):
|
||||
return f"Vertex AI context caching: model {deployment_model} does not support explicit context caching."
|
||||
|
||||
if not cached_messages_end_on_supported_turn(cached_messages):
|
||||
return (
|
||||
"Vertex AI context caching: cached message block ends on a model turn once "
|
||||
|
|
@ -499,9 +507,9 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except _CachedContentsNotServedError:
|
||||
verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url)
|
||||
return messages, optional_params, None
|
||||
for tool_param in _TOOL_PARAMS:
|
||||
optional_params.pop(tool_param, None)
|
||||
if google_cache_name:
|
||||
for tool_param in _TOOL_PARAMS:
|
||||
optional_params.pop(tool_param, None)
|
||||
return non_cached_messages, optional_params, google_cache_name
|
||||
|
||||
## TRANSFORM REQUEST
|
||||
|
|
@ -537,6 +545,17 @@ class ContextCachingEndpoints(VertexBase):
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
if "does not support explicit context caching" in err.response.text or (
|
||||
err.response.status_code == 400
|
||||
and "explicit" in err.response.text.lower()
|
||||
and "caching" in err.response.text.lower()
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"Model %s does not support explicit context caching (%s). Falling back to non-cached request.",
|
||||
model,
|
||||
err.response.text,
|
||||
)
|
||||
return messages, optional_params, None
|
||||
error_code: Final = err.response.status_code
|
||||
raise VertexAIError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException:
|
||||
|
|
@ -552,6 +571,8 @@ class ContextCachingEndpoints(VertexBase):
|
|||
raw_response_cached.get("name"),
|
||||
raw_response_cached.get("expireTime"),
|
||||
)
|
||||
for tool_param in _TOOL_PARAMS:
|
||||
optional_params.pop(tool_param, None)
|
||||
return (
|
||||
non_cached_messages,
|
||||
optional_params,
|
||||
|
|
@ -689,6 +710,17 @@ class ContextCachingEndpoints(VertexBase):
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
if "does not support explicit context caching" in err.response.text or (
|
||||
err.response.status_code == 400
|
||||
and "explicit" in err.response.text.lower()
|
||||
and "caching" in err.response.text.lower()
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"Model %s does not support explicit context caching (%s). Falling back to non-cached request.",
|
||||
model,
|
||||
err.response.text,
|
||||
)
|
||||
return None
|
||||
error_code: Final = err.response.status_code
|
||||
raise VertexAIError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException:
|
||||
|
|
@ -708,6 +740,10 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except _CachedContentsNotServedError:
|
||||
verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url)
|
||||
return messages, optional_params, None
|
||||
|
||||
if cache_name is None:
|
||||
return messages, optional_params, None
|
||||
|
||||
for tool_param in _TOOL_PARAMS:
|
||||
optional_params.pop(tool_param, None)
|
||||
return non_cached_messages, optional_params, cache_name
|
||||
|
|
|
|||
|
|
@ -1262,6 +1262,11 @@ def sync_transform_request_body(
|
|||
|
||||
context_caching_endpoints: Final = ContextCachingEndpoints()
|
||||
|
||||
if context_caching:
|
||||
model_info = litellm_params.get("model_info") or {}
|
||||
if model_info.get("supports_prompt_caching") is False:
|
||||
context_caching = False
|
||||
|
||||
(
|
||||
messages,
|
||||
optional_params,
|
||||
|
|
@ -1398,6 +1403,11 @@ async def async_transform_request_body(
|
|||
|
||||
context_caching_endpoints: Final = ContextCachingEndpoints()
|
||||
|
||||
if context_caching:
|
||||
model_info = litellm_params.get("model_info") or {}
|
||||
if model_info.get("supports_prompt_caching") is False:
|
||||
context_caching = False
|
||||
|
||||
(
|
||||
messages,
|
||||
optional_params,
|
||||
|
|
|
|||
|
|
@ -78192,7 +78192,7 @@
|
|||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
|
|
@ -78252,7 +78252,7 @@
|
|||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
|
|
@ -78310,7 +78310,7 @@
|
|||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
|
|
|
|||
|
|
@ -78192,7 +78192,7 @@
|
|||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
|
|
@ -78252,7 +78252,7 @@
|
|||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
|
|
@ -78310,7 +78310,7 @@
|
|||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,202 @@
|
|||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.anthropic.messages.adapters.transformation import (
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
|
||||
ContextCachingEndpoints,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.transformation import (
|
||||
sync_transform_request_body,
|
||||
)
|
||||
|
||||
|
||||
def test_target_consumes_cache_control_skips_when_model_disables_prompt_caching(local_model_cost_map):
|
||||
"""LiteLLMAnthropicMessagesAdapter should not consume cache_control for models explicitly disabling prompt caching."""
|
||||
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("claude-3-5-sonnet-20241022") is True
|
||||
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("gemini/gemini-2.5-flash") is True
|
||||
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("gemini/gemini-4-argon") is False
|
||||
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("vertex_ai/gemini-4-argon") is False
|
||||
# Bare model name without provider prefix should also be recognized
|
||||
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("gemini-4-argon") is False
|
||||
|
||||
|
||||
def test_check_and_create_cache_skips_when_model_disables_prompt_caching(local_model_cost_map):
|
||||
"""ContextCachingEndpoints should skip explicit cache creation when model has supports_prompt_caching: false."""
|
||||
endpoints = ContextCachingEndpoints()
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hello world test token count " * 1000,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
optional_params = {"tools": [{"type": "function", "function": {"name": "test"}}]}
|
||||
|
||||
res_messages, res_params, cache_id = endpoints.check_and_create_cache(
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test-key",
|
||||
api_base=None,
|
||||
model="gemini/gemini-4-argon",
|
||||
client=MagicMock(),
|
||||
timeout=10,
|
||||
logging_obj=MagicMock(),
|
||||
custom_llm_provider="gemini",
|
||||
vertex_project=None,
|
||||
vertex_location=None,
|
||||
vertex_auth_header=None,
|
||||
)
|
||||
assert cache_id is None
|
||||
# All messages and tools should remain intact
|
||||
assert len(res_messages) == len(messages)
|
||||
assert "tools" in res_params
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_check_and_create_cache_fallback_on_400_explicit_caching_unsupported(local_model_cost_map):
|
||||
"""When Vertex AI returns 400 'does not support explicit context caching', gracefully fallback to non-cached request."""
|
||||
endpoints = ContextCachingEndpoints()
|
||||
endpoints.async_check_cache = AsyncMock(return_value=None)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hello world test token count " * 1000,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
optional_params = {"tools": [{"type": "function", "function": {"name": "test"}}]}
|
||||
|
||||
mock_client = MagicMock(spec=AsyncHTTPHandler)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = json.dumps({
|
||||
"error": {
|
||||
"code": 400,
|
||||
"message": "Model gemini-4-argon does not support explicit context caching.",
|
||||
"status": "INVALID_ARGUMENT",
|
||||
}
|
||||
})
|
||||
mock_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError("Bad Request", request=MagicMock(), response=mock_response)
|
||||
)
|
||||
|
||||
res_messages, res_params, cache_id = await endpoints.async_check_and_create_cache(
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test-key",
|
||||
api_base=None,
|
||||
model="gemini/gemini-2.5-flash",
|
||||
client=mock_client,
|
||||
timeout=10,
|
||||
logging_obj=MagicMock(),
|
||||
custom_llm_provider="gemini",
|
||||
vertex_project=None,
|
||||
vertex_location=None,
|
||||
vertex_auth_header=None,
|
||||
)
|
||||
assert cache_id is None
|
||||
assert mock_client.post.called
|
||||
# Critical checks: messages must not be dropped, tools must not be stripped!
|
||||
assert len(res_messages) == len(messages)
|
||||
assert "tools" in res_params
|
||||
|
||||
|
||||
def test_sync_check_and_create_cache_fallback_on_400_explicit_caching_unsupported(local_model_cost_map):
|
||||
"""Sync check_and_create_cache fallback on 400 explicit caching unsupported."""
|
||||
endpoints = ContextCachingEndpoints()
|
||||
endpoints.check_cache = MagicMock(return_value=None)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hello world test token count " * 1000,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
optional_params = {"tools": [{"type": "function", "function": {"name": "test"}}]}
|
||||
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = json.dumps({
|
||||
"error": {
|
||||
"code": 400,
|
||||
"message": "Model gemini-4-argon does not support explicit context caching.",
|
||||
"status": "INVALID_ARGUMENT",
|
||||
}
|
||||
})
|
||||
mock_client.post.side_effect = httpx.HTTPStatusError(
|
||||
"Bad Request", request=MagicMock(), response=mock_response
|
||||
)
|
||||
|
||||
res_messages, res_params, cache_id = endpoints.check_and_create_cache(
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test-key",
|
||||
api_base=None,
|
||||
model="gemini/gemini-2.5-flash",
|
||||
client=mock_client,
|
||||
timeout=10,
|
||||
logging_obj=MagicMock(),
|
||||
custom_llm_provider="gemini",
|
||||
vertex_project=None,
|
||||
vertex_location=None,
|
||||
vertex_auth_header=None,
|
||||
)
|
||||
assert cache_id is None
|
||||
assert mock_client.post.called
|
||||
# Critical checks: messages must not be dropped, tools must not be stripped!
|
||||
assert len(res_messages) == len(messages)
|
||||
assert "tools" in res_params
|
||||
|
||||
|
||||
def test_sync_transform_request_body_skips_caching_on_model_info_override(monkeypatch):
|
||||
"""sync_transform_request_body should respect supports_prompt_caching: false in model_info."""
|
||||
mock_endpoints = MagicMock()
|
||||
mock_endpoints.check_and_create_cache.return_value = ([], {}, None)
|
||||
monkeypatch.setattr(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.ContextCachingEndpoints",
|
||||
MagicMock(return_value=mock_endpoints),
|
||||
)
|
||||
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
sync_transform_request_body(
|
||||
messages=messages,
|
||||
model="gemini-2.5-flash",
|
||||
custom_llm_provider="gemini",
|
||||
litellm_params={"model_info": {"supports_prompt_caching": False}},
|
||||
optional_params={},
|
||||
gemini_api_key="test-key",
|
||||
api_base=None,
|
||||
client=None,
|
||||
timeout=None,
|
||||
extra_headers=None,
|
||||
logging_obj=MagicMock(),
|
||||
vertex_project=None,
|
||||
vertex_location=None,
|
||||
vertex_auth_header=None,
|
||||
)
|
||||
# context_caching=False must be passed to check_and_create_cache
|
||||
assert mock_endpoints.check_and_create_cache.called
|
||||
_, kwargs = mock_endpoints.check_and_create_cache.call_args
|
||||
assert kwargs.get("context_caching") is False
|
||||
Loading…
Add table
Reference in a new issue