diff --git a/CHANGELOG.md b/CHANGELOG.md index a9ea8bf..44d90d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ - Added `grok-imagine-image-2.0` to the `ImageGenerationModel` known-model type literal - **`xhigh` Reasoning Effort**: Added `"xhigh"` as an accepted `reasoning_effort` value (maps to `EFFORT_XHIGH`; supported by models such as `grok-4.6`) - Added `grok-4.6` to the `ChatModel` known-model type literal +- **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 +- **Image Generation Tool Outputs**: Added `Response.image_outputs` for retrieving images generated by the server-side `image_generation` tool as decoded bytes (`output.image`), along with `mime_type`, `data_url`, `image_uuid`, and the originating `tool_call` - **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": }` 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`). diff --git a/examples/aio/image_generation_tool.py b/examples/aio/image_generation_tool.py new file mode 100644 index 0000000..3e42f50 --- /dev/null +++ b/examples/aio/image_generation_tool.py @@ -0,0 +1,161 @@ +"""Examples for the server-side `image_generation` tool. + +Unlike the standalone Image API (`client.image`, see `image_generation.py`), the +`image_generation` tool lets the model generate and edit images mid-conversation +as part of a single agentic chat request. Generated images are exposed on the +response via `response.image_outputs`. +""" + +import asyncio +import os + +from xai_sdk import AsyncClient +from xai_sdk.chat import Response, image, user +from xai_sdk.tools import get_tool_call_type, image_generation, web_search + + +async def generate_image(client: AsyncClient) -> None: + """Generates an image in a single turn.""" + chat = client.chat.create( + model="grok-4.5", + 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 = await chat.sample() + print(response.content) + _save_images(response, "corgi_surfing") + print(response.server_side_tool_usage) + + +async def edit_input_image(client: AsyncClient) -> None: + """Edits an image attached to the chat context. + + With `action="edit"` (or the default `"auto"`), the model can edit any image + already in the conversation — images attached as input via the `image()` + content helper as well as images it generated earlier. + """ + chat = client.chat.create( + model="grok-4.5", + tools=[image_generation(action="edit")], + ) + chat.append( + user( + "Edit this image so it looks like a watercolor painting.", + image("https://docs.x.ai/assets/api-examples/images/style-realistic.png"), + ) + ) + response = await chat.sample() + print(response.content) + _save_images(response, "watercolor") + + +async def generate_then_edit_image(client: AsyncClient) -> None: + """Generates an image, then edits it in a follow-up turn of the same chat.""" + chat = client.chat.create( + model="grok-4.5", + # Auto mode: exposes both imagine_text_to_image AND imagine_image_to_image. + # (action="generate" would hide the edit tool from every turn of this chat.) + tools=[image_generation()], + ) + + # -- Turn 1: generate ---------------------------------------------------- + chat.append(user("Generate an image of a corgi surfing a big wave, in the style of a Japanese woodblock print")) + response = await chat.sample() + print(response.content) + _save_images(response, "corgi_surfing") + + chat.append(response) + + # -- Turn 2: edit -------------------------------------------------------- + chat.append(user("Edit the image you just generated: make it night time, lit by a full moon")) + response = await chat.sample() + print(response.content) + _save_images(response, "corgi_surfing_night") + + +async def search_and_generate_image(client: AsyncClient) -> 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 = await chat.sample() + print(response.content) + _save_images(response, "uk_forecast_infographic") + print(response.server_side_tool_usage) + + +async def stream_image_generation(client: AsyncClient) -> None: + """Streams a response that generates an image. + + With `include=["verbose_streaming"]`, tool-call activity and text deltas + arrive on the streamed chunks as they happen. The image payload itself + streams as a (large) tool output; the decoded bytes are exposed on the + accumulated response via `response.image_outputs` once the stream ends. + """ + chat = client.chat.create( + model="grok-4.5", + tools=[image_generation()], + include=["verbose_streaming"], + ) + chat.append(user("Generate an image of an origami fox in a paper forest")) + + last_response: Response | None = None + async for response, chunk in chat.stream(): + last_response = response + for tool_call in chunk.tool_calls: + if get_tool_call_type(tool_call) == "image_generation_tool": + print(f"\nGenerating image: {tool_call.function.arguments}") + for tool_output in chunk.tool_outputs: + # The raw payload is a base64 data URL envelope; don't print it. + print(f"\nReceived image payload chunk ({len(tool_output.content)} chars)") + if chunk.content: + print(chunk.content, end="", flush=True) + print() + if last_response is not None: + _save_images(last_response, "origami_fox") + + +def _save_images(response: Response, prefix: str) -> None: + """Saves every image produced by image_generation tool calls in a response.""" + if not response.image_outputs: + print("No images were generated.") + return + for i, output in enumerate(response.image_outputs): + extension = output.mime_type.removeprefix("image/") + filename = f"{prefix}_{i}.{extension}" + with open(filename, "wb") as f: + f.write(output.image) + print(f"Saved {filename} ({output.tool_call.function.name}, image_uuid={output.image_uuid})") + + +async def main() -> None: + client = AsyncClient(api_key=os.getenv("XAI_API_KEY")) + + await generate_image(client) + + # Edit an image attached to the chat context. + # await edit_input_image(client) + + # Multi-turn: generate an image, then edit it in a follow-up turn. + # await generate_then_edit_image(client) + + # Combine web search with image generation. + # await search_and_generate_image(client) + + # Stream tool activity and text deltas while the image is generated. + # await stream_image_generation(client) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/sync/image_generation_tool.py b/examples/sync/image_generation_tool.py new file mode 100644 index 0000000..0863096 --- /dev/null +++ b/examples/sync/image_generation_tool.py @@ -0,0 +1,160 @@ +"""Examples for the server-side `image_generation` tool. + +Unlike the standalone Image API (`client.image`, see `image_generation.py`), the +`image_generation` tool lets the model generate and edit images mid-conversation +as part of a single agentic chat request. Generated images are exposed on the +response via `response.image_outputs`. +""" + +import os + +from xai_sdk import Client +from xai_sdk.chat import Response, image, user +from xai_sdk.tools import get_tool_call_type, image_generation, web_search + + +def generate_image(client: Client) -> None: + """Generates an image in a single turn.""" + chat = client.chat.create( + model="grok-4.5", + 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 edit_input_image(client: Client) -> None: + """Edits an image attached to the chat context. + + With `action="edit"` (or the default `"auto"`), the model can edit any image + already in the conversation — images attached as input via the `image()` + content helper as well as images it generated earlier. + """ + chat = client.chat.create( + model="grok-4.5", + tools=[image_generation(action="edit")], + ) + chat.append( + user( + "Edit this image so it looks like a watercolor painting.", + image("https://docs.x.ai/assets/api-examples/images/style-realistic.png"), + ) + ) + response = chat.sample() + print(response.content) + _save_images(response, "watercolor") + + +def generate_then_edit_image(client: Client) -> None: + """Generates an image, then edits it in a follow-up turn of the same chat.""" + chat = client.chat.create( + model="grok-4.5", + # Auto mode: exposes both imagine_text_to_image AND imagine_image_to_image. + # (action="generate" would hide the edit tool from every turn of this chat.) + tools=[image_generation()], + ) + + # -- Turn 1: generate ---------------------------------------------------- + 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") + + chat.append(response) + + # -- Turn 2: edit -------------------------------------------------------- + chat.append(user("Edit the image you just generated: make it night time, lit by a full moon")) + response = chat.sample() + print(response.content) + _save_images(response, "corgi_surfing_night") + + +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) + _save_images(response, "uk_forecast_infographic") + print(response.server_side_tool_usage) + + +def stream_image_generation(client: Client) -> None: + """Streams a response that generates an image. + + With `include=["verbose_streaming"]`, tool-call activity and text deltas + arrive on the streamed chunks as they happen. The image payload itself + streams as a (large) tool output; the decoded bytes are exposed on the + accumulated response via `response.image_outputs` once the stream ends. + """ + chat = client.chat.create( + model="grok-4.5", + tools=[image_generation()], + include=["verbose_streaming"], + ) + chat.append(user("Generate an image of an origami fox in a paper forest")) + + last_response: Response | None = None + for response, chunk in chat.stream(): + last_response = response + for tool_call in chunk.tool_calls: + if get_tool_call_type(tool_call) == "image_generation_tool": + print(f"\nGenerating image: {tool_call.function.arguments}") + for tool_output in chunk.tool_outputs: + # The raw payload is a base64 data URL envelope; don't print it. + print(f"\nReceived image payload chunk ({len(tool_output.content)} chars)") + if chunk.content: + print(chunk.content, end="", flush=True) + print() + if last_response is not None: + _save_images(last_response, "origami_fox") + + +def _save_images(response: Response, prefix: str) -> None: + """Saves every image produced by image_generation tool calls in a response.""" + if not response.image_outputs: + print("No images were generated.") + return + for i, output in enumerate(response.image_outputs): + extension = output.mime_type.removeprefix("image/") + filename = f"{prefix}_{i}.{extension}" + with open(filename, "wb") as f: + f.write(output.image) + print(f"Saved {filename} ({output.tool_call.function.name}, image_uuid={output.image_uuid})") + + +def main() -> None: + client = Client(api_key=os.getenv("XAI_API_KEY")) + + generate_image(client) + + # Edit an image attached to the chat context. + # edit_input_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) + + # Stream tool activity and text deltas while the image is generated. + # stream_image_generation(client) + + +if __name__ == "__main__": + main() diff --git a/src/xai_sdk/chat.py b/src/xai_sdk/chat.py index c6386b9..c01ee41 100644 --- a/src/xai_sdk/chat.py +++ b/src/xai_sdk/chat.py @@ -1,4 +1,5 @@ import abc +import base64 import datetime import json from collections import Counter, defaultdict @@ -1122,6 +1123,84 @@ def finish_reason(self) -> sample_pb2.FinishReason: return self.proto.finish_reason +class ImageGenerationOutput(ProtoDecorator[chat_pb2.CompletionOutput]): + """A single image produced by the server-side `image_generation` tool. + + Exposes the generated image as decoded `image` bytes (with `mime_type` and + `data_url` describing the payload), the `image_uuid` used to reference the + image in follow-up edit requests, and the completed `tool_call` that + produced it. + """ + + _tool_call: chat_pb2.ToolCall + _data_url: str + _image_uuid: str + _image: bytes + + def __init__(self, proto: chat_pb2.CompletionOutput, tool_call: chat_pb2.ToolCall) -> None: + """Initializes a new instance of the `ImageGenerationOutput` class. + + Args: + proto: The `ROLE_TOOL` completion output carrying the image generation result. + tool_call: The completed image generation tool call that produced this image. + + Raises: + ValueError: If the output content is not a well-formed image generation result + envelope. Completed calls always carry the full envelope in a single chunk, + so this indicates a corrupt payload, not a normal stream state. + """ + super().__init__(proto) + self._tool_call = tool_call + content = proto.message.content + error_message = f"Output content is not a well-formed image generation result envelope: {content[:64]!r}" + try: + envelope = json.loads(content) + except json.JSONDecodeError as e: + raise ValueError(error_message) from e + if not isinstance(envelope, dict) or envelope.get("__type") != "image_generation_result": + raise ValueError(error_message) + data_url = envelope.get("result") + if not isinstance(data_url, str) or not data_url.startswith("data:image/") or "base64," not in data_url: + raise ValueError(error_message) + self._data_url = data_url + image_uuid = envelope.get("image_uuid", "") + self._image_uuid = image_uuid if isinstance(image_uuid, str) else "" + _, encoded = data_url.split("base64,", 1) + self._image = base64.b64decode(encoded) + + @property + def tool_call(self) -> chat_pb2.ToolCall: + """The completed image generation tool call that produced this image. + + `tool_call.function.name` distinguishes text-to-image generations + (`imagine_text_to_image`) from image edits (`imagine_image_to_image`). + """ + return self._tool_call + + @property + def data_url(self) -> str: + """The generated image as a `data:;base64,` URL.""" + return self._data_url + + @property + def mime_type(self) -> str: + """The MIME type of the generated image (e.g. `image/jpeg`).""" + return self._data_url.removeprefix("data:").split(";", 1)[0] + + @property + def image_uuid(self) -> str: + """The short reference ID the model uses to address this image in follow-up tool calls. + + Empty if the server did not report one. + """ + return self._image_uuid + + @property + def image(self) -> bytes: + """The generated image as raw bytes.""" + return self._image + + class _ResponseProtoDecorator(ProtoDecorator[chat_pb2.GetChatCompletionResponse]): def __init__(self, proto: chat_pb2.GetChatCompletionResponse) -> None: """Initialize with proto and content buffers for efficient accumulation.""" @@ -1373,6 +1452,39 @@ def tool_outputs(self) -> Sequence[chat_pb2.CompletionOutput]: """Returns the output entries that contain the tool outputs.""" return [output for output in self.proto.outputs if output.message.role == chat_pb2.MessageRole.ROLE_TOOL] + @property + def image_outputs(self) -> Sequence[ImageGenerationOutput]: + """Returns the images generated by the server-side `image_generation` tool. + + One entry per completed image generation/edit call, in response order. + Failed calls carry an error payload instead of an image and are omitted. + + Safe to access while streaming: the server delivers a completed call and + its full result envelope in the same chunk, so entries appear fully + formed as soon as their tool call completes. A `ValueError` here means a + completed call carried a corrupt envelope (a server bug), not a normal + stream state. + + Example: + ``` + 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) + ``` + """ + image_outputs = [] + for output in self.tool_outputs: + completed_calls = [ + tool_call + for tool_call in output.message.tool_calls + if tool_call.type == chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL + and tool_call.status == chat_pb2.TOOL_CALL_STATUS_COMPLETED + ] + if completed_calls: + image_outputs.append(ImageGenerationOutput(output, completed_calls[0])) + return image_outputs + @property def server_side_tool_usage(self) -> dict[str, int]: """Returns the server side tools used for this response.""" diff --git a/src/xai_sdk/tools.py b/src/xai_sdk/tools.py index a5c8701..363f9df 100644 --- a/src/xai_sdk/tools.py +++ b/src/xai_sdk/tools.py @@ -1,5 +1,5 @@ import datetime -from typing import Optional, Union +from typing import Literal, Optional, Union from google.protobuf.timestamp_pb2 import Timestamp @@ -170,6 +170,38 @@ def code_execution() -> chat_pb2.Tool: return chat_pb2.Tool(code_execution=chat_pb2.CodeExecution()) +def image_generation(action: Optional[Literal["auto", "generate", "edit"]] = None) -> chat_pb2.Tool: + """Creates a server-side tool for image generation, typically used in agentic requests. + + This tool enables the model to generate images from text prompts and to edit images + (both previously generated ones and images provided in the conversation) as part of + generating responses. + + Args: + action: Which image capabilities to expose to the model. One of "auto" (the + default; both generation and editing), "generate" (text-to-image only), or + "edit" (image editing only). Defaults to None, which the server treats + as "auto". + + Returns: + A `chat_pb2.Tool` object configured for image generation. + + Example: + ``` + from xai_sdk.tools import image_generation + + # Create an image generation tool with both generation and editing enabled + tool = image_generation() + + # Restrict the tool to text-to-image generation only + tool = image_generation(action="generate") + ``` + """ + if action is None: + return chat_pb2.Tool(image_generation=chat_pb2.ImageGeneration()) + return chat_pb2.Tool(image_generation=chat_pb2.ImageGeneration(action=action)) + + def collections_search( collection_ids: list[str], limit: Optional[int] = None, @@ -298,6 +330,7 @@ def get_tool_call_type(tool_call: chat_pb2.ToolCall) -> str: Returns: The type of the tool call as a string, valid values are: "client_side_tool", "web_search_tool", - "x_search_tool", "code_execution_tool", "collections_search_tool", "mcp_tool", "attachment_search_tool". + "x_search_tool", "code_execution_tool", "collections_search_tool", "mcp_tool", "attachment_search_tool", + "image_generation_tool". """ return chat_pb2.ToolCallType.Name(tool_call.type).removeprefix("TOOL_CALL_TYPE_").lower() diff --git a/tests/aio/chat_test.py b/tests/aio/chat_test.py index 180d487..e5acaf5 100644 --- a/tests/aio/chat_test.py +++ b/tests/aio/chat_test.py @@ -29,7 +29,7 @@ from xai_sdk.cost import USD_PER_TICK from xai_sdk.proto import chat_pb2, image_pb2, sample_pb2, usage_pb2 from xai_sdk.search import SearchParameters, news_source, rss_source, web_source, x_source -from xai_sdk.tools import code_execution, web_search, x_search +from xai_sdk.tools import code_execution, image_generation, web_search, x_search from .. import server @@ -429,6 +429,56 @@ async def test_agentic_tool_calling_non_streaming(client): assert response.tool_calls[0].function.arguments == '{"query":"What is the weather in London?"}' +@pytest.mark.asyncio(loop_scope="session") +async def test_image_generation_tool_non_streaming(client): + chat = client.chat.create( + "grok-4-fast", + tools=[image_generation()], + ) + chat.append(user("Generate an image of a corgi")) + response = await chat.sample() + + image_outputs = response.image_outputs + assert len(image_outputs) == 1 + output = image_outputs[0] + assert output.tool_call.function.name == "imagine_text_to_image" + assert output.tool_call.status == chat_pb2.TOOL_CALL_STATUS_COMPLETED + assert output.mime_type == "image/jpeg" + assert output.data_url.startswith("data:image/jpeg;base64,") + assert output.image_uuid == server.IMAGE_GENERATION_IMAGE_UUID + assert output.image == server.read_image() + + assert response.content == "Here is your image." + + +@pytest.mark.asyncio(loop_scope="session") +async def test_image_generation_tool_streaming(client): + chat = client.chat.create( + "grok-4-fast", + tools=[image_generation()], + ) + chat.append(user("Generate an image of a corgi")) + + last_response = None + async for response, _chunk in chat.stream(): + # Safe to access mid-stream: a completed call and its full envelope + # always arrive in the same chunk, so entries appear fully formed. + for output in response.image_outputs: + assert output.tool_call.type == chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL + last_response = response + + assert last_response is not None + image_outputs = last_response.image_outputs + assert len(image_outputs) == 1 + output = image_outputs[0] + assert output.tool_call.function.name == "imagine_text_to_image" + assert output.mime_type == "image/jpeg" + assert output.image_uuid == server.IMAGE_GENERATION_IMAGE_UUID + assert output.image == server.read_image() + + assert last_response.content == "Here is your image." + + @pytest.mark.asyncio(loop_scope="session") async def test_structured_output_parse(client): class Weather(BaseModel): @@ -1464,11 +1514,12 @@ def test_chat_create_with_server_side_tools(client: AsyncClient): enable_video_understanding=True, ), code_execution(), + image_generation(), ], ) chat_completion_request = chat.proto - assert len(chat_completion_request.tools) == 3 + assert len(chat_completion_request.tools) == 4 expected_from_date_pb = timestamp_pb2.Timestamp() expected_from_date_pb.FromDatetime(from_date) @@ -1496,9 +1547,12 @@ def test_chat_create_with_server_side_tools(client: AsyncClient): expected_code_execution_tool = chat_pb2.Tool(code_execution=chat_pb2.CodeExecution()) + expected_image_generation_tool = chat_pb2.Tool(image_generation=chat_pb2.ImageGeneration()) + assert chat_completion_request.tools[0] == expected_web_search_tool assert chat_completion_request.tools[1] == expected_x_search_tool assert chat_completion_request.tools[2] == expected_code_execution_tool + assert chat_completion_request.tools[3] == expected_image_generation_tool @pytest.mark.parametrize( diff --git a/tests/chat_test.py b/tests/chat_test.py index c37043e..15ace36 100644 --- a/tests/chat_test.py +++ b/tests/chat_test.py @@ -661,6 +661,35 @@ def test_server_side_tool_image_search_enum(): assert usage_pb2.ServerSideTool.Name(usage_pb2.SERVER_SIDE_TOOL_IMAGE_SEARCH) == "SERVER_SIDE_TOOL_IMAGE_SEARCH" +def test_image_generation_tool(): + """Test that image_generation util function correctly creates an image generation tool.""" + from xai_sdk.tools import image_generation + + tool = image_generation() + assert isinstance(tool, chat_pb2.Tool) + assert tool.HasField("image_generation") + assert not tool.image_generation.HasField("action") + + for action in ("auto", "generate", "edit"): + tool = image_generation(action=action) + assert tool.HasField("image_generation") + assert tool.image_generation.action == action + + +def test_image_generation_tool_call_type(): + """Test that image generation tool calls map to the expected type string.""" + tool_call = chat_pb2.ToolCall(type=chat_pb2.ToolCallType.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL) + assert get_tool_call_type(tool_call) == "image_generation_tool" + + +def test_server_side_tool_image_generation_enum(): + assert usage_pb2.SERVER_SIDE_TOOL_IMAGE_GENERATION == 11 + assert ( + usage_pb2.ServerSideTool.Name(usage_pb2.SERVER_SIDE_TOOL_IMAGE_GENERATION) + == "SERVER_SIDE_TOOL_IMAGE_GENERATION" + ) + + def test_developer_message(): """Test that developer() creates a message with ROLE_DEVELOPER role.""" # Simple string content diff --git a/tests/server.py b/tests/server.py index 95250db..607820e 100644 --- a/tests/server.py +++ b/tests/server.py @@ -3,6 +3,7 @@ import base64 import contextlib import http.server +import json import os.path import threading import time @@ -158,6 +159,112 @@ def _use_server_side_tools(request: chat_pb2.GetCompletionsRequest) -> bool: return any(tool.WhichOneof("tool") != "function" for tool in request.tools) +def _use_image_generation(request: chat_pb2.GetCompletionsRequest) -> bool: + return any(tool.WhichOneof("tool") == "image_generation" for tool in request.tools) + + +IMAGE_GENERATION_IMAGE_UUID = "img01" + + +def image_generation_envelope() -> str: + """The serialized image generation result envelope carried by ROLE_TOOL outputs.""" + encoded = base64.b64encode(read_image()).decode() + return json.dumps( + { + "__type": "image_generation_result", + "result": f"data:image/jpeg;base64,{encoded}", + "image_uuid": IMAGE_GENERATION_IMAGE_UUID, + } + ) + + +def _image_generation_tool_call() -> chat_pb2.ToolCall: + return chat_pb2.ToolCall( + id="test-image-tool-call", + type=chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL, + status=chat_pb2.TOOL_CALL_STATUS_COMPLETED, + function=chat_pb2.FunctionCall( + name="imagine_text_to_image", + arguments='{"prompt":"a corgi"}', + ), + ) + + +def _add_server_side_tool_outputs( + request: chat_pb2.GetCompletionsRequest, response: chat_pb2.GetChatCompletionResponse +) -> None: + """Adds the agentic tool calling outputs to a completion response.""" + if _use_image_generation(request): + _add_image_generation_outputs(response) + return + response.outputs.add( + index=0, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_ASSISTANT, + tool_calls=[ + chat_pb2.ToolCall( + id="test-tool-call", + function=chat_pb2.FunctionCall( + name="web_search", + arguments='{"query":"What is the weather in London?"}', + ), + ) + ], + ), + ) + response.outputs.add( + index=1, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_TOOL, + tool_calls=[ + chat_pb2.ToolCall( + id="test-tool-call", + function=chat_pb2.FunctionCall( + name="web_search", + arguments='{"query":"What is the weather in London?"}', + ), + ) + ], + content="I am tool response", + ), + ) + response.outputs.add( + index=2, + finish_reason=sample_pb2.FinishReason.REASON_STOP, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_ASSISTANT, + content="I am searching.", + ), + ) + + +def _add_image_generation_outputs(response: chat_pb2.GetChatCompletionResponse) -> None: + """Adds the agentic image generation outputs to a completion response.""" + response.outputs.add( + index=0, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_ASSISTANT, + tool_calls=[_image_generation_tool_call()], + ), + ) + response.outputs.add( + index=1, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_TOOL, + tool_calls=[_image_generation_tool_call()], + content=image_generation_envelope(), + ), + ) + response.outputs.add( + index=2, + finish_reason=sample_pb2.FinishReason.REASON_STOP, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_ASSISTANT, + content="Here is your image.", + ), + ) + + class AuthServicer(auth_pb2_grpc.AuthServicer): """A dummy implementation of the Auth service for testing.""" @@ -207,45 +314,7 @@ def GetCompletion(self, request: chat_pb2.GetCompletionsRequest, context: grpc.S for i in range(request.n): if len(request.tools) > 0 and _use_server_side_tools(request): - response.outputs.add( - index=0, - message=chat_pb2.CompletionMessage( - role=chat_pb2.ROLE_ASSISTANT, - tool_calls=[ - chat_pb2.ToolCall( - id="test-tool-call", - function=chat_pb2.FunctionCall( - name="web_search", - arguments='{"query":"What is the weather in London?"}', - ), - ) - ], - ), - ) - response.outputs.add( - index=1, - message=chat_pb2.CompletionMessage( - role=chat_pb2.ROLE_TOOL, - tool_calls=[ - chat_pb2.ToolCall( - id="test-tool-call", - function=chat_pb2.FunctionCall( - name="web_search", - arguments='{"query":"What is the weather in London?"}', - ), - ) - ], - content="I am tool response", - ), - ) - response.outputs.add( - index=2, - finish_reason=sample_pb2.FinishReason.REASON_STOP, - message=chat_pb2.CompletionMessage( - role=chat_pb2.ROLE_ASSISTANT, - content="I am searching.", - ), - ) + _add_server_side_tool_outputs(request, response) elif len(request.tools) > 0: response.outputs.add( finish_reason=sample_pb2.FinishReason.REASON_TOOL_CALLS, @@ -375,8 +444,32 @@ def GetCompletionChunk(self, request: chat_pb2.GetCompletionsRequest, context: g ".", ] + # The backend delivers a completed image generation call and its full + # result envelope in a single ROLE_TOOL chunk. + image_generation_chunks = [ + chat_pb2.CompletionOutputChunk( + delta=chat_pb2.Delta( + role=chat_pb2.ROLE_ASSISTANT, + tool_calls=[_image_generation_tool_call()], + ), + index=0, + ), + chat_pb2.CompletionOutputChunk( + delta=chat_pb2.Delta( + role=chat_pb2.ROLE_TOOL, + content=image_generation_envelope(), + tool_calls=[_image_generation_tool_call()], + ), + index=1, + ), + "Here is", + " your image.", + ] + chunks = normal_chunks if len(request.tools) == 0 else function_call_chunks - if len(request.tools) > 0 and _use_server_side_tools(request): + if len(request.tools) > 0 and _use_image_generation(request): + chunks = image_generation_chunks + elif len(request.tools) > 0 and _use_server_side_tools(request): # Agentic tool calling. chunks = agentic_tool_calling_chunks elif len(request.tools) == 0: diff --git a/tests/sync/chat_test.py b/tests/sync/chat_test.py index 786918c..c2e24d2 100644 --- a/tests/sync/chat_test.py +++ b/tests/sync/chat_test.py @@ -30,7 +30,7 @@ from xai_sdk.proto import chat_pb2, image_pb2, sample_pb2, usage_pb2 from xai_sdk.proto import documents_pb2 as _documents_pb2 from xai_sdk.search import SearchParameters, news_source, rss_source, web_source, x_source -from xai_sdk.tools import code_execution, collections_search, mcp, web_search, x_search +from xai_sdk.tools import code_execution, collections_search, image_generation, mcp, web_search, x_search from .. import server @@ -390,6 +390,104 @@ def test_agentic_tool_calling_non_streaming(client): assert response.tool_calls[0].function.arguments == '{"query":"What is the weather in London?"}' +def test_image_generation_tool_non_streaming(client): + chat = client.chat.create( + "grok-4-fast", + tools=[image_generation()], + ) + chat.append(user("Generate an image of a corgi")) + response = chat.sample() + + image_outputs = response.image_outputs + assert len(image_outputs) == 1 + output = image_outputs[0] + assert output.tool_call.function.name == "imagine_text_to_image" + assert output.tool_call.status == chat_pb2.TOOL_CALL_STATUS_COMPLETED + assert output.mime_type == "image/jpeg" + assert output.data_url.startswith("data:image/jpeg;base64,") + assert output.image_uuid == server.IMAGE_GENERATION_IMAGE_UUID + assert output.image == server.read_image() + + assert response.content == "Here is your image." + + +def test_image_generation_tool_streaming(client): + chat = client.chat.create( + "grok-4-fast", + tools=[image_generation()], + ) + chat.append(user("Generate an image of a corgi")) + + last_response = None + for response, _chunk in chat.stream(): + # Safe to access mid-stream: a completed call and its full envelope + # always arrive in the same chunk, so entries appear fully formed. + for output in response.image_outputs: + assert output.tool_call.type == chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL + last_response = response + + assert last_response is not None + image_outputs = last_response.image_outputs + assert len(image_outputs) == 1 + output = image_outputs[0] + assert output.tool_call.function.name == "imagine_text_to_image" + assert output.mime_type == "image/jpeg" + assert output.image_uuid == server.IMAGE_GENERATION_IMAGE_UUID + assert output.image == server.read_image() + + assert last_response.content == "Here is your image." + + +def test_image_outputs_omits_failed_image_generation_calls(): + response_proto = chat_pb2.GetChatCompletionResponse( + outputs=[ + chat_pb2.CompletionOutput( + index=0, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_TOOL, + tool_calls=[ + chat_pb2.ToolCall( + id="failed-image-call", + type=chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL, + status=chat_pb2.TOOL_CALL_STATUS_FAILED, + error_message="generation failed", + ) + ], + content='{"error": "generation failed"}', + ), + ) + ] + ) + response = Response(response_proto, index=None) + assert response.image_outputs == [] + + +def test_image_outputs_raises_on_malformed_envelope(): + response_proto = chat_pb2.GetChatCompletionResponse( + outputs=[ + chat_pb2.CompletionOutput( + index=0, + message=chat_pb2.CompletionMessage( + role=chat_pb2.ROLE_TOOL, + tool_calls=[ + chat_pb2.ToolCall( + id="completed-image-call", + type=chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL, + status=chat_pb2.TOOL_CALL_STATUS_COMPLETED, + ) + ], + content="not json", + ), + ) + ] + ) + response = Response(response_proto, index=None) + # A completed call always carries the full envelope, so a parse failure is a + # corrupt payload and raises as soon as the outputs are listed. + with pytest.raises(ValueError, match="image generation result envelope"): + _ = response.image_outputs + + def test_structured_output_parse(client: Client): class Weather(BaseModel): city: str @@ -1436,11 +1534,12 @@ def test_chat_create_with_server_side_tools(client: Client): allowed_tool_names=["chat", "completions"], authorization="lin-1234567890", ), + image_generation(action="generate"), ], ) chat_completion_request = chat.proto - assert len(chat_completion_request.tools) == 5 + assert len(chat_completion_request.tools) == 6 expected_from_date_pb = timestamp_pb2.Timestamp() expected_from_date_pb.FromDatetime(from_date) @@ -1487,11 +1586,14 @@ def test_chat_create_with_server_side_tools(client: Client): ) ) + expected_image_generation_tool = chat_pb2.Tool(image_generation=chat_pb2.ImageGeneration(action="generate")) + assert chat_completion_request.tools[0] == expected_web_search_tool assert chat_completion_request.tools[1] == expected_x_search_tool assert chat_completion_request.tools[2] == expected_code_execution_tool assert chat_completion_request.tools[3] == expected_collections_search_tool assert chat_completion_request.tools[4] == expected_mcp_tool + assert chat_completion_request.tools[5] == expected_image_generation_tool @pytest.mark.parametrize(