Add image_generation server-side tool

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [Unreleased]
 ### Added
+- **Image Generation Tool**: Added `xai_sdk.tools.image_generation` helper for the server-side `image_generation` tool, enabling image generation and editing in agentic requests. Accepts an optional `action` parameter (`"auto"`, `"generate"`, or `"edit"`) to control which image capabilities are exposed to the model
 - **Imagine File Storage**: Image and video generation (sync and async) now accept a `storage_options` parameter to persist generated assets to the Files API. It takes a dict with a required `filename` and optional `expires_after` (an `int` in seconds or a `datetime.timedelta`) and `public_url` (`True` to create a public URL with default expiry, or `{"expires_after": <seconds>}` for an independent URL expiry). Image and video responses expose new `file_output`, `storage_error`, `public_url`, and `public_url_error` properties.
 - **File-ID Inputs for Generation**: Image and video generation now accept Files API `file_id` references as inputs alongside URLs/base64 — `image_file_id` / `image_file_ids` for `image.sample()` / `image.sample_batch()`, and `image_file_id` / `video_file_id` / `reference_image_file_ids` for `video.generate()` / `video.extend()` (and the batch `prepare` helpers). URL and file-ID lists may be mixed in the same multi-image request (file IDs are sent first).
 - **Public File URLs**: Added `client.files.create_public_url()` and `client.files.revoke_public_url()` (sync and async) to create and revoke publicly shareable, unauthenticated URLs for stored files. `create_public_url()` accepts an optional `expires_after` (an `int` in seconds or a `datetime.timedelta`).

examples/sync/image_generation_tool.py

@@ -0,0 +1,126 @@
+import base64
+import json
+import os
+
+from xai_sdk import Client
+from xai_sdk.chat import user
+from xai_sdk.proto import chat_pb2
+from xai_sdk.tools import get_tool_call_type, image_generation, web_search
+
+
+def save_images(response, prefix: str) -> None:
+    """Saves every image produced by image_generation tool calls in a response.
+
+    Image generation calls produce ROLE_TOOL outputs. Select the completed ones
+    by the tool-call type and status enums, then read the data URL out of the
+    result envelope. Failed calls carry an error payload instead of an image
+    envelope, so they are skipped.
+    """
+    image_outputs = [
+        output
+        for output in response.tool_outputs
+        if any(
+            tool_call.type == chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL
+            and tool_call.status == chat_pb2.TOOL_CALL_STATUS_COMPLETED
+            for tool_call in output.message.tool_calls
+        )
+    ]
+    for i, output in enumerate(image_outputs):
+        envelope = json.loads(output.message.content)
+        data_url = envelope["result"]  # data:image/jpeg;base64,...
+        mime_type, payload = data_url.removeprefix("data:").split(";base64,", 1)
+        image_bytes = base64.b64decode(payload)
+        filename = f"{prefix}_{i}.{mime_type.removeprefix('image/')}"
+        with open(filename, "wb") as f:
+            f.write(image_bytes)
+        print(f"Saved {filename} ({len(image_bytes)} bytes)")
+
+
+def generate_image(client: Client) -> None:
+    """Generates an image in a single turn."""
+    chat = client.chat.create(
+        model="grok-4.5",
+        # Optionally restrict the tool with image_generation(action="generate")
+        # (text-to-image only) or image_generation(action="edit") (editing only).
+        tools=[image_generation()],
+    )
+
+    chat.append(user("Generate an image of a corgi surfing a big wave, in the style of a Japanese woodblock print"))
+    response = chat.sample()
+
+    print(response.content)
+    save_images(response, "corgi_surfing")
+    print(response.server_side_tool_usage)
+
+
+def generate_then_edit_image(client: Client) -> None:
+    """Generates an image, then edits it in a follow-up turn.
+
+    Storing the first turn (`store_messages=True`) and chaining the second one
+    via `previous_response_id` replays the full agentic state on the server,
+    so the model can edit the image it generated on the previous turn.
+    """
+    chat = client.chat.create(
+        model="grok-4.5",
+        tools=[image_generation(action="generate")],
+        store_messages=True,
+    )
+    chat.append(user("Generate an image of a corgi surfing a big wave, in the style of a Japanese woodblock print"))
+    response = chat.sample()
+
+    print(response.content)
+    save_images(response, "corgi_surfing")
+
+    follow_up_chat = client.chat.create(
+        model="grok-4.5",
+        tools=[image_generation()],
+        store_messages=True,
+        previous_response_id=response.id,
+    )
+    follow_up_chat.append(user("Edit the image you just generated: make it night time, lit by a full moon"))
+    follow_up_response = follow_up_chat.sample()
+
+    print(follow_up_response.content)
+    save_images(follow_up_response, "corgi_surfing_night")
+    print(follow_up_response.server_side_tool_usage)
+
+
+def search_and_generate_image(client: Client) -> None:
+    """Combines web search with image generation in a single agentic request.
+
+    The model first looks up live data with the web_search tool, then renders
+    what it found into a generated image.
+    """
+    chat = client.chat.create(
+        model="grok-4.5",
+        tools=[web_search(), image_generation()],
+    )
+    chat.append(
+        user(
+            "Generate an infographic image based on next week's temperature forecast in the UK, "
+            "with key city icons along with their forecast in the image"
+        )
+    )
+    response = chat.sample()
+
+    print(response.content)
+    for tool_call in response.tool_calls:
+        print(f"Tool call: {get_tool_call_type(tool_call)} ({tool_call.function.name})")
+    save_images(response, "uk_forecast_infographic")
+    print(response.server_side_tool_usage)
+
+
+def main() -> None:
+    client = Client(api_key=os.getenv("XAI_API_KEY"))
+
+    generate_image(client)
+
+    # Multi-turn: generate an image, then edit it in a follow-up turn.
+    generate_then_edit_image(client)
+
+    # Combine web search with image generation.
+    search_and_generate_image(client)
+
+
+if __name__ == "__main__":
+    main()

