litellm/tests/local_testing/test_router_init.py

# 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
from litellm.router_utils.client_initalization_utils import (
    InitalizeOpenAISDKClient,
    OpenAISDKClientInitializationParams,
)

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/chatgpt-v-2",
                    "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/"
            )  # openai python adds the extra /
            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/chatgpt-v-2",
                    "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/chatgpt-v-2",
                    "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/chatgpt-v-2",
                    "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/chatgpt-v-2",
                    "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 "No such organization: org-ikDc4ex8NB" 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 "No such organization: org-ikDc4ex8NB" 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_text_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 "No such organization: org-ikDc4ex8NB" 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/chatgpt-v-2",
                    "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/chatgpt-v-2",
                "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")


def test_should_initialize_sync_client():
    from litellm.types.router import RouterGeneralSettings

    # Test case 1: Router instance is None
    assert InitalizeOpenAISDKClient.should_initialize_sync_client(None) is False

    # Test case 2: Router instance without router_general_settings
    router = Router(model_list=[])
    assert InitalizeOpenAISDKClient.should_initialize_sync_client(router) is True

    # Test case 3: Router instance with async_only_mode = False
    router = Router(
        model_list=[],
        router_general_settings=RouterGeneralSettings(async_only_mode=False),
    )
    assert InitalizeOpenAISDKClient.should_initialize_sync_client(router) is True

    # Test case 4: Router instance with async_only_mode = True
    router = Router(
        model_list=[],
        router_general_settings=RouterGeneralSettings(async_only_mode=True),
    )
    assert InitalizeOpenAISDKClient.should_initialize_sync_client(router) is False

    # Test case 5: Router instance with router_general_settings but without async_only_mode
    router = Router(model_list=[], router_general_settings=RouterGeneralSettings())
    assert InitalizeOpenAISDKClient.should_initialize_sync_client(router) is True

    print("All test cases passed!")


@pytest.mark.parametrize(
    "model_name, custom_llm_provider, expected_result",
    [
        ("gpt-3.5-turbo", None, True),  # OpenAI chat completion model
        ("text-embedding-ada-002", None, True),  # OpenAI embedding model
        ("claude-2", None, False),  # Non-OpenAI model
        ("gpt-3.5-turbo", "azure", True),  # Azure OpenAI
        ("text-davinci-003", "azure_text", True),  # Azure OpenAI
        ("gpt-3.5-turbo", "openai", True),  # OpenAI
        ("custom-model", "custom_openai", True),  # Custom OpenAI compatible
        ("text-davinci-003", "text-completion-openai", True),  # OpenAI text completion
        (
            "ft:gpt-3.5-turbo-0613:my-org:custom-model:7p4lURel",
            None,
            True,
        ),  # Fine-tuned GPT model
        ("mistral-7b", "huggingface", False),  # Non-OpenAI provider
        ("custom-model", "anthropic", False),  # Non-OpenAI compatible provider
    ],
)
def test_should_create_openai_sdk_client_for_model(
    model_name, custom_llm_provider, expected_result
):
    result = InitalizeOpenAISDKClient._should_create_openai_sdk_client_for_model(
        model_name, custom_llm_provider
    )
    assert (
        result == expected_result
    ), f"Failed for model: {model_name}, provider: {custom_llm_provider}"


def test_should_create_openai_sdk_client_for_model_openai_compatible_providers():
    # Test with a known OpenAI compatible provider
    assert InitalizeOpenAISDKClient._should_create_openai_sdk_client_for_model(
        "custom-model", "groq"
    ), "Should return True for OpenAI compatible provider"

    # Add a new compatible provider and test
    litellm.openai_compatible_providers.append("new_provider")
    assert InitalizeOpenAISDKClient._should_create_openai_sdk_client_for_model(
        "custom-model", "new_provider"
    ), "Should return True for newly added OpenAI compatible provider"

    # Clean up
    litellm.openai_compatible_providers.remove("new_provider")


def test_get_client_initialization_params_openai():
    """Test basic OpenAI configuration with direct parameter passing."""
    model = {}
    model_name = "gpt-3.5-turbo"
    custom_llm_provider = None
    litellm_params = {"api_key": "sk-openai-key", "timeout": 30, "max_retries": 3}
    default_api_key = None
    default_api_base = None

    result = InitalizeOpenAISDKClient._get_client_initialization_params(
        model=model,
        model_name=model_name,
        custom_llm_provider=custom_llm_provider,
        litellm_params=litellm_params,
        default_api_key=default_api_key,
        default_api_base=default_api_base,
    )

    assert isinstance(result, OpenAISDKClientInitializationParams)
    assert result.api_key == "sk-openai-key"
    assert result.timeout == 30
    assert result.max_retries == 3
    assert result.model_name == "gpt-3.5-turbo"


def test_get_client_initialization_params_azure():
    """Test Azure OpenAI configuration with specific Azure parameters."""
    model = {}
    model_name = "azure/gpt-4"
    custom_llm_provider = "azure"
    litellm_params = {
        "api_key": "azure-key",
        "api_base": "https://example.azure.openai.com",
        "api_version": "2023-05-15",
    }
    default_api_key = None
    default_api_base = None

    result = InitalizeOpenAISDKClient._get_client_initialization_params(
        model=model,
        model_name=model_name,
        custom_llm_provider=custom_llm_provider,
        litellm_params=litellm_params,
        default_api_key=default_api_key,
        default_api_base=default_api_base,
    )

    assert result.api_key == "azure-key"
    assert result.api_base == "https://example.azure.openai.com"
    assert result.api_version == "2023-05-15"
    assert result.custom_llm_provider == "azure"


