app/dataplane/reverse/protocol/xai_image_edit.py

"""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",
]