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. Takes an optional `action` parameter (`"auto"`, `"generate"`, or `"edit"`) to control which image capabilities are exposed to the model. Also adds the `TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL` tool call type, the `SERVER_SIDE_TOOL_IMAGE_GENERATION` usage enum value, and `num_image_generations` / `num_image_edits` fields on `SamplingUsage`
 - **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/aio/image_generation_tool.py

@@ -0,0 +1,85 @@
+"""Examples for the server-side `image_generation` tool (async client).
+
+Unlike the standalone image API (see `image_generation.py`), the `image_generation`
+tool runs inside an agentic chat request: the model decides when to generate or
+edit images as part of answering the conversation.
+
+The generated image rides on the response as a `ROLE_TOOL` output message whose
+content is a JSON envelope of the form:
+
+    {"__type": "image_generation_result", "result": "data:<mime>;base64,..."}
+"""
+
+import asyncio
+import base64
+import json
+import pathlib
+
+from xai_sdk import AsyncClient
+from xai_sdk.chat import user
+from xai_sdk.proto import chat_pb2
+from xai_sdk.tools import get_tool_call_type, image_generation
+
+
+def save_generated_images(response, output_dir: pathlib.Path) -> list[pathlib.Path]:
+    """Decodes and saves every image the image_generation tool produced."""
+    saved_paths = []
+    tool_outputs = [
+        output
+        for output in response.proto.outputs
+        if output.message.role == chat_pb2.MessageRole.ROLE_TOOL
+        and any(get_tool_call_type(tool_call) == "image_generation_tool" for tool_call in output.message.tool_calls)
+    ]
+    for i, output in enumerate(tool_outputs):
+        envelope = json.loads(output.message.content)
+        if envelope.get("__type") != "image_generation_result":
+            continue
+        # The result is a data URL, e.g. "data:image/jpeg;base64,<payload>".
+        data_url = envelope["result"]
+        header, payload = data_url.split(",", 1)
+        mime_type = header.removeprefix("data:").removesuffix(";base64")
+        extension = mime_type.removeprefix("image/")
+        path = output_dir / f"generated_image_{i}.{extension}"
+        path.write_bytes(base64.b64decode(payload))
+        saved_paths.append(path)
+    return saved_paths
+
+
+async def generate_image_streaming(client: AsyncClient, model: str) -> None:
+    """Asks the model to generate an image, streaming progress as it works."""
+    chat = client.chat.create(
+        model=model,
+        # action="auto" (the default) allows both generation and editing.
+        # Use action="generate" for text-to-image only, or action="edit" for editing only.
+        tools=[image_generation()],
+    )
+    chat.append(user("Generate an image of a red panda drinking boba tea."))
+
+    response = None
+    async for response, chunk in chat.stream():  # noqa: B007
+        for tool_call in chunk.tool_calls:
+            if get_tool_call_type(tool_call) == "image_generation_tool":
+                print(f"Tool call: {tool_call.function.name} with arguments: {tool_call.function.arguments}")
+        if chunk.content:
+            print(chunk.content, end="", flush=True)
+    print()
+
+    assert response is not None
+
+    saved_paths = save_generated_images(response, pathlib.Path.cwd())
+    for path in saved_paths:
+        print(f"Saved generated image to: {path}")
+
+    print(f"Server-side tool usage: {response.server_side_tool_usage}")
+    print(f"Image generations: {response.usage.num_image_generations}")
+    if response.cost_usd is not None:
+        print(f"Cost: ${response.cost_usd:.4f}")
+
+
+async def main() -> None:
+    client = AsyncClient()
+    await generate_image_streaming(client, model="grok-4.20")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/sync/image_generation_tool.py

