mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(zai): address review and CI feedback for protocol passthrough
This commit is contained in:
parent
92c1eebe56
commit
75ffe2112d
4 changed files with 58 additions and 40 deletions
|
|
@ -2,7 +2,7 @@
|
|||
Z.AI Anthropic-compatible messages transformation config.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
|
||||
|
|
@ -32,21 +32,22 @@ class ZAIAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
|
||||
@staticmethod
|
||||
def get_api_base(api_base: str | None = None) -> str:
|
||||
return api_base or get_secret_str("ZAI_ANTHROPIC_API_BASE") or "https://api.z.ai/api/anthropic"
|
||||
return api_base or "https://api.z.ai/api/anthropic"
|
||||
|
||||
def validate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
headers: dict[str, str],
|
||||
model: str,
|
||||
messages: list[Any],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
messages: list[dict[str, object]],
|
||||
optional_params: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | 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 headers and "authorization" not in headers and dynamic_api_key is not None:
|
||||
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:
|
||||
|
|
@ -67,8 +68,8 @@ class ZAIAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
optional_params: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
base_url = self.get_api_base(api_base=api_base).rstrip("/")
|
||||
|
|
|
|||
|
|
@ -28,13 +28,13 @@ class ZAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
headers: dict[str, str],
|
||||
model: str,
|
||||
litellm_params: GenericLiteLLMParams | None,
|
||||
) -> dict:
|
||||
) -> dict[str, str]:
|
||||
litellm_params = litellm_params or GenericLiteLLMParams()
|
||||
|
||||
api_key: Final = litellm_params.api_key or litellm.api_key or get_secret_str("ZAI_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:
|
||||
|
|
@ -44,7 +44,7 @@ class ZAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
litellm_params: dict,
|
||||
litellm_params: dict[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
|
||||
|
|
@ -52,11 +52,7 @@ class ZAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
# the chat-only bases and use the Responses base instead.
|
||||
normalized_api_base = (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 get_secret_str("ZAI_RESPONSES_API_BASE") or "https://api.z.ai/api/v1"
|
||||
)
|
||||
base_url = 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"):
|
||||
|
|
|
|||
|
|
@ -26,8 +26,7 @@ def test_anthropic_provider_keeps_default_config_for_zai_named_model():
|
|||
assert not isinstance(config, ZAIAnthropicMessagesConfig)
|
||||
|
||||
|
||||
def test_zai_anthropic_messages_config_defaults(monkeypatch):
|
||||
monkeypatch.delenv("ZAI_ANTHROPIC_API_BASE", raising=False)
|
||||
def test_zai_anthropic_messages_config_defaults():
|
||||
config = ZAIAnthropicMessagesConfig()
|
||||
|
||||
assert config.custom_llm_provider == "zai"
|
||||
|
|
@ -75,3 +74,23 @@ def test_zai_anthropic_messages_headers_use_zai_key():
|
|||
assert headers["x-api-key"] == "sk-zai"
|
||||
assert headers["anthropic-version"] == "2023-06-01"
|
||||
assert headers["content-type"] == "application/json"
|
||||
|
||||
|
||||
def test_zai_anthropic_messages_respects_existing_case_insensitive_auth_headers():
|
||||
config = ZAIAnthropicMessagesConfig()
|
||||
|
||||
headers, _ = config.validate_anthropic_messages_environment(
|
||||
headers={"Authorization": "Bearer caller-token"},
|
||||
model="glm-5.3",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="sk-zai",
|
||||
api_base="https://api.z.ai/api/anthropic",
|
||||
)
|
||||
|
||||
assert headers == {
|
||||
"Authorization": "Bearer caller-token",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,8 +15,7 @@ def test_zai_provider_uses_responses_api_config():
|
|||
assert config.custom_llm_provider == LlmProviders.ZAI
|
||||
|
||||
|
||||
def test_zai_responses_url_defaults_to_responses_endpoint(monkeypatch):
|
||||
monkeypatch.delenv("ZAI_RESPONSES_API_BASE", raising=False)
|
||||
def test_zai_responses_url_defaults_to_responses_endpoint():
|
||||
config = ZAIResponsesAPIConfig()
|
||||
|
||||
url_cases = {
|
||||
|
|
@ -30,8 +29,7 @@ def test_zai_responses_url_defaults_to_responses_endpoint(monkeypatch):
|
|||
assert config.get_complete_url(api_base=api_base, litellm_params={}) == expected_url
|
||||
|
||||
|
||||
def test_zai_responses_url_ignores_chat_completions_api_base(monkeypatch):
|
||||
monkeypatch.delenv("ZAI_RESPONSES_API_BASE", raising=False)
|
||||
def test_zai_responses_url_ignores_chat_completions_api_base():
|
||||
config = ZAIResponsesAPIConfig()
|
||||
|
||||
chat_bases = (
|
||||
|
|
@ -44,8 +42,7 @@ def test_zai_responses_url_ignores_chat_completions_api_base(monkeypatch):
|
|||
assert config.get_complete_url(api_base=chat_base, litellm_params={}) == "https://api.z.ai/api/v1/responses"
|
||||
|
||||
|
||||
def test_zai_responses_url_keeps_custom_api_base(monkeypatch):
|
||||
monkeypatch.delenv("ZAI_RESPONSES_API_BASE", raising=False)
|
||||
def test_zai_responses_url_keeps_custom_api_base():
|
||||
config = ZAIResponsesAPIConfig()
|
||||
|
||||
assert (
|
||||
|
|
@ -57,19 +54,6 @@ def test_zai_responses_url_keeps_custom_api_base(monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
def test_zai_responses_url_env_overrides_chat_completions_api_base(monkeypatch):
|
||||
monkeypatch.setenv("ZAI_RESPONSES_API_BASE", "https://gateway.example.com/responses-root")
|
||||
config = ZAIResponsesAPIConfig()
|
||||
|
||||
assert (
|
||||
config.get_complete_url(
|
||||
api_base="https://api.z.ai/api/paas/v4",
|
||||
litellm_params={},
|
||||
)
|
||||
== "https://gateway.example.com/responses-root/responses"
|
||||
)
|
||||
|
||||
|
||||
def test_zai_responses_headers_use_bearer_token():
|
||||
config = ZAIResponsesAPIConfig()
|
||||
litellm_params = GenericLiteLLMParams(api_key="sk-zai")
|
||||
|
|
@ -95,3 +79,21 @@ def test_zai_responses_headers_fall_back_to_environment_key(monkeypatch):
|
|||
)
|
||||
|
||||
assert headers["Authorization"] == "Bearer sk-zai-env"
|
||||
|
||||
|
||||
def test_zai_responses_headers_prefer_zai_key_over_global_key(monkeypatch):
|
||||
monkeypatch.setenv("ZAI_API_KEY", "sk-zai-env")
|
||||
original_api_key = litellm.api_key
|
||||
litellm.api_key = "sk-global-other-provider"
|
||||
try:
|
||||
config = ZAIResponsesAPIConfig()
|
||||
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="glm-5.3",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
)
|
||||
finally:
|
||||
litellm.api_key = original_api_key
|
||||
|
||||
assert headers["Authorization"] == "Bearer sk-zai-env"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue