diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py index 6ab5eaf6174..6e46e935463 100644 --- a/tests/litellm_utils_tests/test_litellm_overhead.py +++ b/tests/litellm_utils_tests/test_litellm_overhead.py @@ -23,7 +23,7 @@ import litellm "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 """ @@ -37,20 +37,15 @@ async def test_litellm_overhead(model): ######################################################### # Specific cases for models ######################################################### - if model == "vertex_ai/gemini-1.0-pro-vision-001": + 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) - 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( - **kwargs - ) + + + response = await litellm.acompletion( + **kwargs + ) ######################################################### # End of specific cases for models ######################################################### @@ -92,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()