add testing and remove redundant function

This commit is contained in:
Philip Kiely
2025-08-21 11:25:26 -07:00
parent b8167cf304
commit 5fa4a18073
2 changed files with 54 additions and 11 deletions
-11
View File
@@ -3783,17 +3783,6 @@ def get_optional_params( # noqa: PLR0915
else False
),
)
elif custom_llm_provider == "baseten":
optional_params = litellm.BasetenConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(
drop_params
if drop_params is not None and isinstance(drop_params, bool)
else False
),
)
elif custom_llm_provider == "xai":
optional_params = litellm.XAIChatConfig().map_openai_params(
model=model,
@@ -0,0 +1,54 @@
import os
import pytest
from unittest.mock import patch
from litellm.llms.baseten.chat import BasetenConfig
class TestBasetenRouting:
"""Test Baseten routing logic"""
def test_routing_logic(self):
"""Test routing between Model API and dedicated deployments"""
config = BasetenConfig()
# Dedicated deployment (8-character alphanumeric)
assert config.get_api_base_for_model("abcd1234") == "https://model-abcd1234.api.baseten.co/environments/production/sync/v1"
# Model API (non-8-character)
assert config.get_api_base_for_model("openai/gpt-oss-120b") == "https://inference.baseten.co/v1"
class TestBasetenModelAPI:
"""Test Baseten Model API inference"""
@patch.dict(os.environ, {"BASETEN_API_KEY": "test-key"})
def test_model_api_inference(self):
"""Test Model API inference with basic parameters"""
config = BasetenConfig()
# Test parameter mapping
non_default_params = {
"max_tokens": 100,
"temperature": 0.7,
"top_p": 0.9
}
result = config.map_openai_params(
non_default_params=non_default_params,
optional_params={},
model="openai/gpt-oss-120b",
drop_params=False
)
assert result["max_tokens"] == 100
assert result["temperature"] == 0.7
assert result["top_p"] == 0.9
# Test provider info
api_base, api_key = config._get_openai_compatible_provider_info(None, "test-key")
assert api_base == "https://inference.baseten.co/v1"
assert api_key == "test-key"
if __name__ == "__main__":
pytest.main([__file__])