atlas-iac/services/hermes/scripts/image_broker.py

#!/usr/bin/env python3
"""Private GPT Image broker for isolated Hermes chat tenants."""

from __future__ import annotations

import base64
import hmac
import importlib.util
import json
import os
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


HOST = os.environ.get("HERMES_IMAGE_BROKER_LISTEN_HOST", "0.0.0.0")
# Kubernetes reserves <SERVICE>_PORT for service-link discovery and injects a
# tcp://... value.  Keep listener configuration out of that namespace.
PORT = int(os.environ.get("HERMES_IMAGE_BROKER_LISTEN_PORT", "9002"))


def _relay_token() -> str:
    """Read relay auth without exporting it to provider subprocesses."""
    path = Path(
        os.environ.get(
            "HERMES_IMAGE_BROKER_KEY_FILE", "/runtime-access/chat-relay-key"
        )
    )
    try:
        return path.read_text(encoding="utf-8").strip()
    except OSError:
        return ""


TOKEN = _relay_token()
PROVIDER_PATH = Path(
    os.environ.get(
        "HERMES_IMAGE_PROVIDER_PATH",
        "/opt/hermes/plugins/image_gen/openai-codex/__init__.py",
    )
)
DEFAULT_MODEL = os.environ.get(
    "HERMES_IMAGE_BROKER_DEFAULT_MODEL", "atlas-image-auto-high"
)
LOCAL_URL = os.environ.get(
    "HERMES_LOCAL_IMAGE_URL",
    "http://hermes-local-image.hermes.svc.cluster.local:9004",
).rstrip("/")
POLICY_PATH = Path(
    os.environ.get("HERMES_IMAGE_POLICY_PATH", "/etc/hermes-image-policy/policy.json")
)
MAX_BODY_BYTES = int(os.environ.get("HERMES_IMAGE_BROKER_MAX_BODY", str(96 << 20)))
QUEUE_TIMEOUT_SECONDS = float(
    os.environ.get("HERMES_IMAGE_BROKER_QUEUE_TIMEOUT", "900")
)
ALLOWED_MODELS = {
    "atlas-image-auto-high",
    "gpt-image-2-low",
    "gpt-image-2-medium",
    "gpt-image-2-high",
    "flux-2-klein-4b-local",
}
ALLOWED_ASPECTS = {"landscape", "square", "portrait"}
_generation_lock = threading.Lock()

# This is intentionally a tiny deterministic hard boundary, not a hosted-style
# appearance/race classifier. Operators can add site-specific phrases through
# the ConfigMap, but benign restoration and appearance correction remain valid.
_MINOR_TERMS = {
    "baby",
    "child",
    "children",
    "infant",
    "kid",
    "minor",
    "preteen",
    "toddler",
    "underage",
}
_SEXUAL_TERMS = {
    "erotic",
    "explicit",
    "genital",
    "naked",
    "nude",
    "porn",
    "sexual",
}


def _load_provider() -> Any:
    """Load the pinned Hermes provider without duplicating its API client."""
    spec = importlib.util.spec_from_file_location(
        "atlas_openai_codex_image_provider", PROVIDER_PATH
    )
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot load image provider from {PROVIDER_PATH}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.OpenAICodexImageGenProvider()


_PROVIDER: Any | None = None


def _provider() -> Any:
    """Load the provider once after the sidecar environment is ready."""
    global _PROVIDER
    if _PROVIDER is None:
        _PROVIDER = _load_provider()
    return _PROVIDER


def _authorized(header: str | None) -> bool:
    """Compare the bearer token without leaking timing information."""
    if not TOKEN or not header or not header.startswith("Bearer "):
        return False
    return hmac.compare_digest(header[7:].strip(), TOKEN)


def _policy_error(prompt: str) -> str | None:
    """Return the narrow hard-boundary reason, otherwise allow the request."""
    lowered = prompt.casefold()
    words = {"".join(char for char in word if char.isalnum()) for word in lowered.split()}
    if words & _MINOR_TERMS and words & _SEXUAL_TERMS:
        return "sexual content involving minors is not permitted"
    try:
        configured = json.loads(POLICY_PATH.read_text(encoding="utf-8"))
    except (OSError, ValueError, json.JSONDecodeError):
        configured = {}
    for phrase in configured.get("additional_blocked_phrases") or []:
        if isinstance(phrase, str) and phrase.strip().casefold() in lowered:
            return "request is blocked by the operator image policy"
    return None


def _local_request(payload: dict[str, Any], timeout: float = 1800.0) -> dict[str, Any]:
    """Call the fixed in-cluster FLUX endpoint without forwarding credentials."""
    request = Request(
        f"{LOCAL_URL}/v1/images/generations",
        data=json.dumps(payload).encode("utf-8"),
        method="POST",
        headers={"Content-Type": "application/json", "Accept": "application/json"},
    )
    try:
        with urlopen(request, timeout=timeout) as response:  # noqa: S310 - fixed internal URL
            return json.loads(response.read())
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")[:1000]
        raise RuntimeError(f"local image service returned HTTP {exc.code}: {body}") from exc
    except URLError as exc:
        raise RuntimeError(f"local image service is unavailable: {exc.reason}") from exc


def _local_health() -> dict[str, Any]:
    """Read local availability for the broker health and AUTO router."""
    request = Request(f"{LOCAL_URL}/health", method="GET")
    try:
        with urlopen(request, timeout=5) as response:  # noqa: S310 - fixed internal URL
            return json.loads(response.read())
    except Exception as exc:
        return {"success": False, "available": False, "error": str(exc)}


def _generate_hosted(
    payload: dict[str, Any], model: str, prompt: str, aspect: str
) -> dict[str, Any]:
    """Generate through the owner's hosted GPT Image provider."""
    old_model = os.environ.get("OPENAI_IMAGE_MODEL")
    try:
        os.environ["OPENAI_IMAGE_MODEL"] = model
        result = _provider().generate(
            prompt,
            aspect,
            image_url=payload.get("image_url"),
            reference_image_urls=payload.get("reference_image_urls"),
        )
    finally:
        if old_model is None:
            os.environ.pop("OPENAI_IMAGE_MODEL", None)
        else:
            os.environ["OPENAI_IMAGE_MODEL"] = old_model
    if not result.get("success"):
        return result
    image_path = Path(str(result["image"]))
    try:
        raw = image_path.read_bytes()
    finally:
        image_path.unlink(missing_ok=True)
    response = {key: value for key, value in result.items() if key != "image"}
    response.update(
        {
            "image_b64": base64.b64encode(raw).decode("ascii"),
            "mime_type": "image/png",
            "route": "openai",
        }
    )
    return response


def _generate(payload: dict[str, Any]) -> dict[str, Any]:
    """Route one request and keep generated bytes inside the tenant boundary."""
    prompt = str(payload.get("prompt") or "").strip()
    if not prompt:
        raise ValueError("prompt is required")
    policy_error = _policy_error(prompt)
    if policy_error:
        raise ValueError(policy_error)
    aspect = str(payload.get("aspect_ratio") or "landscape").lower()
    if aspect not in ALLOWED_ASPECTS:
        raise ValueError("aspect_ratio must be landscape, square, or portrait")
    model = str(payload.get("model") or DEFAULT_MODEL)
    if model not in ALLOWED_MODELS:
        raise ValueError("unsupported image quality tier")
    image_url = payload.get("image_url")
    references = payload.get("reference_image_urls")
    if image_url is not None and not isinstance(image_url, str):
        raise ValueError("image_url must be a string")
    if references is not None and not (
        isinstance(references, list)
        and all(isinstance(item, str) for item in references)
    ):
        raise ValueError("reference_image_urls must be a list of strings")

    acquired = _generation_lock.acquire(timeout=QUEUE_TIMEOUT_SECONDS)
    if not acquired:
        raise TimeoutError("image generation queue is busy; retry shortly")
    try:
        if model == "flux-2-klein-4b-local":
            return _local_request(payload)
        if model != "atlas-image-auto-high":
            return _generate_hosted(payload, model, prompt, aspect)

        hosted_error = "hosted image generation was unavailable"
        try:
            result = _generate_hosted(payload, "gpt-image-2-high", prompt, aspect)
            if result.get("success"):
                result["requested_model"] = model
                return result
            hosted_error = str(result.get("error") or hosted_error)
        except Exception as exc:
            hosted_error = f"{type(exc).__name__}: {exc}"
        try:
            result = _local_request({**payload, "model": "flux-2-klein-4b-local"})
        except Exception as exc:
            return {
                "success": False,
                "error": (
                    f"OpenAI image route failed ({hosted_error}); local FLUX fallback "
                    f"also failed ({type(exc).__name__}: {exc})"
                ),
                "error_type": "all_image_routes_failed",
            }
        result["requested_model"] = model
        result["hosted_fallback_reason"] = hosted_error
        return result
    finally:
        _generation_lock.release()


class Handler(BaseHTTPRequestHandler):
    """Small authenticated JSON API; prompts and images are never logged."""

    server_version = "HermesImageBroker/1"

    def _json(self, status: int, value: dict[str, Any]) -> None:
        body = json.dumps(value, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def _check_auth(self) -> bool:
        if _authorized(self.headers.get("Authorization")):
            return True
        self._json(401, {"success": False, "error": "unauthorized"})
        return False

    def do_GET(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if self.path != "/health":
            self._json(404, {"success": False, "error": "not found"})
            return
        if not self._check_auth():
            return
        local = _local_health()
        try:
            openai_available = bool(_provider().is_available())
        except Exception:
            openai_available = False
        self._json(
            200,
            {
                "success": True,
                "provider": "atlas-image-router",
                "available": openai_available or bool(local.get("available")),
                "openai_available": openai_available,
                "local": local,
                "busy": _generation_lock.locked(),
            },
        )

    def do_POST(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if self.path != "/v1/images/generations":
            self._json(404, {"success": False, "error": "not found"})
            return
        if not self._check_auth():
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            length = 0
        if length <= 0 or length > MAX_BODY_BYTES:
            self._json(413, {"success": False, "error": "invalid request size"})
            return
        try:
            payload = json.loads(self.rfile.read(length))
            if not isinstance(payload, dict):
                raise ValueError("JSON object required")
            result = _generate(payload)
            self._json(200 if result.get("success") else 502, result)
        except ValueError as exc:
            self._json(400, {"success": False, "error": str(exc)})
        except TimeoutError as exc:
            self._json(503, {"success": False, "error": str(exc)})
        except Exception as exc:  # keep provider details useful but omit prompts
            self._json(
                502,
                {
                    "success": False,
                    "error": f"image provider failed: {type(exc).__name__}: {exc}",
                },
            )

    def log_message(self, format: str, *args: Any) -> None:
        """Log only method/path/status metadata, never request bodies."""
        print(f"image-broker {self.address_string()} {format % args}", flush=True)


def main() -> None:
    """Serve until Kubernetes terminates the sidecar."""
    if not TOKEN:
        raise SystemExit("runtime relay key is unavailable")
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    server.serve_forever()


if __name__ == "__main__":
    main()