scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

style: apply black formatting to changed files

Browse source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
flex-myeonghyeon 2026-04-13 14:44:44 +09:00
parent e7cece2dc3
commit e9792a2d88
2 changed files with 33 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -93,9 +93,9 @@ 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(

38
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -203,7 +203,10 @@ class TestContextCachingEndpoints:
# Verify cache key was generated with tools, tool_choice, and model
mock_cache_obj.get_cache_key.assert_called_once_with(
messages=cached_messages, tools=self.sample_tools, tool_choice=None, model="gemini-1.5-pro"
messages=cached_messages,
tools=self.sample_tools,
tool_choice=None,
model="gemini-1.5-pro",
)
@pytest.mark.parametrize(
@ -477,7 +480,10 @@ class TestContextCachingEndpoints:
# Verify cache key was generated with tools, tool_choice, and model
mock_cache_obj.get_cache_key.assert_called_once_with(
messages=cached_messages, tools=self.sample_tools, tool_choice=None, model="gemini-1.5-pro"
messages=cached_messages,
tools=self.sample_tools,
tool_choice=None,
model="gemini-1.5-pro",
)
@pytest.mark.asyncio
@ -1428,7 +1434,11 @@ 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",
@ -1441,13 +1451,19 @@ class TestVertexAIGlobalLocation:
# Assert correct URL format for global
expected_url = "https://aiplatform.googleapis.com/v1/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"
assert (
"global-aiplatform" not in url
), "URL should not contain 'global-aiplatform' prefix"
def test_regional_location_url_construction_v1(self):
"""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",
@ -1465,7 +1481,11 @@ 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",
@ -1478,7 +1498,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"
assert url == expected_url, f"Expected {expected_url}, got {url}"
assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix"
assert (
"global-aiplatform" not in url
), "URL should not contain 'global-aiplatform' prefix"
def test_gemini_context_caching_with_custom_api_base_passes_model(self):
"""Gemini context caching with custom api_base must pass model to _check_custom_proxy.
@ -1515,4 +1537,4 @@ class TestVertexAIGlobalLocation:
)
assert "generativelanguage.googleapis.com" in url
assert "cachedContents" in url
assert "cachedContents" in url