Add quality parameter to image generation and regen protos from xai-proto

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 ### Added
+- Added a `quality` parameter (`"low"`, `"medium"`) to image generation (`client.image.sample`, `sample_batch`, and batch `prepare`), mapping to the `GenerateImageRequest.quality` field. When omitted, the default is `"medium"`. Only supported for `grok-imagine-image-2.0`.
+- 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
 - **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.

src/xai_sdk/aio/image.py

@@ -10,6 +10,7 @@
     BaseImageResponse,
     ImageAspectRatio,
     ImageFormat,
+    ImageQuality,
     ImageResolution,
     _make_generate_request,
     _make_span_request_attributes,
@@ -39,6 +40,7 @@ def prepare(
         image_format: Optional[ImageFormat] = None,
         aspect_ratio: Optional[ImageAspectRatio] = None,
         resolution: Optional[ImageResolution] = None,
+        quality: Optional[ImageQuality] = None,
         storage_options: Optional[Union[StorageOptions, image_pb2.StorageOptions]] = None,
     ) -> batch_pb2.BatchRequest:
         """Prepares an image generation request for batch processing.
@@ -73,6 +75,10 @@ def prepare(
             image_format: The format of the image to return ("url" or "base64"). Defaults to "url".
             aspect_ratio: The aspect ratio of the image to generate.
             resolution: The image resolution to generate ("1k" or "2k").
+            quality: Control generation quality with the optional quality
+                parameter. Allowed values are ``"low"`` and ``"medium"``. When
+                omitted, the default is ``"medium"``. The parameter is only
+                supported for ``grok-imagine-image-2.0``.
             storage_options: Persist the result to the Files API. Accepts a dict
                 with a required ``filename`` and optional ``expires_after`` and ``public_url`` keys.
                 Set ``public_url`` to also create a publicly shareable URL.
@@ -124,6 +130,7 @@ def prepare(
             image_format=image_format,
             aspect_ratio=aspect_ratio,
             resolution=resolution,
+            quality=quality,
             storage_options=storage_options,
         )
         return batch_pb2.BatchRequest(
@@ -144,6 +151,7 @@ async def sample(
         image_format: Optional[ImageFormat] = None,
         aspect_ratio: Optional[ImageAspectRatio] = None,
         resolution: Optional[ImageResolution] = None,
+        quality: Optional[ImageQuality] = None,
         storage_options: Optional[Union[StorageOptions, image_pb2.StorageOptions]] = None,
     ) -> "ImageResponse":
         """Samples a single image asynchronously based on the provided prompt.
