[Feat] Add litellm overhead metric for VertexAI (#15040)

* test_litellm_overhead

* vertex track overhead

* fix config.yaml used for testing

* test_litellm_overhead_stream

* add update_response_metadata for caching handler

* Revert "add update_response_metadata for caching handler"

This reverts commit f2a891f2b4.

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

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:

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()