Text Transform

1"""2Text Transform34Two modes, picked by what is connected:5- text mode (no image): apply the prompt to `input_text` and return the6  elaborated text. The messages sent in this mode are a regression contract7  (R-3): they must stay byte-identical to 1.3.0.8- image mode (`image` connected): answer the prompt about the image with a9  vision model ("answer a question about an image"). `input_text`, if also10  given, is appended as context. The image is bounded before it is sent:11  long edge <= 2048 px, encoded size <= 4 MB (eng D3 = A).1213`output_format: json` (both modes) strips code fences, parses the reply and14returns it normalized in `output_text`; an unparseable reply fails the node15with the first 2000 characters of the raw reply.16"""1718from __future__ import annotations1920import asyncio21import base6422import io23import json24import mimetypes25import re26import sys27import traceback28from typing import Any2930from gais import Gais3132INPUT_DIR = "/data/input"3334SYSTEM_PROMPT = (35    "You are a careful text-transformation assistant. "36    "Apply the user's instructions to the provided text and return ONLY the "37    "transformed text — no preamble, no commentary, no surrounding quotes."38)3940IMAGE_SYSTEM_PROMPT = (41    "Answer the instruction about the attached image. "42    "If asked for JSON, return only JSON."43)4445# Image size bound for image mode (eng D3 = A).46MAX_IMAGE_EDGE = 204847MAX_IMAGE_BYTES = 4 * 1024 * 102448JPEG_QUALITY = 9249FALLBACK_JPEG_QUALITY = 8550_JPEG_FORMATS = ("JPEG", "MPO")51# Pillow refuses to decode images above ~179 megapixels (decompression-bomb52# guard). Say so instead of blaming the format.53_TOO_MANY_PIXELS_ERROR = "Image is too large to read (over ~179 megapixels); downscale it and try again"5455JSON_ERROR_EXCERPT_CHARS = 20005657OUTPUT_FORMATS = ("text", "json")5859# Explicit map first: the stdlib mimetypes table varies by platform and misses60# some image types (e.g. .webp on some builds).61_IMAGE_MIME_BY_EXT = {62    "jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png",63    "gif": "image/gif", "webp": "image/webp", "bmp": "image/bmp",64    "tif": "image/tiff", "tiff": "image/tiff", "heic": "image/heic",65    "heif": "image/heif", "avif": "image/avif", "svg": "image/svg+xml",66}6768# Output-language directive (language-persistence workstream). Values are69# platform locale codes; legacy full names are accepted via the allowlist.70# Anything else (wired upstream output, free-form strings) falls back to71# auto: the value is interpolated into the system prompt, so unknown72# strings must never pass through.73LANGUAGE_NAMES = {74    "de": "German", "en": "English", "es": "Spanish", "fr": "French",75    "it": "Italian", "ja": "Japanese", "ko": "Korean", "nl": "Dutch",76    "pt": "Portuguese", "ru": "Russian", "zh": "Chinese",77}787980def resolve_language_name(language):81    """Full display name for an allowlisted language value, else None."""82    if not isinstance(language, str):83        return None84    value = language.strip()85    if not value or value == "auto":86        return None87    if value in LANGUAGE_NAMES:88        return LANGUAGE_NAMES[value]89    if value in LANGUAGE_NAMES.values():90        return value91    print(f"[language] unrecognized value {value!r} ignored - using auto", file=sys.stderr)92    return None939495# ── messages ────────────────────────────────────────────────────────────────969798def build_text_messages(input_text: str, prompt: str, language: Any = "auto") -> list[dict[str, Any]]:99    """Text-mode messages. Regression contract R-3: keep byte-identical."""100    system_prompt = SYSTEM_PROMPT101    lang_name = resolve_language_name(language)102    if lang_name:103        system_prompt += (104            f"\n\nOUTPUT LANGUAGE: {lang_name}. Write the transformed text in "105            f"{lang_name} regardless of the source text's language. This "106            "OVERRIDES any language stated in the instructions or implied by "107            "the source text. Keep proper nouns and verbatim quotes in their "108            "original language."109        )110    user_message = (111        f"<instructions>{prompt}</instructions>\n\n"112        f"<text>{input_text}</text>"113    )114    return [115        {"role": "system", "content": system_prompt},116        {"role": "user", "content": user_message},117    ]118119120def build_image_messages(121    prompt: str, input_text: str, data_url: str, language: Any = "auto",122) -> list[dict[str, Any]]:123    """Image-mode messages: the prompt (+ optional text context) and the image."""124    system_prompt = IMAGE_SYSTEM_PROMPT125    lang_name = resolve_language_name(language)126    if lang_name:127        system_prompt += (128            f"\n\nOUTPUT LANGUAGE: {lang_name}. Write the answer in {lang_name}. "129            "Keep JSON keys, proper nouns and text quoted from the image in "130            "their original form."131        )132    text = prompt + (f"\n\n<text>{input_text}</text>" if input_text else "")133    return [134        {"role": "system", "content": system_prompt},135        {136            "role": "user",137            "content": [138                {"type": "text", "text": text},139                {"type": "image_url", "image_url": {"url": data_url}},140            ],141        },142    ]143144145# ── image input ─────────────────────────────────────────────────────────────146147148def _image_file_name(raw: Any) -> str | None:149    """Upstream file ports deliver a bare filename under /data/input."""150    if isinstance(raw, list):151        items = [item for item in raw if item]152        if len(items) > 1:153            raise ValueError(154                f"Text Transform accepts a single image, got {len(items)}. "155                "Use one Text Transform per image."156            )157        raw = items[0] if items else None158    if not raw or not isinstance(raw, str):159        return None160    base = raw.replace("\\", "/").rsplit("/", 1)[-1].strip()161    if not base or base in (".", ".."):162        return None163    return base164165166def _guess_mime(name: str) -> str:167    ext = name.rsplit(".", 1)[-1].lower() if "." in name else ""168    return (169        _IMAGE_MIME_BY_EXT.get(ext)170        or mimetypes.guess_type(name)[0]171        or "application/octet-stream"172    )173174175def _bounded_size(width: int, height: int) -> tuple[int, int]:176    """Scale so the long edge is <= MAX_IMAGE_EDGE, keeping the aspect ratio."""177    long_edge = max(width, height)178    if long_edge <= MAX_IMAGE_EDGE:179        return width, height180    if width >= height:181        return MAX_IMAGE_EDGE, max(1, round(height * MAX_IMAGE_EDGE / width))182    return max(1, round(width * MAX_IMAGE_EDGE / height)), MAX_IMAGE_EDGE183184185def _to_rgb(img: Any) -> Any:186    """JPEG has no alpha: composite transparent images onto white."""187    if img.mode in ("RGB", "L"):188        return img189    if img.mode in ("RGBA", "LA"):190        from PIL import Image191192        rgba = img.convert("RGBA")193        background = Image.new("RGB", rgba.size, (255, 255, 255))194        background.paste(rgba, mask=rgba.getchannel("A"))195        return background196    return img.convert("RGB")197198199def _encode(img: Any, fmt: str, quality: int = JPEG_QUALITY) -> bytes:200    buf = io.BytesIO()201    if fmt == "JPEG":202        _to_rgb(img).save(buf, format="JPEG", quality=quality)203    else:204        img.save(buf, format="PNG")205    return buf.getvalue()206207208def load_image_as_data_url(path: str, mime: str) -> tuple[str, dict[str, Any]]:209    """Open, orient, bound and encode the image; return (data URL, stats).210211    PNG stays PNG and JPEG stays JPEG; other formats are sent as PNG. If the212    encoded result is still above MAX_IMAGE_BYTES it is re-encoded as JPEG213    q85. Re-encoding also applies the EXIF orientation and drops all EXIF214    (including GPS) before the image leaves the executor.215    """216    from PIL import Image, ImageOps217218    try:219        img = Image.open(path)220    except FileNotFoundError:221        raise ValueError(f"Input image not found: {path.rsplit('/', 1)[-1]}") from None222    except Image.DecompressionBombError:223        raise ValueError(_TOO_MANY_PIXELS_ERROR) from None224    except Exception as e:225        # gstack-shortcut(dec-bf78833f): HEIC needs manual conversion, upgrade when the first real HEIC upload arrives226        raise ValueError(f"Can't read image ({mime}); convert to JPG or PNG") from e227228    try:229        source_format = img.format230        source_size = img.size231        # Pillow opens many phone JPEGs (an MPF marker) as MPO, a JPEG232        # subclass: same reduced decode, same JPEG output.233        is_jpeg = source_format in _JPEG_FORMATS234        if is_jpeg:235            # Decode at a reduced DCT scale (1/2, 1/4, 1/8) that is still at236            # least the target size: a 50 MP photo never fully decodes.237            img.draft(img.mode, _bounded_size(*img.size))238        img.load()239        # Apply EXIF orientation: the re-encoded copy carries no EXIF, so a240        # rotated phone photo would otherwise reach the model sideways.241        ImageOps.exif_transpose(img, in_place=True)242        if img.mode.startswith("I;16"):243            img = img.convert("I")244        if img.mode == "I":245            # 16-bit samples: scale to 8 bits first; a plain convert clips246            # everything above 255, and the model would see a white image.247            img = img.point(lambda v: v * (1 / 256)).convert("L")248        if img.mode not in ("L", "LA", "RGB", "RGBA"):249            has_alpha = "transparency" in img.info or any(b in ("A", "a") for b in img.getbands())250            img = img.convert("RGBA" if has_alpha else "RGB")251    except Image.DecompressionBombError:252        raise ValueError(_TOO_MANY_PIXELS_ERROR) from None253    except Exception as e:254        # gstack-shortcut(dec-bf78833f): HEIC needs manual conversion, upgrade when the first real HEIC upload arrives255        raise ValueError(f"Can't read image ({mime}); convert to JPG or PNG") from e256257    target = _bounded_size(*img.size)258    if target != img.size:259        img = img.resize(target, Image.Resampling.LANCZOS)260261    out_format = "JPEG" if is_jpeg else "PNG"262    data = _encode(img, out_format)263    if len(data) > MAX_IMAGE_BYTES:264        out_format = "JPEG"265        data = _encode(img, "JPEG", quality=FALLBACK_JPEG_QUALITY)266267    out_mime = "image/jpeg" if out_format == "JPEG" else "image/png"268    stats = {269        "source_format": source_format,270        "source_size": source_size,271        "sent_size": img.size,272        "sent_mime": out_mime,273        "sent_bytes": len(data),274    }275    data_url = f"data:{out_mime};base64,{base64.b64encode(data).decode('ascii')}"276    return data_url, stats277278279# ── output ──────────────────────────────────────────────────────────────────280281_FENCE = "```"282_FENCE_LANG_RE = re.compile(r"[\w+-]*")283284285def _strip_code_fences(raw: str) -> str:286    """The body of a reply wrapped in one ``` fence, else of its first fenced287    block, else the reply. Line-based on purpose: a regex over the whole reply288    backtracks quadratically on a long run of spaces."""289    text = raw.strip()290    if len(text) >= 2 * len(_FENCE) and text.startswith(_FENCE) and text.endswith(_FENCE):291        body = text[len(_FENCE):-len(_FENCE)]292        head, newline, rest = body.partition("\n")293        if newline and _FENCE_LANG_RE.fullmatch(head.strip()):294            return rest.strip()295        return body.strip()296    lines = text.split("\n")297    for start, line in enumerate(lines):298        at = line.find(_FENCE)299        if at == -1 or not _FENCE_LANG_RE.fullmatch(line[at + len(_FENCE):].strip()):300            continue301        for end in range(start + 1, len(lines)):302            if lines[end].lstrip().startswith(_FENCE):303                return "\n".join(lines[start + 1:end]).strip()304        break305    return text306307308def _reject_non_json_constant(name: str) -> Any:309    raise ValueError(f"{name} is not valid JSON")310311312def normalize_json_reply(raw: str) -> str:313    """Parse the model reply as JSON and return it re-serialized."""314    try:315        obj = json.loads(_strip_code_fences(raw), parse_constant=_reject_non_json_constant)316        # allow_nan=False also refuses 1e999, which loads as float("inf").317        return json.dumps(obj, ensure_ascii=False, allow_nan=False)318    except ValueError:319        raise ValueError(320            "Model did not return valid JSON: " + raw[:JSON_ERROR_EXCERPT_CHARS]321        ) from None322323324def _output_format(raw: Any) -> str:325    if raw is None or raw == "":326        return "text"327    value = str(raw).strip().lower()328    if value not in OUTPUT_FORMATS:329        raise ValueError(f"output_format must be 'text' or 'json', got {raw!r}")330    return value331332333# ── run ─────────────────────────────────────────────────────────────────────334335336async def _chat(messages: list[dict[str, Any]], model_id: str) -> str:337    response = await Gais.llm.chat_async(messages, model_id=model_id, temperature=0, thinking=False)338    return response.text.strip()339340341def process(inputs: dict[str, Any]) -> dict[str, Any]:342    input_text = inputs.get("input_text")343    prompt = inputs.get("prompt")344    llm_model_id = inputs.get("llmModelId")345    language = inputs.get("language", "auto")346    output_format = _output_format(inputs.get("output_format"))347    image_name = _image_file_name(inputs.get("image"))348349    if not input_text and not image_name:350        raise ValueError(351            "Required input missing: provide 'input_text' or connect an 'image'"352        )353    if not prompt:354        raise ValueError("Required input 'prompt' not provided")355    if not llm_model_id:356        raise ValueError("Required input 'llmModelId' not provided")357358    if image_name:359        mime = _guess_mime(image_name)360        data_url, stats = load_image_as_data_url(f"{INPUT_DIR}/{image_name}", mime)361        print(362            f"[text-transform] mode=image model_id={llm_model_id} "363            f"source={stats['source_format']} {stats['source_size'][0]}x{stats['source_size'][1]} "364            f"sent={stats['sent_mime']} {stats['sent_size'][0]}x{stats['sent_size'][1]} "365            f"bytes={stats['sent_bytes']} context_chars={len(input_text or '')}",366            file=sys.stderr,367        )368        if stats["sent_bytes"] > MAX_IMAGE_BYTES:369            print(370                f"[text-transform] image still {stats['sent_bytes']} bytes after JPEG "371                f"q{FALLBACK_JPEG_QUALITY}; sending anyway",372                file=sys.stderr,373            )374        messages = build_image_messages(prompt, input_text or "", data_url, language)375    else:376        print(377            f"[text-transform] mode=text model_id={llm_model_id} input_chars={len(input_text)}",378            file=sys.stderr,379        )380        messages = build_text_messages(input_text, prompt, language)381382    output_text = asyncio.run(_chat(messages, llm_model_id))383    if output_format == "json":384        output_text = normalize_json_reply(output_text)385386    print(387        f"[text-transform] done output_format={output_format} output_chars={len(output_text)}",388        file=sys.stderr,389    )390    return {"output_text": output_text}391392393def main() -> None:394    try:395        envelope = json.loads(sys.stdin.read() or "{}")396        inputs: dict[str, Any] = envelope.get("inputs", {}) if isinstance(envelope, dict) else {}397        json.dump(process(inputs or {}), sys.stdout)398    except Exception as e:399        error = {400            "error": str(e),401            "errorType": type(e).__name__,402            "traceback": traceback.format_exc(),403        }404        print(json.dumps(error), file=sys.stderr)405        sys.exit(1)406407408if __name__ == "__main__":409    main()410