litellm/tests/local_testing/test_lowest_latency_routing.py

#### What this tests ####
#    This tests the router's ability to pick deployment with lowest latency

import asyncio
import os
import random
import sys
import time
import traceback
from datetime import datetime, timedelta

from dotenv import load_dotenv

load_dotenv()
import copy
import os

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system path
import pytest

import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler

### UNIT TESTS FOR LATENCY ROUTING ###


@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_latency_memory_leak(sync_mode):
    """
    Test to make sure there's no memory leak caused by lowest latency routing

    - make 10 calls -> check memory
    - make 11th call -> no change in memory
    """
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
    model_group = "gpt-3.5-turbo"
    deployment_id = "1234"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 50}}
    time.sleep(5)
    end_time = time.time()
    for _ in range(10):
        if sync_mode:
            lowest_latency_logger.log_success_event(
                response_obj=response_obj,
                kwargs=kwargs,
                start_time=start_time,
                end_time=end_time,
            )
        else:
            await lowest_latency_logger.async_log_success_event(
                response_obj=response_obj,
                kwargs=kwargs,
                start_time=start_time,
                end_time=end_time,
            )
    latency_key = f"{model_group}_map"
    cache_value = copy.deepcopy(
        test_cache.get_cache(key=latency_key)
    )  # MAKE SURE NO MEMORY LEAK IN CACHING OBJECT

    if sync_mode:
        lowest_latency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )
    else:
        await lowest_latency_logger.async_log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )
    new_cache_value = test_cache.get_cache(key=latency_key)
    # Assert that the size of the cache doesn't grow unreasonably
    assert get_size(new_cache_value) <= get_size(
        cache_value
    ), f"Memory leak detected in function call! new_cache size={get_size(new_cache_value)}, old cache size={get_size(cache_value)}"


def get_size(obj, seen=None):
    # From https://goshippo.com/blog/measure-real-size-any-python-object/
    # Recursively finds size of objects
    size = sys.getsizeof(obj)
    if seen is None:
        seen = set()
    obj_id = id(obj)
    if obj_id in seen:
        return 0
    seen.add(obj_id)
    if isinstance(obj, dict):
        size += sum([get_size(v, seen) for v in obj.values()])
        size += sum([get_size(k, seen) for k in obj.keys()])
    elif hasattr(obj, "__dict__"):
        size += get_size(obj.__dict__, seen)
    elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes, bytearray)):
        size += sum([get_size(i, seen) for i in obj])
    return size


def test_latency_updated():
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
    model_group = "gpt-3.5-turbo"
    deployment_id = "1234"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 50}}
    time.sleep(5)
    end_time = time.time()
    lowest_latency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )
    latency_key = f"{model_group}_map"
    assert (
        end_time - start_time
        == test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0]
    )


# test_tpm_rpm_updated()


def test_latency_updated_custom_ttl():
    """
    Invalidate the cached request.

    Test that the cache is empty
    """
    test_cache = DualCache()
    model_list = []
    cache_time = 3
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"ttl": cache_time}
    )
    model_group = "gpt-3.5-turbo"
    deployment_id = "1234"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 50}}
    time.sleep(5)
    end_time = time.time()
    lowest_latency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )
    latency_key = f"{model_group}_map"
    print(f"cache: {test_cache.get_cache(key=latency_key)}")
    assert isinstance(test_cache.get_cache(key=latency_key), dict)
    time.sleep(cache_time)
    assert test_cache.get_cache(key=latency_key) is None


def test_get_available_deployments():
    test_cache = DualCache()
    model_list = [
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "azure/gpt-4.1-mini"},
            "model_info": {"id": "1234"},
        },
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "azure/gpt-4.1-mini"},
            "model_info": {"id": "5678"},
        },
    ]
    lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
    model_group = "gpt-3.5-turbo"
    ## DEPLOYMENT 1 ##
    deployment_id = "1234"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 50}}
    time.sleep(3)
    end_time = time.time()
    lowest_latency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )
    ## DEPLOYMENT 2 ##
    deployment_id = "5678"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 20}}
    time.sleep(2)
    end_time = time.time()
    lowest_latency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )

    ## CHECK WHAT'S SELECTED ##
    print(
        lowest_latency_logger.get_available_deployments(
            model_group=model_group, healthy_deployments=model_list
        )
    )
    assert (
        lowest_latency_logger.get_available_deployments(
            model_group=model_group, healthy_deployments=model_list
        )["model_info"]["id"]
        == "5678"
    )


