scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(token_kiosk): register provider in get_llm_provider and add full url test coverage

Browse source 
Signed-off-by: hgaib <hung.cheng@gaib.ai>
This commit is contained in:
hgaib 2026-08-26 03:57:59 +00:00
parent bebaf7886a
commit dab227866b
2 changed files with 38 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -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")

33
tests/test_litellm/test_token_kiosk.py
View file

@ -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"