# # this tests if the router is initialized correctly # import asyncio # import os # import sys # import time # import traceback # import pytest # sys.path.insert( # 0, os.path.abspath("../..") # ) # Adds the parent directory to the system path # from collections import defaultdict # from concurrent.futures import ThreadPoolExecutor # from dotenv import load_dotenv # import litellm # from litellm import Router # load_dotenv() # # every time we load the router we should have 4 clients: # # Async # # Sync # # Async + Stream # # Sync + Stream # def test_init_clients(): # litellm.set_verbose = True # import logging # from litellm._logging import verbose_router_logger # verbose_router_logger.setLevel(logging.DEBUG) # try: # print("testing init 4 clients with diff timeouts") # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": os.getenv("AZURE_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, # }, # }, # ] # router = Router(model_list=model_list, set_verbose=True) # for elem in router.model_list: # model_id = elem["model_info"]["id"] # assert router.cache.get_cache(f"{model_id}_client") is not None # assert router.cache.get_cache(f"{model_id}_async_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None # # check if timeout for stream/non stream clients is set correctly # async_client = router.cache.get_cache(f"{model_id}_async_client") # stream_async_client = router.cache.get_cache( # f"{model_id}_stream_async_client" # ) # assert async_client.timeout == 0.01 # assert stream_async_client.timeout == 0.000_001 # print(vars(async_client)) # print() # print(async_client._base_url) # assert ( # async_client._base_url # == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/" # ) # assert ( # stream_async_client._base_url # == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/" # ) # print("PASSED !") # except Exception as e: # traceback.print_exc() # pytest.fail(f"Error occurred: {e}") # # test_init_clients() # def test_init_clients_basic(): # litellm.set_verbose = True # try: # print("Test basic client init") # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": os.getenv("AZURE_API_BASE"), # }, # }, # ] # router = Router(model_list=model_list) # for elem in router.model_list: # model_id = elem["model_info"]["id"] # assert router.cache.get_cache(f"{model_id}_client") is not None # assert router.cache.get_cache(f"{model_id}_async_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None # print("PASSED !") # # see if we can init clients without timeout or max retries set # except Exception as e: # traceback.print_exc() # pytest.fail(f"Error occurred: {e}") # # test_init_clients_basic() # def test_init_clients_basic_azure_cloudflare(): # # init azure + cloudflare # # init OpenAI gpt-3.5 # # init OpenAI text-embedding # # init OpenAI comptaible - Mistral/mistral-medium # # init OpenAI compatible - xinference/bge # litellm.set_verbose = True # try: # print("Test basic client init") # model_list = [ # { # "model_name": "azure-cloudflare", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": "https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1", # }, # }, # { # "model_name": "gpt-openai", # "litellm_params": { # "model": "gpt-3.5-turbo", # "api_key": os.getenv("OPENAI_API_KEY"), # }, # }, # { # "model_name": "text-embedding-ada-002", # "litellm_params": { # "model": "text-embedding-ada-002", # "api_key": os.getenv("OPENAI_API_KEY"), # }, # }, # { # "model_name": "mistral", # "litellm_params": { # "model": "mistral/mistral-tiny", # "api_key": os.getenv("MISTRAL_API_KEY"), # }, # }, # { # "model_name": "bge-base-en", # "litellm_params": { # "model": "xinference/bge-base-en", # "api_base": "http://127.0.0.1:9997/v1", # "api_key": os.getenv("OPENAI_API_KEY"), # }, # }, # ] # router = Router(model_list=model_list) # for elem in router.model_list: # model_id = elem["model_info"]["id"] # assert router.cache.get_cache(f"{model_id}_client") is not None # assert router.cache.get_cache(f"{model_id}_async_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None # print("PASSED !") # # see if we can init clients without timeout or max retries set # except Exception as e: # traceback.print_exc() # pytest.fail(f"Error occurred: {e}") # # test_init_clients_basic_azure_cloudflare() # def test_timeouts_router(): # """ # Test the timeouts of the router with multiple clients. This HASas to raise a timeout error # """ # import openai # litellm.set_verbose = True # try: # print("testing init 4 clients with diff timeouts") # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": os.getenv("AZURE_API_BASE"), # "timeout": 0.000001, # "stream_timeout": 0.000_001, # }, # }, # ] # router = Router(model_list=model_list, num_retries=0) # print("PASSED !") # async def test(): # try: # await router.acompletion( # model="gpt-3.5-turbo", # messages=[ # {"role": "user", "content": "hello, write a 20 pg essay"} # ], # ) # except Exception as e: # raise e # asyncio.run(test()) # except openai.APITimeoutError as e: # print( # "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e # ) # print(type(e)) # pass # except Exception as e: # pytest.fail( # f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}" # ) # # test_timeouts_router() # def test_stream_timeouts_router(): # """ # Test the stream timeouts router. See if it selected the correct client with stream timeout # """ # import openai # litellm.set_verbose = True # try: # print("testing init 4 clients with diff timeouts") # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": os.getenv("AZURE_API_BASE"), # "timeout": 200, # regular calls will not timeout, stream calls will # "stream_timeout": 10, # }, # }, # ] # router = Router(model_list=model_list) # print("PASSED !") # data = { # "model": "gpt-3.5-turbo", # "messages": [{"role": "user", "content": "hello, write a 20 pg essay"}], # "stream": True, # } # selected_client = router._get_client( # deployment=router.model_list[0], # kwargs=data, # client_type=None, # ) # print("Select client timeout", selected_client.timeout) # assert selected_client.timeout == 10 # # make actual call # response = router.completion(**data) # for chunk in response: # print(f"chunk: {chunk}") # except openai.APITimeoutError as e: # print( # "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e # ) # print(type(e)) # pass # except Exception as e: # pytest.fail( # f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}" # ) # # test_stream_timeouts_router() # def test_xinference_embedding(): # # [Test Init Xinference] this tests if we init xinference on the router correctly # # [Test Exception Mapping] tests that xinference is an openai comptiable provider # print("Testing init xinference") # print( # "this tests if we create an OpenAI client for Xinference, with the correct API BASE" # ) # model_list = [ # { # "model_name": "xinference", # "litellm_params": { # "model": "xinference/bge-base-en", # "api_base": "os.environ/XINFERENCE_API_BASE", # }, # } # ] # router = Router(model_list=model_list) # print(router.model_list) # print(router.model_list[0]) # assert ( # router.model_list[0]["litellm_params"]["api_base"] == "http://0.0.0.0:9997" # ) # set in env # openai_client = router._get_client( # deployment=router.model_list[0], # kwargs={"input": ["hello"], "model": "xinference"}, # ) # assert openai_client._base_url == "http://0.0.0.0:9997" # assert "xinference" in litellm.openai_compatible_providers # print("passed") # # test_xinference_embedding() # def test_router_init_gpt_4_vision_enhancements(): # try: # # tests base_url set when any base_url with /openai/deployments passed to router # print("Testing Azure GPT_Vision enhancements") # model_list = [ # { # "model_name": "gpt-4-vision-enhancements", # "litellm_params": { # "model": "azure/gpt-4-vision", # "api_key": os.getenv("AZURE_API_KEY"), # "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/", # "dataSources": [ # { # "type": "AzureComputerVision", # "parameters": { # "endpoint": "os.environ/AZURE_VISION_ENHANCE_ENDPOINT", # "key": "os.environ/AZURE_VISION_ENHANCE_KEY", # }, # } # ], # }, # } # ] # router = Router(model_list=model_list) # print(router.model_list) # print(router.model_list[0]) # assert ( # router.model_list[0]["litellm_params"]["base_url"] # == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/" # ) # set in env # assert ( # router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][ # "endpoint" # ] # == os.environ["AZURE_VISION_ENHANCE_ENDPOINT"] # ) # assert ( # router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][ # "key" # ] # == os.environ["AZURE_VISION_ENHANCE_KEY"] # ) # azure_client = router._get_client( # deployment=router.model_list[0], # kwargs={"stream": True, "model": "gpt-4-vision-enhancements"}, # client_type="async", # ) # assert ( # azure_client._base_url # == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/" # ) # print("passed") # except Exception as e: # pytest.fail(f"Error occurred: {e}") # @pytest.mark.parametrize("sync_mode", [True, False]) # @pytest.mark.asyncio # async def test_openai_with_organization(sync_mode): # try: # print("Testing OpenAI with organization") # model_list = [ # { # "model_name": "openai-bad-org", # "litellm_params": { # "model": "gpt-3.5-turbo", # "organization": "org-ikDc4ex8NB", # }, # }, # { # "model_name": "openai-good-org", # "litellm_params": {"model": "gpt-3.5-turbo"}, # }, # ] # router = Router(model_list=model_list) # print(router.model_list) # print(router.model_list[0]) # if sync_mode: # openai_client = router._get_client( # deployment=router.model_list[0], # kwargs={"input": ["hello"], "model": "openai-bad-org"}, # ) # print(vars(openai_client)) # assert openai_client.organization == "org-ikDc4ex8NB" # # bad org raises error # try: # response = router.completion( # model="openai-bad-org", # messages=[{"role": "user", "content": "this is a test"}], # ) # pytest.fail( # "Request should have failed - This organization does not exist" # ) # except Exception as e: # print("Got exception: " + str(e)) # assert "header should match organization for API key" in str( # e # ) or "No such organization" in str(e) # # good org works # response = router.completion( # model="openai-good-org", # messages=[{"role": "user", "content": "this is a test"}], # max_tokens=5, # ) # else: # openai_client = router._get_client( # deployment=router.model_list[0], # kwargs={"input": ["hello"], "model": "openai-bad-org"}, # client_type="async", # ) # print(vars(openai_client)) # assert openai_client.organization == "org-ikDc4ex8NB" # # bad org raises error # try: # response = await router.acompletion( # model="openai-bad-org", # messages=[{"role": "user", "content": "this is a test"}], # ) # pytest.fail( # "Request should have failed - This organization does not exist" # ) # except Exception as e: # print("Got exception: " + str(e)) # assert "header should match organization for API key" in str( # e # ) or "No such organization" in str(e) # # good org works # response = await router.acompletion( # model="openai-good-org", # messages=[{"role": "user", "content": "this is a test"}], # max_tokens=5, # ) # except Exception as e: # pytest.fail(f"Error occurred: {e}") # def test_init_clients_azure_command_r_plus(): # # This tests that the router uses the OpenAI client for Azure/Command-R+ # # For azure/command-r-plus we need to use openai.OpenAI because of how the Azure provider requires requests being sent # litellm.set_verbose = True # import logging # from litellm._logging import verbose_router_logger # verbose_router_logger.setLevel(logging.DEBUG) # try: # print("testing init 4 clients with diff timeouts") # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/command-r-plus", # "api_key": os.getenv("AZURE_COHERE_API_KEY"), # "api_base": os.getenv("AZURE_COHERE_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, # }, # }, # ] # router = Router(model_list=model_list, set_verbose=True) # for elem in router.model_list: # model_id = elem["model_info"]["id"] # async_client = router.cache.get_cache(f"{model_id}_async_client") # stream_async_client = router.cache.get_cache( # f"{model_id}_stream_async_client" # ) # # Assert the Async Clients used are OpenAI clients and not Azure # # For using Azure/Command-R-Plus and Azure/Mistral the clients NEED to be OpenAI clients used # # this is weirdness introduced on Azure's side # assert "openai.AsyncOpenAI" in str(async_client) # assert "openai.AsyncOpenAI" in str(stream_async_client) # print("PASSED !") # except Exception as e: # traceback.print_exc() # pytest.fail(f"Error occurred: {e}") # @pytest.mark.asyncio # async def test_aaaaatext_completion_with_organization(): # try: # print("Testing Text OpenAI with organization") # model_list = [ # { # "model_name": "openai-bad-org", # "litellm_params": { # "model": "text-completion-openai/gpt-3.5-turbo-instruct", # "api_key": os.getenv("OPENAI_API_KEY", None), # "organization": "org-ikDc4ex8NB", # }, # }, # { # "model_name": "openai-good-org", # "litellm_params": { # "model": "text-completion-openai/gpt-3.5-turbo-instruct", # "api_key": os.getenv("OPENAI_API_KEY", None), # "organization": os.getenv("OPENAI_ORGANIZATION", None), # }, # }, # ] # router = Router(model_list=model_list) # print(router.model_list) # print(router.model_list[0]) # openai_client = router._get_client( # deployment=router.model_list[0], # kwargs={"input": ["hello"], "model": "openai-bad-org"}, # ) # print(vars(openai_client)) # assert openai_client.organization == "org-ikDc4ex8NB" # # bad org raises error # try: # response = await router.atext_completion( # model="openai-bad-org", # prompt="this is a test", # ) # pytest.fail("Request should have failed - This organization does not exist") # except Exception as e: # print("Got exception: " + str(e)) # assert "header should match organization for API key" in str( # e # ) or "No such organization" in str(e) # # good org works # response = await router.atext_completion( # model="openai-good-org", # prompt="this is a test", # max_tokens=5, # ) # print("working response: ", response) # except Exception as e: # pytest.fail(f"Error occurred: {e}") # def test_init_clients_async_mode(): # litellm.set_verbose = True # import logging # from litellm._logging import verbose_router_logger # from litellm.types.router import RouterGeneralSettings # verbose_router_logger.setLevel(logging.DEBUG) # try: # print("testing init 4 clients with diff timeouts") # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": os.getenv("AZURE_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, # }, # }, # ] # router = Router( # model_list=model_list, # set_verbose=True, # router_general_settings=RouterGeneralSettings(async_only_mode=True), # ) # for elem in router.model_list: # model_id = elem["model_info"]["id"] # # sync clients not initialized in async_only_mode=True # assert router.cache.get_cache(f"{model_id}_client") is None # assert router.cache.get_cache(f"{model_id}_stream_client") is None # # only async clients initialized in async_only_mode=True # assert router.cache.get_cache(f"{model_id}_async_client") is not None # assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None # except Exception as e: # pytest.fail(f"Error occurred: {e}") # @pytest.mark.parametrize( # "environment,expected_models", # [ # ("development", ["gpt-3.5-turbo"]), # ("production", ["gpt-4", "gpt-3.5-turbo", "gpt-4o"]), # ], # ) # def test_init_router_with_supported_environments(environment, expected_models): # """ # Tests that the correct models are setup on router when LITELLM_ENVIRONMENT is set # """ # os.environ["LITELLM_ENVIRONMENT"] = environment # model_list = [ # { # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", # "api_key": os.getenv("AZURE_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": os.getenv("AZURE_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, # }, # "model_info": {"supported_environments": ["development", "production"]}, # }, # { # "model_name": "gpt-4", # "litellm_params": { # "model": "openai/gpt-4", # "api_key": os.getenv("OPENAI_API_KEY"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, # }, # "model_info": {"supported_environments": ["production"]}, # }, # { # "model_name": "gpt-4o", # "litellm_params": { # "model": "openai/gpt-4o", # "api_key": os.getenv("OPENAI_API_KEY"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, # }, # "model_info": {"supported_environments": ["production"]}, # }, # ] # router = Router(model_list=model_list, set_verbose=True) # _model_list = router.get_model_names() # print("model_list: ", _model_list) # print("expected_models: ", expected_models) # assert set(_model_list) == set(expected_models) # os.environ.pop("LITELLM_ENVIRONMENT")