diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index b53eef19d1e..a237bfca9b8 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/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, + } diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index 069ac5727f6..2429e7c05c3 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/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