diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 005e94ebe82..88c6bb16d6a 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -826,6 +826,11 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, ) = litellm.AIMLChatConfig()._get_openai_compatible_provider_info(api_base, api_key) + elif custom_llm_provider == "token_kiosk": + ( + api_base, + dynamic_api_key, + ) = litellm.TokenKioskConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "wandb": api_base = api_base or get_secret("WANDB_API_BASE") or "https://api.inference.wandb.ai/v1" dynamic_api_key = api_key or get_secret_str("WANDB_API_KEY") diff --git a/tests/test_litellm/test_token_kiosk.py b/tests/test_litellm/test_token_kiosk.py index 24df672b98e..aa252016463 100644 --- a/tests/test_litellm/test_token_kiosk.py +++ b/tests/test_litellm/test_token_kiosk.py @@ -13,6 +13,26 @@ def test_token_kiosk_config_get_complete_url(): ) assert url == "https://agent-router.gaib.ai/v1/chat/completions" + # Test trailing slash + url_slash = config.get_complete_url( + api_base="https://agent-router.gaib.ai/v1/", + api_key="test-key", + model="token_kiosk/claude-3-5-sonnet", + optional_params={}, + litellm_params={}, + ) + assert url_slash == "https://agent-router.gaib.ai/v1/chat/completions" + + # Test full endpoint path already ending with /chat/completions + url_full = config.get_complete_url( + api_base="https://agent-router.gaib.ai/v1/chat/completions", + api_key="test-key", + model="token_kiosk/claude-3-5-sonnet", + optional_params={}, + litellm_params={}, + ) + assert url_full == "https://agent-router.gaib.ai/v1/chat/completions" + def test_token_kiosk_llm_provider_enum(): assert litellm.LlmProviders.TOKEN_KIOSK.value == "token_kiosk" @@ -30,3 +50,16 @@ def test_token_kiosk_get_openai_compatible_provider_info(): def test_token_kiosk_lazy_import(): assert litellm.TokenKioskConfig is TokenKioskConfig + + +def test_token_kiosk_get_llm_provider(): + model, custom_llm_provider, dynamic_api_key, api_base = ( + litellm.get_llm_provider( + model="token_kiosk/claude-3-5-sonnet", + api_key="test-token-kiosk-key", + ) + ) + assert model == "claude-3-5-sonnet" + assert custom_llm_provider == "token_kiosk" + assert dynamic_api_key == "test-token-kiosk-key" + assert api_base == "https://agent-router.gaib.ai/v1"