async def _deploy(lowest_latency_logger, deployment_id, tokens_used, duration):
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": tokens_used}}
    await asyncio.sleep(duration)
    end_time = time.time()
    lowest_latency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )


async def _gather_deploy(all_deploys):
    return await asyncio.gather(*[_deploy(*t) for t in all_deploys])


@pytest.mark.parametrize(
    "ans_rpm", [1, 5]
)  # 1 should produce nothing, 10 should select first
@pytest.mark.flaky(retries=3, delay=1)
def test_get_available_endpoints_tpm_rpm_check_async(ans_rpm):
    """
    Pass in list of 2 valid models

    Update cache with 1 model clearly being at tpm/rpm limit

    assert that only the valid model is returned
    """
    test_cache = DualCache()
    ans = "1234"
    non_ans_rpm = 3
    assert ans_rpm != non_ans_rpm, "invalid test"
    if ans_rpm < non_ans_rpm:
        ans = None
    model_list = [
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "azure/gpt-4.1-mini"},
            "model_info": {"id": "1234", "rpm": ans_rpm},
        },
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "azure/gpt-4.1-mini"},
            "model_info": {"id": "5678", "rpm": non_ans_rpm},
        },
    ]
    lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
    model_group = "gpt-3.5-turbo"
    d1 = [(lowest_latency_logger, "1234", 50, 0.01)] * non_ans_rpm
    d2 = [(lowest_latency_logger, "5678", 50, 0.01)] * non_ans_rpm
    asyncio.run(_gather_deploy([*d1, *d2]))
    time.sleep(3)
    ## CHECK WHAT'S SELECTED ##
    d_ans = lowest_latency_logger.get_available_deployments(
        model_group=model_group, healthy_deployments=model_list
    )
    print(d_ans)
    assert (d_ans and d_ans["model_info"]["id"]) == ans


# test_get_available_endpoints_tpm_rpm_check_async()


@pytest.mark.parametrize(
    "ans_rpm", [1, 5]
)  # 1 should produce nothing, 10 should select first
@pytest.mark.flaky(retries=3, delay=1)
def test_get_available_endpoints_tpm_rpm_check(ans_rpm):
    """
    Pass in list of 2 valid models

    Update cache with 1 model clearly being at tpm/rpm limit

    assert that only the valid model is returned
    """
    test_cache = DualCache()
    ans = "1234"
    non_ans_rpm = 3
    assert ans_rpm != non_ans_rpm, "invalid test"
    if ans_rpm < non_ans_rpm:
        ans = None
    model_list = [
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "azure/gpt-4.1-mini"},
            "model_info": {"id": "1234", "rpm": ans_rpm},
        },
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "azure/gpt-4.1-mini"},
            "model_info": {"id": "5678", "rpm": non_ans_rpm},
        },
    ]
    lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
    model_group = "gpt-3.5-turbo"
    ## DEPLOYMENT 1 ##
    deployment_id = "1234"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    for _ in range(non_ans_rpm):
        start_time = time.time()
        response_obj = {"usage": {"total_tokens": 50}}
        time.sleep(0.01)
        end_time = time.time()
        lowest_latency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )
    ## DEPLOYMENT 2 ##
    deployment_id = "5678"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "gpt-3.5-turbo",
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }
    for _ in range(non_ans_rpm):
        start_time = time.time()
        response_obj = {"usage": {"total_tokens": 20}}
        time.sleep(0.5)
        end_time = time.time()
        lowest_latency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    ## CHECK WHAT'S SELECTED ##
    d_ans = lowest_latency_logger.get_available_deployments(
        model_group=model_group, healthy_deployments=model_list
    )
    print(d_ans)
    assert (d_ans and d_ans["model_info"]["id"]) == ans


def test_router_get_available_deployments():
    """
    Test if routers 'get_available_deployments' returns the fastest deployment
    """
    model_list = [
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-turbo",
                "api_key": "os.environ/AZURE_FRANCE_API_KEY",
                "api_base": "https://openai-france-1234.openai.azure.com",
                "rpm": 1440,
            },
            "model_info": {"id": 1},
        },
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-35-turbo",
                "api_key": "os.environ/AZURE_EUROPE_API_KEY",
                "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
                "rpm": 6,
            },
            "model_info": {"id": 2},
        },
    ]
    router = Router(
        model_list=model_list,
        routing_strategy="latency-based-routing",
        set_verbose=False,
        num_retries=3,
    )  # type: ignore

    ## DEPLOYMENT 1 ##
    deployment_id = 1
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "azure-model",
            },
            "model_info": {"id": 1},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 50}}
    time.sleep(3)
    end_time = time.time()
    router.lowestlatency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )
    ## DEPLOYMENT 2 ##
    deployment_id = 2
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "azure-model",
            },
            "model_info": {"id": 2},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 20}}
    time.sleep(2)
    end_time = time.time()
    router.lowestlatency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )

    ## CHECK WHAT'S SELECTED ##
    # print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
    print(router.get_available_deployment(model="azure-model"))
    assert (
        router.get_available_deployment(model="azure-model")["model_info"]["id"] == "2"
    )


# test_router_get_available_deployments()


@pytest.mark.asyncio
async def test_router_completion_streaming():
    messages = [
        {"role": "user", "content": "Hello, can you generate a 500 words poem?"}
    ]
    model = "azure-model"
    model_list = [
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-turbo",
                "api_key": "os.environ/AZURE_FRANCE_API_KEY",
                "api_base": "https://openai-france-1234.openai.azure.com",
                "rpm": 1440,
                "mock_response": "Hello world",
            },
            "model_info": {"id": 1},
        },
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-35-turbo",
                "api_key": "os.environ/AZURE_EUROPE_API_KEY",
                "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
                "rpm": 6,
                "mock_response": "Hello world",
            },
            "model_info": {"id": 2},
        },
    ]
    router = Router(
        model_list=model_list,
        routing_strategy="latency-based-routing",
        set_verbose=False,
        num_retries=3,
    )  # type: ignore

    ### Make 3 calls, test if 3rd call goes to fastest deployment

    ## CALL 1+2
    tasks = []
    response = None
    final_response = None
    for _ in range(2):
        tasks.append(router.acompletion(model=model, messages=messages))
    response = await asyncio.gather(*tasks)

    if response is not None:
        ## CALL 3
        await asyncio.sleep(1)  # let the cache update happen
        picked_deployment = router.lowestlatency_logger.get_available_deployments(
            model_group=model, healthy_deployments=router.healthy_deployments
        )
        final_response = await router.acompletion(model=model, messages=messages)
        print(f"min deployment id: {picked_deployment}")
        print(f"model id: {final_response._hidden_params['model_id']}")
        assert (
            final_response._hidden_params["model_id"]
            == picked_deployment["model_info"]["id"]
        )


# asyncio.run(test_router_completion_streaming())


@pytest.mark.asyncio
async def test_lowest_latency_routing_with_timeouts():
    """
    PROD Test:
    - Endpoint 1: triggers timeout errors (it takes 10+ seconds to respond)
    - Endpoint 2: Responds in under 1s
    - Run 5 requests to collect data on latency
    - Run Wait till cache is filled with data
    - Run 10 more requests
    - All requests should have been routed to endpoint 2
    """
    import litellm

    litellm.set_verbose = True

    router = Router(
        model_list=[
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/slow-endpoint",
                    "api_base": "https://exampleopenaiendpoint-production-c715.up.railway.app/",  # If you are Krrish, this is OpenAI Endpoint3 on our Railway endpoint :)
                    "api_key": "fake-key",
                },
                "model_info": {"id": "slow-endpoint"},
            },
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/fast-endpoint",
                    "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
                    "api_key": "fake-key",
                },
                "model_info": {"id": "fast-endpoint"},
            },
        ],
        routing_strategy="latency-based-routing",
        set_verbose=True,
        debug_level="DEBUG",
        timeout=1,
    )  # type: ignore

    # make 4 requests
    for _ in range(4):
        try:
            response = await router.acompletion(
                model="azure-model", messages=[{"role": "user", "content": "hello"}]
            )
            print(response)
        except Exception as e:
            print("got exception", e)

    await asyncio.sleep(1)
    print("done sending initial requests to collect latency")
    """
    Note: for debugging
    - By this point: slow-endpoint should have timed out 3-4 times and should be heavily penalized :)
    - The next 10 requests should all be routed to the fast-endpoint
    """

    deployments = {}
    # make 10 requests
    for _ in range(10):
        response = await router.acompletion(
            model="azure-model", messages=[{"role": "user", "content": "hello"}]
        )
        print(response)
        _picked_model_id = response._hidden_params["model_id"]
        if _picked_model_id not in deployments:
            deployments[_picked_model_id] = 1
        else:
            deployments[_picked_model_id] += 1
    print("deployments", deployments)

    # ALL the Requests should have been routed to the fast-endpoint
    assert deployments["fast-endpoint"] == 10