@@ -193,6 +201,10 @@ async def sample(
             - `"1k"`: ~1 megapixel total. Dimensions vary by aspect ratio.
             - `"2k"`: ~4 megapixels total. Dimensions vary by aspect ratio.
             Only supported for grok-imagine models.
+            quality: Control generation quality with the optional quality
+                parameter. Allowed values are ``"low"`` and ``"medium"``. When
+                omitted, the default is ``"medium"``. The parameter is only
+                supported for ``grok-imagine-image-2.0``.
             storage_options: Persist the result to the Files API. Accepts a dict
                 with a required ``filename`` and optional ``expires_after`` and ``public_url`` keys.
                 Set ``public_url`` to also create a publicly shareable URL.
@@ -217,6 +229,7 @@ async def sample(
             image_format=image_format,
             aspect_ratio=aspect_ratio,
             resolution=resolution,
+            quality=quality,
             storage_options=storage_options,
         )
         with tracer.start_as_current_span(
@@ -243,6 +256,7 @@ async def sample_batch(
         image_format: Optional[ImageFormat] = None,
         aspect_ratio: Optional[ImageAspectRatio] = None,
         resolution: Optional[ImageResolution] = None,
+        quality: Optional[ImageQuality] = None,
         storage_options: Optional[Union[StorageOptions, image_pb2.StorageOptions]] = None,
     ) -> Sequence["ImageResponse"]:
         """Samples a batch of images asynchronously based on the provided prompt.
@@ -293,6 +307,10 @@ async def sample_batch(
             - `"1k"`: ~1 megapixel total. Dimensions vary by aspect ratio.
             - `"2k"`: ~4 megapixels total. Dimensions vary by aspect ratio.
             Only supported for grok-imagine models.
+            quality: Control generation quality with the optional quality
+                parameter. Allowed values are ``"low"`` and ``"medium"``. When
+                omitted, the default is ``"medium"``. The parameter is only
+                supported for ``grok-imagine-image-2.0``.
             storage_options: Persist the results to the Files API. Accepts a dict
                 with a required ``filename`` and optional ``expires_after`` and ``public_url`` keys.
                 Set ``public_url`` to also create a publicly shareable URL.
@@ -318,6 +336,7 @@ async def sample_batch(
             image_format=image_format,
             aspect_ratio=aspect_ratio,
             resolution=resolution,
+            quality=quality,
             storage_options=storage_options,
         )
         with tracer.start_as_current_span(

src/xai_sdk/image.py

@@ -9,7 +9,7 @@
 from .meta import ProtoDecorator
 from .proto import image_pb2, image_pb2_grpc, usage_pb2
 from .telemetry import should_disable_sensitive_attributes
-from .types import ImageAspectRatio, ImageFormat, ImageGenerationModel, ImageResolution
+from .types import ImageAspectRatio, ImageFormat, ImageGenerationModel, ImageQuality, ImageResolution
 
 _IMAGE_ASPECT_RATIO_MAP: dict[ImageAspectRatio, image_pb2.ImageAspectRatio] = {
     "1:1": image_pb2.ImageAspectRatio.IMG_ASPECT_RATIO_1_1,
@@ -183,6 +183,7 @@ def _make_generate_request(
     image_format: ImageFormat | None = None,
     aspect_ratio: ImageAspectRatio | None = None,
     resolution: ImageResolution | None = None,
+    quality: ImageQuality | None = None,
     storage_options: Union[StorageOptions, image_pb2.StorageOptions, None] = None,
 ) -> image_pb2.GenerateImageRequest:
     _validate_image_inputs(image_url, image_file_id, image_urls, image_file_ids)
@@ -237,11 +238,24 @@ def _make_generate_request(
         request.aspect_ratio = convert_image_aspect_ratio_to_pb(aspect_ratio)
     if resolution is not None:
         request.resolution = convert_image_resolution_to_pb(resolution)
+    if quality is not None:
+        request.quality = convert_image_quality_to_pb(quality)
     if storage_options is not None:
         request.storage_options.CopyFrom(_resolve_storage_options_pb(storage_options))
     return request
 
 
+def _add_storage_span_attributes(attributes: dict[str, str | int], storage_options: image_pb2.StorageOptions) -> None:
+    """Adds storage-related span attributes for a request that persists its output."""
+    attributes["gen_ai.request.storage"] = True
+    if storage_options.filename:
+        attributes["gen_ai.request.storage.filename"] = storage_options.filename
+    if storage_options.expires_after:
+        attributes["gen_ai.request.storage.expires_after"] = storage_options.expires_after
+    if storage_options.HasField("public_url"):
+        attributes["gen_ai.request.storage.public_url"] = True
+
+
 def _make_span_request_attributes(request: image_pb2.GenerateImageRequest) -> dict[str, str | int]:
     """Creates the image sampling span request attributes."""
     attributes: dict[str, str | int] = {
@@ -260,13 +274,7 @@ def _make_span_request_attributes(request: image_pb2.GenerateImageRequest) -> di
     attributes["gen_ai.prompt"] = request.prompt
 
     if request.HasField("storage_options"):
-        attributes["gen_ai.request.storage"] = True
-        if request.storage_options.filename:
-            attributes["gen_ai.request.storage.filename"] = request.storage_options.filename
-        if request.storage_options.expires_after:
-            attributes["gen_ai.request.storage.expires_after"] = request.storage_options.expires_after
-        if request.storage_options.HasField("public_url"):
-            attributes["gen_ai.request.storage.public_url"] = True
+        _add_storage_span_attributes(attributes, request.storage_options)
 
     if request.HasField("n"):
         attributes["gen_ai.request.image.count"] = request.n
@@ -276,6 +284,10 @@ def _make_span_request_attributes(request: image_pb2.GenerateImageRequest) -> di
         attributes["gen_ai.request.image.resolution"] = (
             image_pb2.ImageResolution.Name(request.resolution).removeprefix("IMG_RESOLUTION_").lower()
         )
+    if request.HasField("quality"):
+        attributes["gen_ai.request.image.quality"] = (
+            image_pb2.ImageQuality.Name(request.quality).removeprefix("IMG_QUALITY_").lower()
+        )
     if request.user:
         attributes["user_id"] = request.user
 
@@ -375,3 +387,14 @@ def convert_image_resolution_to_pb(resolution: ImageResolution) -> image_pb2.Ima
             return image_pb2.ImageResolution.IMG_RESOLUTION_2K
         case _:
             raise ValueError(f"Invalid image resolution {resolution}.")
+
+
+def convert_image_quality_to_pb(quality: ImageQuality) -> image_pb2.ImageQuality:
+    """Converts a string literal representation of an image quality to its protobuf enum variant."""
+    match quality:
+        case "low":
+            return image_pb2.ImageQuality.IMG_QUALITY_LOW
+        case "medium":
+            return image_pb2.ImageQuality.IMG_QUALITY_MEDIUM
+        case _:
+            raise ValueError(f"Invalid image quality {quality}.")

src/xai_sdk/proto/v5/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__ = ()
@@ -130,6 +131,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
@@ -438,22 +440,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")
@@ -524,6 +528,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/chat_pb2_grpc.py

@@ -104,10 +104,10 @@ def DeleteStoredCompletion(self, request, context):
         raise NotImplementedError('Method not implemented!')
 
     def CompactContext(self, request, context):
-        """Compacts a full responses input context and returns a compacted context.
+        """Compacts a full input context and returns a compacted context.
         The client sends the current input items and receives back a compacted
-        set of items (with an opaque compaction summary) suitable for use as
-        the input to the next /v1/responses call.
+        set of items (with an opaque compaction blob) suitable for use as
+        the input to the next request.
         """
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details('Method not implemented!')