src/xai_sdk/proto/v5/chat_pb2.pyi

@@ -73,6 +73,7 @@ class ToolCallType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     TOOL_CALL_TYPE_COLLECTIONS_SEARCH_TOOL: _ClassVar[ToolCallType]
     TOOL_CALL_TYPE_MCP_TOOL: _ClassVar[ToolCallType]
     TOOL_CALL_TYPE_ATTACHMENT_SEARCH_TOOL: _ClassVar[ToolCallType]
+    TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL: _ClassVar[ToolCallType]
 
 class ToolCallStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -128,6 +129,7 @@ TOOL_CALL_TYPE_CODE_EXECUTION_TOOL: ToolCallType
 TOOL_CALL_TYPE_COLLECTIONS_SEARCH_TOOL: ToolCallType
 TOOL_CALL_TYPE_MCP_TOOL: ToolCallType
 TOOL_CALL_TYPE_ATTACHMENT_SEARCH_TOOL: ToolCallType
+TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL: ToolCallType
 TOOL_CALL_STATUS_IN_PROGRESS: ToolCallStatus
 TOOL_CALL_STATUS_COMPLETED: ToolCallStatus
 TOOL_CALL_STATUS_INCOMPLETE: ToolCallStatus
@@ -436,22 +438,24 @@ class ToolChoice(_message.Message):
     def __init__(self, mode: _Optional[_Union[ToolMode, str]] = ..., function_name: _Optional[str] = ...) -> None: ...
 
 class Tool(_message.Message):
-    __slots__ = ("function", "web_search", "x_search", "code_execution", "collections_search", "mcp", "attachment_search")
+    __slots__ = ("function", "web_search", "x_search", "code_execution", "collections_search", "mcp", "attachment_search", "image_generation")
     FUNCTION_FIELD_NUMBER: _ClassVar[int]
     WEB_SEARCH_FIELD_NUMBER: _ClassVar[int]
     X_SEARCH_FIELD_NUMBER: _ClassVar[int]
     CODE_EXECUTION_FIELD_NUMBER: _ClassVar[int]
     COLLECTIONS_SEARCH_FIELD_NUMBER: _ClassVar[int]
     MCP_FIELD_NUMBER: _ClassVar[int]
     ATTACHMENT_SEARCH_FIELD_NUMBER: _ClassVar[int]
+    IMAGE_GENERATION_FIELD_NUMBER: _ClassVar[int]
     function: Function
     web_search: WebSearch
     x_search: XSearch
     code_execution: CodeExecution
     collections_search: CollectionsSearch
     mcp: MCP
     attachment_search: AttachmentSearch
