From 9554dcec8c72c44938fcbe8e77740ef592de9978 Mon Sep 17 00:00:00 2001
From: Esteban Zeller <esteban@serfe.com>
Date: Tue, 24 Feb 2026 11:55:02 -0300
Subject: [PATCH] Enhance tests for map_openai_params functionality

Added tests for mapping OpenAI parameters including handling of tools and tool_choice.
---
 .../test_gradient_ai_chat_transformation.py   | 121 +++++++++++++++++-
 1 file changed, 120 insertions(+), 1 deletion(-)

diff --git a/tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py b/tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
index 66b4b36fcd3..92eb1ef8cc9 100644
--- a/tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
+++ b/tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
@@ -7,6 +7,7 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 
 from litellm.llms.gradient_ai.chat.transformation import GradientAIConfig, GRADIENT_AI_SERVERLESS_ENDPOINT
+from litellm.utils import UnsupportedParamsError
 
 DO_ENDPOINT_PATH = "/api/v1/chat/completions"
 DO_BASE_URL = "https://api.gradient_ai.com"
@@ -88,4 +89,122 @@ def test_get_openai_compatible_provider_info_default(monkeypatch, config):
     monkeypatch.setenv("GRADIENT_AI_API_KEY", "env-key")
     api_base, api_key = config._get_openai_compatible_provider_info(None, None)
     assert api_base == GRADIENT_AI_SERVERLESS_ENDPOINT
-    assert api_key == "env-key"
\ No newline at end of file
+    assert api_key == "env-key"
+
+
+
+def test_map_openai_params_allows_tools_explicitly(config):
+    """
+    Verifies that 'tools', 'tool_choice', and 'parallel_tool_calls' are included
+    in optional_params even if not present in the base supported_openai_params list.
+    """
+    non_default_params = {
+        "temperature": 0.7,
+        "max_tokens": 512,
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_current_weather",
+                    "description": "Get the current weather in a given location",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"}
+                        },
+                        "required": ["location"]
+                    }
+                }
+            }
+        ],
+        "tool_choice": "auto",
+        "parallel_tool_calls": True,
+        "unknown_param": "this should be ignored or raise depending on drop_params"
+    }
+
+    optional_params = {}
+
+    # Case 1: drop_params=False → should raise only on unknown_param, NOT on tools/tool_choice
+    with pytest.raises(UnsupportedParamsError) as exc_info:
+        result = config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params.copy(),
+            model="gradient_ai/test-model",
+            drop_params=False,
+        )
+
+    error_msg = str(exc_info.value)
+    assert "unknown_param" in error_msg
+    assert "tools" not in error_msg          # Important: should NOT complain about tools
+    assert "tool_choice" not in error_msg
+
+    # Case 2: drop_params=True → should pass without raising and preserve tools
+    result = config.map_openai_params(
+        non_default_params=non_default_params,
+        optional_params=optional_params.copy(),
+        model="gradient_ai/test-model",
+        drop_params=True,
+    )
+
+    assert result["temperature"] == 0.7
+    assert result["max_tokens"] == 512
+    assert "tools" in result
+    assert len(result["tools"]) == 1
+    assert result["tools"][0]["type"] == "function"
+    assert result["tool_choice"] == "auto"
+    assert result["parallel_tool_calls"] is True
+    # unknown_param should be dropped
+    assert "unknown_param" not in result
+
+
+def test_map_openai_params_no_tools_if_not_provided(config):
+    """Ensures tools-related keys are not added if not present in non_default_params"""
+    non_default_params = {
+        "temperature": 0.9,
+        "top_p": 0.95,
+    }
+
+    optional_params = {}
+    result = config.map_openai_params(
+        non_default_params=non_default_params,
+        optional_params=optional_params,
+        model="gradient_ai/test-model",
+        drop_params=False,
+    )
+
+    assert "tools" not in result
+    assert "tool_choice" not in result
+    assert "parallel_tool_calls" not in result
+    assert result["temperature"] == 0.9
+
+
+def test_map_openai_params_tool_choice_as_dict(config):
+    """Tests tool_choice in object format (forcing a specific tool)"""
+    non_default_params = {
+        "tool_choice": {
+            "type": "function",
+            "function": {"name": "search_database"}
+        },
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search_database",
+                    "description": "Search the internal database",
+                    "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}
+                }
+            }
+        ]
+    }
+
+    result = config.map_openai_params(
+        non_default_params=non_default_params,
+        optional_params={},
+        model="gradient_ai/test-model",
+        drop_params=True,
+    )
+
+    assert "tool_choice" in result
+    assert result["tool_choice"]["type"] == "function"
+    assert result["tool_choice"]["function"]["name"] == "search_database"
+    assert "tools" in result