From 342ce6ae4a4549cacaf96e47aa80a096ae1035bd Mon Sep 17 00:00:00 2001 From: DewellNaa Date: Thu, 29 Oct 2026 20:30:50 -0700 Subject: [PATCH] feat(zai): support response_format and max_completion_tokens, add glm-5v-turbo pricing (#48335) * feat(zai): fall back to BIGMODEL/ZHIPUAI api keys and add GLM-5.3 aliases and glm-5v-turbo Resolve the zai api key from ZAI_API_KEY, BIGMODEL_API_KEY or ZHIPUAI_API_KEY so deployments predating the zai rename keep working. Accept response_format and max_completion_tokens as supported openai params. Register zai/GLM-5.3 and zai/GLM-5.3-Flash aliases plus the glm-5v-turbo vision model in both the main and backup cost maps * fix(zai): read only ZAI_API_KEY, dropping the BigModel key fallbacks * fix(zai): map max_completion_tokens to max_tokens, correct glm-5v-turbo pricing, drop redundant aliases * fix(zai): reject json_schema response_format, add glm-5v-turbo cache read price * fix(zai): forward every response_format type so allowed_openai_params keeps working * fix(zai): follow json_schema natively on glm-5.3, prompt other models, drop over-limit max_tokens under drop_params * fix(zai): keep BigModel's json_schema as sent and never drop the caller's max_tokens under drop_params * fix(zai): write the prompted JSON schema as JSON so models answer with valid JSON * test(zai): type the zai provider tests and mark their locals Final --------- Co-authored-by: DewellNaa <210961864+DewellNaa@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/factory.py | 34 ++- litellm/llms/bigmodel/chat/transformation.py | 4 + litellm/llms/zai/chat/transformation.py | 94 ++++++++- ...odel_prices_and_context_window_backup.json | 17 ++ model_prices_and_context_window.json | 17 ++ .../llms/bigmodel/test_bigmodel_provider.py | 23 ++ tests/unit/llms/zai/test_zai_provider.py | 197 ++++++++++++++++++ 7 files changed, 364 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index fad4caa59a0..834d26b4fe8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5878,38 +5878,36 @@ def function_call_prompt( ] -def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str: +def response_schema_prompt(model: str, response_schema: Mapping[str, object], llm_provider: str) -> str: """ Decides if a user-defined custom prompt or default needs to be used Returns the prompt str that's passed to the model as a user message """ - custom_prompt_details: dict | None = None - response_schema_as_message: Final = [{"role": "user", "content": f"{response_schema}"}] - if f"{model}/response_schema_prompt" in litellm.custom_prompt_dict: - custom_prompt_details = litellm.custom_prompt_dict[ - f"{model}/response_schema_prompt" - ] # allow user to define custom response schema prompt by model - elif "response_schema_prompt" in litellm.custom_prompt_dict: - custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"] - - if custom_prompt_details is not None: - return render_custom_prompt_template( - template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider - ) - else: + prompts: Final[Mapping[str, Mapping[str, object]]] = litellm.custom_prompt_dict + model_prompt_key: Final = f"{model}/response_schema_prompt" + custom_prompt_details: Final = prompts.get( + model_prompt_key if model_prompt_key in prompts else "response_schema_prompt" + ) + if custom_prompt_details is None: return default_response_schema_prompt(response_schema=response_schema) + return render_custom_prompt_template( + template=custom_prompt_details, + messages=[{"role": "user", "content": f"{response_schema}"}], + model=model, + llm_provider=llm_provider, + ) -def default_response_schema_prompt(response_schema: dict) -> str: +def default_response_schema_prompt(response_schema: Mapping[str, object]) -> str: """ Used if provider/model doesn't support 'response_schema' param. This is the default prompt. Allow user to override this with a custom_prompt. """ prompt_str: Final = f"""Use this JSON schema: - ```json - {response_schema} + ```json + {json.dumps(response_schema, default=str)} ```""" return prompt_str diff --git a/litellm/llms/bigmodel/chat/transformation.py b/litellm/llms/bigmodel/chat/transformation.py index 12d5e912c15..93e9e24fe7b 100644 --- a/litellm/llms/bigmodel/chat/transformation.py +++ b/litellm/llms/bigmodel/chat/transformation.py @@ -26,6 +26,10 @@ class BigModelChatConfig(ZAIChatConfig): def custom_llm_provider(self) -> str: return "bigmodel" + @property + def prompts_unsupported_json_schema(self) -> bool: + return False + def get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]: return get_bigmodel_url(api_base, "chat"), get_bigmodel_api_key(api_key) diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index e0e42d772f8..9e2b97d2bed 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -1,11 +1,43 @@ +from collections.abc import Mapping from typing import Final +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 + +import litellm +from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ChatCompletionUserMessage +from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration from ...openai.chat.gpt_transformation import OpenAIGPTConfig ZAI_API_BASE: Final = "https://api.z.ai/api/paas/v4" +TOKEN_LIMIT_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens"}) + + +def _is_json_object( + value: object, +) -> TypeIs[Mapping[str, object]]: # guard-ok: cost-map rows and request params are JSON + return isinstance(value, Mapping) + + +def _json_field(value: object, key: str) -> object: + return value.get(key) if _is_json_object(value) else None + + +def _exceeds_output_limit(model: str, max_tokens: object) -> bool: + info: Final = lookup_model_info_for_registration(model=model, custom_llm_provider="zai") + limit: Final = None if isinstance(info, ModelInfoLookupMiss) else info.get("max_output_tokens") + return isinstance(max_tokens, int) and isinstance(limit, int) and max_tokens > limit + + +def _schema_without_native_support(model: str, response_format: object) -> Mapping[str, object] | None: + if _json_field(response_format, "type") != "json_schema": + return None + if litellm.supports_response_schema(model=model, custom_llm_provider="zai"): + return None + schema: Final = _json_field(_json_field(response_format, "json_schema"), "schema") + return schema if _is_json_object(schema) else None class ZAIChatConfig(OpenAIGPTConfig): @@ -13,6 +45,10 @@ class ZAIChatConfig(OpenAIGPTConfig): def custom_llm_provider(self) -> str | None: return "zai" + @property + def prompts_unsupported_json_schema(self) -> bool: + return True + def _get_openai_compatible_provider_info( self, api_base: str | None, api_key: str | None ) -> tuple[str | None, str | None]: @@ -40,9 +76,11 @@ class ZAIChatConfig(OpenAIGPTConfig): # GLM/ZAI supports cache_control, so return messages and tools unchanged return messages, tools - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list[str]: base_params: Final = [ "max_tokens", + "max_completion_tokens", + "response_format", "stream", "stream_options", "temperature", @@ -52,8 +90,6 @@ class ZAIChatConfig(OpenAIGPTConfig): "tool_choice", ] - import litellm - try: if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_params.append("thinking") @@ -61,3 +97,53 @@ class ZAIChatConfig(OpenAIGPTConfig): pass return base_params + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: request pipeline adds parameters to the returned body + supported_params: Final = self.get_supported_openai_params(model) + sent: Final = {param: value for param, value in non_default_params.items() if param in supported_params} + max_completion_tokens: Final = sent.get("max_completion_tokens") + falls_back_to_max_tokens: Final = max_completion_tokens is None or ( + (litellm.drop_params or drop_params) and _exceeds_output_limit(model, max_completion_tokens) + ) + max_tokens: Final = sent.get("max_tokens") if falls_back_to_max_tokens else max_completion_tokens + return { + **optional_params, + **{param: value for param, value in sent.items() if param not in TOKEN_LIMIT_PARAMS}, + **({} if max_tokens is None else {"max_tokens": max_tokens}), + } + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: OpenAIGPTConfig.transform_request signature + optional_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature + litellm_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature + headers: dict[str, str], # mutable-ok: OpenAIGPTConfig.transform_request signature + ) -> dict[str, object]: # mutable-ok: inherited contract returns the request body + schema: Final = ( + _schema_without_native_support(model, optional_params.get("response_format")) + if self.prompts_unsupported_json_schema + else None + ) + schema_message: Final = ( + () + if schema is None + else ( + ChatCompletionUserMessage( + role="user", content=response_schema_prompt(model=model, response_schema=schema, llm_provider="zai") + ), + ) + ) + return super().transform_request( + model, + [*messages, *schema_message], + optional_params if schema is None else {**optional_params, "response_format": {"type": "json_object"}}, + litellm_params, + headers, + ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 77ac4f49e8b..96a04c96f44 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54512,6 +54512,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true }, "zai/glm-5.3-flash": { @@ -54526,10 +54527,26 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing", "supports_vision": true }, + "zai/glm-5v-turbo": { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://vercel.com/ai-gateway/models/glm-5v-turbo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "zai/glm-5.1": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 2.6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 77ac4f49e8b..96a04c96f44 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54512,6 +54512,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true }, "zai/glm-5.3-flash": { @@ -54526,10 +54527,26 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing", "supports_vision": true }, + "zai/glm-5v-turbo": { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://vercel.com/ai-gateway/models/glm-5v-turbo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "zai/glm-5.1": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 2.6e-07, diff --git a/tests/unit/llms/bigmodel/test_bigmodel_provider.py b/tests/unit/llms/bigmodel/test_bigmodel_provider.py index cdf98c8f9a4..6abfcb060d0 100644 --- a/tests/unit/llms/bigmodel/test_bigmodel_provider.py +++ b/tests/unit/llms/bigmodel/test_bigmodel_provider.py @@ -197,6 +197,29 @@ def test_chat_does_not_send_the_openai_organization_to_bigmodel(monkeypatch: pyt assert "openai-organization" not in route.calls.last.request.headers +@respx.mock +def test_chat_sends_json_schema_response_format_as_given() -> None: + route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + json={ + "id": "chat_test", + "object": "chat.completion", + "created": 1, + "model": "glm-4.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "{}"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + messages: Final = [{"role": "user", "content": "Ann likes 7"}] + response_format: Final = { + "type": "json_schema", + "json_schema": {"name": "profile", "schema": {"type": "object", "properties": {"n": {"type": "integer"}}}}, + } + _ = litellm.completion(model="bigmodel/glm-4.6", messages=messages, response_format=response_format) + body: Final = json.loads(route.calls.last.request.content) + assert body["response_format"] == response_format + assert body["messages"] == messages + + @pytest.mark.parametrize("is_async", (False, True)) @pytest.mark.asyncio @respx.mock diff --git a/tests/unit/llms/zai/test_zai_provider.py b/tests/unit/llms/zai/test_zai_provider.py index 32849d5eef1..fb126bd8f41 100644 --- a/tests/unit/llms/zai/test_zai_provider.py +++ b/tests/unit/llms/zai/test_zai_provider.py @@ -2,13 +2,27 @@ Tests for Z.AI (Zhipu AI) provider - GLM models """ +import json import math +from collections.abc import Mapping +from typing import Final import pytest +import respx import litellm from litellm import completion from litellm.cost_calculator import cost_per_token +from litellm.llms.zai.chat.transformation import ZAIChatConfig +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +JSON_SCHEMA_RESPONSE_FORMAT: Final = { + "type": "json_schema", + "json_schema": { + "name": "profile", + "schema": {"type": "object", "properties": {"favorite_number": {"type": "integer"}}}, + }, +} @pytest.fixture @@ -99,3 +113,186 @@ 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 + + +def test_zai_sends_response_format_and_maps_max_completion_tokens( + respx_mock: respx.MockRouter, zai_response: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + 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) + + completion( + model="zai/glm-5.3", + messages=[{"role": "user", "content": "Reply in JSON"}], + response_format={"type": "json_object"}, + max_completion_tokens=64, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["response_format"] == {"type": "json_object"} + assert body["max_tokens"] == 64 + assert "max_completion_tokens" not in body + + +@pytest.mark.parametrize( + "token_limits", + [{"max_tokens": 300, "max_completion_tokens": 15}, {"max_completion_tokens": 15, "max_tokens": 300}], +) +def test_zai_max_completion_tokens_wins_over_max_tokens(token_limits: dict[str, int]) -> None: + mapped: Final = ZAIChatConfig().map_openai_params(token_limits, {}, "glm-5.3", drop_params=False) + + assert mapped["max_tokens"] == 15 + assert "max_completion_tokens" not in mapped + + +@pytest.mark.parametrize( + ("over_limit_by", "drop_params", "forwarded"), + [(1, True, False), (1, False, True), (0, True, True)], +) +def test_zai_drop_params_drops_max_completion_tokens_above_model_limit( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + over_limit_by: int, + drop_params: bool, + forwarded: bool, +) -> None: + 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) + requested: Final = litellm.model_cost["zai/glm-4.5-flash"]["max_output_tokens"] + over_limit_by + + completion( + model="zai/glm-4.5-flash", + messages=[{"role": "user", "content": "Count"}], + max_completion_tokens=requested, + drop_params=drop_params, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body.get("max_tokens") == (requested if forwarded else None) + + +@pytest.mark.parametrize("with_over_limit_max_completion_tokens", [False, True]) +def test_zai_drop_params_keeps_max_tokens_above_model_limit( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + with_over_limit_max_completion_tokens: bool, +) -> None: + 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) + over_limit: Final = litellm.model_cost["zai/glm-4.5-flash"]["max_output_tokens"] + 1 + over_limit_max_completion_tokens: Final = ( + {"max_completion_tokens": over_limit + 1} if with_over_limit_max_completion_tokens else {} + ) + + completion( + model="zai/glm-4.5-flash", + messages=[{"role": "user", "content": "Count"}], + max_tokens=over_limit, + drop_params=True, + **over_limit_max_completion_tokens, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["max_tokens"] == over_limit + + +@pytest.mark.parametrize("pass_through", [{}, {"allowed_openai_params": ["response_format"]}]) +def test_zai_forwards_json_schema_to_models_that_enforce_it( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + pass_through: dict[str, list[str]], +) -> None: + 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) + messages: Final = [{"role": "user", "content": "Ann likes 7"}] + + completion(model="zai/glm-5.3", messages=messages, response_format=JSON_SCHEMA_RESPONSE_FORMAT, **pass_through) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT + assert body["messages"] == messages + + +@pytest.mark.parametrize("pass_through", [{}, {"allowed_openai_params": ["response_format"]}]) +def test_zai_puts_json_schema_in_prompt_for_models_that_ignore_it( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + pass_through: dict[str, list[str]], +) -> None: + 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) + + completion( + model="zai/glm-4.6", + messages=[{"role": "user", "content": "Ann likes 7"}], + response_format=JSON_SCHEMA_RESPONSE_FORMAT, + **pass_through, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["response_format"] == {"type": "json_object"} + assert body["messages"][0] == {"role": "user", "content": "Ann likes 7"} + assert body["messages"][-1]["role"] == "user" + prompt: Final = body["messages"][-1]["content"] + rendered_schema: Final = prompt[prompt.index("{") : prompt.rindex("}") + 1] + assert json.loads(rendered_schema) == JSON_SCHEMA_RESPONSE_FORMAT["json_schema"]["schema"] + + +def test_zai_glm_5v_turbo_bills_cached_input_at_cache_read_price(local_model_cost_map: None) -> None: + usage: Final = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800), + ) + + prompt_cost, _ = cost_per_token(model="glm-5v-turbo", custom_llm_provider="zai", usage_object=usage) + + entry: Final = litellm.model_cost["zai/glm-5v-turbo"] + assert math.isclose( + prompt_cost, 200 * entry["input_cost_per_token"] + 800 * entry["cache_read_input_token_cost"] + ) + + +@pytest.mark.parametrize("model", ["glm-5v-turbo", "GLM-5V-Turbo", "GLM-5.3", "GLM-5.3-Flash"]) +def test_zai_cost_resolves_case_insensitively(local_model_cost_map: None, model: str) -> None: + prompt_cost, completion_cost = cost_per_token( + model=model, custom_llm_provider="zai", prompt_tokens=1000, completion_tokens=1000 + ) + + entry: Final = litellm.model_cost[f"zai/{model.lower()}"] + assert math.isclose(prompt_cost, 1000 * entry["input_cost_per_token"]) + assert math.isclose(completion_cost, 1000 * entry["output_cost_per_token"]) + + +def test_zai_glm_5v_turbo_accepts_thinking( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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) + + completion( + model="zai/glm-5v-turbo", + messages=[{"role": "user", "content": "Hello"}], + thinking={"type": "disabled"}, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["thinking"] == {"type": "disabled"}