mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(wandb): gate reasoning effort on model capabilities
This commit is contained in:
Wolfram Ravenwolf
parent
1bf9418917
commit
13ddd1ec64
4 changed files with 219 additions and 37 deletions
|
|
@ -6,12 +6,16 @@ This is OpenAI compatible - no translation needed / occurs
|
|||
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
|
||||
class WandbConfig(OpenAIGPTConfig):
|
||||
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract
|
||||
return super().get_supported_openai_params(model) + ["reasoning_effort"] # mutable-ok: inherited contract
|
||||
supported_params: Final = super().get_supported_openai_params(model)
|
||||
if litellm.supports_reasoning(model=model, custom_llm_provider="wandb"):
|
||||
return supported_params + ["reasoning_effort"] # mutable-ok: inherited contract
|
||||
return supported_params
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -45960,6 +45960,7 @@
|
|||
"output_cost_per_token": 0.0
|
||||
},
|
||||
"wandb/openai/gpt-oss-120b": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
|
|
@ -45970,6 +45971,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/openai/gpt-oss-20b": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
|
|
@ -45980,6 +45982,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/zai-org/GLM-4.5": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
|
|
@ -46008,6 +46011,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
|
|
@ -46064,6 +46068,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V3.1": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 161000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -46074,6 +46079,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-R1-0528": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 161000,
|
||||
"max_input_tokens": 161000,
|
||||
"max_output_tokens": 161000,
|
||||
|
|
@ -55631,6 +55637,7 @@
|
|||
"supports_vision": false
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Flash": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 1048576,
|
||||
"max_input_tokens": 1048576,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
|
|
@ -55643,6 +55650,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Flash-0731": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 1.3e-07,
|
||||
|
|
@ -55655,6 +55663,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Pro": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 1048576,
|
||||
"max_input_tokens": 1048576,
|
||||
"input_cost_per_token": 1.15e-06,
|
||||
|
|
@ -55667,6 +55676,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/google/gemma-4-31B-it": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 1e-07,
|
||||
|
|
@ -55707,6 +55717,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/MiniMaxAI/MiniMax-M3": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 2.3e-07,
|
||||
|
|
@ -55719,6 +55730,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/moonshotai/Kimi-K2.7-Code": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 7.1e-07,
|
||||
|
|
@ -55731,6 +55743,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/moonshotai/Kimi-K2.6": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 6.5e-07,
|
||||
|
|
@ -55743,6 +55756,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 1e-07,
|
||||
|
|
@ -55755,6 +55769,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
|
|
@ -55777,6 +55792,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.8-27B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 4e-07,
|
||||
|
|
@ -55789,6 +55805,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.6-35B-A3B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -55799,6 +55816,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.6-27B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 6e-07,
|
||||
|
|
@ -55811,6 +55829,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.5-35B-A3B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -55830,7 +55849,28 @@
|
|||
"supports_vision": false,
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Pro-0813": {
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"supports_reasoning": true,
|
||||
"input_cost_per_token": 0.00000131,
|
||||
"output_cost_per_token": 0.00000396,
|
||||
"cache_read_input_token_cost": 0.000000044,
|
||||
"supports_prompt_caching": true,
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/ibm-granite/granite-4.2-8b": {
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"supports_reasoning": true,
|
||||
"input_cost_per_token": 0.0000001,
|
||||
"output_cost_per_token": 0.00000015,
|
||||
"cache_read_input_token_cost": 0.00000005,
|
||||
"supports_prompt_caching": true,
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/zai-org/GLM-5.2": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 7.6e-07,
|
||||
|
|
|
|||
|
|
@ -45960,6 +45960,7 @@
|
|||
"output_cost_per_token": 0.0
|
||||
},
|
||||
"wandb/openai/gpt-oss-120b": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
|
|
@ -45970,6 +45971,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/openai/gpt-oss-20b": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
|
|
@ -45980,6 +45982,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/zai-org/GLM-4.5": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
|
|
@ -46008,6 +46011,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
|
|
@ -46064,6 +46068,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V3.1": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 161000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -46074,6 +46079,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-R1-0528": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 161000,
|
||||
"max_input_tokens": 161000,
|
||||
"max_output_tokens": 161000,
|
||||
|
|
@ -55631,6 +55637,7 @@
|
|||
"supports_vision": false
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Flash": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 1048576,
|
||||
"max_input_tokens": 1048576,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
|
|
@ -55643,6 +55650,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Flash-0731": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 1.3e-07,
|
||||
|
|
@ -55655,6 +55663,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Pro": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 1048576,
|
||||
"max_input_tokens": 1048576,
|
||||
"input_cost_per_token": 1.15e-06,
|
||||
|
|
@ -55667,6 +55676,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/google/gemma-4-31B-it": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 1e-07,
|
||||
|
|
@ -55707,6 +55717,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/MiniMaxAI/MiniMax-M3": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 2.3e-07,
|
||||
|
|
@ -55719,6 +55730,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/moonshotai/Kimi-K2.7-Code": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 7.1e-07,
|
||||
|
|
@ -55731,6 +55743,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/moonshotai/Kimi-K2.6": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 6.5e-07,
|
||||
|
|
@ -55743,6 +55756,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 1e-07,
|
||||
|
|
@ -55755,6 +55769,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
|
|
@ -55777,6 +55792,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.8-27B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 4e-07,
|
||||
|
|
@ -55789,6 +55805,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.6-35B-A3B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -55799,6 +55816,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.6-27B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 6e-07,
|
||||
|
|
@ -55811,6 +55829,7 @@
|
|||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/Qwen/Qwen3.5-35B-A3B": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -55830,7 +55849,28 @@
|
|||
"supports_vision": false,
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V4-Pro-0813": {
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"supports_reasoning": true,
|
||||
"input_cost_per_token": 0.00000131,
|
||||
"output_cost_per_token": 0.00000396,
|
||||
"cache_read_input_token_cost": 0.000000044,
|
||||
"supports_prompt_caching": true,
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/ibm-granite/granite-4.2-8b": {
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"supports_reasoning": true,
|
||||
"input_cost_per_token": 0.0000001,
|
||||
"output_cost_per_token": 0.00000015,
|
||||
"cache_read_input_token_cost": 0.00000005,
|
||||
"supports_prompt_caching": true,
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
},
|
||||
"wandb/zai-org/GLM-5.2": {
|
||||
"supports_reasoning": true,
|
||||
"max_tokens": 262144,
|
||||
"max_input_tokens": 262144,
|
||||
"input_cost_per_token": 7.6e-07,
|
||||
|
|
|
|||
|
|
@ -6,30 +6,90 @@ Nebius AI Studio is an OpenAI-compatible provider with minor customizations.
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
from litellm.llms.wandb.chat.transformation import WandbConfig
|
||||
|
||||
|
||||
WANDB_REASONING_MODELS: Final = (
|
||||
"deepseek-ai/DeepSeek-V4-Flash",
|
||||
"deepseek-ai/DeepSeek-V4-Flash-0731",
|
||||
"deepseek-ai/DeepSeek-V4-Pro",
|
||||
"deepseek-ai/DeepSeek-V4-Pro-0813",
|
||||
"deepseek-ai/DeepSeek-V3.1",
|
||||
"google/gemma-4-31B-it",
|
||||
"ibm-granite/granite-4.2-8b",
|
||||
"MiniMaxAI/MiniMax-M3",
|
||||
"moonshotai/Kimi-K2.7-Code",
|
||||
"moonshotai/Kimi-K2.6",
|
||||
"nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B",
|
||||
"nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B",
|
||||
"openai/gpt-oss-120b",
|
||||
"openai/gpt-oss-20b",
|
||||
"Qwen/Qwen3.8-27B",
|
||||
"Qwen/Qwen3.6-35B-A3B",
|
||||
"Qwen/Qwen3.6-27B",
|
||||
"Qwen/Qwen3.5-35B-A3B",
|
||||
"zai-org/GLM-5.2",
|
||||
"moonshotai/Kimi-K2.5",
|
||||
"MiniMaxAI/MiniMax-M2.5",
|
||||
"zai-org/GLM-4.5",
|
||||
"Qwen/Qwen3-235B-A22B-Thinking-2507",
|
||||
"deepseek-ai/DeepSeek-R1-0528",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def wandb_test_config(local_model_cost_map, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "telemetry", False)
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def wandb_request_mock(respx_mock: respx.MockRouter) -> respx.Route:
|
||||
return respx_mock.post("https://api.inference.wandb.ai/v1/chat/completions").respond(
|
||||
json={
|
||||
"id": "chatcmpl-123",
|
||||
"object": "chat.completion",
|
||||
"created": 1677652288,
|
||||
"model": "test-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Done"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
|
||||
class TestWandbConfig:
|
||||
"""Test class for WandB Inference functionality"""
|
||||
|
||||
def test_map_openai_params_preserves_reasoning_effort(self):
|
||||
supported_params = litellm.get_supported_openai_params(model="wandb/openai/gpt-oss-20b")
|
||||
@pytest.mark.parametrize("model", WANDB_REASONING_MODELS)
|
||||
def test_map_openai_params_preserves_reasoning_effort(self, wandb_test_config, model: str):
|
||||
assert litellm.model_cost[f"wandb/{model}"].get("supports_reasoning") is True
|
||||
supported_params = litellm.get_supported_openai_params(model=f"wandb/{model}")
|
||||
assert supported_params is not None
|
||||
assert "reasoning_effort" in supported_params
|
||||
|
||||
result = WandbConfig().map_openai_params(
|
||||
non_default_params={"reasoning_effort": "medium"},
|
||||
non_default_params={"reasoning_effort": "medium", "max_completion_tokens": 64},
|
||||
optional_params={},
|
||||
model="openai/gpt-oss-20b",
|
||||
model=model,
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
assert result == {"reasoning_effort": "medium"}
|
||||
assert result == {"reasoning_effort": "medium", "max_tokens": 64}
|
||||
|
||||
def test_default_api_base(self):
|
||||
"""Test that default API base is used when none is provided"""
|
||||
|
|
@ -155,40 +215,78 @@ class TestWandbConfig:
|
|||
assert "Hey from LiteLLM" in content
|
||||
|
||||
@pytest.mark.respx()
|
||||
def test_wandb_completion_preserves_reasoning_effort_with_drop_params(self, respx_mock, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
api_base = "https://api.inference.wandb.ai/v1"
|
||||
request_mock = respx_mock.post(f"{api_base}/chat/completions").respond(
|
||||
json={
|
||||
"id": "chatcmpl-123",
|
||||
"object": "chat.completion",
|
||||
"created": 1677652288,
|
||||
"model": "gpt-oss-20b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Done"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 1,
|
||||
"completion_tokens": 1,
|
||||
"total_tokens": 2,
|
||||
},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,effort",
|
||||
tuple((model, "medium") for model in WANDB_REASONING_MODELS)
|
||||
+ (
|
||||
("Qwen/Qwen3.8-27B", "low"),
|
||||
("Qwen/Qwen3.8-27B", "xhigh"),
|
||||
),
|
||||
)
|
||||
def test_wandb_completion_preserves_reasoning_effort_with_drop_params(
|
||||
self, wandb_test_config, wandb_request_mock: respx.Route, model: str, effort: str
|
||||
):
|
||||
completion(
|
||||
model="wandb/openai/gpt-oss-20b",
|
||||
model=f"wandb/{model}",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_key="fake-wandb-key",
|
||||
api_base=api_base,
|
||||
reasoning_effort="medium",
|
||||
api_base="https://api.inference.wandb.ai/v1",
|
||||
reasoning_effort=effort,
|
||||
max_completion_tokens=64,
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
request_body = json.loads(request_mock.calls[0].request.content)
|
||||
assert request_body["reasoning_effort"] == "medium"
|
||||
assert wandb_request_mock.call_count == 1
|
||||
request_body = json.loads(wandb_request_mock.calls[0].request.content)
|
||||
assert request_body["model"] == model
|
||||
assert request_body["reasoning_effort"] == effort
|
||||
assert request_body["max_tokens"] == 64
|
||||
assert "max_completion_tokens" not in request_body
|
||||
|
||||
@pytest.mark.respx(assert_all_called=False)
|
||||
@pytest.mark.parametrize("drop_params", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"model,explicit_false",
|
||||
[
|
||||
("meta-llama/Llama-3.1-8B-Instruct", False),
|
||||
("unknown-model", False),
|
||||
("openai/gpt-oss-20b", True),
|
||||
],
|
||||
)
|
||||
def test_wandb_completion_without_reasoning_support(
|
||||
self,
|
||||
wandb_test_config,
|
||||
wandb_request_mock: respx.Route,
|
||||
respx_mock: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
model: str,
|
||||
explicit_false: bool,
|
||||
drop_params: bool,
|
||||
):
|
||||
with monkeypatch.context() as context:
|
||||
if explicit_false:
|
||||
context.setitem(litellm.model_cost[f"wandb/{model}"], "supports_reasoning", False)
|
||||
|
||||
kwargs = {
|
||||
"model": f"wandb/{model}",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"api_key": "fake-wandb-key",
|
||||
"api_base": "https://api.inference.wandb.ai/v1",
|
||||
"reasoning_effort": "medium",
|
||||
"drop_params": drop_params,
|
||||
}
|
||||
if not drop_params:
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match="reasoning_effort"):
|
||||
completion(**kwargs)
|
||||
assert len(respx_mock.calls) == 0
|
||||
return
|
||||
|
||||
completion(**kwargs)
|
||||
assert wandb_request_mock.call_count == 1
|
||||
request_body = json.loads(wandb_request_mock.calls[0].request.content)
|
||||
assert request_body["model"] == model
|
||||
assert "reasoning_effort" not in request_body
|
||||
|
||||
supported_params = litellm.get_supported_openai_params(model=f"wandb/{model}")
|
||||
assert supported_params is not None
|
||||
assert "reasoning_effort" not in supported_params
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue