mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
* test: move the unit half of 126 mixed legacy files into tests/unit * test: restore litellm globals that moved tests set * test: finalize migration test cleanup * test: restore original bodies of moved legacy tests The move into tests/unit had rewritten 612 test bodies, and some of the rewrites dropped assertions. Each moved test now carries its original body from the legacy file, with only the imports, helpers, fake provider credentials and monkeypatched env it needs to run under tests/unit test_timeout_streaming goes back to tests/local_testing because it needs the fake OpenAI endpoint server. The image payload fixture moves with its only user, and two tests that leaked global state (a registered model cost entry and queued logging tasks) are now isolated * test: drop module imports shadowed by restored local imports * test: assert on LiteLLM output in no-assertion moved tests and isolate leaks Twenty no-assertion candidates get one assertion on the value LiteLLM returns, with the original lines unchanged. Four tests go back to their legacy files because they only check types or imports, write into the working directory, or cannot assert without a body change Two moved tests leaked globals into later tests in the same worker, so monkeypatch fixtures now restore the retry-after header parser and the end user cost tracking flags * test: drain queued logging tasks before the Phoenix span test The moved Phoenix test counted spans from logging tasks that earlier tests had queued, so the drain fixture moves to tests/unit/conftest.py and both it and the Datadog batch test use it. test_factory_function goes back to its legacy file because its returned wrapper calls the real Assistants API and cannot be asserted on without a body change --------- Co-authored-by: yuneng <yuneng@berri.ai>
982 lines
29 KiB
Python
982 lines
29 KiB
Python
#### What this tests ####
|
|
# This tests litellm router
|
|
|
|
import asyncio
|
|
import os
|
|
import time
|
|
import traceback
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
from pydantic import BaseModel
|
|
|
|
import litellm
|
|
import litellm.types
|
|
import litellm.types.router
|
|
from litellm import Router
|
|
from litellm.types.router import DeploymentTypedDict
|
|
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
|
|
|
load_dotenv()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_router_deployment_typing():
|
|
deployment_typed_dict = DeploymentTypedDict(
|
|
model_name="hi", litellm_params={"model": "hello-world"}
|
|
)
|
|
for value in deployment_typed_dict.items():
|
|
assert not isinstance(value, BaseModel)
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_router_provider_wildcard_routing_regex():
|
|
"""
|
|
Pass list of orgs in 1 model definition,
|
|
expect a unique deployment for each to be created
|
|
"""
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "openai/fo::*:static::*",
|
|
"litellm_params": {
|
|
"model": "openai/fo::*:static::*",
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "openai/foo3::hello::*",
|
|
"litellm_params": {
|
|
"model": "openai/foo3::hello::*",
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
},
|
|
},
|
|
]
|
|
)
|
|
|
|
print("router model list = ", router.get_model_list())
|
|
|
|
response1 = await router.acompletion(
|
|
model="openai/fo::anything-can-be-here::static::anything-can-be-here",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
)
|
|
|
|
print("response 1 = ", response1)
|
|
|
|
response2 = await router.acompletion(
|
|
model="openai/foo3::hello::static::anything-can-be-here",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
)
|
|
|
|
print("response 2 = ", response2)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_router_sensitive_keys():
|
|
try:
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-3.5-turbo", # openai model name
|
|
"litellm_params": { # params for litellm completion/embedding call
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": "special-key",
|
|
},
|
|
"model_info": {"id": 12345},
|
|
},
|
|
],
|
|
)
|
|
except Exception as e:
|
|
print(f"error msg - {str(e)}")
|
|
if "special-key" in str(e):
|
|
pytest.fail("router error leaked the api key")
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [False, True])
|
|
@pytest.mark.asyncio
|
|
async def test_router_retries(sync_mode):
|
|
"""
|
|
- make sure retries work as expected
|
|
"""
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "bad-key"},
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-4.1-nano",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
]
|
|
|
|
router = Router(model_list=model_list, num_retries=2)
|
|
|
|
if sync_mode:
|
|
router.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "Hey, how's it going?"}],
|
|
)
|
|
else:
|
|
response = await router.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "Hey, how's it going?"}],
|
|
)
|
|
|
|
print(response.choices[0].message)
|
|
|
|
|
|
def test_exception_raising():
|
|
# this tests if the router raises an exception when invalid params are set
|
|
# in this test both deployments have bad keys - Keep this test. It validates if the router raises the most recent exception
|
|
litellm.set_verbose = True
|
|
import openai
|
|
|
|
try:
|
|
print("testing if router raises an exception")
|
|
old_api_key = os.environ["AZURE_AI_API_KEY"]
|
|
os.environ["AZURE_AI_API_KEY"] = ""
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo", # openai model name
|
|
"litellm_params": { # params for litellm completion/embedding call
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": "bad-key",
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
|
},
|
|
"tpm": 240000,
|
|
"rpm": 1800,
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo", # openai model name
|
|
"litellm_params": { #
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "bad-key",
|
|
},
|
|
"tpm": 240000,
|
|
"rpm": 1800,
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
redis_host=os.getenv("REDIS_HOST"),
|
|
redis_password=os.getenv("REDIS_PASSWORD"),
|
|
redis_port=int(os.getenv("REDIS_PORT")),
|
|
routing_strategy="simple-shuffle",
|
|
set_verbose=False,
|
|
num_retries=1,
|
|
) # type: ignore
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "hello this request will fail"}],
|
|
)
|
|
os.environ["AZURE_AI_API_KEY"] = old_api_key
|
|
pytest.fail(f"Should have raised an Auth Error")
|
|
except openai.AuthenticationError:
|
|
print(
|
|
"Test Passed: Caught an OPENAI AUTH Error, Good job. This is what we needed!"
|
|
)
|
|
os.environ["AZURE_AI_API_KEY"] = old_api_key
|
|
router.reset()
|
|
except Exception as e:
|
|
os.environ["AZURE_AI_API_KEY"] = old_api_key
|
|
print("Got unexpected exception on router!", e)
|
|
|
|
|
|
# test_exception_raising()
|
|
|
|
|
|
def test_reading_key_from_model_list():
|
|
# [PROD TEST CASE]
|
|
# this tests if the router can read key from model list and make completion call, and completion + stream call. This is 90% of the router use case
|
|
# DO NOT REMOVE THIS TEST. It's an IMP ONE. Speak to Ishaan, if you are tring to remove this
|
|
litellm.set_verbose = False
|
|
import openai
|
|
|
|
try:
|
|
print("testing if router raises an exception")
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo", # openai model name
|
|
"litellm_params": { # params for litellm completion/embedding call
|
|
"model": "gpt-4.1-nano",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
"tpm": 240000,
|
|
"rpm": 1800,
|
|
}
|
|
]
|
|
|
|
router = Router(
|
|
model_list=model_list,
|
|
redis_host=os.getenv("REDIS_HOST"),
|
|
redis_password=os.getenv("REDIS_PASSWORD"),
|
|
redis_port=int(os.getenv("REDIS_PORT")),
|
|
routing_strategy="simple-shuffle",
|
|
set_verbose=True,
|
|
num_retries=1,
|
|
) # type: ignore
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "hello this request will fail"}],
|
|
)
|
|
print("\n response", response)
|
|
str_response = response.choices[0].message.content
|
|
print("\n str_response", str_response)
|
|
assert len(str_response) > 0
|
|
|
|
print("\n Testing streaming response")
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "hello this request will fail"}],
|
|
stream=True,
|
|
)
|
|
completed_response = ""
|
|
for chunk in response:
|
|
if chunk is not None:
|
|
print(chunk)
|
|
completed_response += chunk.choices[0].delta.content or ""
|
|
print("\n completed_response", completed_response)
|
|
assert len(completed_response) > 0
|
|
print("\n Passed Streaming")
|
|
router.reset()
|
|
except Exception as e:
|
|
print(f"FAILED TEST")
|
|
pytest.fail(f"Got unexpected exception on router! - {e}")
|
|
|
|
|
|
# test_reading_key_from_model_list()
|
|
|
|
|
|
def test_call_one_endpoint():
|
|
# [PROD TEST CASE]
|
|
# user passes one deployment they want to call on the router, we call the specified one
|
|
# this test makes a completion calls azure/gpt-4.1-mini, it should work
|
|
try:
|
|
print("Testing calling a specific deployment")
|
|
old_api_key = os.environ["AZURE_AI_API_KEY"]
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo", # openai model name
|
|
"litellm_params": { # params for litellm completion/embedding call
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": old_api_key,
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
|
},
|
|
"tpm": 240000,
|
|
"rpm": 1800,
|
|
},
|
|
{
|
|
"model_name": "text-embedding-ada-002",
|
|
"litellm_params": {
|
|
"model": "azure/text-embedding-ada-002",
|
|
"api_key": os.environ["AZURE_AI_API_KEY"],
|
|
"api_base": os.environ["AZURE_AI_API_BASE"],
|
|
},
|
|
"tpm": 100000,
|
|
"rpm": 10000,
|
|
},
|
|
]
|
|
litellm.set_verbose = True
|
|
router = Router(
|
|
model_list=model_list,
|
|
routing_strategy="simple-shuffle",
|
|
set_verbose=True,
|
|
num_retries=1,
|
|
) # type: ignore
|
|
old_api_base = os.environ.pop("AZURE_AI_API_BASE", None)
|
|
|
|
async def call_azure_completion():
|
|
response = await router.acompletion(
|
|
model="azure/gpt-4.1-mini",
|
|
messages=[{"role": "user", "content": "hello this request will pass"}],
|
|
specific_deployment=True,
|
|
)
|
|
print("\n response", response)
|
|
|
|
async def call_azure_embedding():
|
|
response = await router.aembedding(
|
|
model="azure/text-embedding-ada-002",
|
|
input=["good morning from litellm"],
|
|
specific_deployment=True,
|
|
)
|
|
|
|
print("\n response", response)
|
|
|
|
asyncio.run(call_azure_completion())
|
|
asyncio.run(call_azure_embedding())
|
|
|
|
os.environ["AZURE_AI_API_BASE"] = old_api_base
|
|
os.environ["AZURE_AI_API_KEY"] = old_api_key
|
|
except Exception as e:
|
|
print(f"FAILED TEST")
|
|
pytest.fail(f"Got unexpected exception on router! - {e}")
|
|
|
|
|
|
# test_call_one_endpoint()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
async def test_async_router_context_window_fallback(sync_mode):
|
|
"""
|
|
- Give a gpt-4 model group with different context windows (8192k vs. 128k)
|
|
- Send a 10k prompt
|
|
- Assert it works
|
|
"""
|
|
import os
|
|
|
|
from large_text import text
|
|
|
|
litellm.set_verbose = False
|
|
litellm.turn_on_debug()
|
|
|
|
print(f"len(text): {len(text)}")
|
|
try:
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-4", # openai model name
|
|
"litellm_params": { # params for litellm completion/embedding call
|
|
"model": "gpt-4",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"api_base": os.getenv("OPENAI_API_BASE"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-4-turbo", # openai model name
|
|
"litellm_params": { # params for litellm completion/embedding call
|
|
"model": "gpt-4-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
]
|
|
|
|
router = Router(model_list=model_list, set_verbose=True, context_window_fallbacks=[{"gpt-4": ["gpt-4-turbo"]}], num_retries=0) # type: ignore
|
|
if sync_mode is False:
|
|
response = await router.acompletion(
|
|
model="gpt-4",
|
|
messages=[
|
|
{"role": "system", "content": text * 2},
|
|
{"role": "user", "content": "Who was Alexander?"},
|
|
],
|
|
)
|
|
|
|
print(f"response: {response}")
|
|
assert "gpt-4-turbo" in response.model
|
|
else:
|
|
response = router.completion(
|
|
model="gpt-4",
|
|
messages=[
|
|
{"role": "system", "content": text * 2},
|
|
{"role": "user", "content": "Who was Alexander?"},
|
|
],
|
|
)
|
|
assert "gpt-4-turbo" in response.model
|
|
except Exception as e:
|
|
pytest.fail(f"Got unexpected exception on router! - {str(e)}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
### FUNCTION CALLING
|
|
|
|
|
|
def test_function_calling():
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
"tpm": 100000,
|
|
"rpm": 10000,
|
|
},
|
|
]
|
|
|
|
messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
|
|
functions = [
|
|
{
|
|
"name": "get_current_weather",
|
|
"description": "Get the current weather in a given location",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "The city and state, e.g. San Francisco, CA",
|
|
},
|
|
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo", messages=messages, functions=functions
|
|
)
|
|
router.reset()
|
|
print(response)
|
|
|
|
|
|
# test_acompletion_on_router()
|
|
|
|
|
|
def test_function_calling_on_router():
|
|
try:
|
|
litellm.set_verbose = True
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
]
|
|
function1 = [
|
|
{
|
|
"name": "get_current_weather",
|
|
"description": "Get the current weather in a given location",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "The city and state, e.g. San Francisco, CA",
|
|
},
|
|
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
}
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
redis_host=os.getenv("REDIS_HOST"),
|
|
redis_password=os.getenv("REDIS_PASSWORD"),
|
|
redis_port=os.getenv("REDIS_PORT"),
|
|
)
|
|
messages = [{"role": "user", "content": "what's the weather in boston"}]
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo", messages=messages, functions=function1
|
|
)
|
|
print(f"final returned response: {response}")
|
|
router.reset()
|
|
assert isinstance(response["choices"][0]["message"]["function_call"], dict)
|
|
except Exception as e:
|
|
print(f"An exception occurred: {e}")
|
|
|
|
|
|
# test_function_calling_on_router()
|
|
|
|
|
|
### IMAGE GENERATION
|
|
# asyncio.run(test_aimg_gen_on_router())
|
|
|
|
|
|
# test_img_gen_on_router()
|
|
###
|
|
|
|
|
|
# test_aembedding_on_router()
|
|
|
|
|
|
# test_azure_embedding_on_router()
|
|
|
|
|
|
# test_bedrock_on_router()
|
|
|
|
|
|
# test openai-compatible endpoint
|
|
# asyncio.run(test_mistral_on_router())
|
|
|
|
|
|
def test_openai_completion_on_router():
|
|
# [PROD Use Case] - Makes an acompletion call + async acompletion call, and sync acompletion call, sync completion + stream
|
|
# 4 LLM API calls made here. If it fails, add retries. Do not remove this test.
|
|
litellm.set_verbose = True
|
|
print("\n Testing OpenAI on router\n")
|
|
try:
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
|
|
async def test():
|
|
response = await router.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "hello from litellm test",
|
|
}
|
|
],
|
|
)
|
|
print(response)
|
|
assert len(response.choices[0].message.content) > 0
|
|
|
|
print("\n streaming + acompletion test")
|
|
response = await router.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": f"hello from litellm test {time.time()}",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
complete_response = ""
|
|
print(response)
|
|
# if you want to see all the attributes and methods
|
|
async for chunk in response:
|
|
print(chunk)
|
|
complete_response += chunk.choices[0].delta.content or ""
|
|
print("\n complete response: ", complete_response)
|
|
assert len(complete_response) > 0
|
|
|
|
asyncio.run(test())
|
|
print("\n Testing Sync completion calls \n")
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "hello from litellm test2",
|
|
}
|
|
],
|
|
)
|
|
print(response)
|
|
assert len(response.choices[0].message.content) > 0
|
|
|
|
print("\n streaming + completion test")
|
|
response = router.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "hello from litellm test3",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
complete_response = ""
|
|
print(response)
|
|
for chunk in response:
|
|
print(chunk)
|
|
complete_response += chunk.choices[0].delta.content or ""
|
|
print("\n complete response: ", complete_response)
|
|
assert len(complete_response) > 0
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
# test_openai_completion_on_router()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# test_reading_keys_os_environ()
|
|
|
|
|
|
|
|
|
|
# test_reading_openai_keys_os_environ()
|
|
|
|
|
|
def test_router_anthropic_key_dynamic():
|
|
anthropic_api_key = os.environ.pop("ANTHROPIC_API_KEY")
|
|
model_list = [
|
|
{
|
|
"model_name": "anthropic-claude",
|
|
"litellm_params": {
|
|
"model": os.environ.get(
|
|
"CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"
|
|
),
|
|
"api_key": anthropic_api_key,
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
messages = [{"role": "user", "content": "Hey, how's it going?"}]
|
|
router.completion(model="anthropic-claude", messages=messages)
|
|
os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
|
|
|
|
|
|
def test_router_timeout():
|
|
litellm.set_verbose = True
|
|
import logging
|
|
|
|
from litellm._logging import verbose_logger
|
|
|
|
verbose_logger.setLevel(logging.DEBUG)
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "openai/slow-endpoint",
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
router = Router(model_list=model_list)
|
|
messages = [{"role": "user", "content": "Hey, how's it going?"}]
|
|
start_time = time.time()
|
|
try:
|
|
res = router.completion(
|
|
model="gpt-3.5-turbo", messages=messages, timeout=0.5
|
|
)
|
|
print(res)
|
|
pytest.fail("this should have timed out")
|
|
except litellm.exceptions.Timeout as e:
|
|
print("got timeout exception")
|
|
print(e)
|
|
print(vars(e))
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_amoderation():
|
|
model_list = [
|
|
{
|
|
"model_name": "openai-moderations",
|
|
"litellm_params": {
|
|
"model": "omni-moderation-latest",
|
|
"api_key": os.getenv("OPENAI_API_KEY", None),
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
## Test 1: user facing function
|
|
result = await router.amoderation(
|
|
model="omni-moderation-latest", input="this is valid good text"
|
|
)
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_text_completion_client():
|
|
# This tests if we re-use the Async OpenAI client
|
|
# This test fails when we create a new Async OpenAI client per request
|
|
try:
|
|
model_list = [
|
|
{
|
|
"model_name": "fake-openai-endpoint",
|
|
"litellm_params": {
|
|
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
|
|
"api_key": os.getenv("OPENAI_API_KEY", None),
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
},
|
|
}
|
|
]
|
|
router = Router(model_list=model_list, debug_level="DEBUG", set_verbose=True)
|
|
tasks = []
|
|
for _ in range(300):
|
|
tasks.append(
|
|
router.atext_completion(
|
|
model="fake-openai-endpoint",
|
|
prompt="hello from litellm test",
|
|
)
|
|
)
|
|
|
|
# Execute all coroutines concurrently
|
|
responses = await asyncio.gather(*tasks)
|
|
print(responses)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_response() -> litellm.ModelResponse:
|
|
return litellm.ModelResponse(
|
|
**{
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1699896916,
|
|
"model": "gpt-3.5-turbo-0125",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_abc123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_weather",
|
|
"arguments": '{\n"location": "Boston, MA"\n}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
"logprobs": None,
|
|
"finish_reason": "tool_calls",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10},
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_model_usage(mock_response):
|
|
"""
|
|
Test if tracking used model tpm works as expected
|
|
"""
|
|
model = "my-fake-model"
|
|
model_tpm = 100
|
|
setattr(
|
|
mock_response,
|
|
"usage",
|
|
litellm.Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10),
|
|
)
|
|
|
|
print(f"mock_response: {mock_response}")
|
|
model_tpm = 100
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"tpm": model_tpm,
|
|
"mock_response": mock_response,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
allowed_fails = 1 # allow for changing b/w minutes
|
|
|
|
for _ in range(2):
|
|
try:
|
|
_ = await llm_router.acompletion(
|
|
model=model, messages=[{"role": "user", "content": "Hey!"}]
|
|
)
|
|
await asyncio.sleep(3)
|
|
|
|
initial_usage_tuple = await llm_router.get_model_group_usage(
|
|
model_group=model
|
|
)
|
|
initial_usage = initial_usage_tuple[0]
|
|
|
|
# completion call - 10 tokens
|
|
_ = await llm_router.acompletion(
|
|
model=model, messages=[{"role": "user", "content": "Hey!"}]
|
|
)
|
|
|
|
await asyncio.sleep(3)
|
|
updated_usage_tuple = await llm_router.get_model_group_usage(
|
|
model_group=model
|
|
)
|
|
updated_usage = updated_usage_tuple[0]
|
|
|
|
assert updated_usage == initial_usage + 10 # type: ignore
|
|
break
|
|
except Exception as e:
|
|
if allowed_fails > 0:
|
|
print(
|
|
f"Decrementing allowed_fails: {allowed_fails}.\nReceived error - {str(e)}"
|
|
)
|
|
allowed_fails -= 1
|
|
else:
|
|
print(f"allowed_fails: {allowed_fails}")
|
|
raise e
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_router_cooldown_api_connection_error():
|
|
from litellm.router_utils.cooldown_handlers import _is_cooldown_required
|
|
|
|
with pytest.raises(litellm.APIConnectionError) as exc_info:
|
|
_ = litellm.completion(
|
|
model="vertex_ai/gemini-1.5-pro",
|
|
messages=[{"role": "admin", "content": "Fail on this!"}],
|
|
)
|
|
e = exc_info.value
|
|
assert (
|
|
_is_cooldown_required(
|
|
litellm_router_instance=Router(),
|
|
model_id="",
|
|
exception_status=e.code,
|
|
exception_str=str(e),
|
|
)
|
|
is False
|
|
)
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gemini-1.5-pro",
|
|
"litellm_params": {"model": "vertex_ai/gemini-1.5-pro"},
|
|
}
|
|
]
|
|
)
|
|
|
|
try:
|
|
router.completion(
|
|
model="gemini-1.5-pro",
|
|
messages=[{"role": "admin", "content": "Fail on this!"}],
|
|
)
|
|
except litellm.APIConnectionError:
|
|
pass
|
|
|
|
|
|
def test_router_correctly_reraise_error():
|
|
"""
|
|
User feedback: There is a problem with my messages array, but the error exception thrown is a Rate Limit error.
|
|
```
|
|
Rate Limit: Error code: 429 - {'error': {'message': 'No deployments available for selected model, Try again in 60 seconds. Passed model=gemini-2.5-flash-lite..
|
|
```
|
|
What they want? Propagation of the real error.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gemini-1.5-pro",
|
|
"litellm_params": {
|
|
"model": "vertex_ai/gemini-1.5-pro",
|
|
"mock_response": "litellm.RateLimitError",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
try:
|
|
router.completion(
|
|
model="gemini-1.5-pro",
|
|
messages=[{"role": "admin", "content": "Fail on this!"}],
|
|
)
|
|
except litellm.RateLimitError:
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# @pytest.mark.parametrize("on_error", [True, False])
|
|
# @pytest.mark.asyncio
|
|
# async def test_router_response_headers(on_error):
|
|
# router = Router(
|
|
# model_list=[
|
|
# {
|
|
# "model_name": "gpt-3.5-turbo",
|
|
# "litellm_params": {
|
|
# "model": "azure/gpt-4.1-mini",
|
|
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
|
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
|
# "tpm": 100000,
|
|
# "rpm": 100000,
|
|
# },
|
|
# },
|
|
# {
|
|
# "model_name": "gpt-3.5-turbo",
|
|
# "litellm_params": {
|
|
# "model": "azure/gpt-4.1-mini",
|
|
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
|
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
|
# "tpm": 500,
|
|
# "rpm": 500,
|
|
# },
|
|
# },
|
|
# ]
|
|
# )
|
|
|
|
# response = await router.acompletion(
|
|
# model="gpt-3.5-turbo",
|
|
# messages=[{"role": "user", "content": "Hello world!"}],
|
|
# mock_testing_rate_limit_error=on_error,
|
|
# )
|
|
|
|
# response_headers = response._hidden_params["additional_headers"]
|
|
|
|
# print(response_headers)
|
|
|
|
# assert response_headers["x-ratelimit-limit-requests"] == 100500
|
|
# assert int(response_headers["x-ratelimit-remaining-requests"]) > 0
|
|
# assert response_headers["x-ratelimit-limit-tokens"] == 100500
|
|
# assert int(response_headers["x-ratelimit-remaining-tokens"]) > 0
|
|
|
|
|
|
def test_router_completion_with_model_id():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {"model": "gpt-3.5-turbo"},
|
|
"model_info": {"id": "123"},
|
|
}
|
|
]
|
|
)
|
|
|
|
with patch.object(
|
|
router, "routing_strategy_pre_call_checks"
|
|
) as mock_pre_call_checks:
|
|
router.completion(model="123", messages=[{"role": "user", "content": "hi"}])
|
|
mock_pre_call_checks.assert_not_called()
|