@pytest.mark.asyncio
async def test_lowest_latency_routing_first_pick():
    """
    PROD Test:
    - When all deployments are latency=0, it should randomly pick a deployment
    - IT SHOULD NEVER PICK THE Very First deployment everytime all deployment latencies are 0
    - This ensures that after the ttl window resets it randomly picks a deployment
    """
    import litellm

    litellm.set_verbose = True

    router = Router(
        model_list=[
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/fast-endpoint",
                    "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
                    "api_key": "fake-key",
                },
                "model_info": {"id": "fast-endpoint"},
            },
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/fast-endpoint-2",
                    "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
                    "api_key": "fake-key",
                },
                "model_info": {"id": "fast-endpoint-2"},
            },
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/fast-endpoint-2",
                    "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
                    "api_key": "fake-key",
                },
                "model_info": {"id": "fast-endpoint-3"},
            },
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/fast-endpoint-2",
                    "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
                    "api_key": "fake-key",
                },
                "model_info": {"id": "fast-endpoint-4"},
            },
        ],
        routing_strategy="latency-based-routing",
        routing_strategy_args={"ttl": 0.0000000001},
        set_verbose=True,
        debug_level="DEBUG",
    )  # type: ignore

    deployments = {}
    for _ in range(10):
        response = await router.acompletion(
            model="azure-model", messages=[{"role": "user", "content": "hello"}]
        )
        print(response)
        _picked_model_id = response._hidden_params["model_id"]
        if _picked_model_id not in deployments:
            deployments[_picked_model_id] = 1
        else:
            deployments[_picked_model_id] += 1
        await asyncio.sleep(0.000000000005)

    print("deployments", deployments)

    # assert that len(deployments) >1
    assert len(deployments) > 1


@pytest.mark.parametrize("buffer", [0, 1])
@pytest.mark.asyncio
async def test_lowest_latency_routing_buffer(buffer):
    """
    Allow shuffling calls within a certain latency buffer
    """
    model_list = [
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-turbo",
                "api_key": "os.environ/AZURE_FRANCE_API_KEY",
                "api_base": "https://openai-france-1234.openai.azure.com",
                "rpm": 1440,
            },
            "model_info": {"id": 1},
        },
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-35-turbo",
                "api_key": "os.environ/AZURE_EUROPE_API_KEY",
                "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
                "rpm": 6,
            },
            "model_info": {"id": 2},
        },
    ]
    router = Router(
        model_list=model_list,
        routing_strategy="latency-based-routing",
        set_verbose=False,
        num_retries=3,
        routing_strategy_args={"lowest_latency_buffer": buffer},
    )  # type: ignore

    ## DEPLOYMENT 1 ##
    deployment_id = 1
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "azure-model",
            },
            "model_info": {"id": 1},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 50}}
    time.sleep(3)
    end_time = time.time()
    router.lowestlatency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )
    ## DEPLOYMENT 2 ##
    deployment_id = 2
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "azure-model",
            },
            "model_info": {"id": 2},
        }
    }
    start_time = time.time()
    response_obj = {"usage": {"total_tokens": 20}}
    time.sleep(2)
    end_time = time.time()
    router.lowestlatency_logger.log_success_event(
        response_obj=response_obj,
        kwargs=kwargs,
        start_time=start_time,
        end_time=end_time,
    )

    ## CHECK WHAT'S SELECTED ##
    # print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
    selected_deployments = {}
    for _ in range(50):
        print(router.get_available_deployment(model="azure-model"))
        selected_deployments[
            router.get_available_deployment(model="azure-model")["model_info"]["id"]
        ] = 1

    if buffer == 0:
        assert len(selected_deployments.keys()) == 1
    else:
        assert len(selected_deployments.keys()) == 2


