diff --git a/litellm/__init__.py b/litellm/__init__.py index 038787f5cc..aa8357b6c7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -60,6 +60,7 @@ from litellm.constants import ( empower_models, together_ai_models, baseten_models, + WANDB_MODELS, REPEATED_STREAMING_CHUNK_LIMIT, request_timeout, open_ai_embedding_models, @@ -242,6 +243,7 @@ novita_api_key: Optional[str] = None snowflake_key: Optional[str] = None gradient_ai_api_key: Optional[str] = None nebius_key: Optional[str] = None +wandb_key: Optional[str] = None heroku_key: Optional[str] = None cometapi_key: Optional[str] = None ovhcloud_key: Optional[str] = None @@ -524,6 +526,7 @@ cometapi_models: Set = set() oci_models: Set = set() vercel_ai_gateway_models: Set = set() volcengine_models: Set = set() +wandb_models: Set = set(WANDB_MODELS) ovhcloud_models: Set = set() ovhcloud_embedding_models: Set = set() @@ -740,6 +743,8 @@ def add_known_models(): oci_models.add(key) elif value.get("litellm_provider") == "volcengine": volcengine_models.add(key) + elif value.get("litellm_provider") == "wandb": + wandb_models.add(key) elif value.get("litellm_provider") == "ovhcloud": ovhcloud_models.add(key) elif value.get("litellm_provider") == "ovhcloud-embedding-models": @@ -838,6 +843,7 @@ model_list = list( | heroku_models | vercel_ai_gateway_models | volcengine_models + | wandb_models | ovhcloud_models ) @@ -920,6 +926,7 @@ models_by_provider: dict = { "cometapi": cometapi_models, "oci": oci_models, "volcengine": volcengine_models, + "wandb": wandb_models, "ovhcloud": ovhcloud_models | ovhcloud_embedding_models, } @@ -1259,6 +1266,7 @@ from .llms.watsonx.chat.transformation import IBMWatsonXChatConfig from .llms.watsonx.embed.transformation import IBMWatsonXEmbeddingConfig from .llms.github_copilot.chat.transformation import GithubCopilotConfig from .llms.nebius.chat.transformation import NebiusConfig +from .llms.wandb.chat.transformation import WandbConfig from .llms.dashscope.chat.transformation import DashScopeChatConfig from .llms.moonshot.chat.transformation import MoonshotChatConfig from .llms.v0.chat.transformation import V0ChatConfig diff --git a/litellm/constants.py b/litellm/constants.py index dd61ba3fa8..005eb2bb6d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -313,6 +313,7 @@ LITELLM_CHAT_PROVIDERS = [ "morph", "lambda_ai", "vercel_ai_gateway", + "wandb", "ovhcloud", ] @@ -448,6 +449,7 @@ openai_compatible_endpoints: List = [ "https://api.lambda.ai/v1", "https://api.hyperbolic.xyz/v1", "https://ai-gateway.vercel.sh/v1", + "https://api.inference.wandb.ai/v1", ] @@ -492,6 +494,7 @@ openai_compatible_providers: List = [ "hyperbolic", "vercel_ai_gateway", "aiml", + "wandb", ] openai_text_completion_compatible_providers: List = ( [ # providers that support `/v1/completions` @@ -507,6 +510,7 @@ openai_text_completion_compatible_providers: List = ( "v0", "lambda_ai", "hyperbolic", + "wandb", ] ) _openai_like_providers: List = [ @@ -757,6 +761,38 @@ nebius_embedding_models: set = set( ] ) +WANDB_MODELS: set = set( + [ + # openai models + "openai/gpt-oss-120b", + "openai/gpt-oss-20b", + + # zai-org models + "zai-org/GLM-4.5", + + # Qwen models + "Qwen/Qwen3-235B-A22B-Instruct-2507", + "Qwen/Qwen3-Coder-480B-A35B-Instruct", + "Qwen/Qwen3-235B-A22B-Thinking-2507", + + # moonshotai + "moonshotai/Kimi-K2-Instruct", + + # meta models + "meta-llama/Llama-3.1-8B-Instruct", + "meta-llama/Llama-3.3-70B-Instruct", + "meta-llama/Llama-4-Scout-17B-16E-Instruct", + + # deepseek-ai + "deepseek-ai/DeepSeek-V3.1", + "deepseek-ai/DeepSeek-R1-0528", + "deepseek-ai/DeepSeek-V3-0324", + + # microsoft + "microsoft/Phi-4-mini-instruct", + ] +) + BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[ "cohere", "anthropic", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 414ccb7ab8..69c996d813 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -252,6 +252,9 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "https://ai-gateway.vercel.sh/v1": custom_llm_provider = "vercel_ai_gateway" dynamic_api_key = get_secret_str("VERCEL_AI_GATEWAY_API_KEY") + elif endpoint == "https://api.inference.wandb.ai/v1": + custom_llm_provider = "wandb" + dynamic_api_key = get_secret_str("WANDB_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception( @@ -773,6 +776,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ) = litellm.AIMLChatConfig()._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" + ) # type: ignore + dynamic_api_key = api_key or get_secret_str("WANDB_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(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 d77f53bd79..06e650f938 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -149,6 +149,9 @@ def get_supported_openai_params( # noqa: PLR0915 elif custom_llm_provider == "nebius": if request_type == "chat_completion": return litellm.NebiusConfig().get_supported_openai_params(model=model) + elif custom_llm_provider == "wandb": + if request_type == "chat_completion": + return litellm.WandbConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "replicate": return litellm.ReplicateConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "huggingface": diff --git a/litellm/llms/wandb/__init__.py b/litellm/llms/wandb/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/llms/wandb/chat/__init__.py b/litellm/llms/wandb/chat/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/llms/wandb/chat/transformation.py b/litellm/llms/wandb/chat/transformation.py new file mode 100644 index 0000000000..1cb2ab492b --- /dev/null +++ b/litellm/llms/wandb/chat/transformation.py @@ -0,0 +1,27 @@ +""" +Wandb Chat Completions API - Transformation + +This is OpenAI compatible - no translation needed / occurs +""" + +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + + +class WandbConfig(OpenAIGPTConfig): + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + map max_completion_tokens param to max_tokens + """ + supported_openai_params = self.get_supported_openai_params(model=model) + for param, value in non_default_params.items(): + if param == "max_completion_tokens": + optional_params["max_tokens"] = value + elif param in supported_openai_params: + optional_params[param] = value + return optional_params diff --git a/litellm/main.py b/litellm/main.py index 908c364220..5493c7e34e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1981,6 +1981,7 @@ def completion( # type: ignore # noqa: PLR0915 or custom_llm_provider == "openai" or custom_llm_provider == "together_ai" or custom_llm_provider == "nebius" + or custom_llm_provider == "wandb" or custom_llm_provider in litellm.openai_compatible_providers or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo ): # allow user to make an openai call with a custom base @@ -4400,6 +4401,27 @@ def embedding( # noqa: PLR0915 or "api.studio.nebius.ai/v1" ) + response = openai_chat_completions.embedding( + model=model, + input=input, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + ) + elif custom_llm_provider == "wandb": + api_key = api_key or litellm.api_key or get_secret_str("WANDB_API_KEY") + api_base = ( + api_base + or litellm.api_base + or get_secret_str("WANDB_API_BASE") + or "https://api.inference.wandb.ai/v1" + ) + response = openai_chat_completions.embedding( model=model, input=input, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 25f353b3bc..02d2166da2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2397,6 +2397,7 @@ class LlmProviders(str, Enum): AUTO_ROUTER = "auto_router" VERCEL_AI_GATEWAY = "vercel_ai_gateway" DOTPROMPT = "dotprompt" + WANDB = "wandb" OVHCLOUD = "ovhcloud" diff --git a/litellm/utils.py b/litellm/utils.py index e4cbaec706..515b6990ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3275,6 +3275,7 @@ def pre_process_optional_params( and custom_llm_provider != "openrouter" and custom_llm_provider != "vercel_ai_gateway" and custom_llm_provider != "nebius" + and custom_llm_provider != "wandb" and custom_llm_provider not in litellm.openai_compatible_providers ): if custom_llm_provider == "ollama": @@ -4446,6 +4447,9 @@ def get_api_key(llm_provider: str, dynamic_api_key: Optional[str]): # nebius elif llm_provider == "nebius": api_key = api_key or litellm.nebius_key or get_secret("NEBIUS_API_KEY") + # wandb + elif llm_provider == "wandb": + api_key = api_key or litellm.wandb_key or get_secret("WANDB_API_KEY") return api_key @@ -5530,6 +5534,11 @@ def validate_environment( # noqa: PLR0915 keys_in_environment = True else: missing_keys.append("NEBIUS_API_KEY") + elif custom_llm_provider == "wandb": + if "WANDB_API_KEY" in os.environ: + keys_in_environment = True + else: + missing_keys.append("WANDB_API_KEY") elif custom_llm_provider == "dashscope": if "DASHSCOPE_API_KEY" in os.environ: keys_in_environment = True @@ -5644,6 +5653,11 @@ def validate_environment( # noqa: PLR0915 keys_in_environment = True else: missing_keys.append("NEBIUS_API_KEY") + elif model in litellm.wandb_models: + if "WANDB_API_KEY" in os.environ: + keys_in_environment = True + else: + missing_keys.append("WANDB_API_KEY") def filter_missing_keys(keys: List[str], exclude_pattern: str) -> List[str]: """Filter out keys that contain the exclude_pattern (case insensitive).""" @@ -7046,6 +7060,8 @@ class ProviderConfigManager: return litellm.NovitaConfig() elif litellm.LlmProviders.NEBIUS == provider: return litellm.NebiusConfig() + elif litellm.LlmProviders.WANDB == provider: + return litellm.WandbConfig() elif litellm.LlmProviders.DASHSCOPE == provider: return litellm.DashScopeChatConfig() elif litellm.LlmProviders.MOONSHOT == provider: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 43ea3af320..aa30efd606 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20943,6 +20943,132 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "wandb/openai/gpt-oss-120b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 0.015, + "output_cost_per_token": 0.06, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/openai/gpt-oss-20b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 0.005, + "output_cost_per_token": 0.02, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/zai-org/GLM-4.5": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 0.055, + "output_cost_per_token": 0.2, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/Qwen/Qwen3-235B-A22B-Instruct-2507": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.01, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/Qwen/Qwen3-Coder-480B-A35B-Instruct": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 0.1, + "output_cost_per_token": 0.15, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.01, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/moonshotai/Kimi-K2-Instruct": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.135, + "output_cost_per_token": 0.4, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/meta-llama/Llama-3.1-8B-Instruct": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.022, + "output_cost_per_token": 0.022, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/deepseek-ai/DeepSeek-V3.1": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.055, + "output_cost_per_token": 0.165, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/deepseek-ai/DeepSeek-R1-0528": { + "max_tokens": 161000, + "max_input_tokens": 161000, + "max_output_tokens": 161000, + "input_cost_per_token": 0.135, + "output_cost_per_token": 0.54, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/deepseek-ai/DeepSeek-V3-0324": { + "max_tokens": 161000, + "max_input_tokens": 161000, + "max_output_tokens": 161000, + "input_cost_per_token": 0.114, + "output_cost_per_token": 0.275, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/meta-llama/Llama-3.3-70B-Instruct": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.071, + "output_cost_per_token": 0.071, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "max_tokens": 64000, + "max_input_tokens": 64000, + "max_output_tokens": 64000, + "input_cost_per_token": 0.017, + "output_cost_per_token": 0.066, + "litellm_provider": "wandb", + "mode": "chat" + }, + "wandb/microsoft/Phi-4-mini-instruct": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.008, + "output_cost_per_token": 0.035, + "litellm_provider": "wandb", + "mode": "chat" + }, "watsonx/ibm/granite-3-8b-instruct": { "input_cost_per_token": 0.0002, "litellm_provider": "watsonx", diff --git a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py b/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py new file mode 100644 index 0000000000..ef7bb0e44f --- /dev/null +++ b/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py @@ -0,0 +1,146 @@ +""" +Unit tests for WandB Inference configuration. + +These tests validate the WandbInferenceConfig class which extends OpenAIGPTConfig. +Nebius AI Studio is an OpenAI-compatible provider with minor customizations. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +import pytest + +import litellm +from litellm import completion +from litellm.llms.wandb.chat.transformation import WandbConfig + + +class TestWandbConfig: + """Test class for WandB Inference functionality""" + + def test_default_api_base(self): + """Test that default API base is used when none is provided""" + config = WandbConfig() + headers = {} + api_key = "fake-wandb-key" + + # Call validate_environment without specifying api_base + result = config.validate_environment( + headers=headers, + model="wandb/openai/gpt-oss-20b", + messages=[{"role": "user", "content": "Hey"}], + optional_params={}, + litellm_params={}, + api_key=api_key, + api_base=None, # Not providing api_base + ) + + # Verify headers are still set correctly + assert result["Authorization"] == f"Bearer {api_key}" + assert result["Content-Type"] == "application/json" + + # We can't directly test the api_base value here since validate_environment + # only returns the headers, but we can verify it doesn't raise an exception + # which would happen if api_base handling was incorrect + + @pytest.mark.respx() + def test_wandb_completion_mock(self, respx_mock): + """ + Mock test for WandB Inference completion using the model format from docs. + This test mocks the actual HTTP request to test the integration properly. + """ + + litellm.disable_aiohttp_transport = ( + True # since this uses respx, we need to set use_aiohttp_transport to False + ) + + # Set up environment variables for the test + api_key = "fake-wandb-key" + api_base = "https://api.inference.wandb.ai/v1" + model = "wandb/openai/gpt-oss-20b" + model_name = "gpt-oss-20b" # The actual model name without provider prefix + + # Mock the HTTP request to the WandB Inference API + respx_mock.post(f"{api_base}/chat/completions").respond( + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model_name, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": '```python\nprint("Hey from LiteLLM!")\n```\n\nThis simple Python code prints a greeting message from LiteLLM.', + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + status_code=200, + ) + + # Make the actual API call through LiteLLM + response = completion( + model=model, + messages=[ + {"role": "user", "content": "write code for saying hey from LiteLLM"} + ], + api_key=api_key, + api_base=api_base, + ) + + # Verify response structure + assert response is not None + + # If response is a streaming wrapper, extract the first chunk for assertions + # This handles both streaming and non-streaming responses + # For streaming, response is typically an iterator yielding (event, data) tuples + if hasattr(response, "__iter__") and not hasattr(response, "choices"): + # Streaming response: get the first chunk + first_chunk = next(iter(response)) + # first_chunk is likely a tuple: (event, data) + # Try to extract the data part + if isinstance(first_chunk, tuple) and len(first_chunk) == 2: + data = first_chunk[1] + else: + data = first_chunk + + # The data object should have .choices[0] with .delta or .message + choices = getattr(data, "choices", None) + assert choices is not None + assert len(choices) > 0 + choice = choices[0] + # For streaming, content may be in .delta or .message + content = None + if hasattr(choice, "delta") and hasattr(choice.delta, "content"): + content = choice.delta.content + elif hasattr(choice, "message") and hasattr(choice.message, "content"): + content = choice.message.content + assert content is not None + assert "```python" in content + assert "Hey from LiteLLM" in content + else: + # Non-streaming response + choices = getattr(response, "choices", None) + assert choices is not None + assert len(choices) > 0 + choice = choices[0] + message = getattr(choice, "message", None) + assert message is not None + content = getattr(message, "content", None) + assert content is not None + + # Check for specific content in the response + assert "```python" in content + assert "Hey from LiteLLM" in content