def test_get_client_initialization_params_environment_variable_parsing():
    """Test parsing of environment variables for configuration."""
    os.environ["UNIQUE_OPENAI_API_KEY"] = "env-openai-key"
    os.environ["UNIQUE_TIMEOUT"] = "45"

    model = {}
    model_name = "gpt-4"
    custom_llm_provider = None
    litellm_params = {
        "api_key": "os.environ/UNIQUE_OPENAI_API_KEY",
        "timeout": "os.environ/UNIQUE_TIMEOUT",
        "organization": "os.environ/UNIQUE_ORG_ID",
    }
    default_api_key = None
    default_api_base = None

    result = InitalizeOpenAISDKClient._get_client_initialization_params(
        model=model,
        model_name=model_name,
        custom_llm_provider=custom_llm_provider,
        litellm_params=litellm_params,
        default_api_key=default_api_key,
        default_api_base=default_api_base,
    )

    assert result.api_key == "env-openai-key"
    assert result.timeout == 45.0
    assert result.organization is None  # Since ORG_ID is not set in the environment


def test_get_client_initialization_params_azure_ai_studio_mistral():
    """
    Test configuration for Azure AI Studio Mistral model.

    - /v1/ is added to the api_base if it is not present
    - custom_llm_provider is set to openai (Azure AI Studio Mistral models need to use OpenAI route)
    """

    model = {}
    model_name = "azure/mistral-large-latest"
    custom_llm_provider = "azure"
    litellm_params = {
        "api_key": "azure-key",
        "api_base": "https://example.azure.openai.com",
    }
    default_api_key = None
    default_api_base = None

    result = InitalizeOpenAISDKClient._get_client_initialization_params(
        model,
        model_name,
        custom_llm_provider,
        litellm_params,
        default_api_key,
        default_api_base,
    )

    assert result.custom_llm_provider == "openai"
    assert result.model_name == "mistral-large-latest"
    assert result.api_base == "https://example.azure.openai.com/v1/"


def test_get_client_initialization_params_default_values():
    """
    Test use of default values when specific parameters are not provided.

    This is used typically for OpenAI compatible providers - example Together AI

    """
    model = {}
    model_name = "together/meta-llama-3.1-8b-instruct"
    custom_llm_provider = None
    litellm_params = {}
    default_api_key = "together-api-key"
    default_api_base = "https://together.xyz/api.openai.com"

    result = InitalizeOpenAISDKClient._get_client_initialization_params(
        model=model,
        model_name=model_name,
        custom_llm_provider=custom_llm_provider,
        litellm_params=litellm_params,
        default_api_key=default_api_key,
        default_api_base=default_api_base,
    )

    assert result.api_key == "together-api-key"
    assert result.api_base == "https://together.xyz/api.openai.com"
    assert result.timeout == litellm.request_timeout
    assert result.max_retries == 0


def test_get_client_initialization_params_all_env_vars():
    # Set up environment variables
    os.environ["TEST_API_KEY"] = "test-api-key"
    os.environ["TEST_API_BASE"] = "https://test.openai.com"
    os.environ["TEST_API_VERSION"] = "2023-05-15"
    os.environ["TEST_TIMEOUT"] = "30"
    os.environ["TEST_STREAM_TIMEOUT"] = "60"
    os.environ["TEST_MAX_RETRIES"] = "3"
    os.environ["TEST_ORGANIZATION"] = "test-org"

    model = {}
    model_name = "gpt-4"
    custom_llm_provider = None
    litellm_params = {
        "api_key": "os.environ/TEST_API_KEY",
        "api_base": "os.environ/TEST_API_BASE",
        "api_version": "os.environ/TEST_API_VERSION",
        "timeout": "os.environ/TEST_TIMEOUT",
        "stream_timeout": "os.environ/TEST_STREAM_TIMEOUT",
        "max_retries": "os.environ/TEST_MAX_RETRIES",
        "organization": "os.environ/TEST_ORGANIZATION",
    }
    default_api_key = None
    default_api_base = None

    result = InitalizeOpenAISDKClient._get_client_initialization_params(
        model=model,
        model_name=model_name,
        custom_llm_provider=custom_llm_provider,
        litellm_params=litellm_params,
        default_api_key=default_api_key,
        default_api_base=default_api_base,
    )

    assert isinstance(result, OpenAISDKClientInitializationParams)
    assert result.api_key == "test-api-key"
    assert result.api_base == "https://test.openai.com"
    assert result.api_version == "2023-05-15"
    assert result.timeout == 30.0
    assert result.stream_timeout == 60.0
    assert result.max_retries == 3
    assert result.organization == "test-org"
    assert result.model_name == "gpt-4"
    assert result.custom_llm_provider is None

    # Clean up environment variables
    for key in [
        "TEST_API_KEY",
        "TEST_API_BASE",
        "TEST_API_VERSION",
        "TEST_TIMEOUT",
        "TEST_STREAM_TIMEOUT",
        "TEST_MAX_RETRIES",
        "TEST_ORGANIZATION",
    ]:
        os.environ.pop(key)