diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index dc3a6cf15e5..9354ae6e67a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3,7 +3,6 @@ ## Initial implementation - covers gemini + image gen calls import json import time -from litellm._uuid import uuid from copy import deepcopy from functools import partial from typing import ( @@ -25,6 +24,7 @@ import litellm import litellm.litellm_core_utils import litellm.litellm_core_utils.litellm_logging from litellm import verbose_logger +from litellm._uuid import uuid from litellm.constants import ( DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, @@ -32,8 +32,8 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH, - DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE, + DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -1596,7 +1596,7 @@ async def make_call( ) try: - response = await client.post(api_base, headers=headers, data=data, stream=True) + response = await client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) response.raise_for_status() except httpx.HTTPStatusError as e: exception_string = str(await e.response.aread()) @@ -1643,7 +1643,7 @@ def make_sync_call( if client is None: client = HTTPHandler() # Create a new client if none provided - response = client.post(api_base, headers=headers, data=data, stream=True) + response = client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -1842,7 +1842,7 @@ class VertexLLM(VertexBase): try: response = await client.post( - api_base, headers=headers, json=cast(dict, request_body) + api_base, headers=headers, json=cast(dict, request_body), logging_obj=logging_obj ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: @@ -2045,7 +2045,7 @@ class VertexLLM(VertexBase): client = client try: - response = client.post(url=url, headers=headers, json=data) # type: ignore + response = client.post(url=url, headers=headers, json=data, logging_obj=logging_obj) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 5445ba3e2b2..60eef9604e1 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -23,6 +23,9 @@ model_list: litellm_params: model: gemini/* api_key: os.environ/GEMINI_API_KEY + - model_name: vertex_ai/* + litellm_params: + model: vertex_ai/* guardrails: diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py index 8d0bdf313dd..6e46e935463 100644 --- a/tests/litellm_utils_tests/test_litellm_overhead.py +++ b/tests/litellm_utils_tests/test_litellm_overhead.py @@ -20,23 +20,35 @@ import litellm "openai/gpt-4o", "openai/self_hosted", "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0", + "vertex_ai/gemini-1.0-pro-vision-001", ], ) -async def test_litellm_overhead(model): +async def test_litellm_overhead_non_streaming(model): + """ + - Test we can see the litellm overhead and that it is less than 40% of the total request time + """ litellm._turn_on_debug() start_time = datetime.now() - if model == "openai/self_hosted": - response = await litellm.acompletion( - model=model, - messages=[{"role": "user", "content": "Hello, world!"}], - api_base="https://exampleopenaiendpoint-production.up.railway.app/", - ) - else: - response = await litellm.acompletion( - model=model, - messages=[{"role": "user", "content": "Hello, world!"}], - ) + kwargs ={ + "messages": [{"role": "user", "content": "Hello, world!"}], + "model": model + } + ######################################################### + # Specific cases for models + ######################################################### + if model == "vertex_ai/gemini-1.0-pro-vision-001" or model == "openai/self_hosted": + kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/" + # warmup call for auth validation on vertex_ai models + await litellm.acompletion(**kwargs) + + + response = await litellm.acompletion( + **kwargs + ) + ######################################################### + # End of specific cases for models + ######################################################### end_time = datetime.now() total_time_ms = (end_time - start_time).total_seconds() * 1000 print(response) @@ -75,19 +87,22 @@ async def test_litellm_overhead_stream(model): litellm._turn_on_debug() start_time = datetime.now() + kwargs ={ + "messages": [{"role": "user", "content": "Hello, world!"}], + "model": model, + "stream": True, + } + ######################################################### + # Specific cases for models + ######################################################### if model == "openai/self_hosted": - response = await litellm.acompletion( - model=model, - messages=[{"role": "user", "content": "Hello, world!"}], - api_base="https://exampleopenaiendpoint-production.up.railway.app/", - stream=True, - ) - else: - response = await litellm.acompletion( - model=model, - messages=[{"role": "user", "content": "Hello, world!"}], - stream=True, - ) + kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/" + # warmup call for auth validation on vertex_ai models + await litellm.acompletion(**kwargs) + + response = await litellm.acompletion( + **kwargs + ) async for chunk in response: print()