@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_lowest_latency_routing_time_to_first_token(sync_mode):
    """
    If a deployment has
    - a fast time to first token
    - slow latency/output token

    test if:
    - for streaming, the deployment with fastest time to first token is picked
    - for non-streaming, fastest overall deployment is picked
    """
    model_list = [
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-turbo",
                "api_key": "os.environ/AZURE_FRANCE_API_KEY",
                "api_base": "https://openai-france-1234.openai.azure.com",
            },
            "model_info": {"id": 1},
        },
        {
            "model_name": "azure-model",
            "litellm_params": {
                "model": "azure/gpt-35-turbo",
                "api_key": "os.environ/AZURE_EUROPE_API_KEY",
                "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
            },
            "model_info": {"id": 2},
        },
    ]
    router = Router(
        model_list=model_list,
        routing_strategy="latency-based-routing",
        set_verbose=False,
        num_retries=3,
    )  # type: ignore
    ## DEPLOYMENT 1 ##
    deployment_id = 1
    start_time = datetime.now()
    one_second_later = start_time + timedelta(seconds=1)

    # Compute 3 seconds after the current time
    three_seconds_later = start_time + timedelta(seconds=3)
    four_seconds_later = start_time + timedelta(seconds=4)

    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "azure-model",
            },
            "model_info": {"id": 1},
        },
        "stream": True,
        "completion_start_time": one_second_later,
    }

    response_obj = litellm.ModelResponse(
        usage=litellm.Usage(completion_tokens=50, total_tokens=50)
    )
    end_time = four_seconds_later

    if sync_mode:
        router.lowestlatency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )
    else:
        await router.lowestlatency_logger.async_log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )
    ## DEPLOYMENT 2 ##
    deployment_id = 2
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": "azure-model",
            },
            "model_info": {"id": 2},
        },
        "stream": True,
        "completion_start_time": three_seconds_later,
    }
    response_obj = litellm.ModelResponse(
        usage=litellm.Usage(completion_tokens=50, total_tokens=50)
    )
    end_time = three_seconds_later
    if sync_mode:
        router.lowestlatency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )
    else:
        await router.lowestlatency_logger.async_log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    """
    TESTING

    - expect deployment 1 to be picked for streaming
    - expect deployment 2 to be picked for non-streaming
    """
    # print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
    selected_deployments = {}
    for _ in range(3):
        print(router.get_available_deployment(model="azure-model"))
        ## for non-streaming
        selected_deployments[
            router.get_available_deployment(model="azure-model")["model_info"]["id"]
        ] = 1

    assert len(selected_deployments.keys()) == 1
    assert "2" in list(selected_deployments.keys())

    selected_deployments = {}
    for _ in range(50):
        print(router.get_available_deployment(model="azure-model"))
        ## for non-streaming
        selected_deployments[
            router.get_available_deployment(
                model="azure-model", request_kwargs={"stream": True}
            )["model_info"]["id"]
        ] = 1

    assert len(selected_deployments.keys()) == 1
    assert "1" in list(selected_deployments.keys())


def test_latency_list_trimming_discards_oldest_entry():
    """
    When the latency list reaches max_latency_list_size, the oldest entry is
    discarded to make room for new entries. The newest entry is appended at
    the end of the list.
    """
    max_size = 3
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
    )

    model_group = "gpt-3.5-turbo"
    deployment_id = "test-deployment"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": model_group,
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }

    # With 1 completion token, the logged latency value equals the raw
    # response time, so we can use distinct, identifiable values.
    latencies_to_add = []
    for i in range(max_size + 1):  # One more than max to trigger trimming
        start_time = time.time()
        response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
        expected_latency = float(i + 1)  # 1.0, 2.0, 3.0, 4.0
        end_time = start_time + expected_latency
        latencies_to_add.append(expected_latency)

        lowest_latency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    latency_key = f"{model_group}_map"
    cached_data = test_cache.get_cache(key=latency_key)
    latency_list = cached_data[deployment_id]["latency"]

    assert (
        len(latency_list) == max_size
    ), f"Expected {max_size} entries, got {len(latency_list)}"

    newest_latency = latencies_to_add[-1]  # 4.0
    oldest_latency = latencies_to_add[0]  # 1.0
    tolerance = 0.1

    # Newest entry is at the end of the list.
    assert (
        abs(latency_list[-1] - newest_latency) < tolerance
    ), f"Newest latency {newest_latency} should be at end, got {latency_list[-1]}"

    # Oldest entry is no longer in the list.
    for latency in latency_list:
        assert (
            abs(latency - oldest_latency) > tolerance
        ), f"Oldest latency {oldest_latency} should have been discarded, found {latency}"


