From 781426707ebee85f26f14b8abd3f65debd911457 Mon Sep 17 00:00:00 2001 From: Andrew Zheng Date: Sat, 25 Jul 2026 09:55:57 -0700 Subject: [PATCH] feat: expose request_id on ImageResponse Capture the x-request-id from the gRPC response initial metadata in image.sample() / image.sample_batch() (sync and aio) and expose it as ImageResponse.request_id, so callers can reference the server-side request in bug reports. Mirrors the request-id extraction the chat client uses internally. Co-authored-by: Cursor --- CHANGELOG.md | 1 + src/xai_sdk/aio/image.py | 13 +++++++++---- src/xai_sdk/image.py | 36 +++++++++++++++++++++++++++++++++++- src/xai_sdk/sync/image.py | 11 +++++++---- tests/aio/image_test.py | 21 +++++++++++++++++++++ tests/server.py | 3 +++ tests/sync/image_test.py | 19 +++++++++++++++++++ 7 files changed, 95 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 64d0412..abdb028 100644 --- a/CHANGELOG.md +++ b/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": }` 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/src/xai_sdk/aio/image.py b/src/xai_sdk/aio/image.py index 6beb03a..621fcb7 100644 --- a/src/xai_sdk/aio/image.py +++ b/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 diff --git a/src/xai_sdk/image.py b/src/xai_sdk/image.py index 54ab127..e58d633 100644 --- a/src/xai_sdk/image.py +++ b/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.""" diff --git a/src/xai_sdk/sync/image.py b/src/xai_sdk/sync/image.py index f25d333..30e946f 100644 --- a/src/xai_sdk/sync/image.py +++ b/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 diff --git a/tests/aio/image_test.py b/tests/aio/image_test.py index 33bbed0..0b25444 100644 --- a/tests/aio/image_test.py +++ b/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() diff --git a/tests/server.py b/tests/server.py index 95250db..400fdf5 100644 --- a/tests/server.py +++ b/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( diff --git a/tests/sync/image_test.py b/tests/sync/image_test.py index b26d50d..c11df96 100644 --- a/tests/sync/image_test.py +++ b/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()