feat: support zai as inference provider (#3395)

* feat: support zai as inference provider

* update

* update

* update

* update

Author: tomsun28

src/huggingface_hub/inference/_client.py

@@ -130,7 +130,7 @@ class InferenceClient:
             Note: for better compatibility with OpenAI's client, `model` has been aliased as `base_url`. Those 2
             arguments are mutually exclusive. If a URL is passed as `model` or `base_url` for chat completion, the `(/v1)/chat/completions` suffix path will be appended to the URL.
         provider (`str`, *optional*):
-            Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"featherless-ai"`, `"fireworks-ai"`, `"groq"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"nscale"`, `"openai"`, `publicai`, `"replicate"`, `"sambanova"`, `"scaleway"` or `"together"`.
+            Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"featherless-ai"`, `"fireworks-ai"`, `"groq"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"nscale"`, `"openai"`, `publicai`, `"replicate"`, `"sambanova"`, `"scaleway"`, `"together"` or `"zai-org"`.
             Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers.
             If model is a URL or `base_url` is passed, then `provider` is not used.
         token (`str`, *optional*):

src/huggingface_hub/inference/_generated/_async_client.py

@@ -118,7 +118,7 @@ class AsyncInferenceClient:
             Note: for better compatibility with OpenAI's client, `model` has been aliased as `base_url`. Those 2
             arguments are mutually exclusive. If a URL is passed as `model` or `base_url` for chat completion, the `(/v1)/chat/completions` suffix path will be appended to the URL.
         provider (`str`, *optional*):
-            Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"featherless-ai"`, `"fireworks-ai"`, `"groq"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"nscale"`, `"openai"`, `publicai`, `"replicate"`, `"sambanova"`, `"scaleway"` or `"together"`.
+            Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"featherless-ai"`, `"fireworks-ai"`, `"groq"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"nscale"`, `"openai"`, `publicai`, `"replicate"`, `"sambanova"`, `"scaleway"`, `"together"` or `"zai-org"`.
             Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers.
             If model is a URL or `base_url` is passed, then `provider` is not used.
         token (`str`, *optional*):

src/huggingface_hub/inference/_providers/__init__.py

@@ -41,6 +41,7 @@
 from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
 from .scaleway import ScalewayConversationalTask, ScalewayFeatureExtractionTask
 from .together import TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask
+from .zai_org import ZaiConversationalTask
 
 
 logger = logging.get_logger(__name__)
@@ -65,6 +66,7 @@
     "sambanova",
     "scaleway",
     "together",
+    "zai-org",
 ]
 
 PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]]
@@ -170,6 +172,9 @@
         "conversational": TogetherConversationalTask(),
         "text-generation": TogetherTextGenerationTask(),
     },
+    "zai-org": {
+        "conversational": ZaiConversationalTask(),
+    },
 }
 
 

src/huggingface_hub/inference/_providers/_common.py

@@ -35,6 +35,7 @@
     "sambanova": {},
     "scaleway": {},
     "together": {},
+    "zai-org": {},
 }
 
 

src/huggingface_hub/inference/_providers/zai_org.py

@@ -0,0 +1,17 @@
+from typing import Any, Dict
+
+from huggingface_hub.inference._providers._common import BaseConversationalTask
+
+
+class ZaiConversationalTask(BaseConversationalTask):
+    def __init__(self):
+        super().__init__(provider="zai-org", base_url="https://api.z.ai")
+
+    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict[str, Any]:
+        headers = super()._prepare_headers(headers, api_key)
+        headers["Accept-Language"] = "en-US,en"
+        headers["x-source-channel"] = "hugging_face"
+        return headers
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        return "/api/paas/v4/chat/completions"

tests/test_inference_providers.py

@@ -53,6 +53,7 @@
 from huggingface_hub.inference._providers.sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
 from huggingface_hub.inference._providers.scaleway import ScalewayConversationalTask, ScalewayFeatureExtractionTask
 from huggingface_hub.inference._providers.together import TogetherTextToImageTask
+from huggingface_hub.inference._providers.zai_org import ZaiConversationalTask
 
 from .testing_utils import assert_in_logs
 
@@ -1412,6 +1413,28 @@ def test_text_to_image_get_response(self):
         assert response == b"image_bytes"
 
 
+class TestZaiProvider:
+    def test_prepare_route(self):
+        helper = ZaiConversationalTask()
+        route = helper._prepare_route("test-model", "zai_token")
+        assert route == "/api/paas/v4/chat/completions"
+
+    def test_prepare_headers(self):
+        helper = ZaiConversationalTask()
+        headers = helper._prepare_headers({}, "test_key")
+        assert headers["Accept-Language"] == "en-US,en"
+
+    def test_prepare_url(self):
+        helper = ZaiConversationalTask()
+        assert helper.task == "conversational"
+        url = helper._prepare_url("zai_token", "test-model")
+        assert url == "https://api.z.ai/api/paas/v4/chat/completions"
+
+        # Test with HF token (should route through HF proxy)
+        url = helper._prepare_url("hf_token", "test-model")
+        assert url.startswith("https://router.huggingface.co/zai-org")
+
+
 class TestBaseConversationalTask:
     def test_prepare_route(self):
         helper = BaseConversationalTask(provider="test-provider", base_url="https://api.test.com")