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

Litellm dev 06 23 2025 p1 (#11989)

Browse Source 
* fix(litellm_logging.py): fix using router model id for logging calls

Fixes https://github.com/BerriAI/litellm/issues/11975#issuecomment-2995882238

* test(test_litellm_logging.py): add unit test for custom price tracking

* fix(vertex_ai/): don't send invalid format parameter to vertex

causes calls to fail

* fix(vertex_ai_context_caching.py): if cached content present and tools in message, cache tools as well

gemini throws errors if tools passed in alongside cached content

* test: add unit tests

* fix: fix linting errors

* test: test_vertex_ai_common_utils.py

update test

* fix(streaming_handler.py): unset response cost when creating model response
This commit is contained in:
Krish Dholakia
2025-06-23 22:33:06 -07:00
committed by GitHub
parent 3f53c2eb18
commit a89397a798
12 changed files with 826 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/tutorials/claude_responses_api.md
+1 -1
View File
@@ -47,7 +47,7 @@ curl -X POST http://0.0.0.0:4000/v1/messages \
- Setup environment variables
```bash
export ANTHROPIC_API_BASE="http://0.0.0.0:4000"
export ANTHROPIC_BASE_URL="http://0.0.0.0:4000"
export ANTHROPIC_API_KEY="sk-1234" # replace with your LiteLLM key
```
litellm/cost_calculator.py
+3 -3
View File
@@ -666,9 +666,9 @@ def completion_cost( # noqa: PLR0915
or isinstance(completion_response, dict)
): # tts returns a custom class
if isinstance(completion_response, dict):
usage_obj: Optional[
Union[dict, Usage]
] = completion_response.get("usage", {})
usage_obj: Optional[Union[dict, Usage]] = (
completion_response.get("usage", {})
)
else:
usage_obj = getattr(completion_response, "usage", {})
if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects(
litellm/litellm_core_utils/litellm_logging.py
+12
View File
@@ -1083,6 +1083,18 @@ class Logging(LiteLLMLoggingBaseClass):
used for consistent cost calculation across response headers + logging integrations.
"""
if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"):
hidden_params = getattr(result, "_hidden_params", {})
if (
"response_cost" in hidden_params
and hidden_params["response_cost"] is not None
): # use cost if already calculated
return hidden_params["response_cost"]
elif (
router_model_id is None and "model_id" in hidden_params
): # use model_id if not already set
router_model_id = hidden_params["model_id"]
## RESPONSE COST ##
custom_pricing = use_custom_pricing_for_model(
litellm_params=(
litellm/litellm_core_utils/streaming_handler.py
+10 -9
View File
@@ -85,9 +85,9 @@ class CustomStreamWrapper:
self.system_fingerprint: Optional[str] = None
self.received_finish_reason: Optional[str] = None
self.intermittent_finish_reason: Optional[
str
] = None # finish reasons that show up mid-stream
self.intermittent_finish_reason: Optional[str] = (
None # finish reasons that show up mid-stream
)
self.special_tokens = [
"<|assistant|>",
"<|system|>",
@@ -643,6 +643,7 @@ class CustomStreamWrapper:
model_response._hidden_params = {
**model_response._hidden_params,
**self._hidden_params,
"response_cost": None,
}
if (
@@ -1322,9 +1323,9 @@ class CustomStreamWrapper:
_json_delta = delta.model_dump()
print_verbose(f"_json_delta: {_json_delta}")
if "role" not in _json_delta or _json_delta["role"] is None:
_json_delta[
"role"
] = "assistant" # mistral's api returns role as None
_json_delta["role"] = (
"assistant" # mistral's api returns role as None
)
if "tool_calls" in _json_delta and isinstance(
_json_delta["tool_calls"], list
):
@@ -1715,9 +1716,9 @@ class CustomStreamWrapper:
chunk = next(self.completion_stream)
if chunk is not None and chunk != b"":
print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}")
processed_chunk: Optional[
ModelResponseStream
] = self.chunk_creator(chunk=chunk)
processed_chunk: Optional[ModelResponseStream] = (
self.chunk_creator(chunk=chunk)
)
print_verbose(
f"PROCESSED CHUNK POST CHUNK CREATOR: {processed_chunk}"
)
litellm/llms/vertex_ai/common_utils.py
+6
View File
@@ -204,6 +204,7 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
add_object_type(parameters)
# Postprocessing
# Filter out fields that don't exist in Schema
parameters = filter_schema_fields(parameters, valid_schema_fields)
if add_property_ordering:
@@ -318,6 +319,11 @@ def filter_schema_fields(
k: filter_schema_fields(v, valid_fields, processed)
for k, v in value.items()
}
elif key == "format":
if value in {"enum", "date-time"}:
result[key] = value
else:
continue
elif key == "items" and isinstance(value, dict):
result[key] = filter_schema_fields(value, valid_fields, processed)
elif key == "anyOf" and isinstance(value, list):
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
+41 -18
View File
@@ -205,6 +205,7 @@ class ContextCachingEndpoints(VertexBase):
def check_and_create_cache(
self,
messages: List[AllMessageValues], # receives openai format messages
optional_params: dict, # cache the tools if present, in case cache content exists in messages
api_key: str,
api_base: Optional[str],
model: str,
@@ -213,7 +214,7 @@ class ContextCachingEndpoints(VertexBase):
logging_obj: Logging,
extra_headers: Optional[dict] = None,
cached_content: Optional[str] = None,
) -> Tuple[List[AllMessageValues], Optional[str]]:
) -> Tuple[List[AllMessageValues], dict, Optional[str]]:
"""
Receives
- messages: List of dict - messages in the openai format
@@ -225,7 +226,16 @@ class ContextCachingEndpoints(VertexBase):
Follows - https://ai.google.dev/api/caching#request-body
"""
if cached_content is not None:
return messages, cached_content
return messages, optional_params, cached_content
cached_messages, non_cached_messages = separate_cached_messages(
messages=messages
)
if len(cached_messages) == 0:
return messages, optional_params, None
tools = optional_params.pop("tools", None)
## AUTHORIZATION ##
token, url = self._get_token_and_url_context_caching(
@@ -252,15 +262,10 @@ class ContextCachingEndpoints(VertexBase):
else:
client = client
cached_messages, non_cached_messages = separate_cached_messages(
messages=messages
)
if len(cached_messages) == 0:
return messages, None
## CHECK IF CACHED ALREADY
generated_cache_key = local_cache_obj.get_cache_key(messages=cached_messages)
generated_cache_key = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools
)
google_cache_name = self.check_cache(
cache_key=generated_cache_key,
client=client,
@@ -270,7 +275,7 @@ class ContextCachingEndpoints(VertexBase):
logging_obj=logging_obj,
)
if google_cache_name:
return non_cached_messages, google_cache_name
return non_cached_messages, optional_params, google_cache_name
## TRANSFORM REQUEST
cached_content_request_body = (
@@ -279,6 +284,8 @@ class ContextCachingEndpoints(VertexBase):
)
)
cached_content_request_body["tools"] = tools
## LOGGING
logging_obj.pre_call(
input=messages,
@@ -305,11 +312,16 @@ class ContextCachingEndpoints(VertexBase):
cached_content_response_obj = VertexAICachedContentResponseObject(
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
)
return (non_cached_messages, cached_content_response_obj["name"])
return (
non_cached_messages,
optional_params,
cached_content_response_obj["name"],
)
async def async_check_and_create_cache(
self,
messages: List[AllMessageValues], # receives openai format messages
optional_params: dict, # cache the tools if present, in case cache content exists in messages
api_key: str,
api_base: Optional[str],
model: str,
@@ -318,7 +330,7 @@ class ContextCachingEndpoints(VertexBase):
logging_obj: Logging,
extra_headers: Optional[dict] = None,
cached_content: Optional[str] = None,
) -> Tuple[List[AllMessageValues], Optional[str]]:
) -> Tuple[List[AllMessageValues], dict, Optional[str]]:
"""
Receives
- messages: List of dict - messages in the openai format
@@ -330,14 +342,16 @@ class ContextCachingEndpoints(VertexBase):
Follows - https://ai.google.dev/api/caching#request-body
"""
if cached_content is not None:
return messages, cached_content
return messages, optional_params, cached_content
cached_messages, non_cached_messages = separate_cached_messages(
messages=messages
)
if len(cached_messages) == 0:
return messages, None
return messages, optional_params, None
tools = optional_params.pop("tools", None)
## AUTHORIZATION ##
token, url = self._get_token_and_url_context_caching(
@@ -362,7 +376,9 @@ class ContextCachingEndpoints(VertexBase):
client = client
## CHECK IF CACHED ALREADY
generated_cache_key = local_cache_obj.get_cache_key(messages=cached_messages)
generated_cache_key = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools
)
google_cache_name = await self.async_check_cache(
cache_key=generated_cache_key,
client=client,
@@ -371,8 +387,9 @@ class ContextCachingEndpoints(VertexBase):
api_base=api_base,
logging_obj=logging_obj,
)
if google_cache_name:
return non_cached_messages, google_cache_name
return non_cached_messages, optional_params, google_cache_name
## TRANSFORM REQUEST
cached_content_request_body = (
@@ -381,6 +398,8 @@ class ContextCachingEndpoints(VertexBase):
)
)
cached_content_request_body["tools"] = tools
## LOGGING
logging_obj.pre_call(
input=messages,
@@ -407,7 +426,11 @@ class ContextCachingEndpoints(VertexBase):
cached_content_response_obj = VertexAICachedContentResponseObject(
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
)
return (non_cached_messages, cached_content_response_obj["name"])
return (
non_cached_messages,
optional_params,
cached_content_response_obj["name"],
)
def get_cache(self):
pass
litellm/llms/vertex_ai/gemini/transformation.py
+16 -11
View File
@@ -1,5 +1,5 @@
"""
Transformation logic from OpenAI format to Gemini format.
Transformation logic from OpenAI format to Gemini format.
Why separate file? Make it easy to see how transformation works
"""
@@ -402,16 +402,19 @@ def sync_transform_request_body(
context_caching_endpoints = ContextCachingEndpoints()
if gemini_api_key is not None:
messages, cached_content = context_caching_endpoints.check_and_create_cache(
messages=messages,
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,
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 = optional_params.pop("cached_content", None)
@@ -446,9 +449,11 @@ async def async_transform_request_body(
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,
litellm/proxy/_new_secret_config.yaml
+2 -117
View File
@@ -1,119 +1,4 @@
model_list:
- model_name: codex-mini
- model_name: gemini-2.5-pro
litellm_params:
model: codex-mini-latest
api_key: os.environ/OPENAI_API_KEY
- model_name: bedrock/*
litellm_params:
model: bedrock/*
- model_name: eu.anthropic.claude-3-5-sonnet-20240620-v1:0
litellm_params:
model: eu.anthropic.claude-3-5-sonnet-20240620-v1:0
- model_name: "gpt-4o-mini-openai"
litellm_params:
model: gpt-4o-mini
api_key: os.environ/OPENAI_API_KEY
model_info:
access_groups: ["beta-models"] # 👈 Model Access Group
- model_name: azure_ai/Phi-3-medium
litellm_params:
model: azure_ai/Phi-3-medium
api_key: os.environ/AZURE_AI_PHI_3_MEDIUM_API_KEY
api_base: os.environ/AZURE_AI_PHI_3_MEDIUM_API_BASE
- model_name: "bedrock-nova"
litellm_params:
model: us.amazon.nova-pro-v1:0
- model_name: openrouter_model
litellm_params:
model: openrouter/openrouter_model
api_key: os.environ/OPENROUTER_API_KEY
api_base: http://0.0.0.0:8090
- model_name: dall-e-3-azure
litellm_params:
model: azure/dall-e-3-test
api_version: "2023-12-01-preview"
api_base: os.environ/AZURE_SWEDEN_API_BASE
api_key: os.environ/AZURE_SWEDEN_API_KEY
model_info:
input_cost_per_pixel: 10
- model_name: "claude-3-7-sonnet"
litellm_params:
model: databricks/databricks-claude-3-7-sonnet
api_key: os.environ/DATABRICKS_API_KEY
api_base: os.environ/DATABRICKS_API_BASE
- model_name: "gpt-4.1"
litellm_params:
model: azure/gpt-4.1
api_key: os.environ/AZURE_API_KEY_REALTIME
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/
- model_name: "xai/*"
litellm_params:
model: xai/*
api_key: os.environ/XAI_API_KEY
- model_name: "text-embedding-ada-002"
litellm_params:
model: text-embedding-ada-002
api_key: os.environ/OPENAI_API_KEY
- model_name: gemini/*
litellm_params:
model: gemini/*
- model_name: llama-qwen
litellm_params:
model: ollama/qwen2:0.5b
model_info:
input_cost_per_token: 0.75
output_cost_per_token: 3
- model_name: gpt-image-1
litellm_params:
model: gpt-image-1
api_key: os.environ/OPENAI_API_KEY
# drop_params: true
- model_name: "gpt-4o-batch"
litellm_params:
model: azure/gpt-4o-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
model_info:
id: my-general-azure-deployment
mode: batch
- model_name: "gpt-4o-batch"
litellm_params:
model: azure/gpt-4o-mini
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com
api_key: 04d22fb7e9ad4d9c8afe7c6abf97a6fc
model_info:
id: my-unique-azure-deployment
mode: batch
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
- model_name: "anthropic-claude-vertex"
litellm_params:
model: vertex_ai/claude-3-5-sonnet@20240620
vertex_project: internal-litellm-local-dev
- model_name: "openai-custom/*"
litellm_params:
model: "openai/*"
api_key: os.environ/OPENAI_API_KEY_TEST
- model_name: "anthropic-claude"
litellm_params:
model: "anthropic/claude-3-5-sonnet-latest"
api_key: os.environ/ANTHROPIC_API_KEY
general_settings:
store_model_in_db: true
store_prompts_in_spend_logs: true
token_rate_limit_type: "output"
# master_key: os.environ/PROXY_MASTER_KEY
litellm_settings:
# cache: true
# cache_params:
# type: redis
# ttl: 600
# password: os.environ/REDIS_PASSWORD
# supported_call_types: ["acompletion", "completion"]
callbacks: ["prometheus", "langfuse"]
model: gemini/gemini-2.5-pro
litellm/types/llms/vertex_ai.py
+1 -1
View File
@@ -89,7 +89,7 @@ class SystemInstructions(TypedDict):
class Schema(TypedDict, total=False):
type: Literal["STRING", "INTEGER", "BOOLEAN", "NUMBER", "ARRAY", "OBJECT"]
format: str
format: Literal["enum", "date-time"]
title: str
description: str
nullable: bool
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+36
View File
@@ -178,3 +178,39 @@ def test_get_request_tags():
assert "test-tag" in tags
assert "User-Agent: litellm" in tags
assert "User-Agent: litellm/0.1.0" in tags
def test_response_cost_calculator_with_response_cost_in_hidden_params(logging_obj):
from litellm import Router
from litellm.litellm_core_utils.litellm_logging import Logging
router = Router(
model_list=[
{
"model_name": "DeepSeek-R1",
"litellm_params": {
"model": "together_ai/deepseek-ai/DeepSeek-R1",
},
"model_info": {
"access_groups": ["agent-models"],
"supports_tool_choice": True,
"supports_function_calling": True,
"input_cost_per_token": 100,
"output_cost_per_token": 100,
},
}
]
)
mock_response = router.completion(
model="DeepSeek-R1",
messages=[{"role": "user", "content": "Hey"}],
mock_response="Hello, world!",
)
response_cost = logging_obj._response_cost_calculator(
result=mock_response,
)
assert response_cost is not None
assert response_cost > 100
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
+643
View File
@@ -0,0 +1,643 @@
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
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
ContextCachingEndpoints,
)
class TestContextCachingEndpoints:
"""Test class for ContextCachingEndpoints methods"""
def setup_method(self):
"""Setup for each test method"""
self.context_caching = ContextCachingEndpoints()
self.mock_logging = MagicMock(spec=Logging)
self.mock_client = MagicMock(spec=HTTPHandler)
self.mock_async_client = MagicMock(spec=AsyncHTTPHandler)
# Sample messages for testing
self.sample_messages = [
{
"role": "system",
"content": "You are a helpful assistant",
"cache_control": {"type": "ephemeral"},
},
{"role": "user", "content": "Hello, how are you?"},
]
# Sample tools for testing
self.sample_tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
},
},
}
]
self.sample_optional_params = {"tools": self.sample_tools.copy()}
@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
):
"""Test check_and_create_cache when cached_content is provided"""
# Setup
cached_content = "cached_content_123"
optional_params = self.sample_optional_params.copy()
# Execute
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,
cached_content=cached_content,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == self.sample_messages
assert returned_params == optional_params
assert returned_cache == cached_content
# Verify mocks weren't called since we short-circuited
mock_separate.assert_not_called()
mock_cache_obj.get_cache_key.assert_not_called()
@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):
"""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()
# Execute
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,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == self.sample_messages
assert returned_params == optional_params
assert returned_cache is None
@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
):
"""Test check_and_create_cache when existing cache is found"""
# Setup
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)
mock_cache_obj.get_cache_key.return_value = "test_cache_key"
mock_check_cache.return_value = "existing_cache_name"
optional_params = self.sample_optional_params.copy()
# Execute
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,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == non_cached_messages
assert returned_params == optional_params
assert returned_cache == "existing_cache_name"
# Verify cache key was generated with tools
mock_cache_obj.get_cache_key.assert_called_once_with(
messages=cached_messages, tools=self.sample_tools
)
@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, "check_cache")
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
def test_check_and_create_cache_create_new_cache(
self,
mock_get_token_url,
mock_check_cache,
mock_transform,
mock_cache_obj,
mock_separate,
):
"""Test check_and_create_cache when creating new cache"""
# Setup
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
mock_separate.return_value = (cached_messages, non_cached_messages)
mock_cache_obj.get_cache_key.return_value = "test_cache_key"
mock_check_cache.return_value = None # No existing cache
mock_get_token_url.return_value = ("token", "https://test-url.com")
mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
# Mock successful HTTP response
mock_response = MagicMock()
mock_response.json.return_value = {
"name": "new_cache_name",
"model": "gemini-1.5-pro",
}
self.mock_client.post.return_value = mock_response
optional_params = self.sample_optional_params.copy()
# Execute
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,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == non_cached_messages
assert returned_params == optional_params
assert returned_cache == "new_cache_name"
# Verify HTTP request was made
self.mock_client.post.assert_called_once()
call_args = self.mock_client.post.call_args
assert "tools" in call_args.kwargs["json"]
assert call_args.kwargs["json"]["tools"] == self.sample_tools
@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(
self, mock_get_token_url, mock_check_cache, mock_cache_obj, mock_separate
):
"""Test check_and_create_cache handles HTTP errors properly"""
# Setup
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
mock_separate.return_value = (cached_messages, non_cached_messages)
mock_cache_obj.get_cache_key.return_value = "test_cache_key"
mock_check_cache.return_value = None
mock_get_token_url.return_value = ("token", "https://test-url.com")
# Mock HTTP error
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = "Bad Request"
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()
# Execute and Assert
with pytest.raises(VertexAIError) as exc_info:
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,
)
assert exc_info.value.status_code == 400
assert "Bad Request" in str(exc_info.value.message)
@pytest.mark.asyncio
@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
):
"""Test async_check_and_create_cache when cached_content is provided"""
# Setup
cached_content = "cached_content_123"
optional_params = self.sample_optional_params.copy()
# Execute
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,
cached_content=cached_content,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == self.sample_messages
assert returned_params == optional_params
assert returned_cache == cached_content
@pytest.mark.asyncio
@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):
"""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()
# Execute
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,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == self.sample_messages
assert returned_params == optional_params
assert returned_cache is None
@pytest.mark.asyncio
@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
):
"""Test async_check_and_create_cache when existing cache is found"""
# Setup
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
mock_separate.return_value = (cached_messages, non_cached_messages)
mock_cache_obj.get_cache_key.return_value = "test_cache_key"
mock_async_check_cache.return_value = "existing_cache_name"
optional_params = self.sample_optional_params.copy()
# Execute
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,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == non_cached_messages
assert returned_params == optional_params
assert returned_cache == "existing_cache_name"
# Verify cache key was generated with tools
mock_cache_obj.get_cache_key.assert_called_once_with(
messages=cached_messages, tools=self.sample_tools
)
@pytest.mark.asyncio
@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"
)
async def test_async_check_and_create_cache_create_new_cache(
self,
mock_get_client,
mock_get_token_url,
mock_async_check_cache,
mock_transform,
mock_cache_obj,
mock_separate,
):
"""Test async_check_and_create_cache when creating new cache"""
# Setup
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
mock_separate.return_value = (cached_messages, non_cached_messages)
mock_cache_obj.get_cache_key.return_value = "test_cache_key"
mock_async_check_cache.return_value = None
mock_get_token_url.return_value = ("token", "https://test-url.com")
mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
# Mock successful HTTP response
mock_response = MagicMock()
mock_response.json.return_value = {
"name": "new_cache_name",
"model": "gemini-1.5-pro",
}
self.mock_async_client.post = AsyncMock(return_value=mock_response)
optional_params = self.sample_optional_params.copy()
# Execute
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,
)
# Assert
messages, returned_params, returned_cache = result
assert messages == non_cached_messages
assert returned_params == optional_params
assert returned_cache == "new_cache_name"
# Verify HTTP request was made
self.mock_async_client.post.assert_called_once()
call_args = self.mock_async_client.post.call_args
assert "tools" in call_args.kwargs["json"]
assert call_args.kwargs["json"]["tools"] == self.sample_tools
@pytest.mark.asyncio
@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"
)
async def test_async_check_and_create_cache_timeout_error(
self,
mock_get_client,
mock_get_token_url,
mock_async_check_cache,
mock_cache_obj,
mock_separate,
):
"""Test async_check_and_create_cache handles timeout errors properly"""
# Setup
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
mock_separate.return_value = (cached_messages, non_cached_messages)
mock_cache_obj.get_cache_key.return_value = "test_cache_key"
mock_async_check_cache.return_value = None
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")
)
optional_params = self.sample_optional_params.copy()
# Execute and Assert
with pytest.raises(VertexAIError) as exc_info:
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,
)
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):
"""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
non_cached_messages = [self.sample_messages[1]] # User message
mock_separate.return_value = (cached_messages, non_cached_messages)
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
# 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"
):
# Execute
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,
)
# Assert tools were popped from optional_params
assert "tools" not in optional_params
# 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):
"""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"
) as mock_separate:
mock_separate.return_value = (
[],
self.sample_messages,
) # No cached messages
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
# Execute
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,
)
# Assert tools were NOT popped from optional_params (early return)
assert "tools" in optional_params
assert optional_params["tools"] == original_tools
@pytest.mark.asyncio
async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(
self,
):
"""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"
) as mock_separate:
mock_separate.return_value = (
[],
self.sample_messages,
) # No cached messages
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
# Execute
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,
)
# Assert tools were NOT popped from optional_params (early return)
assert "tools" in optional_params
assert optional_params["tools"] == original_tools
@pytest.mark.asyncio
async def test_async_check_and_create_cache_tools_popped_from_optional_params(self):
"""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
non_cached_messages = [self.sample_messages[1]] # User message
mock_separate.return_value = (cached_messages, non_cached_messages)
optional_params = self.sample_optional_params.copy()
original_tools = optional_params["tools"].copy()
# 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"
):
# Execute
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,
)
# Assert tools were popped from optional_params
assert "tools" not in optional_params
# But original tools should still be available for comparison
assert original_tools == self.sample_tools
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
+55 -5
View File
@@ -237,11 +237,7 @@ def test_build_vertex_schema():
},
"recursion_limit": {"type": "integer"},
"configurable": {"type": "object"},
"run_id": {
"anyOf": [
{"format": "uuid", "type": "string", "nullable": True}
]
},
"run_id": {"anyOf": [{"type": "string", "nullable": True}]},
},
"type": "object",
},
@@ -627,3 +623,57 @@ def test_get_vertex_region_global_only_model(
assert result == expected_region
mock_is_global_only.assert_called_once_with("test-model")
def test_vertex_filter_format_uri():
import json
from litellm.llms.vertex_ai.common_utils import filter_schema_fields
parameters = {
"type": "object",
"properties": {
"url": {
"type": "string",
"format": "uri",
"description": "The URL to fetch content from",
},
"prompt": {
"type": "string",
"description": "The prompt to run on the fetched content",
},
},
"required": ["url", "prompt"],
"$schema": "http://json-schema.org/draft-07/schema#",
}
valid_schema_fields = {
"minLength",
"nullable",
"maxItems",
"required",
"default",
"items",
"propertyOrdering",
"maximum",
"properties",
"anyOf",
"description",
"minProperties",
"minimum",
"minItems",
"maxProperties",
"title",
"pattern",
"example",
"format",
"enum",
"maxLength",
"type",
}
new_parameters = filter_schema_fields(
schema_dict=parameters,
valid_fields=valid_schema_fields,
)
assert "uri" not in json.dumps(new_parameters)