@@ -0,0 +1,118 @@
+"""Examples for the server-side `image_generation` tool.
+
+Unlike the standalone image API (see `image_generation.py`), the `image_generation`
+tool runs inside an agentic chat request: the model decides when to generate or
+edit images as part of answering the conversation.
+
+The generated image rides on the response as a `ROLE_TOOL` output message whose
+content is a JSON envelope of the form:
+
+    {"__type": "image_generation_result", "result": "data:<mime>;base64,..."}
+"""
+
+import base64
+import json
+import pathlib
+
+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
+
+
+def save_generated_images(response, output_dir: pathlib.Path) -> list[pathlib.Path]:
+    """Decodes and saves every image the image_generation tool produced."""
+    saved_paths = []
+    tool_outputs = [
+        output
+        for output in response.proto.outputs
+        if output.message.role == chat_pb2.MessageRole.ROLE_TOOL
+        and any(get_tool_call_type(tool_call) == "image_generation_tool" for tool_call in output.message.tool_calls)
+    ]
+    for i, output in enumerate(tool_outputs):
+        envelope = json.loads(output.message.content)
+        if envelope.get("__type") != "image_generation_result":
+            continue
+        # The result is a data URL, e.g. "data:image/jpeg;base64,<payload>".
+        data_url = envelope["result"]
+        header, payload = data_url.split(",", 1)
+        mime_type = header.removeprefix("data:").removesuffix(";base64")
+        extension = mime_type.removeprefix("image/")
+        path = output_dir / f"generated_image_{i}.{extension}"
+        path.write_bytes(base64.b64decode(payload))
+        saved_paths.append(path)
+    return saved_paths
+
+
+def generate_image(client: Client, model: str) -> None:
+    """Asks the model to generate an image and saves the result to disk."""
+    chat = client.chat.create(
+        model=model,
+        # action="auto" (the default) allows both generation and editing.
+        # Use action="generate" for text-to-image only, or action="edit" for editing only.
+        tools=[image_generation()],
+    )
+    chat.append(user("Generate an image of a red panda drinking boba tea."))
+
+    response = chat.sample()
+
+    for tool_call in response.tool_calls:
+        if get_tool_call_type(tool_call) == "image_generation_tool":
+            print(f"Tool call: {tool_call.function.name} with arguments: {tool_call.function.arguments}")
+
+    print(f"Response: {response.content}")
+
+    saved_paths = save_generated_images(response, pathlib.Path.cwd())
+    for path in saved_paths:
+        print(f"Saved generated image to: {path}")
+
+    print(f"Server-side tool usage: {response.server_side_tool_usage}")
+    print(f"Image generations: {response.usage.num_image_generations}")
+    if response.cost_usd is not None:
+        print(f"Cost: ${response.cost_usd:.4f}")
+
+
+def generate_then_edit(client: Client, model: str) -> None:
+    """Generates an image on turn 1, then edits it on turn 2 via `previous_response_id`."""
+    chat = client.chat.create(
+        model=model,
+        tools=[image_generation(action="generate")],
+        store_messages=True,
+    )
+    chat.append(user("Generate an image of a lighthouse on a cliff."))
+    response = chat.sample()
+    print(f"Turn 1: {response.content}")
+
+    # The follow-up chat replays the stored conversation, including the generated
+    # image, so the model can edit it.
+    follow_up_chat = client.chat.create(
+        model=model,
+        tools=[image_generation()],
+        store_messages=True,
+        previous_response_id=response.id,
+    )
+    follow_up_chat.append(user("Edit the image you just generated: make it nighttime with a full moon."))
+    follow_up_response = follow_up_chat.sample()
+
+    for tool_call in follow_up_response.tool_calls:
+        if get_tool_call_type(tool_call) == "image_generation_tool":
+            # Generation calls are named `imagine_text_to_image`, edits `imagine_image_to_image`.
+            print(f"Tool call: {tool_call.function.name}")
+
+    print(f"Turn 2: {follow_up_response.content}")
+    print(f"Image edits: {follow_up_response.usage.num_image_edits}")
+
+    saved_paths = save_generated_images(follow_up_response, pathlib.Path.cwd())
+    for path in saved_paths:
+        print(f"Saved edited image to: {path}")
+
+
+def main() -> None:
+    client = Client()
+
+    generate_image(client, model="grok-4.20")
+    # generate_then_edit(client, model="grok-4.20")
+
+
+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]