tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-22 22:24:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14416 from xprilion/wandb-inference

Browse Source 
Add W&B Inference to LiteLLM
This commit is contained in:
Krish Dholakia
2025-09-22 16:36:46 -07:00
committed by GitHub
parent fa20abfe48 0f1de92e19
commit 390e47cbf3
12 changed files with 395 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/__init__.py
+8
View File
@@ -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
litellm/constants.py
+36
View File
@@ -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",
litellm/litellm_core_utils/get_llm_provider_logic.py
+10
View File
@@ -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))
litellm/litellm_core_utils/get_supported_openai_params.py
+3
View File
@@ -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":
litellm/llms/wandb/__init__.py
View File
litellm/llms/wandb/chat/__init__.py
View File
litellm/llms/wandb/chat/transformation.py
+27
View File
@@ -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
litellm/main.py
+22
View File
@@ -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,
litellm/types/utils.py
+1
View File
@@ -2397,6 +2397,7 @@ class LlmProviders(str, Enum):
AUTO_ROUTER = "auto_router"
VERCEL_AI_GATEWAY = "vercel_ai_gateway"
DOTPROMPT = "dotprompt"
WANDB = "wandb"
OVHCLOUD = "ovhcloud"
litellm/utils.py
+16
View File
@@ -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:
model_prices_and_context_window.json
+126
View File
@@ -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",
tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py
+146
View File
@@ -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