test_litellm_overhead

tests/litellm_utils_tests/test_litellm_overhead.py

@@ -20,12 +20,27 @@ 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):
+    """
+    - 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()
+    kwargs ={
+        "messages": [{"role": "user", "content": "Hello, world!"}],
+        "model": model
+    }
+    #########################################################
+    # Specific cases for models
+    #########################################################
+    if model == "vertex_ai/gemini-1.0-pro-vision-001":
+        kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/"
+        # warmup call for auth validation on vertex_ai models
+        await litellm.acompletion(**kwargs)
     if model == "openai/self_hosted":
         response = await litellm.acompletion(
             model=model,
@@ -34,9 +49,11 @@ async def test_litellm_overhead(model):
         )
     else:
         response = await litellm.acompletion(
-            model=model,
-            messages=[{"role": "user", "content": "Hello, world!"}],
+            **kwargs
         )
+    #########################################################
+    # End of specific cases for models
+    #########################################################
     end_time = datetime.now()
     total_time_ms = (end_time - start_time).total_seconds() * 1000
     print(response)