"""XAI image-edit protocol — payload builder and SSE field extractors.""" from typing import Any from app.platform.config.snapshot import get_config IMAGE_EDIT_MODEL_NAME = "imagine-image-edit" IMAGE_EDIT_MODEL_KIND = "imagine" IMAGE_POST_MEDIA_TYPE = "MEDIA_POST_TYPE_IMAGE" IMAGE_EDIT_GENERATION_COUNT = 2 def build_image_edit_payload( *, prompt: str, image_references: list[str], parent_post_id: str, ) -> dict[str, Any]: """Build the JSON payload for imagine image-edit chat requests.""" cfg = get_config() return { "temporary": cfg.get_bool("features.temporary", True), "modelName": IMAGE_EDIT_MODEL_NAME, "message": prompt, "enableImageGeneration": True, "returnImageBytes": False, "returnRawGrokInXaiRequest": False, "enableImageStreaming": True, "imageGenerationCount": IMAGE_EDIT_GENERATION_COUNT, "forceConcise": False, "enableSideBySide": True, "sendFinalMetadata": True, "isReasoning": False, "disableTextFollowUps": True, "responseMetadata": { "modelConfigOverride": { "modelMap": { "imageEditModel": IMAGE_EDIT_MODEL_KIND, "imageEditModelConfig": { "imageReferences": image_references, "parentPostId": parent_post_id, }, } } }, "disableMemory": not cfg.get_bool("features.memory", False), "forceSideBySide": False, } def extract_streaming_response(data: dict[str, Any]) -> dict[str, Any] | None: """Return ``response.streamingImageGenerationResponse`` when present.""" result = data.get("result") if not isinstance(result, dict): return None response = result.get("response") if not isinstance(response, dict): return None stream = response.get("streamingImageGenerationResponse") return stream if isinstance(stream, dict) else None def extract_model_response_urls(data: dict[str, Any]) -> list[str]: """Return fallback ``modelResponse.generatedImageUrls`` values.""" result = data.get("result") if not isinstance(result, dict): return [] response = result.get("response") if not isinstance(response, dict): return [] model_response = response.get("modelResponse") if not isinstance(model_response, dict): return [] urls = model_response.get("generatedImageUrls") if not isinstance(urls, list): return [] return [url for url in urls if isinstance(url, str) and url] def extract_model_response_file_attachments(data: dict[str, Any]) -> list[str]: """Return fallback ``modelResponse.fileAttachments`` asset IDs.""" result = data.get("result") if not isinstance(result, dict): return [] response = result.get("response") if not isinstance(response, dict): return [] model_response = response.get("modelResponse") if not isinstance(model_response, dict): return [] attachments = model_response.get("fileAttachments") if not isinstance(attachments, list): return [] return [attachment for attachment in attachments if isinstance(attachment, str) and attachment] __all__ = [ "IMAGE_EDIT_MODEL_NAME", "IMAGE_EDIT_MODEL_KIND", "IMAGE_POST_MEDIA_TYPE", "IMAGE_EDIT_GENERATION_COUNT", "build_image_edit_payload", "extract_streaming_response", "extract_model_response_urls", "extract_model_response_file_attachments", ]