mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat: add Wallaby (OpenAI-compatible provider) — kimi-k3
- provider registry, pricing (root + backup cost maps, byte-synced) - openai_compatible_endpoints base URL for bare-URL autodetect - LlmProviders.WALLABY enum, dashboard provider_create_fields entry - provider_endpoints_support registration (root + backup) - tests: tests/unit/llms/openai_like/test_wallaby_provider.py (9 tests)
This commit is contained in:
Dawncoral
parent
22b36cbcf6
commit
c53b3c9044
9 changed files with 246 additions and 0 deletions
|
|
@ -949,6 +949,7 @@ openai_compatible_endpoints: Final[list] = [
|
|||
"https://api.sailresearch.com/v1",
|
||||
"https://api.cognition.ai/v1",
|
||||
"https://api.scx.ai/v1",
|
||||
"https://api.wallabytoken.com/v1",
|
||||
"https://gigachat.devices.sberbank.ru/api/v1",
|
||||
]
|
||||
|
||||
|
|
@ -1022,6 +1023,7 @@ openai_compatible_providers: Final[list] = [
|
|||
"meta", # Meta Model API (Muse Spark) - JSON-configured provider
|
||||
"cognition",
|
||||
"scx-ai",
|
||||
"wallaby",
|
||||
"sail",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -206,5 +206,10 @@
|
|||
"api_key_env": "SAIL_API_KEY",
|
||||
"api_base_env": "SAIL_API_BASE",
|
||||
"supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"]
|
||||
},
|
||||
"wallaby": {
|
||||
"base_url": "https://api.wallabytoken.com/v1",
|
||||
"api_key_env": "WALLABY_API_KEY",
|
||||
"api_base_env": "WALLABY_API_BASE"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -77910,5 +77910,21 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"wallaby/kimi-k3": {
|
||||
"cache_read_input_token_cost": 2.7e-07,
|
||||
"input_cost_per_token": 2.7e-06,
|
||||
"litellm_provider": "wallaby",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.35e-05,
|
||||
"source": "https://wallabytoken.com/pricing.json",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2316,6 +2316,24 @@
|
|||
"rerank": true
|
||||
}
|
||||
},
|
||||
"wallaby": {
|
||||
"display_name": "Wallaby (`wallaby`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/wallaby",
|
||||
"endpoints": {
|
||||
"chat_completions": true,
|
||||
"messages": false,
|
||||
"responses": false,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"moderations": false,
|
||||
"batches": false,
|
||||
"rerank": false,
|
||||
"a2a": false,
|
||||
"interactions": false
|
||||
}
|
||||
},
|
||||
"wandb": {
|
||||
"display_name": "WandB Inference (`wandb`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/wandb_inference",
|
||||
|
|
|
|||
|
|
@ -3059,6 +3059,34 @@
|
|||
],
|
||||
"default_model_placeholder": "scx-ai/GLM-5.2"
|
||||
},
|
||||
{
|
||||
"provider": "WALLABY",
|
||||
"provider_display_name": "Wallaby",
|
||||
"litellm_provider": "wallaby",
|
||||
"credential_fields": [
|
||||
{
|
||||
"key": "api_base",
|
||||
"label": "API Base",
|
||||
"placeholder": "https://api.wallabytoken.com/v1",
|
||||
"tooltip": null,
|
||||
"required": false,
|
||||
"field_type": "text",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
},
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "API Key",
|
||||
"placeholder": null,
|
||||
"tooltip": null,
|
||||
"required": true,
|
||||
"field_type": "password",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
}
|
||||
],
|
||||
"default_model_placeholder": "wallaby/kimi-k3"
|
||||
},
|
||||
{
|
||||
"provider": "Snowflake",
|
||||
"provider_display_name": "Snowflake",
|
||||
|
|
|
|||
|
|
@ -4139,6 +4139,7 @@ class LlmProviders(str, Enum):
|
|||
PINSTRIPES = "pinstripes"
|
||||
COGNITION = "cognition"
|
||||
SCX_AI = "scx-ai"
|
||||
WALLABY = "wallaby"
|
||||
DARKBLOOM = "darkbloom"
|
||||
META = "meta"
|
||||
SAIL = "sail"
|
||||
|
|
|
|||
|
|
@ -77910,5 +77910,21 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"wallaby/kimi-k3": {
|
||||
"cache_read_input_token_cost": 2.7e-07,
|
||||
"input_cost_per_token": 2.7e-06,
|
||||
"litellm_provider": "wallaby",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.35e-05,
|
||||
"source": "https://wallabytoken.com/pricing.json",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2662,6 +2662,24 @@
|
|||
"rerank": true
|
||||
}
|
||||
},
|
||||
"wallaby": {
|
||||
"display_name": "Wallaby (`wallaby`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/wallaby",
|
||||
"endpoints": {
|
||||
"chat_completions": true,
|
||||
"messages": false,
|
||||
"responses": false,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"moderations": false,
|
||||
"batches": false,
|
||||
"rerank": false,
|
||||
"a2a": false,
|
||||
"interactions": false
|
||||
}
|
||||
},
|
||||
"wandb": {
|
||||
"display_name": "WandB Inference (`wandb`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/wandb_inference",
|
||||
|
|
|
|||
142
tests/unit/llms/openai_like/test_wallaby_provider.py
Normal file
142
tests/unit/llms/openai_like/test_wallaby_provider.py
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
"""
|
||||
Tests for Wallaby provider configuration and integration.
|
||||
"""
|
||||
|
||||
import litellm
|
||||
|
||||
|
||||
class TestWallabyProviderConfig:
|
||||
def test_wallaby_in_provider_list(self):
|
||||
from litellm import LlmProviders
|
||||
|
||||
assert hasattr(LlmProviders, "WALLABY")
|
||||
assert LlmProviders.WALLABY.value == "wallaby"
|
||||
assert "wallaby" in litellm.provider_list
|
||||
|
||||
def test_wallaby_json_config_exists(self):
|
||||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
|
||||
assert JSONProviderRegistry.exists("wallaby")
|
||||
|
||||
wallaby = JSONProviderRegistry.get("wallaby")
|
||||
assert wallaby is not None
|
||||
assert wallaby.base_url == "https://api.wallabytoken.com/v1"
|
||||
assert wallaby.api_key_env == "WALLABY_API_KEY"
|
||||
|
||||
def test_wallaby_in_openai_compatible_providers(self):
|
||||
from litellm.constants import openai_compatible_providers
|
||||
|
||||
assert "wallaby" in openai_compatible_providers
|
||||
|
||||
def test_wallaby_provider_resolution(self):
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
model, provider, api_key, api_base = get_llm_provider(
|
||||
model="wallaby/kimi-k3",
|
||||
custom_llm_provider=None,
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
assert model == "kimi-k3"
|
||||
assert provider == "wallaby"
|
||||
assert api_base == "https://api.wallabytoken.com/v1"
|
||||
|
||||
def test_wallaby_api_base_override(self):
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
model, provider, api_key, api_base = get_llm_provider(
|
||||
model="wallaby/kimi-k3",
|
||||
custom_llm_provider=None,
|
||||
api_base="https://custom.wallabytoken.com/v1",
|
||||
api_key="sk-test",
|
||||
)
|
||||
|
||||
assert provider == "wallaby"
|
||||
assert api_base == "https://custom.wallabytoken.com/v1"
|
||||
assert api_key == "sk-test"
|
||||
|
||||
def test_wallaby_url_autodetection(self):
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
model, provider, api_key, api_base = get_llm_provider(
|
||||
model="kimi-k3",
|
||||
custom_llm_provider=None,
|
||||
api_base="https://api.wallabytoken.com/v1",
|
||||
api_key=None,
|
||||
)
|
||||
assert provider == "wallaby"
|
||||
assert api_base == "https://api.wallabytoken.com/v1"
|
||||
|
||||
def test_wallaby_router_config(self):
|
||||
from litellm import Router
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "wallaby-chat",
|
||||
"litellm_params": {
|
||||
"model": "wallaby/kimi-k3",
|
||||
"api_key": "test-key",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert len(router.model_list) == 1
|
||||
assert router.model_list[0]["model_name"] == "wallaby-chat"
|
||||
|
||||
|
||||
class TestWallabyModelMetadata:
|
||||
WALLABY_MODELS = ("wallaby/kimi-k3",)
|
||||
|
||||
@staticmethod
|
||||
def _load(path_parts):
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
json_path = Path(__file__).parents[4].joinpath(*path_parts)
|
||||
with open(json_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
def test_wallaby_models_synced_to_backup(self):
|
||||
model_cost = self._load(("model_prices_and_context_window.json",))
|
||||
backup = self._load(("litellm", "model_prices_and_context_window_backup.json"))
|
||||
for model in self.WALLABY_MODELS:
|
||||
assert model in backup, f"{model} missing from backup json"
|
||||
assert backup[model] == model_cost[model], (
|
||||
f"{model} differs between root and backup json"
|
||||
)
|
||||
|
||||
|
||||
class TestWallabyDashboardRegistration:
|
||||
@staticmethod
|
||||
def _provider_create_fields():
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import litellm
|
||||
|
||||
path = (
|
||||
Path(litellm.__file__).parent
|
||||
/ "proxy"
|
||||
/ "public_endpoints"
|
||||
/ "provider_create_fields.json"
|
||||
)
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
def test_wallaby_is_selectable_in_the_add_model_form(self):
|
||||
entries = [
|
||||
e for e in self._provider_create_fields() if e["litellm_provider"] == "wallaby"
|
||||
]
|
||||
assert len(entries) == 1, "wallaby must appear exactly once in provider_create_fields.json"
|
||||
|
||||
entry = entries[0]
|
||||
assert entry["provider"] == "WALLABY"
|
||||
assert entry["provider_display_name"] == "Wallaby"
|
||||
assert entry["default_model_placeholder"].startswith("wallaby/")
|
||||
|
||||
fields = {f["key"]: f for f in entry["credential_fields"]}
|
||||
assert fields["api_key"]["required"] is True
|
||||
assert fields["api_key"]["field_type"] == "password"
|
||||
Loading…
Add table
Reference in a new issue