diff --git a/docs/my-website/docs/tutorials/claude_responses_api.md b/docs/my-website/docs/tutorials/claude_responses_api.md index d95f75c7d7..9fe92cf11c 100644 --- a/docs/my-website/docs/tutorials/claude_responses_api.md +++ b/docs/my-website/docs/tutorials/claude_responses_api.md @@ -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 ``` diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 30658cb091..b55da01aff 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 72589d57e3..c608dfc8d8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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=( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 079cef4631..529e7c8399 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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}" ) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 23facabbf8..cceac0ea79 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -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): diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 5cfb9141a5..33a480aa6b 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -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 diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 39edb9642e..85e3f15364 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -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, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index f9c4fe15e1..07e0a96f55 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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 \ No newline at end of file diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index ded64cee6c..1c3eb49fb0 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -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 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 9d95749827..6687fbd0c1 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py new file mode 100644 index 0000000000..ef3404353a --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index b43e5a113c..9ebb9004b7 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -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)