Skip to content

Add image_generation server-side tool - #181

Closed
yanxi0830 wants to merge 2 commits into
xai-org:mainfrom
yanxi0830:feature/add-image-generation-tool
Closed

Add image_generation server-side tool#181
yanxi0830 wants to merge 2 commits into
xai-org:mainfrom
yanxi0830:feature/add-image-generation-tool

Conversation

@yanxi0830

@yanxi0830 yanxi0830 commented Jul 14, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds the server-side image_generation tool to the SDK, mirroring the tool now supported by the xAI API (following the web_search / x_search / code_execution agentic-tool pattern):

  • xai_sdk.tools.image_generation(): new helper that returns a chat_pb2.Tool for image generation and editing in agentic requests. Takes an optional action parameter — "auto" (default; generate + edit), "generate" (text-to-image only), or "edit" (image editing only). Omitting action leaves the field unset so the server applies its default.
  • Regenerated protos (v5 & v6):
    • chat.proto: ImageGeneration message (optional string action = 1), image_generation = 10 in the Tool oneof, and TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL = 10.
    • usage.proto: SERVER_SIDE_TOOL_IMAGE_GENERATION = 11 and num_image_generations / num_image_edits fields on SamplingUsage for per-call generation/edit billing visibility.
  • get_tool_call_type() docstring now lists image_generation_tool as a possible return value.
  • Examples: examples/sync/image_generation_tool.py and examples/aio/image_generation_tool.py showing generation, decoding the ROLE_TOOL result envelope ({"__type": "image_generation_result", "result": "data:<mime>;base64,..."}) to save the image to disk, streaming, and multi-turn editing via previous_response_id.

Usage:

from xai_sdk.tools import image_generation

chat = client.chat.create("grok-4", tools=[image_generation()])
# or restrict to text-to-image only:
chat = client.chat.create("grok-4", tools=[image_generation(action="generate")])

Test plan

  • uv run ruff format --check passes
  • uv run ruff check passes
  • uv run pyright passes (0 errors)
  • uv run pytest -n auto — all 770 tests pass
  • New unit tests: image_generation() default and per-action proto construction, TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL mapping via get_tool_call_type, SERVER_SIDE_TOOL_IMAGE_GENERATION enum value, SamplingUsage image fields, and sync/async chat.create(tools=[...]) round-trips
  • Verified the regenerated v5 protos import and work under protobuf==5.29.4, and v6 under protobuf==6.31.1
  • Live end-to-end verification against the production API: tool call surfaces as imagine_text_to_image / imagine_image_to_image with TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL, the ROLE_TOOL envelope decodes to a valid image, usage reports SERVER_SIDE_TOOL_IMAGE_GENERATION with num_image_generations / num_image_edits populated, and both example scripts (sync generate, multi-turn edit, async streaming) run successfully

Adds the xai_sdk.tools.image_generation() helper for the server-side
image_generation tool, enabling image generation and editing in agentic
requests. The optional action parameter ("auto", "generate", or "edit")
controls which image capabilities are exposed to the model.

Regenerates the v5/v6 chat and usage protos to add the ImageGeneration
tool message, TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL,
SERVER_SIDE_TOOL_IMAGE_GENERATION, and the num_image_generations /
num_image_edits usage fields.

Co-authored-by: Cursor <cursoragent@cursor.com>
@yanxi0830
yanxi0830 requested a review from a team as a code owner July 14, 2026 04:59
Adds sync and async examples showing how to use the server-side
image_generation tool in agentic requests: generating an image,
decoding the ROLE_TOOL result envelope to save it to disk, streaming,
and multi-turn editing via previous_response_id.

Co-authored-by: Cursor <cursoragent@cursor.com>
@yanxi0830 yanxi0830 closed this Jul 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant