365 lines
13 KiB
Python
365 lines
13 KiB
Python
"""Hermes image backend that keeps owner OAuth outside family chat pods."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from urllib.error import HTTPError, URLError
|
||
from urllib.request import Request, urlopen
|
||
|
||
from agent.image_gen_provider import (
|
||
DEFAULT_ASPECT_RATIO,
|
||
ImageGenProvider,
|
||
error_response,
|
||
normalize_reference_images,
|
||
resolve_aspect_ratio,
|
||
save_b64_image,
|
||
success_response,
|
||
)
|
||
|
||
|
||
BROKER_URL = os.environ.get(
|
||
"HERMES_IMAGE_BROKER_URL",
|
||
"http://hermes-image-broker.hermes.svc.cluster.local:9002",
|
||
).rstrip("/")
|
||
DEFAULT_MODEL = "atlas-image-auto-high"
|
||
ROUTE_MODELS = {
|
||
"auto": "atlas-image-auto-high",
|
||
"hosted": "gpt-image-2-high",
|
||
"local": "flux-2-klein-4b-local",
|
||
}
|
||
MODELS = {
|
||
"atlas-image-auto-high": (
|
||
"Auto (GPT Image quality, local fallback)",
|
||
"auto-high",
|
||
"Hosted quality first; FLUX fallback when hosted generation fails",
|
||
),
|
||
"gpt-image-2-low": ("GPT Image 2 (Fast)", "low"),
|
||
"gpt-image-2-medium": ("GPT Image 2 (Balanced)", "medium"),
|
||
"gpt-image-2-high": ("GPT Image 2 (Highest quality)", "high"),
|
||
"flux-2-klein-4b-local": (
|
||
"Local FLUX 2 Klein 4B",
|
||
"local-high",
|
||
"Private local generation and editing on the shared RTX 3080",
|
||
),
|
||
}
|
||
MAX_INPUT_BYTES = 25 << 20
|
||
|
||
IMAGE_GENERATE_PARAMETERS = {
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt": {
|
||
"type": "string",
|
||
"description": "Detailed description of the image or requested edit.",
|
||
},
|
||
"aspect_ratio": {
|
||
"type": "string",
|
||
"enum": ["landscape", "square", "portrait"],
|
||
"default": DEFAULT_ASPECT_RATIO,
|
||
"description": "Requested output aspect ratio.",
|
||
},
|
||
"image_url": {
|
||
"type": "string",
|
||
"description": "Optional uploaded image URL or private local path to edit.",
|
||
},
|
||
"reference_image_urls": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": "Optional reference images for the generation or edit.",
|
||
},
|
||
},
|
||
"required": ["prompt"],
|
||
}
|
||
|
||
LOCAL_IMAGE_SCHEMA = {
|
||
"name": "image_generate_local",
|
||
"description": (
|
||
"Generate or edit an image only with the private local FLUX model on "
|
||
"Atlas hardware. Use this tool whenever the user says local, private, "
|
||
"FLUX, on my hardware, or otherwise explicitly rejects a hosted image "
|
||
"provider. Never substitute the generic image_generate tool for an "
|
||
"explicit local request. The backend is already provisioned; never ask "
|
||
"for an endpoint or tell the user to install Diffusers or ComfyUI."
|
||
),
|
||
"parameters": IMAGE_GENERATE_PARAMETERS,
|
||
}
|
||
|
||
HOSTED_IMAGE_SCHEMA = {
|
||
"name": "image_generate_hosted",
|
||
"description": (
|
||
"Generate or edit an image only with hosted OpenAI GPT Image at the "
|
||
"highest configured quality. Use this tool when the user explicitly "
|
||
"asks for OpenAI, GPT Image, or hosted image generation. Do not use it "
|
||
"when the user explicitly requests local or private generation."
|
||
),
|
||
"parameters": IMAGE_GENERATE_PARAMETERS,
|
||
}
|
||
|
||
|
||
def _broker_key() -> str:
|
||
"""Load the internal relay key from process env or the private .env."""
|
||
value = os.environ.get("HERMES_IMAGE_BROKER_KEY", "").strip()
|
||
if value:
|
||
return value
|
||
home = Path(os.environ.get("HERMES_HOME", "/opt/data"))
|
||
try:
|
||
for line in (home / ".env").read_text(encoding="utf-8").splitlines():
|
||
if line.startswith("HERMES_IMAGE_BROKER_KEY="):
|
||
return line.split("=", 1)[1].strip()
|
||
except OSError:
|
||
pass
|
||
return ""
|
||
|
||
|
||
def _request(
|
||
path: str, payload: dict[str, Any] | None, timeout: float
|
||
) -> dict[str, Any]:
|
||
"""Call the internal broker with a bounded authenticated JSON request."""
|
||
key = _broker_key()
|
||
if not key:
|
||
raise RuntimeError("private image broker key is unavailable")
|
||
data = None if payload is None else json.dumps(payload).encode("utf-8")
|
||
request = Request(
|
||
f"{BROKER_URL}{path}",
|
||
data=data,
|
||
method="GET" if data is None else "POST",
|
||
headers={
|
||
"Authorization": f"Bearer {key}",
|
||
"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"image broker returned HTTP {exc.code}: {body}") from exc
|
||
except URLError as exc:
|
||
raise RuntimeError(f"image broker is unavailable: {exc.reason}") from exc
|
||
|
||
|
||
def _local_image_data_url(value: str) -> str:
|
||
"""Convert a private uploaded image to a data URL for broker transport."""
|
||
candidate = value.strip()
|
||
lowered = candidate.lower()
|
||
if lowered.startswith(("http://", "https://", "data:image/")):
|
||
return candidate
|
||
try:
|
||
from agent.file_safety import get_read_block_error
|
||
|
||
blocked = get_read_block_error(candidate)
|
||
if blocked:
|
||
raise ValueError(blocked)
|
||
except ImportError:
|
||
pass
|
||
path = Path(os.path.expanduser(candidate)).resolve()
|
||
raw = path.read_bytes()
|
||
if not raw or len(raw) > MAX_INPUT_BYTES:
|
||
raise ValueError("reference image must be between 1 byte and 25 MiB")
|
||
if raw.startswith(b"\x89PNG\r\n\x1a\n"):
|
||
mime = "image/png"
|
||
elif raw.startswith(b"\xff\xd8\xff"):
|
||
mime = "image/jpeg"
|
||
elif raw.startswith((b"GIF87a", b"GIF89a")):
|
||
mime = "image/gif"
|
||
elif raw.startswith(b"RIFF") and raw[8:12] == b"WEBP":
|
||
mime = "image/webp"
|
||
else:
|
||
raise ValueError("reference image must be PNG, JPEG, GIF, or WebP")
|
||
return f"data:{mime};base64,{base64.b64encode(raw).decode('ascii')}"
|
||
|
||
|
||
class AtlasBrokerImageProvider(ImageGenProvider):
|
||
"""High-quality GPT Image generation through the owner-only broker."""
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "atlas-broker"
|
||
|
||
@property
|
||
def display_name(self) -> str:
|
||
return "Atlas private image studio"
|
||
|
||
def is_available(self) -> bool:
|
||
try:
|
||
result = _request("/health", None, 5.0)
|
||
return bool(result.get("success") and result.get("available"))
|
||
except Exception:
|
||
return False
|
||
|
||
def list_models(self) -> list[dict[str, Any]]:
|
||
models = []
|
||
for model, values in MODELS.items():
|
||
display, _quality, *strengths = values
|
||
local = model == "flux-2-klein-4b-local"
|
||
models.append(
|
||
{
|
||
"id": model,
|
||
"display": display,
|
||
"speed": "~2–15min" if local else "~30s–3min",
|
||
"strengths": strengths[0]
|
||
if strengths
|
||
else "OpenAI GPT Image 2 quality",
|
||
"price": "local compute" if local else "included account capacity",
|
||
}
|
||
)
|
||
return models
|
||
|
||
def default_model(self) -> str:
|
||
return DEFAULT_MODEL
|
||
|
||
def capabilities(self) -> dict[str, Any]:
|
||
return {"modalities": ["text", "image"], "max_reference_images": 16}
|
||
|
||
def generate(
|
||
self,
|
||
prompt: str,
|
||
aspect_ratio: str = DEFAULT_ASPECT_RATIO,
|
||
*,
|
||
image_url: str | None = None,
|
||
reference_image_urls: list[str] | None = None,
|
||
**kwargs: Any,
|
||
) -> dict[str, Any]:
|
||
prompt = (prompt or "").strip()
|
||
aspect = resolve_aspect_ratio(aspect_ratio)
|
||
if not prompt:
|
||
return error_response(
|
||
error="Prompt is required",
|
||
error_type="invalid_argument",
|
||
provider=self.name,
|
||
aspect_ratio=aspect,
|
||
)
|
||
model = str(kwargs.get("model") or DEFAULT_MODEL)
|
||
if model not in MODELS:
|
||
model = DEFAULT_MODEL
|
||
try:
|
||
primary = _local_image_data_url(image_url) if image_url else None
|
||
references = [
|
||
_local_image_data_url(item)
|
||
for item in (normalize_reference_images(reference_image_urls) or [])[
|
||
:16
|
||
]
|
||
]
|
||
response = _request(
|
||
"/v1/images/generations",
|
||
{
|
||
"prompt": prompt,
|
||
"aspect_ratio": aspect,
|
||
"model": model,
|
||
"image_url": primary,
|
||
"reference_image_urls": references,
|
||
},
|
||
1200.0,
|
||
)
|
||
if not response.get("success"):
|
||
return error_response(
|
||
error=str(response.get("error") or "image generation failed"),
|
||
error_type=str(response.get("error_type") or "provider_error"),
|
||
provider=self.name,
|
||
model=model,
|
||
prompt=prompt,
|
||
aspect_ratio=aspect,
|
||
)
|
||
image = save_b64_image(str(response["image_b64"]), prefix=f"atlas_{model}")
|
||
return success_response(
|
||
image=str(image),
|
||
model=str(response.get("model") or model),
|
||
prompt=prompt,
|
||
aspect_ratio=aspect,
|
||
provider=self.name,
|
||
modality="image" if primary or references else "text",
|
||
extra={
|
||
"quality": response.get("quality", MODELS[model][1]),
|
||
"size": response.get("size"),
|
||
"route": response.get("route"),
|
||
"requested_model": response.get("requested_model", model),
|
||
"hosted_fallback_reason": response.get("hosted_fallback_reason"),
|
||
"input_image_count": response.get(
|
||
"input_image_count", len(references) + bool(primary)
|
||
),
|
||
},
|
||
)
|
||
except Exception as exc:
|
||
return error_response(
|
||
error=f"Private image studio failed: {exc}",
|
||
error_type="broker_error",
|
||
provider=self.name,
|
||
model=model,
|
||
prompt=prompt,
|
||
aspect_ratio=aspect,
|
||
)
|
||
|
||
|
||
def _handle_image_generate(args: dict[str, Any], route: str) -> str:
|
||
"""Dispatch one image request through a fixed route."""
|
||
prompt = str(args.get("prompt") or "").strip()
|
||
if not prompt:
|
||
return json.dumps(
|
||
{
|
||
"success": False,
|
||
"image": None,
|
||
"error": "prompt is required for image generation",
|
||
"error_type": "invalid_argument",
|
||
}
|
||
)
|
||
model = ROUTE_MODELS.get(route)
|
||
if model is None: # pragma: no cover - routes are fixed by the wrappers
|
||
raise ValueError(f"unsupported image route: {route}")
|
||
provider = AtlasBrokerImageProvider()
|
||
result = provider.generate(
|
||
prompt,
|
||
str(args.get("aspect_ratio") or DEFAULT_ASPECT_RATIO),
|
||
model=model,
|
||
image_url=args.get("image_url"),
|
||
reference_image_urls=args.get("reference_image_urls"),
|
||
)
|
||
if isinstance(result, dict):
|
||
result.setdefault("requested_route", route)
|
||
return json.dumps(result, ensure_ascii=False)
|
||
|
||
|
||
def _handle_local_image(args: dict[str, Any], **_kwargs: Any) -> str:
|
||
"""Generate an image without allowing hosted provider substitution."""
|
||
return _handle_image_generate(args, "local")
|
||
|
||
|
||
def _handle_hosted_image(args: dict[str, Any], **_kwargs: Any) -> str:
|
||
"""Generate an image through the highest-quality hosted route."""
|
||
return _handle_image_generate(args, "hosted")
|
||
|
||
|
||
def _image_tool_available() -> bool:
|
||
"""Expose the tool when at least one broker route is healthy."""
|
||
return AtlasBrokerImageProvider().is_available()
|
||
|
||
|
||
def register(ctx: Any) -> None:
|
||
"""Register the broker and explicit local/hosted image tools."""
|
||
ctx.register_image_gen_provider(AtlasBrokerImageProvider())
|
||
ctx.register_tool(
|
||
name="image_generate_local",
|
||
toolset="image_gen",
|
||
schema=LOCAL_IMAGE_SCHEMA,
|
||
handler=_handle_local_image,
|
||
check_fn=_image_tool_available,
|
||
requires_env=[],
|
||
is_async=False,
|
||
description=LOCAL_IMAGE_SCHEMA["description"],
|
||
emoji="🎨",
|
||
)
|
||
ctx.register_tool(
|
||
name="image_generate_hosted",
|
||
toolset="image_gen",
|
||
schema=HOSTED_IMAGE_SCHEMA,
|
||
handler=_handle_hosted_image,
|
||
check_fn=_image_tool_available,
|
||
requires_env=[],
|
||
is_async=False,
|
||
description=HOSTED_IMAGE_SCHEMA["description"],
|
||
emoji="🎨",
|
||
)
|