feat: Add Zhipu (智谱) image generation provider

This commit is contained in:
Jiajun Xie
2026-05-23 17:06:36 +08:00
committed by Xubin Ren
parent c0d4f012c8
commit 3e6f9907fe
4 changed files with 294 additions and 3 deletions
+100
View File
@@ -18,6 +18,7 @@ from nanobot.providers.image_generation import (
OpenAIImageGenerationClient,
OpenRouterImageGenerationClient,
StepFunImageGenerationClient,
ZhipuImageGenerationClient,
)
PNG_BYTES = (
@@ -1027,3 +1028,102 @@ async def test_openai_no_images_raises() -> None:
with pytest.raises(ImageGenerationError, match="returned no images"):
await client.generate(prompt="draw", model="dall-e-3")
# ---------------------------------------------------------------------------
# Zhipu
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_zhipu_image_generation_payload_and_response() -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = ZhipuImageGenerationClient(
api_key="sk-zhipu-test",
api_base="https://open.bigmodel.cn/api/paas/v4",
extra_headers={"X-Test": "1"},
extra_body={"watermark_enabled": False},
client=fake, # type: ignore[arg-type]
)
response = await client.generate(
prompt="a sunset over the ocean",
model="glm-image",
aspect_ratio="16:9",
image_size="2K",
)
assert response.images[0].startswith("data:image/png;base64,")
call = fake.calls[0]
assert call["url"] == "https://open.bigmodel.cn/api/paas/v4/images/generations"
assert call["headers"]["Authorization"] == "Bearer sk-zhipu-test"
assert call["headers"]["X-Test"] == "1"
body = call["json"]
assert body["model"] == "glm-image"
assert body["prompt"] == "a sunset over the ocean"
assert body["size"] == "1728x960"
assert body["watermark_enabled"] is False
@pytest.mark.asyncio
async def test_zhipu_image_generation_with_explicit_size() -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = ZhipuImageGenerationClient(
api_key="sk-zhipu-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(
prompt="a cat",
model="cogview-4",
image_size="1024x1024",
)
body = fake.calls[0]["json"]
assert body["size"] == "1024x1024"
@pytest.mark.asyncio
async def test_zhipu_image_generation_downloads_url_response() -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = ZhipuImageGenerationClient(
api_key="sk-zhipu-test",
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="glm-image")
assert response.images[0].startswith("data:image/png;base64,")
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
@pytest.mark.asyncio
async def test_zhipu_image_generation_requires_api_key() -> None:
client = ZhipuImageGenerationClient(api_key=None)
with pytest.raises(ImageGenerationError, match="API key"):
await client.generate(prompt="draw", model="glm-image")
@pytest.mark.asyncio
async def test_zhipu_image_generation_no_images_raises() -> None:
fake = FakeClient(FakeResponse({"data": [{"text": "sorry"}]}))
client = ZhipuImageGenerationClient(api_key="sk-zhipu-test", client=fake) # type: ignore[arg-type]
with pytest.raises(ImageGenerationError, match="returned no images"):
await client.generate(prompt="draw", model="glm-image")
@pytest.mark.asyncio
async def test_zhipu_image_generation_rejects_reference_images() -> None:
client = ZhipuImageGenerationClient(api_key="sk-zhipu-test")
with pytest.raises(ImageGenerationError, match="reference images"):
await client.generate(
prompt="edit this",
model="glm-image",
reference_images=["ref.png"],
)