mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(gemini): pop tool_choice from optional_params when creating cached content
When using context caching with Gemini, `check_and_create_cache()` correctly pops `tools` from `optional_params` and includes them in the cached content. However, `tool_choice` was not being popped, causing it to remain as `toolConfig` in the final `GenerateContent` request alongside `cachedContent`. The Gemini API rejects requests that include both `cachedContent` and `toolConfig` with error 400: "CachedContent can not be used with GenerateContent request setting system_instruction, tools or tool_config". This fix pops `tool_choice` from `optional_params` (both sync and async paths) and includes it as `toolConfig` in the cached content request body, matching the existing pattern for `tools`. Verified against the Gemini API: the CachedContent creation endpoint accepts `toolConfig` alongside `tools` and `system_instruction`.
This commit is contained in:
Enzo Moreira
parent
c08fb82cae
commit
501b8fa445
2 changed files with 247 additions and 268 deletions
|
|
@ -26,9 +26,7 @@ from .transformation import (
|
|||
transform_openai_messages_to_gemini_context_caching,
|
||||
)
|
||||
|
||||
local_cache_obj = Cache(
|
||||
type=LiteLLMCacheType.LOCAL
|
||||
) # only used for calling 'get_cache_key' function
|
||||
local_cache_obj = Cache(type=LiteLLMCacheType.LOCAL) # only used for calling 'get_cache_key' function
|
||||
|
||||
MAX_PAGINATION_PAGES = 100 # Reasonable upper bound for pagination
|
||||
|
||||
|
|
@ -64,9 +62,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
if custom_llm_provider == "gemini":
|
||||
auth_header = None
|
||||
endpoint = "cachedContents"
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}?key={}".format(
|
||||
endpoint, gemini_api_key
|
||||
)
|
||||
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"
|
||||
|
|
@ -93,9 +89,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
model=model,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_api_version="v1beta1"
|
||||
if custom_llm_provider == "vertex_ai_beta"
|
||||
else "v1",
|
||||
vertex_api_version="v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1",
|
||||
)
|
||||
|
||||
def check_cache(
|
||||
|
|
@ -161,9 +155,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 403:
|
||||
return None
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code, message=e.response.text
|
||||
)
|
||||
raise VertexAIError(status_code=e.response.status_code, message=e.response.text)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
|
|
@ -255,9 +247,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 403:
|
||||
return None
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code, message=e.response.text
|
||||
)
|
||||
raise VertexAIError(status_code=e.response.status_code, message=e.response.text)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
|
|
@ -316,9 +306,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
if cached_content is not None:
|
||||
return messages, optional_params, cached_content
|
||||
|
||||
cached_messages, non_cached_messages = separate_cached_messages(
|
||||
messages=messages
|
||||
)
|
||||
cached_messages, non_cached_messages = separate_cached_messages(messages=messages)
|
||||
|
||||
if len(cached_messages) == 0:
|
||||
return messages, optional_params, None
|
||||
|
|
@ -338,6 +326,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
return messages, optional_params, None
|
||||
|
||||
tools = optional_params.pop("tools", None)
|
||||
tool_choice = optional_params.pop("tool_choice", None)
|
||||
|
||||
## AUTHORIZATION ##
|
||||
token, url = self._get_token_and_url_context_caching(
|
||||
|
|
@ -369,9 +358,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
client = client
|
||||
|
||||
## CHECK IF CACHED ALREADY
|
||||
generated_cache_key = local_cache_obj.get_cache_key(
|
||||
messages=cached_messages, tools=tools, model=model
|
||||
)
|
||||
generated_cache_key = local_cache_obj.get_cache_key(messages=cached_messages, tools=tools, model=model)
|
||||
google_cache_name = self.check_cache(
|
||||
cache_key=generated_cache_key,
|
||||
client=client,
|
||||
|
|
@ -389,18 +376,18 @@ class ContextCachingEndpoints(VertexBase):
|
|||
return non_cached_messages, optional_params, google_cache_name
|
||||
|
||||
## TRANSFORM REQUEST
|
||||
cached_content_request_body = (
|
||||
transform_openai_messages_to_gemini_context_caching(
|
||||
model=model,
|
||||
messages=cached_messages,
|
||||
cache_key=generated_cache_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
cached_content_request_body = transform_openai_messages_to_gemini_context_caching(
|
||||
model=model,
|
||||
messages=cached_messages,
|
||||
cache_key=generated_cache_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
cached_content_request_body["tools"] = tools
|
||||
if tool_choice is not None:
|
||||
cached_content_request_body["toolConfig"] = tool_choice
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
@ -415,7 +402,9 @@ class ContextCachingEndpoints(VertexBase):
|
|||
|
||||
try:
|
||||
response = client.post(
|
||||
url=url, headers=headers, json=cached_content_request_body # type: ignore
|
||||
url=url,
|
||||
headers=headers,
|
||||
json=cached_content_request_body, # type: ignore
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
|
|
@ -464,9 +453,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
if cached_content is not None:
|
||||
return messages, optional_params, cached_content
|
||||
|
||||
cached_messages, non_cached_messages = separate_cached_messages(
|
||||
messages=messages
|
||||
)
|
||||
cached_messages, non_cached_messages = separate_cached_messages(messages=messages)
|
||||
|
||||
if len(cached_messages) == 0:
|
||||
return messages, optional_params, None
|
||||
|
|
@ -486,6 +473,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
return messages, optional_params, None
|
||||
|
||||
tools = optional_params.pop("tools", None)
|
||||
tool_choice = optional_params.pop("tool_choice", None)
|
||||
|
||||
## AUTHORIZATION ##
|
||||
token, url = self._get_token_and_url_context_caching(
|
||||
|
|
@ -507,16 +495,12 @@ class ContextCachingEndpoints(VertexBase):
|
|||
headers.update(extra_headers)
|
||||
|
||||
if client is None or not isinstance(client, AsyncHTTPHandler):
|
||||
client = get_async_httpx_client(
|
||||
params={"timeout": timeout}, llm_provider=litellm.LlmProviders.VERTEX_AI
|
||||
)
|
||||
client = get_async_httpx_client(params={"timeout": timeout}, llm_provider=litellm.LlmProviders.VERTEX_AI)
|
||||
else:
|
||||
client = client
|
||||
|
||||
## CHECK IF CACHED ALREADY
|
||||
generated_cache_key = local_cache_obj.get_cache_key(
|
||||
messages=cached_messages, tools=tools, model=model
|
||||
)
|
||||
generated_cache_key = local_cache_obj.get_cache_key(messages=cached_messages, tools=tools, model=model)
|
||||
google_cache_name = await self.async_check_cache(
|
||||
cache_key=generated_cache_key,
|
||||
client=client,
|
||||
|
|
@ -535,18 +519,18 @@ class ContextCachingEndpoints(VertexBase):
|
|||
return non_cached_messages, optional_params, google_cache_name
|
||||
|
||||
## TRANSFORM REQUEST
|
||||
cached_content_request_body = (
|
||||
transform_openai_messages_to_gemini_context_caching(
|
||||
model=model,
|
||||
messages=cached_messages,
|
||||
cache_key=generated_cache_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
cached_content_request_body = transform_openai_messages_to_gemini_context_caching(
|
||||
model=model,
|
||||
messages=cached_messages,
|
||||
cache_key=generated_cache_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
cached_content_request_body["tools"] = tools
|
||||
if tool_choice is not None:
|
||||
cached_content_request_body["toolConfig"] = tool_choice
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
@ -561,7 +545,9 @@ class ContextCachingEndpoints(VertexBase):
|
|||
|
||||
try:
|
||||
response = await client.post(
|
||||
url=url, headers=headers, json=cached_content_request_body # type: ignore
|
||||
url=url,
|
||||
headers=headers,
|
||||
json=cached_content_request_body, # type: ignore
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
|
|
|
|||
|
|
@ -1,14 +1,11 @@
|
|||
import os
|
||||
import sys
|
||||
from typing import List
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
|
@ -69,18 +66,10 @@ class TestContextCachingEndpoints:
|
|||
"""Teardown for each test method"""
|
||||
self._token_check_patcher.stop()
|
||||
|
||||
@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"
|
||||
)
|
||||
@patch(
|
||||
"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, custom_llm_provider
|
||||
):
|
||||
@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")
|
||||
@patch("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, custom_llm_provider):
|
||||
"""Test check_and_create_cache when cached_content is provided"""
|
||||
# Setup
|
||||
cached_content = "cached_content_123"
|
||||
|
|
@ -115,15 +104,9 @@ 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, custom_llm_provider
|
||||
):
|
||||
@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, 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
|
||||
|
|
@ -153,15 +136,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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj")
|
||||
@patch.object(ContextCachingEndpoints, "check_cache")
|
||||
def test_check_and_create_cache_existing_cache_found(
|
||||
self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider
|
||||
|
|
@ -206,15 +183,9 @@ class TestContextCachingEndpoints:
|
|||
messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
|
||||
)
|
||||
|
||||
@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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj")
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
|
||||
)
|
||||
|
|
@ -281,15 +252,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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj")
|
||||
@patch.object(ContextCachingEndpoints, "check_cache")
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_and_create_cache_http_error(
|
||||
|
|
@ -314,9 +279,7 @@ class TestContextCachingEndpoints:
|
|||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
http_error = httpx.HTTPStatusError(
|
||||
"Error", request=MagicMock(), response=mock_response
|
||||
)
|
||||
http_error = httpx.HTTPStatusError("Error", request=MagicMock(), response=mock_response)
|
||||
self.mock_client.post.side_effect = http_error
|
||||
|
||||
optional_params = self.sample_optional_params.copy()
|
||||
|
|
@ -344,15 +307,9 @@ class TestContextCachingEndpoints:
|
|||
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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("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, custom_llm_provider
|
||||
):
|
||||
|
|
@ -387,15 +344,9 @@ 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, custom_llm_provider
|
||||
):
|
||||
@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, custom_llm_provider):
|
||||
"""Test async_check_and_create_cache when no cached messages are found"""
|
||||
# Setup
|
||||
mock_separate.return_value = ([], self.sample_messages)
|
||||
|
|
@ -426,15 +377,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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj")
|
||||
@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, custom_llm_provider
|
||||
|
|
@ -480,23 +425,15 @@ 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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj")
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "async_check_cache")
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client"
|
||||
)
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client")
|
||||
async def test_async_check_and_create_cache_create_new_cache(
|
||||
self,
|
||||
mock_get_client,
|
||||
|
|
@ -560,20 +497,12 @@ 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"
|
||||
)
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
|
||||
)
|
||||
@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")
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj")
|
||||
@patch.object(ContextCachingEndpoints, "async_check_cache")
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
@patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client"
|
||||
)
|
||||
@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client")
|
||||
async def test_async_check_and_create_cache_timeout_error(
|
||||
self,
|
||||
mock_get_client,
|
||||
|
|
@ -594,9 +523,7 @@ class TestContextCachingEndpoints:
|
|||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
||||
# Mock timeout error
|
||||
self.mock_async_client.post = AsyncMock(
|
||||
side_effect=httpx.TimeoutException("Timeout")
|
||||
)
|
||||
self.mock_async_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
|
||||
|
||||
optional_params = self.sample_optional_params.copy()
|
||||
test_project = "test_project"
|
||||
|
|
@ -622,20 +549,14 @@ class TestContextCachingEndpoints:
|
|||
assert exc_info.value.status_code == 408
|
||||
assert "Timeout error occurred" in str(exc_info.value.message)
|
||||
|
||||
@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
|
||||
):
|
||||
@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"
|
||||
) as mock_separate:
|
||||
# Mock to return cached messages so tools get popped
|
||||
cached_messages = [
|
||||
self.sample_messages[0]
|
||||
] # System message with cache_control
|
||||
cached_messages = [self.sample_messages[0]] # System message with cache_control
|
||||
non_cached_messages = [self.sample_messages[1]] # User message
|
||||
mock_separate.return_value = (cached_messages, non_cached_messages)
|
||||
|
||||
|
|
@ -645,9 +566,7 @@ class TestContextCachingEndpoints:
|
|||
test_location = "test_location"
|
||||
|
||||
# Mock the check_cache to return existing cache so we don't make HTTP calls
|
||||
with patch.object(
|
||||
self.context_caching, "check_cache", return_value="existing_cache"
|
||||
):
|
||||
with patch.object(self.context_caching, "check_cache", return_value="existing_cache"):
|
||||
# Execute
|
||||
result = self.context_caching.check_and_create_cache(
|
||||
messages=self.sample_messages,
|
||||
|
|
@ -670,12 +589,8 @@ class TestContextCachingEndpoints:
|
|||
# But original tools should still be available for comparison
|
||||
assert original_tools == self.sample_tools
|
||||
|
||||
@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
|
||||
):
|
||||
@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"
|
||||
|
|
@ -711,12 +626,8 @@ 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, custom_llm_provider
|
||||
):
|
||||
@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, custom_llm_provider):
|
||||
"""Test that tools are NOT popped from optional_params in async version when there are no cached messages"""
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
|
||||
|
|
@ -752,20 +663,14 @@ 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_popped_from_optional_params(
|
||||
self, custom_llm_provider
|
||||
):
|
||||
@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"
|
||||
) as mock_separate:
|
||||
# Mock to return cached messages so tools get popped
|
||||
cached_messages = [
|
||||
self.sample_messages[0]
|
||||
] # System message with cache_control
|
||||
cached_messages = [self.sample_messages[0]] # System message with cache_control
|
||||
non_cached_messages = [self.sample_messages[1]] # User message
|
||||
mock_separate.return_value = (cached_messages, non_cached_messages)
|
||||
|
||||
|
|
@ -775,9 +680,7 @@ class TestContextCachingEndpoints:
|
|||
test_location = "test_location"
|
||||
|
||||
# Mock the async_check_cache to return existing cache so we don't make HTTP calls
|
||||
with patch.object(
|
||||
self.context_caching, "async_check_cache", return_value="existing_cache"
|
||||
):
|
||||
with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"):
|
||||
# Execute
|
||||
result = await self.context_caching.async_check_and_create_cache(
|
||||
messages=self.sample_messages,
|
||||
|
|
@ -800,15 +703,137 @@ class TestContextCachingEndpoints:
|
|||
# But original tools should still be available for comparison
|
||||
assert original_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"
|
||||
)
|
||||
def test_check_and_create_cache_skips_when_below_min_tokens(
|
||||
self, mock_separate, custom_llm_provider
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
def test_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider):
|
||||
"""Test that tool_choice is popped from optional_params when there are cached messages.
|
||||
|
||||
Gemini rejects requests that include both cachedContent and toolConfig.
|
||||
tool_choice must be moved into the cached content alongside tools.
|
||||
"""
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
|
||||
) as mock_separate:
|
||||
cached_messages = [self.sample_messages[0]]
|
||||
non_cached_messages = [self.sample_messages[1]]
|
||||
mock_separate.return_value = (cached_messages, non_cached_messages)
|
||||
|
||||
optional_params = self.sample_optional_params.copy()
|
||||
optional_params["tool_choice"] = "auto"
|
||||
|
||||
with patch.object(self.context_caching, "check_cache", return_value="existing_cache"):
|
||||
result = self.context_caching.check_and_create_cache(
|
||||
messages=self.sample_messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
model="gemini-1.5-pro",
|
||||
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 "tool_choice" not in optional_params
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(self, custom_llm_provider):
|
||||
"""Test that tool_choice is NOT popped when there are no cached messages"""
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
|
||||
) as mock_separate:
|
||||
mock_separate.return_value = ([], self.sample_messages)
|
||||
|
||||
optional_params = self.sample_optional_params.copy()
|
||||
optional_params["tool_choice"] = "auto"
|
||||
|
||||
result = self.context_caching.check_and_create_cache(
|
||||
messages=self.sample_messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
model="gemini-1.5-pro",
|
||||
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 "tool_choice" in optional_params
|
||||
assert optional_params["tool_choice"] == "auto"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider):
|
||||
"""Test that tool_choice is 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"
|
||||
) as mock_separate:
|
||||
cached_messages = [self.sample_messages[0]]
|
||||
non_cached_messages = [self.sample_messages[1]]
|
||||
mock_separate.return_value = (cached_messages, non_cached_messages)
|
||||
|
||||
optional_params = self.sample_optional_params.copy()
|
||||
optional_params["tool_choice"] = "auto"
|
||||
|
||||
with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"):
|
||||
result = await self.context_caching.async_check_and_create_cache(
|
||||
messages=self.sample_messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
model="gemini-1.5-pro",
|
||||
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 "tool_choice" not in optional_params
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(
|
||||
self, custom_llm_provider
|
||||
):
|
||||
"""Test that tool_choice is NOT popped in async version when there are no cached messages"""
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
|
||||
) as mock_separate:
|
||||
mock_separate.return_value = ([], self.sample_messages)
|
||||
|
||||
optional_params = self.sample_optional_params.copy()
|
||||
optional_params["tool_choice"] = "auto"
|
||||
|
||||
result = await self.context_caching.async_check_and_create_cache(
|
||||
messages=self.sample_messages,
|
||||
optional_params=optional_params,
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
model="gemini-1.5-pro",
|
||||
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 "tool_choice" in optional_params
|
||||
assert optional_params["tool_choice"] == "auto"
|
||||
|
||||
@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_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider):
|
||||
"""Test that context caching is skipped when cached content is below 1024 tokens.
|
||||
|
||||
Gemini requires a minimum of 1024 tokens for context caching. If the cached
|
||||
|
|
@ -855,16 +880,10 @@ class TestContextCachingEndpoints:
|
|||
# Restart the patcher so teardown_method can stop it cleanly
|
||||
self._token_check_patcher.start()
|
||||
|
||||
@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"
|
||||
)
|
||||
@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")
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_check_and_create_cache_skips_when_below_min_tokens(
|
||||
self, mock_separate, custom_llm_provider
|
||||
):
|
||||
async def test_async_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider):
|
||||
"""Test that async context caching is skipped when cached content is below 1024 tokens."""
|
||||
# Stop the default mock so the real token count check runs
|
||||
self._token_check_patcher.stop()
|
||||
|
|
@ -917,13 +936,9 @@ class TestCheckCachePagination:
|
|||
self.mock_client = MagicMock(spec=HTTPHandler)
|
||||
self.mock_async_client = MagicMock(spec=AsyncHTTPHandler)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_finds_cache_on_second_page(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
def test_check_cache_pagination_finds_cache_on_second_page(self, mock_get_token_url, custom_llm_provider):
|
||||
"""Test that check_cache correctly handles pagination and finds cache on second page"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
|
@ -972,13 +987,9 @@ class TestCheckCachePagination:
|
|||
second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"]
|
||||
assert "pageToken=token_page_2" in second_call_url
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_stops_when_no_next_token(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
def test_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider):
|
||||
"""Test that check_cache stops pagination when no nextPageToken is present"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
|
@ -1013,13 +1024,9 @@ class TestCheckCachePagination:
|
|||
assert result is None
|
||||
assert self.mock_client.get.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_multiple_pages(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
def test_check_cache_pagination_multiple_pages(self, mock_get_token_url, custom_llm_provider):
|
||||
"""Test that check_cache correctly iterates through multiple pages"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
|
@ -1064,9 +1071,7 @@ class TestCheckCachePagination:
|
|||
assert self.mock_client.get.call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_finds_cache_on_second_page(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
|
|
@ -1096,9 +1101,7 @@ class TestCheckCachePagination:
|
|||
}
|
||||
|
||||
# Setup mock async client to return different responses
|
||||
self.mock_async_client.get = AsyncMock(
|
||||
side_effect=[first_page_response, second_page_response]
|
||||
)
|
||||
self.mock_async_client.get = AsyncMock(side_effect=[first_page_response, second_page_response])
|
||||
|
||||
# Execute
|
||||
result = await self.context_caching.async_check_cache(
|
||||
|
|
@ -1122,13 +1125,9 @@ class TestCheckCachePagination:
|
|||
assert "pageToken=token_page_2" in second_call_url
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_stops_when_no_next_token(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
async def test_async_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider):
|
||||
"""Test that async_check_cache stops pagination when no nextPageToken is present"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
|
@ -1163,13 +1162,9 @@ class TestCheckCachePagination:
|
|||
assert result is None
|
||||
assert self.mock_async_client.get.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_max_pages_limit(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
def test_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider):
|
||||
"""Test that pagination stops after MAX_PAGINATION_PAGES iterations"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
|
@ -1179,17 +1174,13 @@ class TestCheckCachePagination:
|
|||
def create_page_response(page_num):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}
|
||||
],
|
||||
"cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}],
|
||||
"nextPageToken": f"token_page_{page_num + 1}",
|
||||
}
|
||||
return response
|
||||
|
||||
# Create MAX_PAGINATION_PAGES responses, each with a nextPageToken
|
||||
self.mock_client.get.side_effect = [
|
||||
create_page_response(i) for i in range(MAX_PAGINATION_PAGES)
|
||||
]
|
||||
self.mock_client.get.side_effect = [create_page_response(i) for i in range(MAX_PAGINATION_PAGES)]
|
||||
|
||||
# Execute
|
||||
result = self.context_caching.check_cache(
|
||||
|
|
@ -1211,13 +1202,9 @@ class TestCheckCachePagination:
|
|||
assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_max_pages_limit(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
async def test_async_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider):
|
||||
"""Test that async pagination stops after MAX_PAGINATION_PAGES iterations"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
|
|
@ -1227,9 +1214,7 @@ class TestCheckCachePagination:
|
|||
def create_page_response(page_num):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}
|
||||
],
|
||||
"cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}],
|
||||
"nextPageToken": f"token_page_{page_num + 1}",
|
||||
}
|
||||
return response
|
||||
|
|
@ -1267,7 +1252,9 @@ class TestVertexAIGlobalLocation:
|
|||
caching = ContextCachingEndpoints()
|
||||
|
||||
# Mock the _check_custom_proxy to return the URL unchanged
|
||||
with patch.object(caching, '_check_custom_proxy', side_effect=lambda **kwargs: (kwargs.get('auth_header'), kwargs.get('url'))):
|
||||
with patch.object(
|
||||
caching, "_check_custom_proxy", side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url"))
|
||||
):
|
||||
auth_header, url = caching._get_token_and_url_context_caching(
|
||||
gemini_api_key=None,
|
||||
custom_llm_provider="vertex_ai",
|
||||
|
|
@ -1286,7 +1273,9 @@ class TestVertexAIGlobalLocation:
|
|||
"""Test that regional location uses correct URL (with location prefix) for v1 API."""
|
||||
caching = ContextCachingEndpoints()
|
||||
|
||||
with patch.object(caching, '_check_custom_proxy', side_effect=lambda **kwargs: (kwargs.get('auth_header'), kwargs.get('url'))):
|
||||
with patch.object(
|
||||
caching, "_check_custom_proxy", side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url"))
|
||||
):
|
||||
auth_header, url = caching._get_token_and_url_context_caching(
|
||||
gemini_api_key=None,
|
||||
custom_llm_provider="vertex_ai",
|
||||
|
|
@ -1304,7 +1293,9 @@ class TestVertexAIGlobalLocation:
|
|||
"""Test that global location uses correct URL for v1beta1 API."""
|
||||
caching = ContextCachingEndpoints()
|
||||
|
||||
with patch.object(caching, '_check_custom_proxy', side_effect=lambda **kwargs: (kwargs.get('auth_header'), kwargs.get('url'))):
|
||||
with patch.object(
|
||||
caching, "_check_custom_proxy", side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url"))
|
||||
):
|
||||
auth_header, url = caching._get_token_and_url_context_caching(
|
||||
gemini_api_key=None,
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
|
|
@ -1315,7 +1306,9 @@ class TestVertexAIGlobalLocation:
|
|||
)
|
||||
|
||||
# Assert correct URL format for global with beta API
|
||||
expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents"
|
||||
expected_url = (
|
||||
"https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents"
|
||||
)
|
||||
assert url == expected_url, f"Expected {expected_url}, got {url}"
|
||||
assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix"
|
||||
|
||||
|
|
@ -1354,4 +1347,4 @@ class TestVertexAIGlobalLocation:
|
|||
)
|
||||
|
||||
assert "generativelanguage.googleapis.com" in url
|
||||
assert "cachedContents" in url
|
||||
assert "cachedContents" in url
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue