import time
from abc import ABC
from typing import Any
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
from huggingface_hub.utils import get_session
_PROVIDER = "zai-org"
_BASE_URL = "https://api.z.ai"
_POLLING_INTERVAL = 5 # seconds
_MAX_POLL_ATTEMPTS = 60
class ZaiTask(TaskProviderHelper, ABC):
def __init__(self, task: str):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
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
class ZaiConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
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"
class ZaiTextToImageTask(ZaiTask):
"""Text-to-image task for ZAI provider using async API."""
def __init__(self):
super().__init__("text-to-image")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/api/paas/v4/async/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> dict | None:
width = parameters.pop("width", None)
height = parameters.pop("height", None)
size = None
if width is not None and height is not None:
size = f"{width}x{height}"
payload: dict[str, Any] = {
"model": provider_mapping_info.provider_id,
"prompt": inputs,
}
if size is not None:
payload["size"] = size
payload.update(filter_none(parameters))
return payload
def get_response(
self,
response: bytes | dict,
request_params: RequestParameters | None = None,
) -> Any:
"""Handle async response by polling for results."""
response_dict = _as_dict(response)
task_id = response_dict.get("id")
if task_id is None:
raise ValueError("No task_id in response from ZAI API")
task_status = response_dict.get("task_status")
if task_status == "FAIL":
raise ValueError(f"ZAI image generation failed for request {task_id}")
if task_status == "PROCESSING" and request_params is not None:
return self._poll_for_result(task_id, request_params)
return self._extract_image(response_dict)
def _poll_for_result(self, task_id: str, request_params: RequestParameters) -> bytes:
"""Poll the async-result endpoint until completion."""
session = get_session()
base_url = request_params.url.rsplit("/api/paas/v4/async/images/generations", 1)[0]
poll_url = f"{base_url}/api/paas/v4/async-result/{task_id}"
for _ in range(_MAX_POLL_ATTEMPTS):
poll_response = session.get(poll_url, headers=request_params.headers)
poll_response.raise_for_status()
result = poll_response.json()
task_status = result.get("task_status")
if task_status == "SUCCESS":
return self._extract_image(result)
elif task_status == "FAIL":
raise ValueError(f"Zai text-to-image generation failed for request {task_id}")
time.sleep(_POLLING_INTERVAL)
raise ValueError(
f"Timed out while waiting for the result from Zai API - aborting after {_MAX_POLL_ATTEMPTS} attempts"
)
def _extract_image(self, result: dict) -> bytes:
"""Extract and download the image from the result."""
image_result = result.get("image_result")
if not image_result or not isinstance(image_result, list) or len(image_result) == 0:
raise ValueError("No image_result in response from ZAI API")
image_url = image_result[0].get("url")
if not image_url:
raise ValueError("No image URL in response from ZAI API")
session = get_session()
image_response = session.get(image_url)
image_response.raise_for_status()
return image_response.content