@pytest.mark.asyncio
async def test_latency_list_trimming_discards_oldest_entry_async():
    """
    Async counterpart: the oldest entry is discarded when the latency list is
    trimmed.
    """
    max_size = 3
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
    )

    model_group = "gpt-3.5-turbo"
    deployment_id = "test-deployment"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": model_group,
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }

    latencies_to_add = []
    for i in range(max_size + 1):
        start_time = time.time()
        response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
        expected_latency = float(i + 1)
        end_time = start_time + expected_latency
        latencies_to_add.append(expected_latency)

        await lowest_latency_logger.async_log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    latency_key = f"{model_group}_map"
    cached_data = await test_cache.async_get_cache(key=latency_key)
    latency_list = cached_data[deployment_id]["latency"]

    assert len(latency_list) == max_size

    newest_latency = latencies_to_add[-1]
    oldest_latency = latencies_to_add[0]
    tolerance = 0.1

    assert (
        abs(latency_list[-1] - newest_latency) < tolerance
    ), f"Newest latency {newest_latency} should be at end of list"

    for latency in latency_list:
        assert (
            abs(latency - oldest_latency) > tolerance
        ), f"Oldest latency {oldest_latency} should have been discarded"


def test_ttft_list_trimming_discards_oldest_entry():
    """
    The time_to_first_token list trims the oldest entry when full, matching
    the behavior of the latency list.
    """
    max_size = 3
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
    )

    model_group = "gpt-3.5-turbo"
    deployment_id = "test-deployment"

    ttft_values = []
    for i in range(max_size + 1):
        start_time = time.time()
        expected_ttft = float(i + 1) * 0.1  # 0.1, 0.2, 0.3, 0.4
        completion_start_time = start_time + expected_ttft
        end_time = start_time + float(i + 1)
        ttft_values.append(expected_ttft)

        kwargs = {
            "litellm_params": {
                "metadata": {
                    "model_group": model_group,
                    "deployment": "azure/gpt-4.1-mini",
                },
                "model_info": {"id": deployment_id},
            },
            "stream": True,
            "completion_start_time": completion_start_time,
        }
        # TTFT is only recorded when response_obj is a ModelResponse.
        response_obj = litellm.ModelResponse(
            usage=litellm.Usage(completion_tokens=1, total_tokens=1)
        )

        lowest_latency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    latency_key = f"{model_group}_map"
    cached_data = test_cache.get_cache(key=latency_key)
    ttft_list = cached_data[deployment_id].get("time_to_first_token", [])

    assert (
        len(ttft_list) == max_size
    ), f"Expected {max_size} entries, got {len(ttft_list)}"

    newest_ttft = ttft_values[-1]
    oldest_ttft = ttft_values[0]
    tolerance = 0.05

    assert (
        abs(ttft_list[-1] - newest_ttft) < tolerance
    ), f"Newest TTFT {newest_ttft} should be at end of list"

    for ttft in ttft_list:
        assert (
            abs(ttft - oldest_ttft) > tolerance
        ), f"Oldest TTFT {oldest_ttft} should have been discarded"


@pytest.mark.asyncio
async def test_timeout_penalty_discards_oldest_entry():
    """
    Timeout penalties (1000.0) are appended to the latency list and, when the
    list is full, the oldest entry is discarded.
    """
    max_size = 3
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
    )

    model_group = "gpt-3.5-turbo"
    deployment_id = "test-deployment"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": model_group,
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }

    # Fill the list with max_size normal latency entries first.
    for i in range(max_size):
        start_time = time.time()
        response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
        end_time = start_time + float(i + 1)

        await lowest_latency_logger.async_log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    # Trigger a timeout failure: this appends 1000.0 and should discard the
    # oldest normal entry (1.0).
    timeout_kwargs = {
        **kwargs,
        "exception": litellm.Timeout(
            message="Request timed out", model="test-model", llm_provider="test"
        ),
    }

    await lowest_latency_logger.async_log_failure_event(
        kwargs=timeout_kwargs,
        response_obj=None,
        start_time=time.time(),
        end_time=time.time() + 30,
    )

    latency_key = f"{model_group}_map"
    cached_data = await test_cache.async_get_cache(key=latency_key)
    latency_list = cached_data[deployment_id]["latency"]

    assert len(latency_list) == max_size

    # Timeout penalty is the newest entry.
    assert (
        latency_list[-1] == 1000.0
    ), f"Timeout penalty should be at end of list, got {latency_list[-1]}"

    # Oldest normal entry (1.0) has been discarded.
    tolerance = 0.1
    for latency in latency_list[:-1]:
        assert (
            abs(latency - 1.0) > tolerance
        ), f"Oldest latency 1.0 should have been discarded, found {latency}"


