From 9bcb006794c058eb89dfa271a6e34e10c6fa778a Mon Sep 17 00:00:00 2001 From: Abhi Ram Salammagari Date: Sat, 30 May 2026 00:36:40 -0700 Subject: [PATCH] Add kluster.ai as an OpenAI-compatible chat provider kluster.ai exposes an OpenAI-compatible inference API. This wires it in as a first-class provider: a KlusterAIConfig subclass of OpenAIGPTConfig, registration in the provider enum and OpenAI-compatible provider/endpoint lists, and api_base/api_key resolution (default https://api.kluster.ai/v1, overridable via KLUSTER_AI_API_BASE). Adds mocked unit tests covering config precedence and end-to-end provider resolution. --- litellm/__init__.py | 9 ++ litellm/_lazy_imports_registry.py | 5 + litellm/constants.py | 3 + .../get_llm_provider_logic.py | 10 ++ .../get_supported_openai_params.py | 2 + .../llms/kluster_ai/chat/transformation.py | 56 ++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 4 + .../test_kluster_ai_chat_transformation.py | 101 ++++++++++++++++++ 9 files changed, 191 insertions(+) create mode 100644 litellm/llms/kluster_ai/chat/transformation.py create mode 100644 tests/test_litellm/llms/kluster_ai/chat/test_kluster_ai_chat_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 56d516536e8..77519ef2548 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -616,6 +616,7 @@ snowflake_models: Set = set() gradient_ai_models: Set = set() llama_models: Set = set() nscale_models: Set = set() +kluster_ai_models: Set = set() nebius_models: Set = set() nebius_embedding_models: Set = set() aiml_models: Set = set() @@ -805,6 +806,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): llama_models.add(key) elif value.get("litellm_provider") == "nscale": nscale_models.add(key) + elif value.get("litellm_provider") == "kluster_ai": + kluster_ai_models.add(key) elif value.get("litellm_provider") == "azure_ai": azure_ai_models.add(key) elif value.get("litellm_provider") == "voyage": @@ -1009,6 +1012,7 @@ model_list = list( | llama_models | featherless_ai_models | nscale_models + | kluster_ai_models | deepgram_models | elevenlabs_models | dashscope_models @@ -1107,6 +1111,7 @@ models_by_provider: dict = { "gradient_ai": gradient_ai_models, "meta_llama": llama_models, "nscale": nscale_models, + "kluster_ai": kluster_ai_models, "featherless_ai": featherless_ai_models, "deepgram": deepgram_models, "elevenlabs": elevenlabs_models, @@ -1787,6 +1792,9 @@ if TYPE_CHECKING: PerplexityChatConfig as _PerplexityChatConfig, ) from .llms.nscale.chat.transformation import NscaleConfig as _NscaleConfig + from .llms.kluster_ai.chat.transformation import ( + KlusterAIConfig as _KlusterAIConfig, + ) from .llms.watsonx.chat.transformation import ( IBMWatsonXChatConfig as _IBMWatsonXChatConfig, ) @@ -1821,6 +1829,7 @@ if TYPE_CHECKING: AzureOpenAIO1Config: Type[_AzureOpenAIO1Config] PerplexityChatConfig: Type[_PerplexityChatConfig] NscaleConfig: Type[_NscaleConfig] + KlusterAIConfig: Type[_KlusterAIConfig] IBMWatsonXChatConfig: Type[_IBMWatsonXChatConfig] IBMWatsonXAIConfig: Type[_IBMWatsonXAIConfig] LiteLLMProxyChatConfig: Type[_LiteLLMProxyChatConfig] diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 17eb6609292..8851746751f 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -285,6 +285,7 @@ LLM_CONFIG_NAMES = ( "LMStudioChatConfig", "LmStudioEmbeddingConfig", "NscaleConfig", + "KlusterAIConfig", "PerplexityChatConfig", "AzureOpenAIO1Config", "IBMWatsonXAIConfig", @@ -1088,6 +1089,10 @@ _LLM_CONFIGS_IMPORT_MAP = { "LmStudioEmbeddingConfig", ), "NscaleConfig": (".llms.nscale.chat.transformation", "NscaleConfig"), + "KlusterAIConfig": ( + ".llms.kluster_ai.chat.transformation", + "KlusterAIConfig", + ), "PerplexityChatConfig": ( ".llms.perplexity.chat.transformation", "PerplexityChatConfig", diff --git a/litellm/constants.py b/litellm/constants.py index ae98b37d6e6..344f61be1c6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -611,6 +611,7 @@ LITELLM_CHAT_PROVIDERS = [ "meta_llama", "featherless_ai", "nscale", + "kluster_ai", "nebius", "dashscope", "moonshot", @@ -766,6 +767,7 @@ openai_compatible_endpoints: List = [ "api.llama.com/compat/v1/", "api.featherless.ai/v1", "inference.api.nscale.com/v1", + "api.kluster.ai/v1", "api.studio.nebius.ai/v1", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "https://api.moonshot.ai/v1", @@ -826,6 +828,7 @@ openai_compatible_providers: List = [ "chutes", # Chutes - JSON-configured provider "featherless_ai", "nscale", + "kluster_ai", "nebius", "dashscope", "moonshot", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index ba6d438f16c..0d43b0a4e99 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -311,6 +311,9 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == litellm.NscaleConfig.API_BASE_URL: custom_llm_provider = "nscale" dynamic_api_key = litellm.NscaleConfig.get_api_key() + elif endpoint == "api.kluster.ai/v1": + custom_llm_provider = "kluster_ai" + dynamic_api_key = litellm.KlusterAIConfig.get_api_key() elif endpoint == "dashscope-intl.aliyuncs.com/compatible-mode/v1": custom_llm_provider = "dashscope" dynamic_api_key = get_secret_str("DASHSCOPE_API_KEY") @@ -881,6 +884,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ) = litellm.NscaleConfig()._get_openai_compatible_provider_info( api_base=api_base, api_key=api_key ) + elif custom_llm_provider == "kluster_ai": + ( + api_base, + dynamic_api_key, + ) = litellm.KlusterAIConfig()._get_openai_compatible_provider_info( + api_base=api_base, api_key=api_key + ) elif custom_llm_provider == "heroku": ( api_base, diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index b8cdc8210fc..d6212b57994 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -261,6 +261,8 @@ def get_supported_openai_params( # noqa: PLR0915 return litellm.PerplexityChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "nscale": return litellm.NscaleConfig().get_supported_openai_params(model=model) + elif custom_llm_provider == "kluster_ai": + return litellm.KlusterAIConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "anyscale": return [ "temperature", diff --git a/litellm/llms/kluster_ai/chat/transformation.py b/litellm/llms/kluster_ai/chat/transformation.py new file mode 100644 index 00000000000..63b012bf7e1 --- /dev/null +++ b/litellm/llms/kluster_ai/chat/transformation.py @@ -0,0 +1,56 @@ +from typing import Optional, Tuple + +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.secret_managers.main import get_secret_str + + +class KlusterAIConfig(OpenAIGPTConfig): + """ + Reference: kluster.ai is OpenAI compatible. + + API Key: KLUSTER_AI_API_KEY + Default API Base: https://api.kluster.ai/v1 + """ + + API_BASE_URL = "https://api.kluster.ai/v1" + + @property + def custom_llm_provider(self) -> Optional[str]: + return "kluster_ai" + + @staticmethod + def get_api_key(api_key: Optional[str] = None) -> Optional[str]: + return api_key or get_secret_str("KLUSTER_AI_API_KEY") + + @staticmethod + def get_api_base(api_base: Optional[str] = None) -> Optional[str]: + return ( + api_base + or get_secret_str("KLUSTER_AI_API_BASE") + or KlusterAIConfig.API_BASE_URL + ) + + def _get_openai_compatible_provider_info( + self, api_base: Optional[str], api_key: Optional[str] + ) -> Tuple[Optional[str], Optional[str]]: + resolved_api_base = KlusterAIConfig.get_api_base(api_base) + resolved_api_key = KlusterAIConfig.get_api_key(api_key) + return resolved_api_base, resolved_api_key + + def get_supported_openai_params(self, model: str) -> list: + return [ + "max_tokens", + "n", + "temperature", + "top_p", + "stream", + "logprobs", + "top_logprobs", + "frequency_penalty", + "presence_penalty", + "response_format", + "stop", + "logit_bias", + "tools", + "tool_choice", + ] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8f471b62b5e..8020488d733 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3335,6 +3335,7 @@ class LlmProviders(str, Enum): GRADIENT_AI = "gradient_ai" LLAMA = "meta_llama" NSCALE = "nscale" + KLUSTER_AI = "kluster_ai" PG_VECTOR = "pg_vector" S3_VECTORS = "s3_vectors" HELICONE = "helicone" diff --git a/litellm/utils.py b/litellm/utils.py index 5a9dccc089e..d6a6a3ac538 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8331,6 +8331,10 @@ class ProviderConfigManager: ), LlmProviders.GRADIENT_AI: (lambda: litellm.GradientAIConfig(), False), LlmProviders.NSCALE: (lambda: litellm.NscaleConfig(), False), + LlmProviders.KLUSTER_AI: ( + lambda: litellm.KlusterAIConfig(), + False, + ), LlmProviders.HEROKU: (lambda: litellm.HerokuChatConfig(), False), LlmProviders.OCI: (lambda: litellm.OCIChatConfig(), False), LlmProviders.HYPERBOLIC: (lambda: litellm.HyperbolicChatConfig(), False), diff --git a/tests/test_litellm/llms/kluster_ai/chat/test_kluster_ai_chat_transformation.py b/tests/test_litellm/llms/kluster_ai/chat/test_kluster_ai_chat_transformation.py new file mode 100644 index 00000000000..bfc775e13ff --- /dev/null +++ b/tests/test_litellm/llms/kluster_ai/chat/test_kluster_ai_chat_transformation.py @@ -0,0 +1,101 @@ +import os +import sys +from unittest.mock import patch + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm import get_llm_provider, get_supported_openai_params +from litellm.llms.kluster_ai.chat.transformation import KlusterAIConfig + + +class TestKlusterAIConfig: + def setup_method(self): + self.config = KlusterAIConfig() + + def test_custom_llm_provider(self): + assert self.config.custom_llm_provider == "kluster_ai" + + def test_get_api_key(self): + assert self.config.get_api_key("test-key") == "test-key" + + with patch( + "litellm.llms.kluster_ai.chat.transformation.get_secret_str", + return_value="env-key", + ): + assert self.config.get_api_key() == "env-key" + + with patch.dict(os.environ, {"KLUSTER_AI_API_KEY": "env-key"}, clear=False): + assert self.config.get_api_key() == "env-key" + + def test_get_api_base_precedence(self): + # Explicit argument wins over everything. + assert ( + self.config.get_api_base("https://custom-base.com/v1") + == "https://custom-base.com/v1" + ) + + # KLUSTER_AI_API_BASE override. + with patch( + "litellm.llms.kluster_ai.chat.transformation.get_secret_str", + return_value="https://proxy.internal/v1", + ): + assert self.config.get_api_base() == "https://proxy.internal/v1" + + # Falls back to the default endpoint. + with patch( + "litellm.llms.kluster_ai.chat.transformation.get_secret_str", + return_value=None, + ): + assert self.config.get_api_base() == KlusterAIConfig.API_BASE_URL + assert KlusterAIConfig.API_BASE_URL == "https://api.kluster.ai/v1" + + def test_get_openai_compatible_provider_info(self): + with patch.dict(os.environ, {"KLUSTER_AI_API_KEY": "sk-secret"}, clear=False): + api_base, api_key = self.config._get_openai_compatible_provider_info( + api_base=None, api_key=None + ) + assert api_base == "https://api.kluster.ai/v1" + assert api_key == "sk-secret" + + def test_supported_params_include_tools(self): + params = self.config.get_supported_openai_params( + model="klusterai/Meta-Llama-3.1-405B-Instruct-Turbo" + ) + for expected in ("temperature", "stream", "tools", "tool_choice"): + assert expected in params + + +class TestKlusterAIProviderResolution: + def test_get_llm_provider_resolves_prefixed_model(self): + with patch.dict(os.environ, {"KLUSTER_AI_API_KEY": "sk-secret"}, clear=False): + model, provider, api_key, api_base = get_llm_provider( + model="kluster_ai/Meta-Llama-3.1-405B-Instruct-Turbo" + ) + assert model == "Meta-Llama-3.1-405B-Instruct-Turbo" + assert provider == "kluster_ai" + assert api_key == "sk-secret" + assert api_base == "https://api.kluster.ai/v1" + + def test_get_llm_provider_detects_provider_from_api_base(self): + _, provider, _, _ = get_llm_provider( + model="Meta-Llama-3.1-405B-Instruct-Turbo", + api_base="https://api.kluster.ai/v1", + api_key="sk-secret", + ) + assert provider == "kluster_ai" + + def test_get_supported_openai_params_routes_to_config(self): + params = get_supported_openai_params( + model="Meta-Llama-3.1-405B-Instruct-Turbo", + custom_llm_provider="kluster_ai", + ) + assert params is not None + assert "tools" in params + + def test_provider_registered_in_enum_and_lists(self): + from litellm.types.utils import LlmProviders + + assert LlmProviders.KLUSTER_AI.value == "kluster_ai" + assert "kluster_ai" in litellm.openai_compatible_providers + assert "api.kluster.ai/v1" in litellm.openai_compatible_endpoints