-    def __init__(self, function: _Optional[_Union[Function, _Mapping]] = ..., web_search: _Optional[_Union[WebSearch, _Mapping]] = ..., x_search: _Optional[_Union[XSearch, _Mapping]] = ..., code_execution: _Optional[_Union[CodeExecution, _Mapping]] = ..., collections_search: _Optional[_Union[CollectionsSearch, _Mapping]] = ..., mcp: _Optional[_Union[MCP, _Mapping]] = ..., attachment_search: _Optional[_Union[AttachmentSearch, _Mapping]] = ...) -> None: ...
+    image_generation: ImageGeneration
+    def __init__(self, function: _Optional[_Union[Function, _Mapping]] = ..., web_search: _Optional[_Union[WebSearch, _Mapping]] = ..., x_search: _Optional[_Union[XSearch, _Mapping]] = ..., code_execution: _Optional[_Union[CodeExecution, _Mapping]] = ..., collections_search: _Optional[_Union[CollectionsSearch, _Mapping]] = ..., mcp: _Optional[_Union[MCP, _Mapping]] = ..., attachment_search: _Optional[_Union[AttachmentSearch, _Mapping]] = ..., image_generation: _Optional[_Union[ImageGeneration, _Mapping]] = ...) -> None: ...
 
 class MCP(_message.Message):
     __slots__ = ("server_label", "server_description", "server_url", "allowed_tool_names", "authorization", "extra_headers")
@@ -522,6 +526,12 @@ class CodeExecution(_message.Message):
     __slots__ = ()
     def __init__(self) -> None: ...
 
+class ImageGeneration(_message.Message):
+    __slots__ = ("action",)
+    ACTION_FIELD_NUMBER: _ClassVar[int]
+    action: str
+    def __init__(self, action: _Optional[str] = ...) -> None: ...
+
 class CollectionsSearch(_message.Message):
     __slots__ = ("collection_ids", "limit", "instructions", "hybrid_retrieval", "semantic_retrieval", "keyword_retrieval")
     COLLECTION_IDS_FIELD_NUMBER: _ClassVar[int]

src/xai_sdk/proto/v5/usage_pb2.pyi

