From 5fa4a180736e796124478507ea7be7bd4e69206d Mon Sep 17 00:00:00 2001
From: Philip Kiely <philip.kiely@baseten.co>
Date: Thu, 21 Aug 2025 11:25:26 -0700
Subject: [PATCH] add testing and remove redundant function

---
 litellm/utils.py                              | 11 ----
 .../baseten/chat/test_baseten_completions.py  | 54 +++++++++++++++++++
 2 files changed, 54 insertions(+), 11 deletions(-)
 create mode 100644 tests/test_litellm/llms/baseten/chat/test_baseten_completions.py

diff --git a/litellm/utils.py b/litellm/utils.py
index 9aea4372cb..647c03f908 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -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,
diff --git a/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py b/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py
new file mode 100644
index 0000000000..9420149a8e
--- /dev/null
+++ b/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py
@@ -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__])