def test_list_order_preserved_after_multiple_trims():
    """
    After many trims, the list still holds the most recent `max_size` entries
    in insertion order (oldest at index 0, newest at index -1).
    """
    max_size = 3
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
    )

    model_group = "gpt-3.5-turbo"
    deployment_id = "test-deployment"
    kwargs = {
        "litellm_params": {
            "metadata": {
                "model_group": model_group,
                "deployment": "azure/gpt-4.1-mini",
            },
            "model_info": {"id": deployment_id},
        }
    }

    # Add 10 entries (7 more than max) to trigger multiple trims.
    all_latencies = []
    for i in range(10):
        start_time = time.time()
        response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
        expected_latency = float(i + 1)
        end_time = start_time + expected_latency
        all_latencies.append(expected_latency)

        lowest_latency_logger.log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    latency_key = f"{model_group}_map"
    cached_data = test_cache.get_cache(key=latency_key)
    latency_list = cached_data[deployment_id]["latency"]

    assert len(latency_list) == max_size

    # After inserting 1..10 with max_size=3, the list should be [8, 9, 10].
    expected_remaining = all_latencies[-max_size:]
    tolerance = 0.1

    for i, expected in enumerate(expected_remaining):
        assert (
            abs(latency_list[i] - expected) < tolerance
        ), f"At index {i}, expected ~{expected}, got {latency_list[i]}"


@pytest.mark.asyncio
async def test_ttft_list_trimming_discards_oldest_entry_async():
    """
    Async counterpart: the time_to_first_token list trims the oldest entry
    when full. Exercises the async_log_success_event TTFT path, which only
    runs when response_obj is a ModelResponse and the call is marked as
    streaming with a completion_start_time.
    """
    max_size = 3
    test_cache = DualCache()
    lowest_latency_logger = LowestLatencyLoggingHandler(
        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
    )

    model_group = "gpt-3.5-turbo"
    deployment_id = "test-deployment"

    ttft_values = []
    for i in range(max_size + 1):
        start_time = time.time()
        expected_ttft = float(i + 1) * 0.1  # 0.1, 0.2, 0.3, 0.4
        completion_start_time = start_time + expected_ttft
        end_time = start_time + float(i + 1)
        ttft_values.append(expected_ttft)

        kwargs = {
            "litellm_params": {
                "metadata": {
                    "model_group": model_group,
                    "deployment": "azure/gpt-4.1-mini",
                },
                "model_info": {"id": deployment_id},
            },
            "stream": True,
            "completion_start_time": completion_start_time,
        }
        response_obj = litellm.ModelResponse(
            usage=litellm.Usage(completion_tokens=1, total_tokens=1)
        )

        await lowest_latency_logger.async_log_success_event(
            response_obj=response_obj,
            kwargs=kwargs,
            start_time=start_time,
            end_time=end_time,
        )

    latency_key = f"{model_group}_map"
    cached_data = await test_cache.async_get_cache(key=latency_key)
    ttft_list = cached_data[deployment_id].get("time_to_first_token", [])

    assert (
        len(ttft_list) == max_size
    ), f"Expected {max_size} entries, got {len(ttft_list)}"

    newest_ttft = ttft_values[-1]
    oldest_ttft = ttft_values[0]
    tolerance = 0.05

    assert (
        abs(ttft_list[-1] - newest_ttft) < tolerance
    ), f"Newest TTFT {newest_ttft} should be at end of list"

    for ttft in ttft_list:
        assert (
            abs(ttft - oldest_ttft) > tolerance
        ), f"Oldest TTFT {oldest_ttft} should have been discarded"