@@ -24,6 +24,7 @@ class ServerSideTool(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     SERVER_SIDE_TOOL_MCP: _ClassVar[ServerSideTool]
     SERVER_SIDE_TOOL_ATTACHMENT_SEARCH: _ClassVar[ServerSideTool]
     SERVER_SIDE_TOOL_IMAGE_SEARCH: _ClassVar[ServerSideTool]
+    SERVER_SIDE_TOOL_IMAGE_GENERATION: _ClassVar[ServerSideTool]
 SERVICE_TIER_UNSPECIFIED: ServiceTier
 SERVICE_TIER_DEFAULT: ServiceTier
 SERVICE_TIER_PRIORITY: ServiceTier
@@ -37,6 +38,7 @@ SERVER_SIDE_TOOL_COLLECTIONS_SEARCH: ServerSideTool
 SERVER_SIDE_TOOL_MCP: ServerSideTool
 SERVER_SIDE_TOOL_ATTACHMENT_SEARCH: ServerSideTool
 SERVER_SIDE_TOOL_IMAGE_SEARCH: ServerSideTool
+SERVER_SIDE_TOOL_IMAGE_GENERATION: ServerSideTool
 
 class SamplingUsage(_message.Message):
     __slots__ = ("completion_tokens", "reasoning_tokens", "prompt_tokens", "total_tokens", "prompt_text_tokens", "cached_prompt_text_tokens", "prompt_image_tokens", "num_sources_used", "server_side_tools_used", "cost_in_usd_ticks")

src/xai_sdk/proto/v6/chat_pb2.pyi

@@ -74,6 +74,7 @@ class ToolCallType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     TOOL_CALL_TYPE_COLLECTIONS_SEARCH_TOOL: _ClassVar[ToolCallType]
     TOOL_CALL_TYPE_MCP_TOOL: _ClassVar[ToolCallType]
     TOOL_CALL_TYPE_ATTACHMENT_SEARCH_TOOL: _ClassVar[ToolCallType]
+    TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL: _ClassVar[ToolCallType]
 
 class ToolCallStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -129,6 +130,7 @@ TOOL_CALL_TYPE_CODE_EXECUTION_TOOL: ToolCallType
 TOOL_CALL_TYPE_COLLECTIONS_SEARCH_TOOL: ToolCallType
 TOOL_CALL_TYPE_MCP_TOOL: ToolCallType
 TOOL_CALL_TYPE_ATTACHMENT_SEARCH_TOOL: ToolCallType
+TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL: ToolCallType
 TOOL_CALL_STATUS_IN_PROGRESS: ToolCallStatus
 TOOL_CALL_STATUS_COMPLETED: ToolCallStatus
 TOOL_CALL_STATUS_INCOMPLETE: ToolCallStatus
@@ -437,22 +439,24 @@ class ToolChoice(_message.Message):
     def __init__(self, mode: _Optional[_Union[ToolMode, str]] = ..., function_name: _Optional[str] = ...) -> None: ...
 
 class Tool(_message.Message):
-    __slots__ = ("function", "web_search", "x_search", "code_execution", "collections_search", "mcp", "attachment_search")
+    __slots__ = ("function", "web_search", "x_search", "code_execution", "collections_search", "mcp", "attachment_search", "image_generation")
     FUNCTION_FIELD_NUMBER: _ClassVar[int]
     WEB_SEARCH_FIELD_NUMBER: _ClassVar[int]
     X_SEARCH_FIELD_NUMBER: _ClassVar[int]
     CODE_EXECUTION_FIELD_NUMBER: _ClassVar[int]
     COLLECTIONS_SEARCH_FIELD_NUMBER: _ClassVar[int]
     MCP_FIELD_NUMBER: _ClassVar[int]
     ATTACHMENT_SEARCH_FIELD_NUMBER: _ClassVar[int]
+    IMAGE_GENERATION_FIELD_NUMBER: _ClassVar[int]
     function: Function
     web_search: WebSearch
     x_search: XSearch
     code_execution: CodeExecution
     collections_search: CollectionsSearch
     mcp: MCP
     attachment_search: AttachmentSearch
-    def __init__(self, function: _Optional[_Union[Function, _Mapping]] = ..., web_search: _Optional[_Union[WebSearch, _Mapping]] = ..., x_search: _Optional[_Union[XSearch, _Mapping]] = ..., code_execution: _Optional[_Union[CodeExecution, _Mapping]] = ..., collections_search: _Optional[_Union[CollectionsSearch, _Mapping]] = ..., mcp: _Optional[_Union[MCP, _Mapping]] = ..., attachment_search: _Optional[_Union[AttachmentSearch, _Mapping]] = ...) -> None: ...
+    image_generation: ImageGeneration
+    def __init__(self, function: _Optional[_Union[Function, _Mapping]] = ..., web_search: _Optional[_Union[WebSearch, _Mapping]] = ..., x_search: _Optional[_Union[XSearch, _Mapping]] = ..., code_execution: _Optional[_Union[CodeExecution, _Mapping]] = ..., collections_search: _Optional[_Union[CollectionsSearch, _Mapping]] = ..., mcp: _Optional[_Union[MCP, _Mapping]] = ..., attachment_search: _Optional[_Union[AttachmentSearch, _Mapping]] = ..., image_generation: _Optional[_Union[ImageGeneration, _Mapping]] = ...) -> None: ...
 
 class MCP(_message.Message):
     __slots__ = ("server_label", "server_description", "server_url", "allowed_tool_names", "authorization", "extra_headers")
@@ -523,6 +527,12 @@ class CodeExecution(_message.Message):
     __slots__ = ()
     def __init__(self) -> None: ...
 
+class ImageGeneration(_message.Message):
+    __slots__ = ("action",)
+    ACTION_FIELD_NUMBER: _ClassVar[int]
+    action: str
+    def __init__(self, action: _Optional[str] = ...) -> None: ...
+
 class CollectionsSearch(_message.Message):
     __slots__ = ("collection_ids", "limit", "instructions", "hybrid_retrieval", "semantic_retrieval", "keyword_retrieval")
     COLLECTION_IDS_FIELD_NUMBER: _ClassVar[int]