Merge ae77461294 into 1c61c2606e

litellm/llms/zai/chat/transformation.py

@@ -1,3 +1,5 @@
+from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Final
 
 from litellm.secret_managers.main import get_secret_str
@@ -6,6 +8,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
 from ...openai.chat.gpt_transformation import OpenAIGPTConfig
 
 ZAI_API_BASE: Final = "https://api.z.ai/api/paas/v4"
+ZAI_REASONING_PARAMS: Final = frozenset(("thinking", "reasoning_effort"))
 
 
 class ZAIChatConfig(OpenAIGPTConfig):
@@ -49,8 +52,40 @@ class ZAIChatConfig(OpenAIGPTConfig):
 
         try:
             if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
-                base_params.append("thinking")
+                return [*base_params, *sorted(ZAI_REASONING_PARAMS)]  # mutable-ok: base class returns a list
         except Exception:
             pass
 
         return base_params
+
+    def _map_openai_params(
+        self,
+        non_default_params: dict[str, object],
+        optional_params: dict[str, object],
+        model: str,
+        drop_params: bool,
+    ) -> dict[str, object]:
+        supported_openai_params: Final = frozenset(self.get_supported_openai_params(model))
+        reasoning_params: Final = MappingProxyType(
+            {k: v for k, v in non_default_params.items() if k in ZAI_REASONING_PARAMS and k in supported_openai_params}
+        )
+        passthrough_params: Final = MappingProxyType(
+            {
+                k: v
+                for k, v in non_default_params.items()
+                if k in supported_openai_params and k not in ZAI_REASONING_PARAMS
+            }
+        )
+        if not reasoning_params:
+            return {**optional_params, **passthrough_params}  # mutable-ok: base class returns a dict
+        existing: Final = optional_params.get("extra_body")
+        existing_body: Final = existing if isinstance(existing, Mapping) else MappingProxyType({})
+        extra_body: Final = {  # mutable-ok: the OpenAI SDK json-encodes extra_body from a plain dict
+            **existing_body,
+            **reasoning_params,
+        }
+        return {  # mutable-ok: base class returns a dict
+            **optional_params,
+            **passthrough_params,
+            "extra_body": extra_body,
+        }

tests/test_litellm/llms/zai/test_zai_provider.py

@@ -2,6 +2,7 @@
 Tests for Z.AI (Zhipu AI) provider - GLM models
 """
 
+import json
 import math
 
 import pytest
@@ -89,9 +90,7 @@ async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
     monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
 
-    respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
-        json=zai_response
-    )
+    respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
 
     response = await litellm.acompletion(
         model="zai/glm-4.6",
@@ -115,9 +114,7 @@ def test_zai_sync_completion(respx_mock, zai_response, monkeypatch):
     monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
 
-    respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
-        json=zai_response
-    )
+    respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
 
     response = completion(
         model="zai/glm-4.6",
@@ -127,3 +124,112 @@ def test_zai_sync_completion(respx_mock, zai_response, monkeypatch):
 
     assert response.choices[0].message.content == "Hello! How can I help you today?"
     assert response.usage.total_tokens == 25
+
+
+@pytest.fixture
+def zai_thinking_response():
+    return {
+        "id": "chatcmpl-zai-thinking",
+        "object": "chat.completion",
+        "created": 1700000000,
+        "model": "glm-4.7",
+        "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
+        "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7},
+    }
+
+
+def _captured_body(respx_mock) -> dict:
+    assert len(respx_mock.calls) == 1
+    return json.loads(respx_mock.calls[0].request.content.decode("utf-8"))
+
+
+def test_reasoning_params_supported_on_reasoning_models(local_model_cost_map):
+    from litellm.llms.zai.chat.transformation import ZAIChatConfig
+
+    params = ZAIChatConfig().get_supported_openai_params(model="glm-4.7")
+    assert "thinking" in params
+    assert "reasoning_effort" in params
+
+
+def test_reasoning_params_not_supported_without_reasoning_flag(monkeypatch):
+    from litellm.llms.zai.chat.transformation import ZAIChatConfig
+
+    monkeypatch.setattr(litellm, "supports_reasoning", lambda **kwargs: False)
+    params = ZAIChatConfig().get_supported_openai_params(model="glm-4-32b-0414-128k")
+    assert "thinking" not in params
+    assert "reasoning_effort" not in params
+
+
+def test_thinking_and_reasoning_effort_move_into_extra_body(local_model_cost_map):
+    from litellm.llms.zai.chat.transformation import ZAIChatConfig
+
+    result = ZAIChatConfig()._map_openai_params(
+        non_default_params={"max_tokens": 100, "thinking": {"type": "disabled"}, "reasoning_effort": "low"},
+        optional_params={},
+        model="glm-4.7",
+        drop_params=False,
+    )
+    assert result["max_tokens"] == 100
+    assert "thinking" not in result
+    assert "reasoning_effort" not in result
+    assert result["extra_body"] == {"thinking": {"type": "disabled"}, "reasoning_effort": "low"}
+
+
+def test_caller_optional_params_and_extra_body_are_not_mutated(local_model_cost_map):
+    from litellm.llms.zai.chat.transformation import ZAIChatConfig
+
+    caller_extra_body = {"already_here": True}
+    caller_optional_params = {"stream": False, "extra_body": caller_extra_body}
+    result = ZAIChatConfig()._map_openai_params(
+        non_default_params={"max_tokens": 100, "thinking": {"type": "disabled"}},
+        optional_params=caller_optional_params,
+        model="glm-4.7",
+        drop_params=False,
+    )
+    assert result == {
+        "stream": False,
+        "max_tokens": 100,
+        "extra_body": {"already_here": True, "thinking": {"type": "disabled"}},
+    }
+    assert result is not caller_optional_params
+    assert caller_optional_params == {"stream": False, "extra_body": {"already_here": True}}
+    assert caller_optional_params["extra_body"] is caller_extra_body
+    assert caller_extra_body == {"already_here": True}
+
+
+def test_reasoning_params_dropped_for_non_reasoning_model(monkeypatch):
+    from litellm.utils import get_optional_params
+
+    monkeypatch.setattr(litellm, "supports_reasoning", lambda **kwargs: False)
+    result = get_optional_params(
+        model="glm-4-32b-0414-128k",
+        custom_llm_provider="zai",
+        thinking={"type": "disabled"},
+        reasoning_effort="low",
+        drop_params=True,
+    )
+    assert "thinking" not in result
+    assert "reasoning_effort" not in result
+    assert "thinking" not in result.get("extra_body", {})
+    assert "reasoning_effort" not in result.get("extra_body", {})
+
+
+@pytest.mark.asyncio
+async def test_thinking_and_reasoning_effort_reach_http_body(
+    respx_mock, zai_thinking_response, monkeypatch, local_model_cost_map
+):
+    monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_thinking_response)
+
+    await litellm.acompletion(
+        model="zai/glm-4.7",
+        messages=[{"role": "user", "content": "hi"}],
+        thinking={"type": "disabled"},
+        reasoning_effort="low",
+    )
+
+    body = _captured_body(respx_mock)
+    assert body["thinking"] == {"type": "disabled"}
+    assert body["reasoning_effort"] == "low"
+    assert "extra_body" not in body