214 lines
7.4 KiB
Python
214 lines
7.4 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
|
|
|
|
|
|
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"))
|
|
TOKEN = os.environ.get("HERMES_IMAGE_BROKER_KEY", "").strip()
|
|
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", "gpt-image-2-high"
|
|
)
|
|
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 = {
|
|
"gpt-image-2-low",
|
|
"gpt-image-2-medium",
|
|
"gpt-image-2-high",
|
|
}
|
|
ALLOWED_ASPECTS = {"landscape", "square", "portrait"}
|
|
_generation_lock = threading.Lock()
|
|
|
|
|
|
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 _generate(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Generate once, return bytes, and remove the broker-side cache copy."""
|
|
prompt = str(payload.get("prompt") or "").strip()
|
|
if not prompt:
|
|
raise ValueError("prompt is required")
|
|
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")
|
|
old_model = os.environ.get("OPENAI_IMAGE_MODEL")
|
|
try:
|
|
os.environ["OPENAI_IMAGE_MODEL"] = model
|
|
result = _provider().generate(
|
|
prompt,
|
|
aspect,
|
|
image_url=image_url,
|
|
reference_image_urls=references,
|
|
)
|
|
finally:
|
|
if old_model is None:
|
|
os.environ.pop("OPENAI_IMAGE_MODEL", None)
|
|
else:
|
|
os.environ["OPENAI_IMAGE_MODEL"] = old_model
|
|
_generation_lock.release()
|
|
|
|
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",
|
|
}
|
|
)
|
|
return response
|
|
|
|
|
|
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
|
|
self._json(
|
|
200,
|
|
{
|
|
"success": True,
|
|
"provider": "openai-codex",
|
|
"available": bool(_provider().is_available()),
|
|
"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("HERMES_IMAGE_BROKER_KEY is required")
|
|
server = ThreadingHTTPServer((HOST, PORT), Handler)
|
|
server.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|