mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fa2c8984ba
* 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>
506 lines
19 KiB
Python
506 lines
19 KiB
Python
# Tests for router.get_available_deployment
|
|
# specifically test if it can pick the correct LLM when rpm/tpm set
|
|
# These are fast Tests, and make no API calls
|
|
import os
|
|
import traceback
|
|
from collections import defaultdict
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
|
|
import litellm
|
|
from litellm import Router
|
|
|
|
load_dotenv()
|
|
|
|
def test_weighted_selection_router():
|
|
# this tests if load balancing works based on the provided rpms in the router
|
|
# it's a fast test, only tests get_available_deployment
|
|
# users can pass rpms as a litellm_param
|
|
try:
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"rpm": 6,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"rpm": 1440,
|
|
},
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
)
|
|
selection_counts = defaultdict(int)
|
|
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
for _ in range(1000):
|
|
selected_model = router.get_available_deployment("gpt-3.5-turbo")
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["azure/gpt-4.1-mini"] / total_requests > 0.89
|
|
), f"Assertion failed: 'azure/gpt-4.1-mini' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_weighted_selection_router()
|
|
|
|
def test_weighted_selection_router_tpm():
|
|
# this tests if load balancing works based on the provided tpms in the router
|
|
# it's a fast test, only tests get_available_deployment
|
|
# users can pass rpms as a litellm_param
|
|
try:
|
|
print("\ntest weighted selection based on TPM\n")
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"tpm": 5,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"tpm": 90,
|
|
},
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
)
|
|
selection_counts = defaultdict(int)
|
|
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
for _ in range(1000):
|
|
selected_model = router.get_available_deployment("gpt-3.5-turbo")
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["azure/gpt-4.1-mini"] / total_requests > 0.89
|
|
), f"Assertion failed: 'azure/gpt-4.1-mini' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_weighted_selection_router_tpm()
|
|
|
|
def test_weighted_selection_router_tpm_as_router_param():
|
|
# this tests if load balancing works based on the provided tpms in the router
|
|
# it's a fast test, only tests get_available_deployment
|
|
# users can pass rpms as a litellm_param
|
|
try:
|
|
print("\ntest weighted selection based on TPM\n")
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
"tpm": 5,
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
},
|
|
"tpm": 90,
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
)
|
|
selection_counts = defaultdict(int)
|
|
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
for _ in range(1000):
|
|
selected_model = router.get_available_deployment("gpt-3.5-turbo")
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["azure/gpt-4.1-mini"] / total_requests > 0.89
|
|
), f"Assertion failed: 'azure/gpt-4.1-mini' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_weighted_selection_router_tpm_as_router_param()
|
|
|
|
def test_weighted_selection_router_rpm_as_router_param():
|
|
# this tests if load balancing works based on the provided tpms in the router
|
|
# it's a fast test, only tests get_available_deployment
|
|
# users can pass rpms as a litellm_param
|
|
try:
|
|
print("\ntest weighted selection based on RPM\n")
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
"rpm": 5,
|
|
"tpm": 5,
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
},
|
|
"rpm": 90,
|
|
"tpm": 90,
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
)
|
|
selection_counts = defaultdict(int)
|
|
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
for _ in range(1000):
|
|
selected_model = router.get_available_deployment("gpt-3.5-turbo")
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["azure/gpt-4.1-mini"] / total_requests > 0.89
|
|
), f"Assertion failed: 'azure/gpt-4.1-mini' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_weighted_selection_router_tpm_as_router_param()
|
|
|
|
def test_weighted_selection_router_no_rpm_set():
|
|
# this tests if we can do selection when no rpm is provided too
|
|
# it's a fast test, only tests get_available_deployment
|
|
# users can pass rpms as a litellm_param
|
|
try:
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"rpm": 6,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"rpm": 1440,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "claude-1",
|
|
"litellm_params": {
|
|
"model": "bedrock/claude1.2",
|
|
"rpm": 1440,
|
|
},
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
)
|
|
selection_counts = defaultdict(int)
|
|
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
for _ in range(1000):
|
|
selected_model = router.get_available_deployment("claude-1")
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["bedrock/claude1.2"] / total_requests == 1
|
|
), f"Assertion failed: Selection counts {selection_counts}"
|
|
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_weighted_selection_router_no_rpm_set()
|
|
|
|
def test_model_group_aliases():
|
|
try:
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"tpm": 1,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"tpm": 99,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "claude-1",
|
|
"litellm_params": {
|
|
"model": "bedrock/claude1.2",
|
|
"tpm": 1,
|
|
},
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
model_group_alias={
|
|
"gpt-4": "gpt-3.5-turbo"
|
|
}, # gpt-4 requests sent to gpt-3.5-turbo
|
|
)
|
|
|
|
# test that gpt-4 requests are sent to gpt-3.5-turbo
|
|
for _ in range(20):
|
|
selected_model = router.get_available_deployment("gpt-4")
|
|
print("\n selected model", selected_model)
|
|
selected_model_name = selected_model.get("model_name")
|
|
if selected_model_name != "gpt-3.5-turbo":
|
|
pytest.fail(
|
|
f"Selected model {selected_model_name} is not gpt-3.5-turbo"
|
|
)
|
|
|
|
# test that
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
selection_counts = defaultdict(int)
|
|
for _ in range(1000):
|
|
selected_model = router.get_available_deployment("gpt-3.5-turbo")
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["azure/gpt-4.1-mini"] / total_requests > 0.89
|
|
), f"Assertion failed: 'azure/gpt-4.1-mini' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_model_group_aliases()
|
|
|
|
@pytest.mark.flaky(retries=3, delay=2)
|
|
def test_usage_based_routing():
|
|
"""
|
|
in this test we, have a model group with two models in it, model-a and model-b.
|
|
Then at some point, we exceed the TPM limit (set in the litellm_params)
|
|
for model-a only; but for model-b we are still under the limit
|
|
"""
|
|
try:
|
|
|
|
def get_azure_params(deployment_name: str):
|
|
params = {
|
|
"model": f"azure/{deployment_name}",
|
|
"api_key": os.environ["AZURE_API_KEY"],
|
|
"api_version": os.environ["AZURE_API_VERSION"],
|
|
"api_base": "https://fake-api.openai.com/v1",
|
|
}
|
|
return params
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "azure/gpt-4",
|
|
"litellm_params": get_azure_params("chatgpt-low-tpm"),
|
|
"tpm": 100,
|
|
},
|
|
{
|
|
"model_name": "azure/gpt-4",
|
|
"litellm_params": get_azure_params("chatgpt-high-tpm"),
|
|
"tpm": 1000,
|
|
},
|
|
]
|
|
|
|
router = Router(
|
|
model_list=model_list,
|
|
set_verbose=True,
|
|
debug_level="DEBUG",
|
|
routing_strategy="usage-based-routing",
|
|
redis_host=os.environ["REDIS_HOST"],
|
|
redis_port=os.environ["REDIS_PORT"],
|
|
)
|
|
|
|
messages = [
|
|
{"content": "Tell me a joke.", "role": "user"},
|
|
]
|
|
|
|
selection_counts = defaultdict(int)
|
|
for _ in range(25):
|
|
response = router.completion(
|
|
model="azure/gpt-4",
|
|
messages=messages,
|
|
timeout=5,
|
|
mock_response="good morning",
|
|
)
|
|
|
|
# print("response", response)
|
|
|
|
selection_counts[response["model"]] += 1
|
|
|
|
print("selection counts", selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
# Assert that 'chatgpt-low-tpm' has more than 2 requests
|
|
assert (
|
|
selection_counts["chatgpt-low-tpm"] > 2
|
|
), f"Assertion failed: 'chatgpt-low-tpm' does not have more than 2 request in the weighted load balancer. Selection counts {selection_counts}"
|
|
|
|
# Assert that 'chatgpt-high-tpm' has about 70% of the total requests [DO NOT MAKE THIS LOWER THAN 70%]
|
|
assert (
|
|
selection_counts["chatgpt-high-tpm"] / total_requests > 0.70
|
|
), f"Assertion failed: 'chatgpt-high-tpm' does not have about 80% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
"""
|
|
Test async router get deployment (Simpl-shuffle)
|
|
"""
|
|
|
|
rpm_list = [[None, None], [6, 1440]]
|
|
tpm_list = [[None, None], [6, 1440]]
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"rpm_list, tpm_list",
|
|
[(rpm, tpm) for rpm in rpm_list for tpm in tpm_list],
|
|
)
|
|
async def test_weighted_selection_router_async(rpm_list, tpm_list):
|
|
# this tests if load balancing works based on the provided rpms in the router
|
|
# it's a fast test, only tests get_available_deployment
|
|
# users can pass rpms as a litellm_param
|
|
try:
|
|
litellm.set_verbose = False
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"rpm": rpm_list[0],
|
|
"tpm": tpm_list[0],
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"api_key": os.getenv("AZURE_API_KEY"),
|
|
"api_base": os.getenv("AZURE_API_BASE"),
|
|
"api_version": os.getenv("AZURE_API_VERSION"),
|
|
"rpm": rpm_list[1],
|
|
"tpm": tpm_list[1],
|
|
},
|
|
},
|
|
]
|
|
router = Router(
|
|
model_list=model_list,
|
|
)
|
|
selection_counts = defaultdict(int)
|
|
|
|
# call get_available_deployment 1k times, it should pick azure/gpt-4.1-mini about 90% of the time
|
|
for _ in range(1000):
|
|
selected_model = await router.async_get_available_deployment(
|
|
"gpt-3.5-turbo", request_kwargs={}
|
|
)
|
|
selected_model_id = selected_model["litellm_params"]["model"]
|
|
selected_model_name = selected_model_id
|
|
selection_counts[selected_model_name] += 1
|
|
print(selection_counts)
|
|
|
|
total_requests = sum(selection_counts.values())
|
|
|
|
if rpm_list[0] is not None or tpm_list[0] is not None:
|
|
# Assert that 'azure/gpt-4.1-mini' has about 90% of the total requests
|
|
assert (
|
|
selection_counts["azure/gpt-4.1-mini"] / total_requests > 0.89
|
|
), f"Assertion failed: 'azure/gpt-4.1-mini' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
|
|
else:
|
|
# Assert both are used
|
|
assert selection_counts["azure/gpt-4.1-mini"] > 0
|
|
assert selection_counts["gpt-3.5-turbo"] > 0
|
|
router.reset()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|