tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-16 08:17:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #15226 from otaviofbrito/chore/vertex-ai-context-caching

Browse Source 
Chore/vertex ai context caching
This commit is contained in:
Krish Dholakia
2025-10-06 20:05:24 -07:00
committed by GitHub
parent 6b4415684e b2a734c1f8
commit 8b357c24c2
7 changed files with 508 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/vertex.md
+150 -6
View File
@@ -191,7 +191,7 @@ print(json.loads(completion.choices[0].message.content))
model_list:
- model_name: gemini-2.5-pro
litellm_params:
model: vertex_ai/gemini-1.5-pro
model: vertex_ai/gemini-2.5-pro
vertex_project: "project-id"
vertex_location: "us-central1"
vertex_credentials: "/path/to/service_account.json" # [OPTIONAL] Do this OR `!gcloud auth application-default login` - run this to add vertex credentials to your env
@@ -277,7 +277,7 @@ except JSONSchemaValidationError as e:
model_list:
- model_name: gemini-2.5-pro
litellm_params:
model: vertex_ai/gemini-1.5-pro
model: vertex_ai/gemini-2.5-pro
vertex_project: "project-id"
vertex_location: "us-central1"
vertex_credentials: "/path/to/service_account.json" # [OPTIONAL] Do this OR `!gcloud auth application-default login` - run this to add vertex credentials to your env
@@ -981,11 +981,155 @@ curl http://0.0.0.0:4000/v1/chat/completions \
### **Context Caching**
Use Vertex AI context caching is supported by calling provider api directly. (Unified Endpoint support coming soon.).
#### Unified Endpoint
Use Vertex AI context caching in the same way as [**Google AI Studio - Context Caching**](../providers/gemini.md#context-caching)
##### Example usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
for _ in range(2):
resp = completion(
model="vertex_ai/gemini-2.5-pro",
messages=[
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement" * 4000,
"cache_control": {"type": "ephemeral"}, # 👈 KEY CHANGE
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
}]
)
print(resp.usage) # 👈 2nd usage block will be less, since cached tokens used
```
</TabItem>
<TabItem value="sdk-ttl" label="SDK with Custom TTL">
```python
from litellm import completion
# Cache for 2 hours (7200 seconds)
resp = completion(
model="vertex_ai/gemini-2.5-pro",
messages=[
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement" * 4000,
"cache_control": {
"type": "ephemeral",
"ttl": "7200s" # 👈 Cache for 2 hours
},
}
],
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {
"type": "ephemeral",
"ttl": "3600s" # 👈 This TTL will be ignored (first one is used)
},
}
],
}
]
)
print(resp.usage)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
1. Setup config.yaml
```yaml
model_list:
- model_name: gemini-2.5-pro
litellm_params:
model: vertex_ai/gemini-2.5-pro
vertex_project: "project-id"
vertex_location: "us-central1"
vertex_credentials: "/path/to/service_account.json"
```
2. Start proxy
```bash
litellm --config /path/to/config.yaml
```
3. Test it!
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "gemini-2.5-flash",
"messages": [
{
"role": "system",
"content": [
{
"type": "text",
"text": "Long cache message (must be >= 1024 tokens)",
"cache_control": {
"type": "ephemeral",
"ttl": "7200s"
}
}
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "What is the text about?"
}
]
}
]
}'
```
#### Calling provider api directly
[**Go straight to provider**](../pass_through/vertex_ai.md#context-caching)
#### 1. Create the Cache
##### 1. Create the Cache
First, create the cache by sending a `POST` request to the `cachedContents` endpoint via the LiteLLM proxy.
@@ -1011,7 +1155,7 @@ curl http://0.0.0.0:4000/vertex_ai/v1/projects/{project_id}/locations/{location}
</TabItem>
</Tabs>
#### 2. Get the Cache Name from the Response
##### 2. Get the Cache Name from the Response
Vertex AI will return a response containing the `name` of the cached content. This name is the identifier for your cached data.
@@ -1030,7 +1174,7 @@ Vertex AI will return a response containing the `name` of the cached content. Th
}
```
#### 3. Use the Cached Content
##### 3. Use the Cached Content
Use the `name` from the response as `cachedContent` or `cached_content` in subsequent API calls to reuse the cached information. This is passed in the body of your request to `/chat/completions`.
litellm/llms/vertex_ai/context_caching/transformation.py
+14 -4
View File
@@ -5,7 +5,7 @@ Why separate file? Make it easy to see how transformation works
"""
import re
from typing import List, Optional, Tuple
from typing import List, Optional, Tuple, Literal
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.vertex_ai import CachedContentRequestBody
@@ -155,13 +155,18 @@ def separate_cached_messages(
def transform_openai_messages_to_gemini_context_caching(
model: str, messages: List[AllMessageValues], cache_key: str
model: str,
messages: List[AllMessageValues],
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
cache_key: str,
vertex_project: Optional[str],
vertex_location: Optional[str],
) -> CachedContentRequestBody:
# Extract TTL from cached messages BEFORE system message transformation
ttl = extract_ttl_from_cached_messages(messages)
supports_system_message = get_supports_system_message(
model=model, custom_llm_provider="gemini"
model=model, custom_llm_provider=custom_llm_provider
)
transformed_system_messages, new_messages = _transform_system_message(
@@ -170,9 +175,14 @@ def transform_openai_messages_to_gemini_context_caching(
transformed_messages = _gemini_convert_messages_with_history(messages=new_messages)
model_name = "models/{}".format(model)
if custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta":
model_name = f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/{model_name}"
data = CachedContentRequestBody(
contents=transformed_messages,
model="models/{}".format(model),
model=model_name,
displayName=cache_key,
)
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
+64 -9
View File
@@ -41,8 +41,11 @@ class ContextCachingEndpoints(VertexBase):
def _get_token_and_url_context_caching(
self,
gemini_api_key: Optional[str],
custom_llm_provider: Literal["gemini"],
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
api_base: Optional[str],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
) -> Tuple[Optional[str], str]:
"""
Internal function. Returns the token and url for the call.
@@ -58,9 +61,15 @@ class ContextCachingEndpoints(VertexBase):
url = "https://generativelanguage.googleapis.com/v1beta/{}?key={}".format(
endpoint, gemini_api_key
)
elif custom_llm_provider == "vertex_ai":
auth_header = vertex_auth_header
endpoint = "cachedContents"
url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}"
else:
raise NotImplementedError
auth_header = vertex_auth_header
endpoint = "cachedContents"
url = f"https://{vertex_location}-aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}"
return self._check_custom_proxy(
api_base=api_base,
@@ -80,6 +89,10 @@ class ContextCachingEndpoints(VertexBase):
api_key: str,
api_base: Optional[str],
logging_obj: Logging,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
) -> Optional[str]:
"""
Checks if content already cached.
@@ -94,8 +107,11 @@ class ContextCachingEndpoints(VertexBase):
_, url = self._get_token_and_url_context_caching(
gemini_api_key=api_key,
custom_llm_provider="gemini",
custom_llm_provider=custom_llm_provider,
api_base=api_base,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header
)
try:
## LOGGING
@@ -145,6 +161,10 @@ class ContextCachingEndpoints(VertexBase):
api_key: str,
api_base: Optional[str],
logging_obj: Logging,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str]
) -> Optional[str]:
"""
Checks if content already cached.
@@ -159,8 +179,11 @@ class ContextCachingEndpoints(VertexBase):
_, url = self._get_token_and_url_context_caching(
gemini_api_key=api_key,
custom_llm_provider="gemini",
custom_llm_provider=custom_llm_provider,
api_base=api_base,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header
)
try:
## LOGGING
@@ -212,6 +235,10 @@ class ContextCachingEndpoints(VertexBase):
client: Optional[HTTPHandler],
timeout: Optional[Union[float, httpx.Timeout]],
logging_obj: Logging,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
extra_headers: Optional[dict] = None,
cached_content: Optional[str] = None,
) -> Tuple[List[AllMessageValues], dict, Optional[str]]:
@@ -240,8 +267,11 @@ class ContextCachingEndpoints(VertexBase):
## AUTHORIZATION ##
token, url = self._get_token_and_url_context_caching(
gemini_api_key=api_key,
custom_llm_provider="gemini",
custom_llm_provider=custom_llm_provider,
api_base=api_base,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header
)
headers = {
@@ -273,6 +303,10 @@ class ContextCachingEndpoints(VertexBase):
api_key=api_key,
api_base=api_base,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header
)
if google_cache_name:
return non_cached_messages, optional_params, google_cache_name
@@ -280,7 +314,12 @@ class ContextCachingEndpoints(VertexBase):
## TRANSFORM REQUEST
cached_content_request_body = (
transform_openai_messages_to_gemini_context_caching(
model=model, messages=cached_messages, cache_key=generated_cache_key
model=model,
messages=cached_messages,
cache_key=generated_cache_key,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
)
@@ -328,6 +367,10 @@ class ContextCachingEndpoints(VertexBase):
client: Optional[AsyncHTTPHandler],
timeout: Optional[Union[float, httpx.Timeout]],
logging_obj: Logging,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
extra_headers: Optional[dict] = None,
cached_content: Optional[str] = None,
) -> Tuple[List[AllMessageValues], dict, Optional[str]]:
@@ -356,8 +399,11 @@ class ContextCachingEndpoints(VertexBase):
## AUTHORIZATION ##
token, url = self._get_token_and_url_context_caching(
gemini_api_key=api_key,
custom_llm_provider="gemini",
custom_llm_provider=custom_llm_provider,
api_base=api_base,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header
)
headers = {
@@ -386,6 +432,10 @@ class ContextCachingEndpoints(VertexBase):
api_key=api_key,
api_base=api_base,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header
)
if google_cache_name:
@@ -394,7 +444,12 @@ class ContextCachingEndpoints(VertexBase):
## TRANSFORM REQUEST
cached_content_request_body = (
transform_openai_messages_to_gemini_context_caching(
model=model, messages=cached_messages, cache_key=generated_cache_key
model=model,
messages=cached_messages,
cache_key=generated_cache_key,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
)
litellm/llms/vertex_ai/gemini/transformation.py
+47 -46
View File
@@ -514,34 +514,35 @@ def sync_transform_request_body(
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
litellm_params: dict,
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
) -> RequestBody:
from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints
context_caching_endpoints = ContextCachingEndpoints()
if gemini_api_key is not None:
(
messages,
optional_params,
cached_content,
) = context_caching_endpoints.check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key=gemini_api_key,
api_base=api_base,
model=model,
client=client,
timeout=timeout,
extra_headers=extra_headers,
cached_content=optional_params.pop("cached_content", None),
logging_obj=logging_obj,
)
else: # [TODO] implement context caching for gemini as well
cached_content = None
if "cached_content" in optional_params:
cached_content = optional_params.pop("cached_content")
elif "cachedContent" in optional_params:
cached_content = optional_params.pop("cachedContent")
(
messages,
optional_params,
cached_content,
) = context_caching_endpoints.check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key=gemini_api_key or "dummy",
api_base=api_base,
model=model,
client=client,
timeout=timeout,
extra_headers=extra_headers,
cached_content=optional_params.pop("cached_content", None),
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
)
return _transform_request_body(
messages=messages,
@@ -565,34 +566,34 @@ async def async_transform_request_body(
logging_obj: litellm.litellm_core_utils.litellm_logging.Logging, # type: ignore
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
litellm_params: dict,
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
) -> RequestBody:
from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints
context_caching_endpoints = ContextCachingEndpoints()
if gemini_api_key is not None:
(
messages,
optional_params,
cached_content,
) = await context_caching_endpoints.async_check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key=gemini_api_key,
api_base=api_base,
model=model,
client=client,
timeout=timeout,
extra_headers=extra_headers,
cached_content=optional_params.pop("cached_content", None),
logging_obj=logging_obj,
)
else: # [TODO] implement context caching for gemini as well
cached_content = None
if "cached_content" in optional_params:
cached_content = optional_params.pop("cached_content")
elif "cachedContent" in optional_params:
cached_content = optional_params.pop("cachedContent")
(
messages,
optional_params,
cached_content,
) = await context_caching_endpoints.async_check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key=gemini_api_key or "dummy",
api_base=api_base,
model=model,
client=client,
timeout=timeout,
extra_headers=extra_headers,
cached_content=optional_params.pop("cached_content", None),
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
)
return _transform_request_body(
messages=messages,
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+18 -3
View File
@@ -1792,7 +1792,6 @@ class VertexLLM(VertexBase):
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
) -> CustomStreamWrapper:
request_body = await async_transform_request_body(**data) # type: ignore
should_use_v1beta1_features = self.is_using_v1beta1_features(
optional_params=optional_params
@@ -1826,6 +1825,13 @@ class VertexLLM(VertexBase):
litellm_params=litellm_params,
)
request_body = await async_transform_request_body(
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header) # type: ignore
## LOGGING
logging_obj.pre_call(
input=messages,
@@ -1913,7 +1919,12 @@ class VertexLLM(VertexBase):
litellm_params=litellm_params,
)
request_body = await async_transform_request_body(**data) # type: ignore
request_body = await async_transform_request_body(
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header) # type: ignore
_async_client_params = {}
if timeout:
_async_client_params["timeout"] = timeout
@@ -2088,7 +2099,11 @@ class VertexLLM(VertexBase):
)
## TRANSFORMATION ##
data = sync_transform_request_body(**transform_request_params)
data = sync_transform_request_body(
**transform_request_params,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header)
## LOGGING
logging_obj.pre_call(
tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py
+60 -12
View File
@@ -191,7 +191,8 @@ class TestTTLExtraction:
class TestTransformationWithTTL:
"""Test the complete transformation with TTL support"""
def test_transform_with_valid_ttl(self):
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
def test_transform_with_valid_ttl(self, custom_llm_provider):
"""Test transformation includes TTL when provided"""
messages = [
{
@@ -205,19 +206,32 @@ class TestTransformationWithTTL:
]
}
]
vertex_location="test_location"
vertex_project="test_project"
result = transform_openai_messages_to_gemini_context_caching(
model="gemini-1.5-pro",
messages=messages,
cache_key="test-cache-key"
cache_key="test-cache-key",
custom_llm_provider=custom_llm_provider,
vertex_location="test_location",
vertex_project="test_project"
)
assert "ttl" in result
assert result["ttl"] == "3600s"
assert result["model"] == "models/gemini-1.5-pro"
if custom_llm_provider == "gemini":
assert result["model"] == "models/gemini-1.5-pro"
else:
assert result["model"] == f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/gemini-1.5-pro"
assert result["displayName"] == "test-cache-key"
def test_transform_without_ttl(self):
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
def test_transform_without_ttl(self, custom_llm_provider):
"""Test transformation without TTL"""
messages = [
{
@@ -231,18 +245,30 @@ class TestTransformationWithTTL:
]
}
]
vertex_location="test_location"
vertex_project="test_project"
result = transform_openai_messages_to_gemini_context_caching(
model="gemini-1.5-pro",
messages=messages,
cache_key="test-cache-key"
cache_key="test-cache-key",
custom_llm_provider=custom_llm_provider,
vertex_location=vertex_location,
vertex_project=vertex_project
)
assert "ttl" not in result
assert result["model"] == "models/gemini-1.5-pro"
if custom_llm_provider == "gemini":
assert result["model"] == "models/gemini-1.5-pro"
else:
assert result["model"] == f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/gemini-1.5-pro"
assert result["displayName"] == "test-cache-key"
def test_transform_with_invalid_ttl(self):
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
def test_transform_with_invalid_ttl(self, custom_llm_provider):
"""Test transformation with invalid TTL (should be ignored)"""
messages = [
{
@@ -256,18 +282,29 @@ class TestTransformationWithTTL:
]
}
]
vertex_location="test_location"
vertex_project="test_project"
result = transform_openai_messages_to_gemini_context_caching(
model="gemini-1.5-pro",
messages=messages,
cache_key="test-cache-key"
cache_key="test-cache-key",
custom_llm_provider=custom_llm_provider,
vertex_location=vertex_location,
vertex_project=vertex_project
)
assert "ttl" not in result
assert result["model"] == "models/gemini-1.5-pro"
if custom_llm_provider == "gemini":
assert result["model"] == "models/gemini-1.5-pro"
else:
assert result["model"] == f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/gemini-1.5-pro"
assert result["displayName"] == "test-cache-key"
def test_transform_with_system_message_and_ttl(self):
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
def test_transform_with_system_message_and_ttl(self, custom_llm_provider):
"""Test transformation with system message and TTL"""
messages = [
{
@@ -290,17 +327,28 @@ class TestTransformationWithTTL:
]
}
]
vertex_location="test_location"
vertex_project="test_project"
result = transform_openai_messages_to_gemini_context_caching(
model="gemini-1.5-pro",
messages=messages,
cache_key="test-cache-key"
cache_key="test-cache-key",
custom_llm_provider=custom_llm_provider,
vertex_location=vertex_location,
vertex_project=vertex_project
)
assert "ttl" in result
assert result["ttl"] == "7200s"
assert "system_instruction" in result
assert result["model"] == "models/gemini-1.5-pro"
if custom_llm_provider == "gemini":
assert result["model"] == "models/gemini-1.5-pro"
else:
assert result["model"] == f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/gemini-1.5-pro"
assert result["displayName"] == "test-cache-key"
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
+155 -11
View File
@@ -55,6 +55,9 @@ class TestContextCachingEndpoints:
self.sample_optional_params = {"tools": self.sample_tools.copy()}
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -62,12 +65,14 @@ class TestContextCachingEndpoints:
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
)
def test_check_and_create_cache_with_cached_content(
self, mock_cache_obj, mock_separate
self, mock_cache_obj, mock_separate, custom_llm_provider
):
"""Test check_and_create_cache when cached_content is provided"""
# Setup
cached_content = "cached_content_123"
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = self.context_caching.check_and_create_cache(
@@ -80,6 +85,10 @@ class TestContextCachingEndpoints:
timeout=30.0,
logging_obj=self.mock_logging,
cached_content=cached_content,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -92,14 +101,21 @@ class TestContextCachingEndpoints:
mock_separate.assert_not_called()
mock_cache_obj.get_cache_key.assert_not_called()
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
def test_check_and_create_cache_no_cached_messages(self, mock_separate):
def test_check_and_create_cache_no_cached_messages(
self, mock_separate, custom_llm_provider
):
"""Test check_and_create_cache when no cached messages are found"""
# Setup
mock_separate.return_value = ([], self.sample_messages) # No cached messages
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = self.context_caching.check_and_create_cache(
@@ -111,6 +127,10 @@ class TestContextCachingEndpoints:
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -119,6 +139,9 @@ class TestContextCachingEndpoints:
assert returned_params == optional_params
assert returned_cache is None
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -127,7 +150,7 @@ class TestContextCachingEndpoints:
)
@patch.object(ContextCachingEndpoints, "check_cache")
def test_check_and_create_cache_existing_cache_found(
self, mock_check_cache, mock_cache_obj, mock_separate
self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider
):
"""Test check_and_create_cache when existing cache is found"""
# Setup
@@ -139,6 +162,8 @@ class TestContextCachingEndpoints:
mock_check_cache.return_value = "existing_cache_name"
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = self.context_caching.check_and_create_cache(
@@ -150,6 +175,10 @@ class TestContextCachingEndpoints:
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -163,6 +192,9 @@ class TestContextCachingEndpoints:
messages=cached_messages, tools=self.sample_tools
)
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -181,6 +213,7 @@ class TestContextCachingEndpoints:
mock_transform,
mock_cache_obj,
mock_separate,
custom_llm_provider,
):
"""Test check_and_create_cache when creating new cache"""
# Setup
@@ -203,6 +236,8 @@ class TestContextCachingEndpoints:
self.mock_client.post.return_value = mock_response
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = self.context_caching.check_and_create_cache(
@@ -214,6 +249,10 @@ class TestContextCachingEndpoints:
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -228,6 +267,9 @@ class TestContextCachingEndpoints:
assert "tools" in call_args.kwargs["json"]
assert call_args.kwargs["json"]["tools"] == self.sample_tools
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -237,7 +279,12 @@ class TestContextCachingEndpoints:
@patch.object(ContextCachingEndpoints, "check_cache")
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
def test_check_and_create_cache_http_error(
self, mock_get_token_url, mock_check_cache, mock_cache_obj, mock_separate
self,
mock_get_token_url,
mock_check_cache,
mock_cache_obj,
mock_separate,
custom_llm_provider,
):
"""Test check_and_create_cache handles HTTP errors properly"""
# Setup
@@ -259,6 +306,8 @@ class TestContextCachingEndpoints:
self.mock_client.post.side_effect = http_error
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute and Assert
with pytest.raises(VertexAIError) as exc_info:
@@ -271,12 +320,19 @@ class TestContextCachingEndpoints:
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
assert exc_info.value.status_code == 400
assert "Bad Request" in str(exc_info.value.message)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -284,12 +340,14 @@ class TestContextCachingEndpoints:
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
)
async def test_async_check_and_create_cache_with_cached_content(
self, mock_cache_obj, mock_separate
self, mock_cache_obj, mock_separate, custom_llm_provider
):
"""Test async_check_and_create_cache when cached_content is provided"""
# Setup
cached_content = "cached_content_123"
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = await self.context_caching.async_check_and_create_cache(
@@ -302,6 +360,10 @@ class TestContextCachingEndpoints:
timeout=30.0,
logging_obj=self.mock_logging,
cached_content=cached_content,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -311,14 +373,21 @@ class TestContextCachingEndpoints:
assert returned_cache == cached_content
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
async def test_async_check_and_create_cache_no_cached_messages(self, mock_separate):
async def test_async_check_and_create_cache_no_cached_messages(
self, mock_separate, custom_llm_provider
):
"""Test async_check_and_create_cache when no cached messages are found"""
# Setup
mock_separate.return_value = ([], self.sample_messages)
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = await self.context_caching.async_check_and_create_cache(
@@ -330,6 +399,10 @@ class TestContextCachingEndpoints:
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -339,6 +412,9 @@ class TestContextCachingEndpoints:
assert returned_cache is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -347,7 +423,7 @@ class TestContextCachingEndpoints:
)
@patch.object(ContextCachingEndpoints, "async_check_cache")
async def test_async_check_and_create_cache_existing_cache_found(
self, mock_async_check_cache, mock_cache_obj, mock_separate
self, mock_async_check_cache, mock_cache_obj, mock_separate, custom_llm_provider
):
"""Test async_check_and_create_cache when existing cache is found"""
# Setup
@@ -359,6 +435,8 @@ class TestContextCachingEndpoints:
mock_async_check_cache.return_value = "existing_cache_name"
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = await self.context_caching.async_check_and_create_cache(
@@ -370,6 +448,10 @@ class TestContextCachingEndpoints:
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -384,6 +466,9 @@ class TestContextCachingEndpoints:
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -406,6 +491,7 @@ class TestContextCachingEndpoints:
mock_transform,
mock_cache_obj,
mock_separate,
custom_llm_provider,
):
"""Test async_check_and_create_cache when creating new cache"""
# Setup
@@ -428,6 +514,8 @@ class TestContextCachingEndpoints:
self.mock_async_client.post = AsyncMock(return_value=mock_response)
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = await self.context_caching.async_check_and_create_cache(
@@ -439,6 +527,10 @@ class TestContextCachingEndpoints:
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert
@@ -454,6 +546,9 @@ class TestContextCachingEndpoints:
assert call_args.kwargs["json"]["tools"] == self.sample_tools
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@@ -472,6 +567,7 @@ class TestContextCachingEndpoints:
mock_async_check_cache,
mock_cache_obj,
mock_separate,
custom_llm_provider,
):
"""Test async_check_and_create_cache handles timeout errors properly"""
# Setup
@@ -489,6 +585,8 @@ class TestContextCachingEndpoints:
)
optional_params = self.sample_optional_params.copy()
test_project = "test_project"
test_location = "test_location"
# Execute and Assert
with pytest.raises(VertexAIError) as exc_info:
@@ -501,12 +599,21 @@ class TestContextCachingEndpoints:
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
assert exc_info.value.status_code == 408
assert "Timeout error occurred" in str(exc_info.value.message)
def test_check_and_create_cache_tools_popped_from_optional_params(self):
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
def test_check_and_create_cache_tools_popped_from_optional_params(
self, custom_llm_provider
):
"""Test that tools are properly popped from optional_params when there are cached messages"""
with patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
@@ -520,6 +627,8 @@ class TestContextCachingEndpoints:
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
test_project = "test_project"
test_location = "test_location"
# Mock the check_cache to return existing cache so we don't make HTTP calls
with patch.object(
@@ -535,6 +644,10 @@ class TestContextCachingEndpoints:
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert tools were popped from optional_params
@@ -543,7 +656,12 @@ class TestContextCachingEndpoints:
# But original tools should still be available for comparison
assert original_tools == self.sample_tools
def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(self):
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(
self, custom_llm_provider
):
"""Test that tools are NOT popped from optional_params when there are no cached messages"""
with patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
@@ -555,6 +673,8 @@ class TestContextCachingEndpoints:
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = self.context_caching.check_and_create_cache(
@@ -566,6 +686,10 @@ class TestContextCachingEndpoints:
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert tools were NOT popped from optional_params (early return)
@@ -573,8 +697,11 @@ class TestContextCachingEndpoints:
assert optional_params["tools"] == original_tools
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(
self,
self, custom_llm_provider
):
"""Test that tools are NOT popped from optional_params in async version when there are no cached messages"""
with patch(
@@ -587,6 +714,8 @@ class TestContextCachingEndpoints:
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
test_project = "test_project"
test_location = "test_location"
# Execute
result = await self.context_caching.async_check_and_create_cache(
@@ -598,6 +727,10 @@ class TestContextCachingEndpoints:
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert tools were NOT popped from optional_params (early return)
@@ -605,7 +738,12 @@ class TestContextCachingEndpoints:
assert optional_params["tools"] == original_tools
@pytest.mark.asyncio
async def test_async_check_and_create_cache_tools_popped_from_optional_params(self):
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
async def test_async_check_and_create_cache_tools_popped_from_optional_params(
self, custom_llm_provider
):
"""Test that tools are properly popped from optional_params in async version when there are cached messages"""
with patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
@@ -619,6 +757,8 @@ class TestContextCachingEndpoints:
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
test_project = "test_project"
test_location = "test_location"
# Mock the async_check_cache to return existing cache so we don't make HTTP calls
with patch.object(
@@ -634,6 +774,10 @@ class TestContextCachingEndpoints:
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider=custom_llm_provider,
vertex_project=test_project,
vertex_location=test_location,
vertex_auth_header="vertext_test_token",
)
# Assert tools were popped from optional_params