From 1b19dea96dab3a422e57c505cd99e09d63865bde Mon Sep 17 00:00:00 2001 From: togear Date: Mon, 7 Sep 2026 19:32:38 +0800 Subject: [PATCH] refactor(zai): satisfy type discipline gates on protocol configs --- litellm/llms/zai/messages/transformation.py | 52 +++++++------------- litellm/llms/zai/responses/transformation.py | 30 ++++++----- 2 files changed, 36 insertions(+), 46 deletions(-) diff --git a/litellm/llms/zai/messages/transformation.py b/litellm/llms/zai/messages/transformation.py index c802798d2d9..91ae5ff70cc 100644 --- a/litellm/llms/zai/messages/transformation.py +++ b/litellm/llms/zai/messages/transformation.py @@ -2,6 +2,7 @@ Z.AI Anthropic-compatible messages transformation config. """ +from collections.abc import Mapping from typing import Final import litellm @@ -36,51 +37,36 @@ class ZAIAnthropicMessagesConfig(AnthropicMessagesConfig): def validate_anthropic_messages_environment( self, - headers: dict[str, str], + headers: Mapping[str, str], model: str, - messages: list[dict[str, object]], - optional_params: dict[str, object], - litellm_params: dict[str, object], + messages: list[Mapping[str, object]], # mutable-ok: matches the pass-through handler's message list contract + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> tuple[dict[str, str], str | None]: - dynamic_api_key: Final = self.get_api_key(api_key=api_key) - header_names: Final = {header_name.lower() for header_name in headers} - - if "x-api-key" not in header_names and "authorization" not in header_names and dynamic_api_key is not None: - headers["x-api-key"] = dynamic_api_key - - if "anthropic-version" not in headers: - headers["anthropic-version"] = "2023-06-01" - if "content-type" not in headers: - headers["content-type"] = "application/json" - - headers = self._update_headers_with_anthropic_beta( + ) -> tuple[dict[str, str], str | None]: # mutable-ok: the handler owns and mutates the returned headers dict + return super().validate_anthropic_messages_environment( headers=headers, + model=model, + messages=messages, optional_params=optional_params, - custom_llm_provider=self.custom_llm_provider or "zai", + litellm_params=litellm_params, + api_key=self.get_api_key(api_key=api_key), + api_base=api_base, ) - return headers, api_base - def get_complete_url( self, api_base: str | None, api_key: str | None, model: str, - optional_params: dict[str, object], - litellm_params: dict[str, object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - base_url = self.get_api_base(api_base=api_base).rstrip("/") + raw_base_url: Final = self.get_api_base(api_base=api_base).rstrip("/") + root_url: Final = raw_base_url.removesuffix("/v1/messages").removesuffix("/v1").removesuffix("/beta") - if base_url.endswith("/v1/messages"): - return base_url - base_url = base_url.removesuffix("/v1/messages") - base_url = base_url.removesuffix("/v1") - base_url = base_url.removesuffix("/beta") - - if not base_url.endswith("/anthropic") and "/anthropic/" not in base_url: - base_url = f"{base_url}/anthropic" - - return f"{base_url}/v1/messages" + if root_url.endswith("/anthropic") or "/anthropic/" in root_url: + return f"{root_url}/v1/messages" + return f"{root_url}/anthropic/v1/messages" diff --git a/litellm/llms/zai/responses/transformation.py b/litellm/llms/zai/responses/transformation.py index 8a00a64ead8..cfcedc2b868 100644 --- a/litellm/llms/zai/responses/transformation.py +++ b/litellm/llms/zai/responses/transformation.py @@ -2,6 +2,7 @@ Z.AI OpenAI-compatible Responses API transformation config. """ +from collections.abc import Mapping from typing import Final import litellm @@ -26,35 +27,38 @@ class ZAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.ZAI + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return api_key or get_secret_str("ZAI_API_KEY") or litellm.api_key + def validate_environment( self, - headers: dict[str, str], + headers: Mapping[str, str], model: str, litellm_params: GenericLiteLLMParams | None, - ) -> dict[str, str]: - litellm_params = litellm_params or GenericLiteLLMParams() + ) -> dict[str, str]: # mutable-ok: the responses handler owns and mutates the returned headers dict + request_api_key: Final = litellm_params.api_key if litellm_params is not None else None + resolved_params: Final = GenericLiteLLMParams(api_key=self.get_api_key(api_key=request_api_key)) - api_key: Final = litellm_params.api_key or get_secret_str("ZAI_API_KEY") or litellm.api_key - - headers.setdefault("Content-Type", "application/json") - if api_key is not None: - headers["Authorization"] = f"Bearer {api_key}" - return headers + return super().validate_environment( + headers=headers, + model=model, + litellm_params=resolved_params, + ) def get_complete_url( self, api_base: str | None, - litellm_params: dict[str, object], + litellm_params: Mapping[str, object], ) -> str: # ``litellm_params.api_base`` can carry the Z.AI chat-completions base # (``/api/paas/v4``) when the generic provider resolver pre-fills it from # the chat config. Z.AI serves Responses on a different base, so ignore # the chat-only bases and use the Responses base instead. - normalized_api_base = (api_base or "").rstrip("/") + normalized_api_base: Final = (api_base or "").rstrip("/") chat_base_passed_in: Final = normalized_api_base.endswith(self._ZAI_CHAT_API_BASE_SUFFIXES) - base_url = api_base if api_base and not chat_base_passed_in else "https://api.z.ai/api/v1" + base_url: Final = normalized_api_base if api_base and not chat_base_passed_in else "https://api.z.ai/api/v1" - base_url = base_url.rstrip("/") if base_url.endswith("/responses"): return base_url return f"{base_url}/responses"