354 lines
12 KiB
Python
354 lines
12 KiB
Python
#!/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()
|