feat: expose request_id on ImageResponse

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [Unreleased]
 ### Added
+- **Image Request IDs**: `ImageResponse` (sync and async) now exposes a `request_id` property containing the `x-request-id` from the gRPC response metadata of `image.sample()` / `image.sample_batch()`. Include it when reporting issues so xAI can look up the request in server logs. Returns `None` if no request ID was captured.
 - **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`).

src/xai_sdk/aio/image.py

@@ -11,6 +11,7 @@
     ImageAspectRatio,
     ImageFormat,
     ImageResolution,
+    _extract_request_id,
     _make_generate_request,
     _make_span_request_attributes,
     _make_span_response_attributes,
@@ -224,8 +225,10 @@ async def sample(
             kind=SpanKind.CLIENT,
             attributes=_make_span_request_attributes(request),
         ) as span:
-            response_pb = await self._stub.GenerateImage(request)
-            image_response = ImageResponse(response_pb, 0)
+            call = self._stub.GenerateImage(request)
+            response_pb = await call
+            request_id = _extract_request_id(await call.initial_metadata())
+            image_response = ImageResponse(response_pb, 0, request_id)
             span.set_attributes(_make_span_response_attributes(request, [image_response]))
             return image_response
 
@@ -325,8 +328,10 @@ async def sample_batch(
             kind=SpanKind.CLIENT,
             attributes=_make_span_request_attributes(request),
         ) as span:
-            response_pb = await self._stub.GenerateImage(request)
-            image_responses = [ImageResponse(response_pb, i) for i in range(n)]
+            call = self._stub.GenerateImage(request)
+            response_pb = await call
+            request_id = _extract_request_id(await call.initial_metadata())
+            image_responses = [ImageResponse(response_pb, i, request_id) for i in range(n)]
             span.set_attributes(_make_span_response_attributes(request, image_responses))
             return image_responses
 

src/xai_sdk/image.py

@@ -38,20 +38,45 @@ def __init__(self, channel: Union[grpc.Channel, grpc.aio.Channel]):
         self._stub = image_pb2_grpc.ImageStub(channel)
 
 
+def _extract_request_id(metadata) -> Optional[str]:
+    """Extracts the x-request-id from gRPC initial metadata.
+
+    Args:
+        metadata: The gRPC initial metadata (result of call.initial_metadata() or await call.initial_metadata()).
+
+    Returns:
+        The x-request-id string if found in initial metadata, None otherwise.
+    """
+    if not metadata:
+        return None
+
+    # Handle sync metadata format (list of tuples)
+    if isinstance(metadata, list | tuple):
+        return dict(metadata).get("x-request-id")
+
+    # Handle async metadata format
+    if hasattr(metadata, "get"):
+        return metadata.get("x-request-id")
+
+    return None
+
+
 class BaseImageResponse(ProtoDecorator[image_pb2.ImageResponse]):
     """Adds auxiliary functions for handling the image response proto."""
 
     _image: image_pb2.GeneratedImage
 
-    def __init__(self, proto: image_pb2.ImageResponse, index: int) -> None:
+    def __init__(self, proto: image_pb2.ImageResponse, index: int, request_id: Optional[str] = None) -> None:
         """Initializes a new instance of the `ImageResponse` class.
 
         Args:
             proto: The proto to wrap.
             index: The index of the image within that proto to expose.
+            request_id: The x-request-id from the gRPC initial metadata, used for tracing and debugging.
         """
         super().__init__(proto)
         self._image = proto.images[index]
+        self._request_id = request_id
 
     @property
     def model(self) -> str:
@@ -63,6 +88,15 @@ def usage(self) -> usage_pb2.SamplingUsage:
         """Token and tool usage for this request."""
         return self._proto.usage
 
+    @property
+    def request_id(self) -> Optional[str]:
+        """Returns the x-request-id from the gRPC initial metadata.
+
+        This is a unique identifier for the request, useful for tracing and debugging.
+        Returns None if no request ID was captured during the request.
+        """
+        return self._request_id
+
     @property
     def cost_usd(self) -> Optional[float]:
         """Cost of the request in USD, or None if the server did not report it."""

src/xai_sdk/sync/image.py

@@ -11,6 +11,7 @@
     ImageAspectRatio,
     ImageFormat,
     ImageResolution,
+    _extract_request_id,
     _make_generate_request,
     _make_span_request_attributes,
     _make_span_response_attributes,
@@ -224,8 +225,9 @@ def sample(
             kind=SpanKind.CLIENT,
             attributes=_make_span_request_attributes(request),
         ) as span:
-            response_pb = self._stub.GenerateImage(request)
-            image_response = ImageResponse(response_pb, 0)
+            response_pb, call = self._stub.GenerateImage.with_call(request)
+            request_id = _extract_request_id(call.initial_metadata())
+            image_response = ImageResponse(response_pb, 0, request_id)
             span.set_attributes(_make_span_response_attributes(request, [image_response]))
             return image_response
 
@@ -325,8 +327,9 @@ def sample_batch(
             kind=SpanKind.CLIENT,
             attributes=_make_span_request_attributes(request),
         ) as span:
-            response_pb = self._stub.GenerateImage(request)
-            image_responses = [ImageResponse(response_pb, i) for i in range(n)]
+            response_pb, call = self._stub.GenerateImage.with_call(request)
+            request_id = _extract_request_id(call.initial_metadata())
+            image_responses = [ImageResponse(response_pb, i, request_id) for i in range(n)]
             span.set_attributes(_make_span_response_attributes(request, image_responses))
             return image_responses
 

tests/aio/image_test.py

@@ -55,6 +55,27 @@ async def test_batch(client: AsyncClient, image_asset: bytes):
         assert image_asset == await r.image
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_sample_returns_request_id(client: AsyncClient):
+    response = await client.image.sample(prompt="foo", model="grok-2-image", image_format="url")
+
+    assert response.request_id == server.IMAGE_REQUEST_ID
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_sample_batch_returns_request_id(client: AsyncClient):
+    responses = await client.image.sample_batch(prompt="foo", model="grok-2-image", n=2, image_format="url")
+
+    assert all(r.request_id == server.IMAGE_REQUEST_ID for r in responses)
+
+
+def test_image_response_request_id_defaults_to_none():
+    proto = image_pb2.ImageResponse(
+        images=[image_pb2.GeneratedImage(url="https://example.com/i.png")],
+    )
+    assert BaseImageResponse(proto, 0).request_id is None
+
+
 @pytest.mark.asyncio(loop_scope="session")
 async def test_sample_passes_aspect_ratio_and_resolution(client: AsyncClient):
     server.clear_last_image_request()

tests/server.py

@@ -48,6 +48,8 @@
 API_KEY = "123"
 MANAGEMENT_API_KEY = "456"
 IMAGE_PATH = "test.jpg"
+# x-request-id returned in the initial metadata of every GenerateImage response.
+IMAGE_REQUEST_ID = "test-image-request-id"
 
 _last_image_request_lock = threading.Lock()
 _last_video_request_lock = threading.Lock()
@@ -683,6 +685,7 @@ def __init__(self, url):
     def GenerateImage(self, request: image_pb2.GenerateImageRequest, context: grpc.ServicerContext):
         _check_auth(context)
         _record_last_image_request(request)
+        context.send_initial_metadata((("x-request-id", IMAGE_REQUEST_ID),))
 
         if request.format == image_pb2.ImageFormat.IMG_FORMAT_URL:
             return image_pb2.ImageResponse(

tests/sync/image_test.py

@@ -51,6 +51,25 @@ def test_batch(client: Client, image_asset: bytes):
         assert image_asset == r.image
 
 
+def test_sample_returns_request_id(client: Client):
+    response = client.image.sample(prompt="foo", model="grok-2-image", image_format="url")
+
+    assert response.request_id == server.IMAGE_REQUEST_ID
+
+
+def test_sample_batch_returns_request_id(client: Client):
+    responses = client.image.sample_batch(prompt="foo", model="grok-2-image", n=2, image_format="url")
+
+    assert all(r.request_id == server.IMAGE_REQUEST_ID for r in responses)
+
+
+def test_image_response_request_id_defaults_to_none():
+    proto = image_pb2.ImageResponse(
+        images=[image_pb2.GeneratedImage(url="https://example.com/i.png")],
+    )
+    assert BaseImageResponse(proto, 0).request_id is None
+
+
 def test_sample_passes_aspect_ratio_and_resolution(client: Client):
     server.clear_last_image_request()