Add image_generation server-side tool

[yanxi0830]

Adds the image_generation server-side (built-in) tool to the SDK, following the web_search / x_search / code_execution agentic-tool pattern, along with a Response.image_outputs property for retrieving the generated images.

Changes

image_generation() tool helper

• tools.py: New image_generation() helper that returns a chat_pb2.Tool. Accepts an optional action parameter — "auto" (default; both generation and editing), "generate" (text-to-image only), or "edit" (image editing only). Also documents "image_generation_tool" as a possible return value of get_tool_call_type().
• Proto updates (v5 and v6): ImageGeneration message and image_generation member in the Tool oneof, TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL in ToolCallType (chat_pb2), and SERVER_SIDE_TOOL_IMAGE_GENERATION in ServerSideTool (usage_pb2) for usage tracking.

Response.image_outputs

Retrieving generated images previously required hand-parsing the ROLE_TOOL output content: JSON-decode the {"__type": "image_generation_result", ...} envelope, split the data URL, and base64-decode the payload. This PR adds an ImageGenerationOutput wrapper and a Response.image_outputs property (mirroring the existing output_files pattern), so callers can simply do:

chat = client.chat.create(model="grok-4.5", tools=[image_generation()])
chat.append(user("Generate an image of a corgi surfing a big wave"))
response = chat.sample()

for i, output in enumerate(response.image_outputs):
    extension = output.mime_type.removeprefix("image/")
    with open(f"image_{i}.{extension}", "wb") as f:
        f.write(output.image)

• Response.image_outputs — one entry per completed image generation/edit call, in response order. Failed calls carry an error payload instead of an image envelope and are omitted; the property returns an empty sequence when no images were generated.
• ImageGenerationOutput — parses the result envelope once and exposes:
  • image: decoded raw bytes
  • mime_type / data_url: self-describing payload metadata
  • image_uuid: reference ID for follow-up edits (empty on older envelopes, matching the server's #[serde(default)])
  • tool_call: the originating completed tool call (imagine_text_to_image vs imagine_image_to_image)

Works for sync and aio clients, and for both sample() and stream() — the accumulated streaming response syncs its content buffers before parsing.

Examples and tests

• examples/{sync,aio}/image_generation_tool.py — single-turn generation, editing an attached input image via the image() content helper, multi-turn generate-then-edit in a single chat via chat.append(response), web_search + image_generation combo, and verbose streaming, all built on the new API (the existing image_generation.py examples only cover the standalone client.image API).
• Fake-server fixtures for the image_generation tool (unary + streaming) and new tests: helper unit tests (default and each action value), tool-call type mapping, usage enum, plus sync/aio × sample/stream image_outputs tests, failed-call omission, and malformed-envelope handling.
• CHANGELOG.md entries.

Usage

from xai_sdk.tools import image_generation

chat = client.chat.create(
    model="grok-4.5",
    tools=[image_generation()],  # or image_generation(action="generate")
)

Test plan

• uv run ruff format --check and uv run ruff check
• uv run pyright — 0 errors
• uv run pytest -n auto — 772 passed
• Verified the regenerated v5 protos import and work under a protobuf 5.29 runtime, and v6 under 6.30

Made with Cursor

[suyanxai]

LGTM