Source code for tools.generate_background

"""GENERATE_BACKGROUND -- Generate + save backgrounds for VN scenes.

Wraps the existing Gemini image generation to also save output
to /home/star/large_files/assets/backgrounds/ so it can be served
by nginx and referenced in scene scripts.

@fire @skull THE CANVAS PAINTS ITSELF.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
from io import BytesIO
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from tool_context import ToolContext

logger = logging.getLogger(__name__)

BACKGROUNDS_DIR = "/home/star/large_files/assets/backgrounds"

TOOL_NAME = "generate_background"
TOOL_DESCRIPTION = (
    "Generate a VN scene background image from a text prompt and save it "
    "to the backgrounds folder for use in compose_scene scripts.\n\n"
    "The image is generated via Gemini, saved to disk, AND sent to the "
    "current channel. Returns the filename for use in scene bg steps.\n\n"
    "Example:\n"
    "  generate_background(prompt='dark classroom at night, anime style, purple lighting')\n"
    "  -> Returns { filename: 'bg_classroom_a3f2.png', path: '/backgrounds/bg_classroom_a3f2.png' }\n\n"
    "Then use in compose_scene:\n"
    "  {'type': 'bg', 'image': 'bg:bg_classroom_a3f2.png'}"
)

TOOL_PARAMETERS = {
    "type": "object",
    "properties": {
        "prompt": {
            "type": "string",
            "description": (
                "Text description of the background to generate. "
                "Be descriptive about mood, lighting, style, and perspective."
            ),
        },
        "name": {
            "type": "string",
            "description": (
                "Short name for the background file (e.g. 'classroom_dark'). "
                "Will be prefixed with 'bg_'. If omitted, auto-generated from hash."
            ),
        },
        "aspect_ratio": {
            "type": "string",
            "description": (
                "Aspect ratio. Default: 16:9 (widescreen). "
                "Supported: 1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9, 21:9."
            ),
        },
    },
    "required": ["prompt"],
}


[docs] async def run( prompt: str, name: str | None = None, aspect_ratio: str = "16:9", ctx: "ToolContext | None" = None, ) -> str: """Generate a VN scene background, persist it, and post it to the channel. Entry point for the ``generate_background`` tool. It borrows the image pipeline from ``tools.generate_image`` -- reusing ``_resolve_api_key`` and ``_call_gemini_native`` -- after wrapping the caller's ``prompt`` in a background-specific instruction (no characters, wide establishing shot) so the result is usable as a ``compose_scene`` backdrop. The returned bytes are normalised to PNG off the event loop via the nested ``_convert_to_png`` and ``asyncio.to_thread``, then the function has three side effects: it writes the PNG to ``BACKGROUNDS_DIR`` on disk (creating the directory if needed) so nginx can serve it under ``/backgrounds/``, uploads the same bytes to the current channel through ``ctx.adapter.send_file``, and appends a record to ``ctx.sent_files``. The deterministic filename embeds a SHA-256 hash of the image (plus a sanitised ``name`` when given). Any failure is caught and returned as a JSON error object. Dispatched by ``tool_loader.py``, which imports this module and resolves ``run`` via ``getattr(module, "run")`` to register it under ``TOOL_NAME`` ("generate_background"). Args: prompt: Text description of the background; mood, lighting, style, and perspective all help. name: Optional short label for the file (prefixed with ``bg_``). When omitted the filename is derived purely from the image hash. aspect_ratio: Output aspect ratio passed through to Gemini (default ``"16:9"``). ctx: The current ``ToolContext``; its ``adapter`` is required to send the file and its ``channel_id`` / ``sent_files`` are used for the upload. When ``ctx`` or ``ctx.adapter`` is missing the call returns an error immediately. Returns: A pretty-printed JSON string with ``success: True`` plus ``filename``, ``path``, ``scene_ref`` (the ``bg:`` reference for scene scripts), ``save_path``, and a usage ``message`` on success; otherwise a JSON object with ``success: False`` and an ``error`` describing the failure. """ if ctx is None or ctx.adapter is None: return json.dumps({"success": False, "error": "No context/adapter available."}) # Import the generate_image internals try: from tools.generate_image import ( _resolve_api_key, _call_gemini_native, IMAGE_GENERATION_SYSTEM_PROMPT, ) except ImportError as e: return json.dumps( {"success": False, "error": f"Cannot import generate_image: {e}"} ) api_key, _ = await _resolve_api_key(ctx) # Enhance prompt for background generation bg_prompt = ( f"Generate a high-quality anime/illustration background scene: {prompt}. " "This is a visual novel background -- no characters, no text, no UI elements. " "Wide establishing shot, atmospheric lighting, rich detail." ) prompt_parts = [{"text": bg_prompt}] try: img_bytes = await _call_gemini_native( prompt_parts, api_key, aspect_ratio, ) if not img_bytes: return json.dumps({"success": False, "error": "No image generated."}) from PIL import Image def _convert_to_png(data: bytes) -> bytes: """Re-encode arbitrary image bytes as PNG. Opens ``data`` with Pillow (``PIL.Image.open`` over an in-memory ``BytesIO``) and re-saves it to a fresh ``BytesIO`` in PNG format, normalising whatever container Gemini returned (often WebP/JPEG) into the PNG the rest of the pipeline expects for hashing, disk persistence under ``BACKGROUNDS_DIR``, and channel upload. This is a synchronous, CPU-bound helper closed over the enclosing ``run`` coroutine; it is invoked off the event loop via ``asyncio.to_thread`` so the decode/encode does not block. It has no side effects of its own (no I/O, network, Redis, or KG access); it purely transforms bytes in memory. No internal callers were found beyond its single ``asyncio.to_thread`` use inside ``run``. Args: data: Raw image bytes returned by the Gemini image model, in any format Pillow can decode. Returns: The same image re-encoded as PNG bytes. Raises: PIL.UnidentifiedImageError: If ``data`` is not a decodable image. Other ``PIL``/``OSError`` exceptions may propagate from decoding or encoding and are caught by ``run``\\ 's outer ``try``/``except``. """ img = Image.open(BytesIO(data)) buf = BytesIO() img.save(buf, format="PNG") return buf.getvalue() png_bytes = await asyncio.to_thread(_convert_to_png, img_bytes) # Generate filename h = hashlib.sha256(png_bytes).hexdigest()[:8] if name: # Sanitize name safe_name = "".join(c if c.isalnum() or c in "_-" else "_" for c in name) fname = f"bg_{safe_name}_{h}.png" else: fname = f"bg_{h}.png" # Save to backgrounds directory os.makedirs(BACKGROUNDS_DIR, exist_ok=True) save_path = os.path.join(BACKGROUNDS_DIR, fname) await asyncio.to_thread(lambda: open(save_path, "wb").write(png_bytes)) logger.info("Background saved: %s (%d bytes)", save_path, len(png_bytes)) # Also send to channel so user can see it file_url = await ctx.adapter.send_file( ctx.channel_id, png_bytes, fname, "image/png", ) ctx.sent_files.append( { "data": png_bytes, "filename": fname, "mimetype": "image/png", "file_url": file_url or "", } ) return json.dumps( { "success": True, "filename": fname, "path": f"/backgrounds/{fname}", "scene_ref": f"bg:{fname}", "save_path": save_path, "message": ( f"Background '{fname}' generated and saved. " f"Use in scene scripts as: {{'type': 'bg', 'image': 'bg:{fname}'}}" ), }, indent=2, ) except Exception as e: logger.error("generate_background error: %s", e, exc_info=True) return json.dumps({"success": False, "error": f"Generation failed: {e}"})