feat(pets): generate a custom pet from a prompt (desktop)
Add an in-app pet generator: describe a creature, get four cheap base-look drafts, pick one, then hatch the full six-state animated atlas and preview every frame before adopting. - backend: agent/pet/generate/ (deterministic atlas assembly ported from the hatch-pet skill, prompt builders, provider wrapper, orchestration) - store.register_local_pet so generated pets install + adopt without a manifest entry - gateway pet.generate (draft variants) + pet.hatch (build preview, not active); adopt via existing pet.select, discard via pet.remove - OpenAI provider: reference-image edit path + opt-in transparent background; imagegen retries without the flag on models that reject it - transparency hardened with a chroma-key cutout pass on base drafts - desktop: Cmd+K "Generate a pet" page (draft grid, retry, animated post-hatch preview, adopt/start-over), egg icon, i18n
This commit is contained in:
Brooklyn Nicholson
parent
25e78f129f
commit
8c48be90ac
29
agent/pet/generate/__init__.py
Normal file
29
agent/pet/generate/__init__.py
Normal file
@ -0,0 +1,29 @@
|
||||
"""Pet generation — base-draft → hatch pipeline.
|
||||
|
||||
Public surface used by the gateway RPCs, the CLI ``hermes pets generate``
|
||||
command, and tests:
|
||||
|
||||
- :func:`generate_base_drafts` / :func:`hatch_pet` — the two-step flow.
|
||||
- :class:`HatchResult`, :class:`GenerationError`.
|
||||
- :mod:`atlas` — deterministic frame extraction + atlas composition/validation.
|
||||
|
||||
Image generation is delegated to the active reference-capable
|
||||
:class:`~agent.image_gen_provider.ImageGenProvider` (OpenAI gpt-image-2 or Krea);
|
||||
atlas assembly is fully deterministic so it's testable without any API calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from agent.pet.generate.imagegen import GenerationError
|
||||
from agent.pet.generate.orchestrate import (
|
||||
HatchResult,
|
||||
generate_base_drafts,
|
||||
hatch_pet,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"GenerationError",
|
||||
"HatchResult",
|
||||
"generate_base_drafts",
|
||||
"hatch_pet",
|
||||
]
|
||||
400
agent/pet/generate/atlas.py
Normal file
400
agent/pet/generate/atlas.py
Normal file
@ -0,0 +1,400 @@
|
||||
"""Deterministic spritesheet assembly — generated row strips → Hermes atlas.
|
||||
|
||||
Image-generation models are good at *drawing* a row of poses but bad at exact
|
||||
grid geometry, so the model never owns the atlas layout: it produces one loose
|
||||
horizontal strip per state, and these deterministic ops slice that strip into
|
||||
clean, centered, transparent ``192x208`` cells and pack them into the sheet our
|
||||
renderer reads.
|
||||
|
||||
The atlas is **Hermes-native**, not the petdex/Codex format. Our renderer
|
||||
(:mod:`agent.pet.render`) keys frames as ``rows = states, cols = frames`` using
|
||||
:data:`agent.pet.constants.STATE_ROWS`, so we emit exactly the six states the
|
||||
engine drives — idle, wave, run, failed, review, jump — left-packed with
|
||||
trailing transparent cells (which the renderer trims). Sheet is
|
||||
``COLUMNS*192 x ROWS*208`` (1152x1248).
|
||||
|
||||
The frame-segmentation, fit-to-cell, and transparency-residue logic is adapted
|
||||
from OpenAI's ``hatch-pet`` skill (openai/skills, Apache-2.0).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
from agent.pet.constants import FRAME_H, FRAME_W
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CELL_WIDTH = FRAME_W
|
||||
CELL_HEIGHT = FRAME_H
|
||||
|
||||
# (state, row index, frame count). Order/row indices MUST match
|
||||
# ``STATE_ROWS`` so the renderer crops the right row for each driven state.
|
||||
# Frame counts are the petdex-ish per-state lengths; the renderer trims any
|
||||
# trailing blank columns, so rows shorter than ``COLUMNS`` just leave the tail
|
||||
# transparent.
|
||||
ROW_SPECS: list[tuple[str, int, int]] = [
|
||||
("idle", 0, 6),
|
||||
("wave", 1, 4),
|
||||
("run", 2, 6),
|
||||
("failed", 3, 6),
|
||||
("review", 4, 6),
|
||||
("jump", 5, 5),
|
||||
]
|
||||
|
||||
ROWS = len(ROW_SPECS)
|
||||
COLUMNS = max(count for _, _, count in ROW_SPECS)
|
||||
ATLAS_WIDTH = COLUMNS * CELL_WIDTH
|
||||
ATLAS_HEIGHT = ROWS * CELL_HEIGHT
|
||||
|
||||
FRAME_COUNTS: dict[str, int] = {state: count for state, _, count in ROW_SPECS}
|
||||
|
||||
# Alpha at/below which a pixel is "background" for component detection.
|
||||
_ALPHA_FLOOR = 16
|
||||
# Cell padding kept around a fitted sprite so poses never touch the edge.
|
||||
_CELL_PAD = 10
|
||||
|
||||
|
||||
# ───────────────────────── background removal ─────────────────────────
|
||||
|
||||
|
||||
def _color_distance(r: int, g: int, b: int, key: tuple[int, int, int]) -> float:
|
||||
return math.sqrt((r - key[0]) ** 2 + (g - key[1]) ** 2 + (b - key[2]) ** 2)
|
||||
|
||||
|
||||
def _has_transparency(image) -> bool:
|
||||
"""True if the strip already carries a real alpha background."""
|
||||
extrema = image.getchannel("A").getextrema()
|
||||
# Min alpha 0 somewhere and a meaningful share of fully-transparent pixels.
|
||||
if extrema[0] > _ALPHA_FLOOR:
|
||||
return False
|
||||
hist = image.getchannel("A").histogram()
|
||||
transparent = sum(hist[: _ALPHA_FLOOR + 1])
|
||||
total = image.width * image.height
|
||||
return transparent > total * 0.05
|
||||
|
||||
|
||||
def _dominant_corner_color(image) -> tuple[int, int, int]:
|
||||
"""Sample the four corners and return the most common opaque color."""
|
||||
from collections import Counter
|
||||
|
||||
w, h = image.width, image.height
|
||||
px = image.load()
|
||||
counter: Counter = Counter()
|
||||
for x, y in ((0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1)):
|
||||
r, g, b, a = px[x, y]
|
||||
if a > _ALPHA_FLOOR:
|
||||
counter[(r, g, b)] += 1
|
||||
if not counter:
|
||||
return (0, 255, 0)
|
||||
return counter.most_common(1)[0][0]
|
||||
|
||||
|
||||
def remove_background(image, *, chroma_key: tuple[int, int, int] | None = None, threshold: float = 110.0):
|
||||
"""Return *image* (RGBA) with its flat background keyed out to transparent.
|
||||
|
||||
If the strip already has a transparent background we leave it alone; else we
|
||||
key out *chroma_key* (or the dominant corner color when not given). This
|
||||
handles both providers that emit transparency natively and those that paint
|
||||
a solid backdrop.
|
||||
"""
|
||||
rgba = image.convert("RGBA")
|
||||
if _has_transparency(rgba):
|
||||
return rgba
|
||||
|
||||
key = chroma_key or _dominant_corner_color(rgba)
|
||||
px = rgba.load()
|
||||
for y in range(rgba.height):
|
||||
for x in range(rgba.width):
|
||||
r, g, b, a = px[x, y]
|
||||
if a > _ALPHA_FLOOR and _color_distance(r, g, b, key) <= threshold:
|
||||
px[x, y] = (0, 0, 0, 0)
|
||||
return rgba
|
||||
|
||||
|
||||
# ───────────────────────── frame extraction ─────────────────────────
|
||||
|
||||
|
||||
def _fit_to_cell(image):
|
||||
"""Crop to content, scale to fit a padded cell, and center on transparent."""
|
||||
from PIL import Image
|
||||
|
||||
target = Image.new("RGBA", (CELL_WIDTH, CELL_HEIGHT), (0, 0, 0, 0))
|
||||
bbox = image.getbbox()
|
||||
if bbox is None:
|
||||
return target
|
||||
|
||||
sprite = image.crop(bbox)
|
||||
max_w = CELL_WIDTH - _CELL_PAD
|
||||
max_h = CELL_HEIGHT - _CELL_PAD
|
||||
scale = min(max_w / sprite.width, max_h / sprite.height, 1.0)
|
||||
if scale != 1.0:
|
||||
sprite = sprite.resize(
|
||||
(max(1, round(sprite.width * scale)), max(1, round(sprite.height * scale))),
|
||||
Image.Resampling.LANCZOS,
|
||||
)
|
||||
left = (CELL_WIDTH - sprite.width) // 2
|
||||
top = (CELL_HEIGHT - sprite.height) // 2
|
||||
target.alpha_composite(sprite, (left, top))
|
||||
return target
|
||||
|
||||
|
||||
def _connected_components(image) -> list[dict]:
|
||||
"""Flood-fill the alpha mask into connected blobs (4-connectivity)."""
|
||||
alpha = image.getchannel("A")
|
||||
w, h = image.size
|
||||
data = alpha.tobytes()
|
||||
visited = bytearray(w * h)
|
||||
out: list[dict] = []
|
||||
|
||||
for start, a in enumerate(data):
|
||||
if a <= _ALPHA_FLOOR or visited[start]:
|
||||
continue
|
||||
stack = [start]
|
||||
visited[start] = 1
|
||||
pixels: list[int] = []
|
||||
min_x = w
|
||||
min_y = h
|
||||
max_x = 0
|
||||
max_y = 0
|
||||
while stack:
|
||||
cur = stack.pop()
|
||||
pixels.append(cur)
|
||||
x = cur % w
|
||||
y = cur // w
|
||||
min_x = min(min_x, x)
|
||||
min_y = min(min_y, y)
|
||||
max_x = max(max_x, x)
|
||||
max_y = max(max_y, y)
|
||||
for nb, ok in (
|
||||
(cur - 1, x > 0),
|
||||
(cur + 1, x + 1 < w),
|
||||
(cur - w, y > 0),
|
||||
(cur + w, y + 1 < h),
|
||||
):
|
||||
if ok and not visited[nb] and data[nb] > _ALPHA_FLOOR:
|
||||
visited[nb] = 1
|
||||
stack.append(nb)
|
||||
out.append(
|
||||
{
|
||||
"pixels": pixels,
|
||||
"area": len(pixels),
|
||||
"bbox": (min_x, min_y, max_x + 1, max_y + 1),
|
||||
"center_x": (min_x + max_x + 1) / 2,
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _group_image(source, components: list[dict], padding: int = 4):
|
||||
from PIL import Image
|
||||
|
||||
w, h = source.size
|
||||
min_x = max(0, min(c["bbox"][0] for c in components) - padding)
|
||||
min_y = max(0, min(c["bbox"][1] for c in components) - padding)
|
||||
max_x = min(w, max(c["bbox"][2] for c in components) + padding)
|
||||
max_y = min(h, max(c["bbox"][3] for c in components) + padding)
|
||||
|
||||
out = Image.new("RGBA", (max_x - min_x, max_y - min_y), (0, 0, 0, 0))
|
||||
src_px = source.load()
|
||||
out_px = out.load()
|
||||
for c in components:
|
||||
for idx in c["pixels"]:
|
||||
x = idx % w
|
||||
y = idx // w
|
||||
out_px[x - min_x, y - min_y] = src_px[x, y]
|
||||
return out
|
||||
|
||||
|
||||
def _component_frames(strip, frame_count: int) -> list | None:
|
||||
"""Segment a strip into *frame_count* sprites by connected components.
|
||||
|
||||
Picks the ``frame_count`` largest blobs as seeds (left→right), attaches
|
||||
smaller blobs to the nearest seed, and returns one fitted cell per group.
|
||||
Returns ``None`` when it can't find enough distinct sprites (caller falls
|
||||
back to equal slicing).
|
||||
"""
|
||||
components = _connected_components(strip)
|
||||
if not components:
|
||||
return None
|
||||
|
||||
largest = max(c["area"] for c in components)
|
||||
seed_threshold = max(120, largest * 0.20)
|
||||
seeds = [c for c in components if c["area"] >= seed_threshold]
|
||||
if len(seeds) < frame_count:
|
||||
seeds = sorted(components, key=lambda c: c["area"], reverse=True)[:frame_count]
|
||||
if len(seeds) < frame_count:
|
||||
return None
|
||||
|
||||
seeds = sorted(
|
||||
sorted(seeds, key=lambda c: c["area"], reverse=True)[:frame_count],
|
||||
key=lambda c: c["center_x"],
|
||||
)
|
||||
seed_ids = {id(s) for s in seeds}
|
||||
groups: list[list[dict]] = [[s] for s in seeds]
|
||||
noise_threshold = max(12, largest * 0.002)
|
||||
for c in components:
|
||||
if id(c) in seed_ids or c["area"] < noise_threshold:
|
||||
continue
|
||||
nearest = min(range(len(seeds)), key=lambda i: abs(seeds[i]["center_x"] - c["center_x"]))
|
||||
groups[nearest].append(c)
|
||||
|
||||
return [_fit_to_cell(_group_image(strip, g)) for g in groups]
|
||||
|
||||
|
||||
def _slot_frames(strip, frame_count: int) -> list:
|
||||
"""Fallback: slice the strip into *frame_count* equal columns."""
|
||||
slot = strip.width / frame_count
|
||||
frames = []
|
||||
for i in range(frame_count):
|
||||
left = round(i * slot)
|
||||
right = round((i + 1) * slot)
|
||||
frames.append(_fit_to_cell(strip.crop((left, 0, right, strip.height))))
|
||||
return frames
|
||||
|
||||
|
||||
def extract_strip_frames(
|
||||
strip,
|
||||
frame_count: int,
|
||||
*,
|
||||
chroma_key: tuple[int, int, int] | None = None,
|
||||
method: str = "auto",
|
||||
) -> list:
|
||||
"""Turn one generated row strip into *frame_count* clean 192x208 cells.
|
||||
|
||||
*strip* is a PIL image (or path). Background is keyed out, then frames are
|
||||
found by connected components (``auto``) with an equal-slot fallback.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
if isinstance(strip, (str, Path)):
|
||||
with Image.open(strip) as opened:
|
||||
strip = opened.convert("RGBA")
|
||||
else:
|
||||
strip = strip.convert("RGBA")
|
||||
|
||||
strip = remove_background(strip, chroma_key=chroma_key)
|
||||
|
||||
if method in ("auto", "components"):
|
||||
frames = _component_frames(strip, frame_count)
|
||||
if frames is not None:
|
||||
return frames
|
||||
if method == "components":
|
||||
raise ValueError(f"could not segment {frame_count} sprites from strip")
|
||||
return _slot_frames(strip, frame_count)
|
||||
|
||||
|
||||
# ───────────────────────── atlas composition ─────────────────────────
|
||||
|
||||
|
||||
def single_frame(image):
|
||||
"""One fitted 192x208 cell from a standalone image (e.g. the base look).
|
||||
|
||||
Used as an idle fallback so a pet always renders even if the idle row
|
||||
generation failed.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
if isinstance(image, (str, Path)):
|
||||
with Image.open(image) as opened:
|
||||
image = opened.convert("RGBA")
|
||||
return _fit_to_cell(remove_background(image))
|
||||
|
||||
|
||||
def _clear_transparent_rgb(image):
|
||||
"""Zero the RGB of fully-transparent pixels (no colored-halo residue)."""
|
||||
from PIL import Image
|
||||
|
||||
rgba = image.convert("RGBA")
|
||||
data = bytearray(rgba.tobytes())
|
||||
for i in range(0, len(data), 4):
|
||||
if data[i + 3] == 0:
|
||||
data[i] = data[i + 1] = data[i + 2] = 0
|
||||
return Image.frombytes("RGBA", rgba.size, bytes(data))
|
||||
|
||||
|
||||
def compose_atlas(frames_by_state: dict[str, list]):
|
||||
"""Pack per-state frame lists into the Hermes atlas (RGBA, residue-cleared).
|
||||
|
||||
Missing/short states leave their trailing cells transparent; extra frames
|
||||
beyond a state's spec are dropped.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
atlas = Image.new("RGBA", (ATLAS_WIDTH, ATLAS_HEIGHT), (0, 0, 0, 0))
|
||||
for state, row, count in ROW_SPECS:
|
||||
frames = frames_by_state.get(state) or []
|
||||
for col, frame in enumerate(frames[:count]):
|
||||
cell = frame.convert("RGBA")
|
||||
if cell.size != (CELL_WIDTH, CELL_HEIGHT):
|
||||
cell = _fit_to_cell(cell)
|
||||
atlas.alpha_composite(cell, (col * CELL_WIDTH, row * CELL_HEIGHT))
|
||||
return _clear_transparent_rgb(atlas)
|
||||
|
||||
|
||||
def atlas_to_webp_bytes(atlas) -> bytes:
|
||||
"""Encode an atlas image to lossless WebP bytes (the on-disk pet format)."""
|
||||
buf = io.BytesIO()
|
||||
atlas.save(buf, format="WEBP", lossless=True, quality=100, method=6, exact=True)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def validate_atlas(atlas) -> dict:
|
||||
"""Check geometry, per-cell occupancy, and transparency invariants.
|
||||
|
||||
Returns ``{ok, width, height, errors, warnings, filled_states}``. Errors are
|
||||
blockers (wrong size, empty used cell, opaque/dirty transparency); warnings
|
||||
are soft (a whole state row blank — generation likely dropped a row).
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
if isinstance(atlas, (str, Path)):
|
||||
with Image.open(atlas) as opened:
|
||||
atlas = opened.convert("RGBA")
|
||||
else:
|
||||
atlas = atlas.convert("RGBA")
|
||||
|
||||
errors: list[str] = []
|
||||
warnings: list[str] = []
|
||||
|
||||
if atlas.size != (ATLAS_WIDTH, ATLAS_HEIGHT):
|
||||
errors.append(f"expected {ATLAS_WIDTH}x{ATLAS_HEIGHT}, got {atlas.width}x{atlas.height}")
|
||||
return {"ok": False, "width": atlas.width, "height": atlas.height, "errors": errors, "warnings": warnings, "filled_states": []}
|
||||
|
||||
filled_states: list[str] = []
|
||||
for state, row, count in ROW_SPECS:
|
||||
row_pixels = 0
|
||||
for col in range(count):
|
||||
left = col * CELL_WIDTH
|
||||
top = row * CELL_HEIGHT
|
||||
cell = atlas.crop((left, top, left + CELL_WIDTH, top + CELL_HEIGHT))
|
||||
nonblank = sum(cell.getchannel("A").histogram()[1:])
|
||||
row_pixels += nonblank
|
||||
if row_pixels > 0:
|
||||
filled_states.append(state)
|
||||
else:
|
||||
warnings.append(f"state '{state}' has no frames")
|
||||
|
||||
if not filled_states:
|
||||
errors.append("atlas is empty — no state produced any frames")
|
||||
|
||||
# Transparent pixels must carry zero RGB (no halo residue).
|
||||
data = atlas.tobytes()
|
||||
residue = 0
|
||||
for i in range(0, len(data), 4):
|
||||
if data[i + 3] == 0 and (data[i] or data[i + 1] or data[i + 2]):
|
||||
residue += 1
|
||||
if residue:
|
||||
errors.append(f"{residue} transparent pixels retain RGB residue")
|
||||
|
||||
return {
|
||||
"ok": not errors,
|
||||
"width": atlas.width,
|
||||
"height": atlas.height,
|
||||
"errors": errors,
|
||||
"warnings": warnings,
|
||||
"filled_states": filled_states,
|
||||
}
|
||||
168
agent/pet/generate/imagegen.py
Normal file
168
agent/pet/generate/imagegen.py
Normal file
@ -0,0 +1,168 @@
|
||||
"""Thin image-generation layer for pet sprites.
|
||||
|
||||
Wraps the active :class:`~agent.image_gen_provider.ImageGenProvider` with the
|
||||
two things sprite generation needs that the agent-facing ``image_generate`` tool
|
||||
doesn't expose: **N variants** (loop) and **reference-image grounding** (so each
|
||||
animation row stays the same character as the chosen base).
|
||||
|
||||
Reference grounding only works on providers that support it — currently OpenAI
|
||||
``gpt-image-2`` (image edits) and Krea (style references). We resolve to one of
|
||||
those and surface a clear, actionable error otherwise rather than silently
|
||||
producing an ungrounded, drifting pet.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Providers that can ground generation on a reference image.
|
||||
_REF_CAPABLE = ("openai", "openai-codex", "krea")
|
||||
|
||||
|
||||
class GenerationError(RuntimeError):
|
||||
"""Raised on any image-generation failure (no provider, API error, IO)."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SpriteProvider:
|
||||
"""Resolved provider plus whether it can take reference images."""
|
||||
|
||||
name: str
|
||||
provider: object
|
||||
supports_references: bool
|
||||
|
||||
|
||||
def _discover() -> None:
|
||||
try:
|
||||
from hermes_cli.plugins import _ensure_plugins_discovered
|
||||
|
||||
_ensure_plugins_discovered()
|
||||
except Exception as exc: # noqa: BLE001 - discovery is best-effort
|
||||
logger.debug("image-gen plugin discovery failed: %s", exc)
|
||||
|
||||
|
||||
def resolve_provider(*, require_references: bool = True) -> SpriteProvider:
|
||||
"""Pick the image provider to use for sprite work.
|
||||
|
||||
Preference: the configured provider when it's reference-capable, else the
|
||||
first available reference-capable provider. With *require_references* off we
|
||||
fall back to any available provider (used for prompt-only base drafts).
|
||||
"""
|
||||
_discover()
|
||||
from agent.image_gen_registry import get_active_provider, get_provider
|
||||
|
||||
# Configured / active provider first.
|
||||
active = None
|
||||
try:
|
||||
active = get_active_provider()
|
||||
except Exception: # noqa: BLE001
|
||||
active = None
|
||||
if active is not None:
|
||||
name = getattr(active, "name", "")
|
||||
if name in _REF_CAPABLE and active.is_available():
|
||||
return SpriteProvider(name=name, provider=active, supports_references=True)
|
||||
|
||||
# Any available reference-capable provider.
|
||||
for name in _REF_CAPABLE:
|
||||
provider = get_provider(name)
|
||||
if provider is not None and provider.is_available():
|
||||
return SpriteProvider(name=name, provider=provider, supports_references=True)
|
||||
|
||||
if not require_references and active is not None and active.is_available():
|
||||
return SpriteProvider(
|
||||
name=getattr(active, "name", "unknown"), provider=active, supports_references=False
|
||||
)
|
||||
|
||||
raise GenerationError(
|
||||
"Pet generation needs a reference-capable image backend. "
|
||||
"Run `hermes tools` → Image Generation → OpenAI (gpt-image-2) and add an "
|
||||
"OpenAI API key (or configure Krea)."
|
||||
)
|
||||
|
||||
|
||||
def _save_local(image_ref: str, *, prefix: str) -> Path:
|
||||
"""Return a local path for *image_ref*, downloading it if it's a URL."""
|
||||
if image_ref.startswith(("http://", "https://")):
|
||||
from agent.image_gen_provider import save_url_image
|
||||
|
||||
return Path(save_url_image(image_ref, prefix=prefix))
|
||||
return Path(image_ref)
|
||||
|
||||
|
||||
def _rejected_background(error: str) -> bool:
|
||||
"""True when a provider error is specifically about the ``background`` param.
|
||||
|
||||
Transparent backgrounds are a per-model capability (e.g. some gpt-image tiers
|
||||
reject ``background=transparent`` outright). We detect that one rejection so
|
||||
we can retry without the flag rather than failing the whole pet — our chroma
|
||||
key pass makes the result transparent regardless.
|
||||
"""
|
||||
lowered = (error or "").lower()
|
||||
return "background" in lowered and ("not supported" in lowered or "transparent" in lowered)
|
||||
|
||||
|
||||
def generate(
|
||||
prompt: str,
|
||||
*,
|
||||
n: int = 1,
|
||||
reference_images: list[Path] | None = None,
|
||||
provider: SpriteProvider | None = None,
|
||||
prefix: str = "pet_gen",
|
||||
) -> list[Path]:
|
||||
"""Generate *n* square sprite images and return their local paths.
|
||||
|
||||
*reference_images* grounds the output on a base image (required for rows).
|
||||
We *ask* for a transparent background, but fall back to an opaque generation
|
||||
(cleaned up downstream by the chroma-key pass) on models that reject the
|
||||
flag. Raises :class:`GenerationError` if nothing usable comes back.
|
||||
"""
|
||||
sprite = provider or resolve_provider(require_references=bool(reference_images))
|
||||
if reference_images and not sprite.supports_references:
|
||||
raise GenerationError(
|
||||
f"image backend '{sprite.name}' cannot use reference images; "
|
||||
"configure OpenAI gpt-image-2 or Krea for pet generation"
|
||||
)
|
||||
|
||||
refs = [str(p) for p in (reference_images or [])]
|
||||
|
||||
def _run(extra: dict) -> tuple[Path | None, str]:
|
||||
kwargs: dict = {"aspect_ratio": "square", **extra}
|
||||
if refs:
|
||||
kwargs["reference_images"] = refs
|
||||
try:
|
||||
result = sprite.provider.generate(prompt, **kwargs)
|
||||
except Exception as exc: # noqa: BLE001 - normalize provider crashes
|
||||
logger.debug("provider.generate crashed: %s", exc)
|
||||
return None, str(exc)
|
||||
if not isinstance(result, dict) or not result.get("success"):
|
||||
return None, (result or {}).get("error", "unknown error") if isinstance(result, dict) else "no result"
|
||||
image_ref = result.get("image")
|
||||
if not image_ref:
|
||||
return None, "provider returned no image"
|
||||
try:
|
||||
return _save_local(str(image_ref), prefix=prefix), ""
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return None, f"could not save generated image: {exc}"
|
||||
|
||||
out: list[Path] = []
|
||||
last_error = ""
|
||||
allow_transparent = True
|
||||
for _ in range(max(1, n)):
|
||||
path, err = _run({"background": "transparent"} if allow_transparent else {})
|
||||
# Model doesn't support the transparent flag → drop it for this and every
|
||||
# remaining variant (no point re-probing a capability we just disproved).
|
||||
if path is None and allow_transparent and _rejected_background(err):
|
||||
allow_transparent = False
|
||||
path, err = _run({})
|
||||
if path is not None:
|
||||
out.append(path)
|
||||
else:
|
||||
last_error = err
|
||||
|
||||
if not out:
|
||||
raise GenerationError(last_error or "image generation produced no output")
|
||||
return out
|
||||
149
agent/pet/generate/orchestrate.py
Normal file
149
agent/pet/generate/orchestrate.py
Normal file
@ -0,0 +1,149 @@
|
||||
"""Pet generation orchestration — the base-draft → hatch flow.
|
||||
|
||||
Two steps, mirroring the UX across every surface:
|
||||
|
||||
1. :func:`generate_base_drafts` — a handful of prompt-only "what should this pet
|
||||
look like" variants. Cheap; the user picks one (or retries for a fresh set).
|
||||
2. :func:`hatch_pet` — takes the chosen base and generates one grounded row
|
||||
strip per Hermes state, slices each into frames, composes the atlas, validates
|
||||
it, and writes the pet into the store.
|
||||
|
||||
Splitting it this way bounds cost (4 cheap base calls per round; the ~6 row
|
||||
calls happen once, on the pet you actually keep) and gives each UI a natural
|
||||
preview/loading point.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
from agent.pet.generate import atlas, imagegen, prompts
|
||||
from agent.pet.generate.imagegen import GenerationError, SpriteProvider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# (event, detail) — e.g. ("row", "idle"), ("compose", ""), ("save", "<slug>").
|
||||
ProgressFn = Callable[[str, str], None]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HatchResult:
|
||||
"""Outcome of a successful :func:`hatch_pet`."""
|
||||
|
||||
slug: str
|
||||
display_name: str
|
||||
spritesheet: Path
|
||||
states: list[str]
|
||||
validation: dict
|
||||
|
||||
|
||||
def _harden_transparency(path: Path) -> Path:
|
||||
"""Key out any solid backdrop the provider painted; save as an RGBA PNG.
|
||||
|
||||
``background=transparent`` is requested on every call, but image models honor
|
||||
it inconsistently — some still paint a flat (often near-white) backdrop. We
|
||||
run the same chroma-key pass the row extractor uses so every base draft the
|
||||
user picks between (and the reference the rows are grounded on) is a clean
|
||||
cutout. Best-effort: a decode failure leaves the original untouched.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
try:
|
||||
with Image.open(path) as opened:
|
||||
keyed = atlas.remove_background(opened.convert("RGBA"))
|
||||
out = path.with_suffix(".png")
|
||||
keyed.save(out, format="PNG")
|
||||
return out
|
||||
except Exception as exc: # noqa: BLE001 - cosmetic; fall back to the raw image
|
||||
logger.debug("base draft transparency hardening failed for %s: %s", path, exc)
|
||||
return path
|
||||
|
||||
|
||||
def generate_base_drafts(
|
||||
concept: str,
|
||||
*,
|
||||
n: int = 4,
|
||||
style: str = "auto",
|
||||
provider: SpriteProvider | None = None,
|
||||
) -> list[Path]:
|
||||
"""Generate *n* candidate base looks for *concept*; returns image paths.
|
||||
|
||||
Each draft is hardened to a transparent cutout (see :func:`_harden_transparency`).
|
||||
"""
|
||||
prompt = prompts.build_base_prompt(concept, style=style)
|
||||
sprite = provider or imagegen.resolve_provider(require_references=False)
|
||||
raw = imagegen.generate(prompt, n=n, provider=sprite, prefix="pet_base")
|
||||
return [_harden_transparency(p) for p in raw]
|
||||
|
||||
|
||||
def hatch_pet(
|
||||
*,
|
||||
base_image: str | Path,
|
||||
slug: str,
|
||||
display_name: str = "",
|
||||
description: str = "",
|
||||
concept: str = "",
|
||||
style: str = "auto",
|
||||
on_progress: ProgressFn | None = None,
|
||||
provider: SpriteProvider | None = None,
|
||||
) -> HatchResult:
|
||||
"""Turn an approved base image into a full, installed Hermes pet.
|
||||
|
||||
Generates a grounded row strip per state, extracts frames, composes +
|
||||
validates the atlas, and registers it. The idle row falls back to the base
|
||||
look so the pet always renders. Raises :class:`GenerationError` on failure.
|
||||
"""
|
||||
base = Path(base_image)
|
||||
if not base.is_file():
|
||||
raise GenerationError(f"base image not found: {base}")
|
||||
|
||||
sprite = provider or imagegen.resolve_provider(require_references=True)
|
||||
progress = on_progress or (lambda *_: None)
|
||||
label = concept or display_name or slug
|
||||
|
||||
frames_by_state: dict[str, list] = {}
|
||||
for state, _row, count in atlas.ROW_SPECS:
|
||||
progress("row", state)
|
||||
row_prompt = prompts.build_row_prompt(state, count, label, style=style)
|
||||
try:
|
||||
strips = imagegen.generate(
|
||||
row_prompt,
|
||||
n=1,
|
||||
reference_images=[base],
|
||||
provider=sprite,
|
||||
prefix=f"pet_row_{state}",
|
||||
)
|
||||
frames_by_state[state] = atlas.extract_strip_frames(strips[0], count, method="auto")
|
||||
except Exception as exc: # noqa: BLE001 - a single row may fail; keep going
|
||||
logger.warning("pet row '%s' failed: %s", state, exc)
|
||||
|
||||
# Idle is the resting state the renderer falls back to — guarantee it.
|
||||
if not frames_by_state.get("idle"):
|
||||
progress("row", "idle-fallback")
|
||||
frames_by_state["idle"] = [atlas.single_frame(base)]
|
||||
|
||||
progress("compose", "")
|
||||
sheet = atlas.compose_atlas(frames_by_state)
|
||||
validation = atlas.validate_atlas(sheet)
|
||||
if not validation["ok"]:
|
||||
raise GenerationError("; ".join(validation["errors"]) or "atlas validation failed")
|
||||
|
||||
from agent.pet import store
|
||||
|
||||
progress("save", slug)
|
||||
pet = store.register_local_pet(
|
||||
sheet,
|
||||
slug=slug,
|
||||
display_name=display_name or slug,
|
||||
description=description,
|
||||
)
|
||||
return HatchResult(
|
||||
slug=pet.slug,
|
||||
display_name=pet.display_name,
|
||||
spritesheet=pet.spritesheet,
|
||||
states=validation["filled_states"],
|
||||
validation=validation,
|
||||
)
|
||||
74
agent/pet/generate/prompts.py
Normal file
74
agent/pet/generate/prompts.py
Normal file
@ -0,0 +1,74 @@
|
||||
"""Prompt builders for pet generation.
|
||||
|
||||
Two prompt shapes: a *base* prompt (prompt-only, produces the canonical look the
|
||||
user picks between) and per-*state* *row* prompts (grounded on the chosen base,
|
||||
produce one horizontal strip of N poses). Prompts stay concise and
|
||||
sprite-production oriented; the identity lock and "one transparent row" framing
|
||||
matter more than flowery description.
|
||||
|
||||
Hermes drives six states (see :data:`agent.pet.generate.atlas.ROW_SPECS`); these
|
||||
mirror that set rather than the petdex/Codex nine.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# What each Hermes state should depict (kept short — these go straight into the
|
||||
# row prompt). Phrased to avoid the common sprite-gen failure modes (detached
|
||||
# effects, motion lines, shadows).
|
||||
STATE_ACTIONS: dict[str, str] = {
|
||||
"idle": "a calm idle loop: subtle breathing, a tiny blink or gentle bob, no big gestures",
|
||||
"wave": "a friendly greeting: raising a paw/hand/limb to wave, clear up-and-down gesture",
|
||||
"run": "focused active work: leaning in, concentrating, busy 'thinking/processing' energy (NOT foot-running)",
|
||||
"failed": "a sad or deflated reaction: slumped, dejected, small frown — readable but not noisy",
|
||||
"review": "careful inspection: a focused lean, head tilt, studying something intently",
|
||||
"jump": "a happy celebration jump: anticipation, lift off the ground, peak, and land",
|
||||
}
|
||||
|
||||
_STYLE_HINTS: dict[str, str] = {
|
||||
"auto": "",
|
||||
"pixel": " Render in clean pixel-art style.",
|
||||
"plush": " Render as a soft plush toy.",
|
||||
"clay": " Render as a claymation / soft 3D clay figure.",
|
||||
"sticker": " Render as a glossy die-cut sticker.",
|
||||
"flat-vector": " Render in flat vector mascot style.",
|
||||
"3d-toy": " Render as a glossy 3D toy.",
|
||||
"painterly": " Render in a soft painterly style.",
|
||||
}
|
||||
|
||||
_BACKGROUND = (
|
||||
"Center one full-body character on a fully transparent background. "
|
||||
"No text, no labels, no shadow, no ground line, no scenery, no frame, no border."
|
||||
)
|
||||
|
||||
|
||||
def style_hint(style: str | None) -> str:
|
||||
return _STYLE_HINTS.get((style or "auto").strip().lower(), "")
|
||||
|
||||
|
||||
def build_base_prompt(concept: str, *, style: str | None = "auto") -> str:
|
||||
"""The base look: a single, clean, centered full-body mascot."""
|
||||
concept = (concept or "a cute friendly mascot creature").strip()
|
||||
return (
|
||||
f"A cute, characterful mascot pet: {concept}. "
|
||||
"Compact, whole-body silhouette that reads clearly at small size, "
|
||||
"appealing face, simple consistent palette. "
|
||||
f"{_BACKGROUND}{style_hint(style)}"
|
||||
)
|
||||
|
||||
|
||||
def build_row_prompt(state: str, frame_count: int, concept: str, *, style: str | None = "auto") -> str:
|
||||
"""A row strip: *frame_count* poses of the SAME character, left→right.
|
||||
|
||||
The attached base image is the identity source of truth; the prompt locks
|
||||
species, palette, face, and props to it.
|
||||
"""
|
||||
action = STATE_ACTIONS.get(state, "a simple idle pose")
|
||||
concept = (concept or "the mascot").strip()
|
||||
return (
|
||||
f"Using the attached reference image as the exact same character "
|
||||
f"(same species, face, colors, markings, proportions, and props), "
|
||||
f"draw a single horizontal strip of {frame_count} animation frames showing {action}. "
|
||||
f"The {frame_count} poses must be evenly spaced left to right, each fully separated "
|
||||
"(not overlapping), same size and baseline, forming a smooth loop. "
|
||||
f"Keep the character identical across all frames. {_BACKGROUND}{style_hint(style)}"
|
||||
)
|
||||
@ -18,6 +18,7 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
@ -173,6 +174,77 @@ def install_pet(slug: str, *, force: bool = False, timeout: float = _DOWNLOAD_TI
|
||||
return pet
|
||||
|
||||
|
||||
def slugify(name: str) -> str:
|
||||
"""Lowercase, hyphenate, and strip a display name into a filesystem slug."""
|
||||
slug = re.sub(r"[^a-z0-9]+", "-", (name or "").strip().lower()).strip("-")
|
||||
return slug or "pet"
|
||||
|
||||
|
||||
def unique_slug(name: str) -> str:
|
||||
"""A :func:`slugify` result that doesn't collide with an existing pet dir."""
|
||||
base = slugify(name)
|
||||
slug = base
|
||||
counter = 2
|
||||
while (pets_dir() / slug).exists():
|
||||
slug = f"{base}-{counter}"
|
||||
counter += 1
|
||||
return slug
|
||||
|
||||
|
||||
def _write_spritesheet(source, dest: Path) -> None:
|
||||
"""Write *source* (PIL image, bytes, or path) as a lossless WebP at *dest*."""
|
||||
if isinstance(source, (bytes, bytearray)):
|
||||
dest.write_bytes(bytes(source))
|
||||
return
|
||||
|
||||
from PIL import Image
|
||||
|
||||
if isinstance(source, (str, Path)):
|
||||
with Image.open(source) as opened:
|
||||
image = opened.convert("RGBA")
|
||||
else:
|
||||
image = source.convert("RGBA")
|
||||
image.save(dest, format="WEBP", lossless=True, quality=100, method=6, exact=True)
|
||||
|
||||
|
||||
def register_local_pet(
|
||||
spritesheet,
|
||||
*,
|
||||
slug: str,
|
||||
display_name: str = "",
|
||||
description: str = "",
|
||||
) -> InstalledPet:
|
||||
"""Write a locally-generated pet into the store and return it.
|
||||
|
||||
*spritesheet* may be a PIL image, raw WebP/PNG bytes, or a path. The pet
|
||||
appears in :func:`installed_pets` immediately, and because :func:`install_pet`
|
||||
returns an already-on-disk pet before consulting the manifest, it can be
|
||||
adopted (``pet.select`` / ``/pet <slug>``) without a manifest entry.
|
||||
"""
|
||||
slug = slugify(slug)
|
||||
directory = pets_dir() / slug
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
sprite_path = directory / "spritesheet.webp"
|
||||
try:
|
||||
_write_spritesheet(spritesheet, sprite_path)
|
||||
except Exception as exc: # noqa: BLE001 - normalize to one error type
|
||||
raise PetStoreError(f"could not write spritesheet for '{slug}': {exc}") from exc
|
||||
|
||||
meta = {
|
||||
"id": slug,
|
||||
"displayName": display_name or slug,
|
||||
"description": description or "",
|
||||
"spritesheetPath": sprite_path.name,
|
||||
"createdBy": "generator",
|
||||
}
|
||||
(directory / "pet.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
|
||||
|
||||
pet = load_pet(slug)
|
||||
if pet is None or not pet.exists:
|
||||
raise PetStoreError(f"register of generated pet '{slug}' did not produce a spritesheet")
|
||||
return pet
|
||||
|
||||
|
||||
_THUMB_FRAME_W = 192
|
||||
_THUMB_FRAME_H = 208
|
||||
_THUMB_W = 96 # rendered ~40px; 2x+ keeps it crisp on HiDPI
|
||||
|
||||
@ -5,6 +5,7 @@ import { useCallback, useEffect, useMemo, useState } from 'react'
|
||||
import { useNavigate } from 'react-router-dom'
|
||||
|
||||
import { HUD_HEADING, HUD_ITEM, HUD_POSITION, HUD_SURFACE, HUD_TEXT } from '@/app/floating-hud'
|
||||
import { useGatewayRequest } from '@/app/gateway/hooks/use-gateway-request'
|
||||
import { setTerminalTakeover } from '@/app/right-sidebar/store'
|
||||
import { Command, CommandEmpty, CommandGroup, CommandInput, CommandItem, CommandList } from '@/components/ui/command'
|
||||
import { KbdCombo } from '@/components/ui/kbd'
|
||||
@ -20,6 +21,7 @@ import {
|
||||
Clock,
|
||||
Cpu,
|
||||
Download,
|
||||
Egg,
|
||||
Globe,
|
||||
type IconComponent,
|
||||
Info,
|
||||
@ -42,6 +44,7 @@ import {
|
||||
import { cn } from '@/lib/utils'
|
||||
import { $commandPaletteOpen, $commandPalettePage, closeCommandPalette, setCommandPaletteOpen } from '@/store/command-palette'
|
||||
import { $bindings } from '@/store/keybinds'
|
||||
import { $petGenStatus, cleanupPetGen, generateDrafts } from '@/store/pet-generate'
|
||||
import { luminance } from '@/themes/color'
|
||||
import { type ThemeMode, useTheme } from '@/themes/context'
|
||||
import { isUserTheme, resolveTheme } from '@/themes/user-themes'
|
||||
@ -63,6 +66,7 @@ import { fieldCopyForSchemaKey } from '../settings/field-copy'
|
||||
import { prettyName } from '../settings/helpers'
|
||||
|
||||
import { MarketplaceThemePage } from './marketplace-theme-page'
|
||||
import { PetGeneratePage } from './pet-generate-page'
|
||||
import { PetInlineToggle, PetPalettePage } from './pet-palette-page'
|
||||
|
||||
interface PaletteItem {
|
||||
@ -89,7 +93,7 @@ interface PaletteGroup {
|
||||
|
||||
// Nested page → its parent, so Back / Esc step up one level instead of closing
|
||||
// the palette. Pages absent here go straight back to the root list.
|
||||
const PAGE_PARENTS: Record<string, string> = { 'install-theme': 'theme' }
|
||||
const PAGE_PARENTS: Record<string, string> = { 'generate-pet': 'pets', 'install-theme': 'theme' }
|
||||
|
||||
/** A nested page reachable from a root item via `to`. */
|
||||
interface PalettePage {
|
||||
@ -210,6 +214,7 @@ export function CommandPalette() {
|
||||
const pendingPage = useStore($commandPalettePage)
|
||||
const bindings = useStore($bindings)
|
||||
const navigate = useNavigate()
|
||||
const { requestGateway } = useGatewayRequest()
|
||||
const { availableThemes, resolvedMode, setMode, setTheme, themeName } = useTheme()
|
||||
const [search, setSearch] = useState('')
|
||||
const [page, setPage] = useState<string | null>(null)
|
||||
@ -245,13 +250,15 @@ export function CommandPalette() {
|
||||
const sessions = useMemo(() => (sessionsQuery.data?.sessions ?? []).map(toSessionEntry), [sessionsQuery.data])
|
||||
const archivedSessions = useMemo(() => (archivedQuery.data?.sessions ?? []).map(toSessionEntry), [archivedQuery.data])
|
||||
|
||||
// Reset the query/sub-page on close so it reopens clean.
|
||||
// Reset the query/sub-page on close so it reopens clean. Cleanup also deletes
|
||||
// a hatched-but-unadopted preview pet so it doesn't linger in the gallery.
|
||||
useEffect(() => {
|
||||
if (!open) {
|
||||
setSearch('')
|
||||
setPage(null)
|
||||
cleanupPetGen(requestGateway)
|
||||
}
|
||||
}, [open])
|
||||
}, [open, requestGateway])
|
||||
|
||||
// Deep-link into a nested page (e.g. `/pet list` → pets picker).
|
||||
useEffect(() => {
|
||||
@ -400,6 +407,13 @@ export function CommandPalette() {
|
||||
keywords: ['pet', 'petdex', 'mascot', 'pets', '/pet', 'paw'],
|
||||
label: cc.pets.title,
|
||||
to: 'pets'
|
||||
},
|
||||
{
|
||||
icon: Egg,
|
||||
id: 'appearance-generate-pet',
|
||||
keywords: ['pet', 'generate', 'create', 'make', 'new pet', 'mascot', 'hatch', 'ai'],
|
||||
label: cc.generatePet.title,
|
||||
to: 'generate-pet'
|
||||
}
|
||||
]
|
||||
},
|
||||
@ -574,6 +588,12 @@ export function CommandPalette() {
|
||||
placeholder: t.commandCenter.pets.placeholder,
|
||||
groups: []
|
||||
},
|
||||
// Server-driven page: describe → draft variants → hatch a custom pet.
|
||||
'generate-pet': {
|
||||
title: t.commandCenter.generatePet.title,
|
||||
placeholder: t.commandCenter.generatePet.placeholder,
|
||||
groups: []
|
||||
},
|
||||
// Server-driven page: items come from the Marketplace, rendered by
|
||||
// <MarketplaceThemePage> (loader + live search + per-row install).
|
||||
'install-theme': {
|
||||
@ -644,6 +664,21 @@ export function CommandPalette() {
|
||||
event.preventDefault()
|
||||
event.stopPropagation()
|
||||
goBack()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// On the generate page, Enter (re)generates from the typed
|
||||
// concept — cmdk has no item to select there, so each Enter,
|
||||
// including a retype after drafts already exist, starts a fresh
|
||||
// round. The page's own Retry/Hatch buttons cover the rest.
|
||||
if (page === 'generate-pet' && event.key === 'Enter' && search.trim()) {
|
||||
const genStatus = $petGenStatus.get()
|
||||
|
||||
if (genStatus !== 'generating' && genStatus !== 'hatching') {
|
||||
event.preventDefault()
|
||||
void generateDrafts(requestGateway, { prompt: search })
|
||||
}
|
||||
}
|
||||
}}
|
||||
onValueChange={setSearch}
|
||||
@ -653,8 +688,10 @@ export function CommandPalette() {
|
||||
/>
|
||||
<CommandList className="dt-portal-scrollbar max-h-[min(20rem,56vh)]">
|
||||
{/* Server-driven pages render their own list; the rest show groups. */}
|
||||
{page === 'pets' ? (
|
||||
<PetPalettePage search={search} />
|
||||
{page === 'generate-pet' ? (
|
||||
<PetGeneratePage search={search} />
|
||||
) : page === 'pets' ? (
|
||||
<PetPalettePage onGenerate={() => { setSearch(''); setPage('generate-pet') }} search={search} />
|
||||
) : page === 'install-theme' ? (
|
||||
<MarketplaceThemePage onPickTheme={setTheme} search={search} />
|
||||
) : (
|
||||
|
||||
275
apps/desktop/src/app/command-palette/pet-generate-page.tsx
Normal file
275
apps/desktop/src/app/command-palette/pet-generate-page.tsx
Normal file
@ -0,0 +1,275 @@
|
||||
/**
|
||||
* Cmd-K → Pets → "Generate" page — describe a pet, pick a draft, hatch it.
|
||||
*
|
||||
* A thin view over the `pet-generate` store. The palette search box doubles as
|
||||
* the concept prompt; this page renders the variant grid, the selection, the
|
||||
* retry/hatch actions, and the loading states. The store owns the two-step
|
||||
* `pet.generate` → `pet.hatch` flow.
|
||||
*/
|
||||
|
||||
import { useStore } from '@nanostores/react'
|
||||
import { useEffect, useState } from 'react'
|
||||
|
||||
import { useGatewayRequest } from '@/app/gateway/hooks/use-gateway-request'
|
||||
import { PetSprite } from '@/components/pet/pet-sprite'
|
||||
import { useI18n } from '@/i18n'
|
||||
import { triggerHaptic } from '@/lib/haptics'
|
||||
import { Check, Egg, Loader2, PawPrint, RefreshCw } from '@/lib/icons'
|
||||
import { cn } from '@/lib/utils'
|
||||
import { closeCommandPalette } from '@/store/command-palette'
|
||||
import { type PetInfo, type PetState } from '@/store/pet'
|
||||
import {
|
||||
$petGenDrafts,
|
||||
$petGenError,
|
||||
$petGenPreview,
|
||||
$petGenSelected,
|
||||
$petGenStatus,
|
||||
adoptHatched,
|
||||
discardHatched,
|
||||
generateDrafts,
|
||||
hatchSelected
|
||||
} from '@/store/pet-generate'
|
||||
|
||||
const VARIANT_COUNT = 4
|
||||
|
||||
// Fixed render scale for the preview so it's a predictable size regardless of
|
||||
// the user's configured `display.pet.scale`.
|
||||
const PREVIEW_SCALE = 0.7
|
||||
|
||||
// States the preview cycles through so the user sees every animation row.
|
||||
const PREVIEW_STATES: PetState[] = ['idle', 'wave', 'run', 'review', 'jump', 'failed']
|
||||
const PREVIEW_STATE_MS = 1500
|
||||
|
||||
interface PetGeneratePageProps {
|
||||
search: string
|
||||
}
|
||||
|
||||
export function PetGeneratePage({ search }: PetGeneratePageProps) {
|
||||
const { t } = useI18n()
|
||||
const copy = t.commandCenter.generatePet
|
||||
const { requestGateway } = useGatewayRequest()
|
||||
|
||||
const status = useStore($petGenStatus)
|
||||
const error = useStore($petGenError)
|
||||
const drafts = useStore($petGenDrafts)
|
||||
const selected = useStore($petGenSelected)
|
||||
const preview = useStore($petGenPreview)
|
||||
const [name, setName] = useState('')
|
||||
|
||||
const prompt = search.trim()
|
||||
const busy = status === 'generating' || status === 'hatching'
|
||||
|
||||
const generate = () => {
|
||||
if (prompt) {
|
||||
void generateDrafts(requestGateway, { prompt })
|
||||
}
|
||||
}
|
||||
|
||||
const hatch = () => {
|
||||
void hatchSelected(requestGateway, { name: name.trim() || prompt, prompt })
|
||||
}
|
||||
|
||||
const adopt = () => {
|
||||
void adoptHatched(requestGateway).then(out => {
|
||||
if (out.ok) {
|
||||
triggerHaptic('crisp')
|
||||
closeCommandPalette()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if (status === 'stale') {
|
||||
return <Status text={copy.staleBackend} tone="error" />
|
||||
}
|
||||
|
||||
// Hatching is slow (several grounded image generations) — own the whole pane.
|
||||
if (status === 'hatching') {
|
||||
return <Status icon={<Loader2 className="size-4 animate-spin" />} text={copy.hatching} />
|
||||
}
|
||||
|
||||
// Preview: play every animation row before the user commits.
|
||||
if ((status === 'preview' || status === 'adopting') && preview) {
|
||||
return (
|
||||
<HatchPreview
|
||||
adopting={status === 'adopting'}
|
||||
error={error}
|
||||
onAdopt={adopt}
|
||||
onDiscard={() => void discardHatched(requestGateway)}
|
||||
pet={preview}
|
||||
/>
|
||||
)
|
||||
}
|
||||
|
||||
const hasDrafts = drafts.length > 0
|
||||
const generating = status === 'generating'
|
||||
const cells = generating ? Array.from({ length: VARIANT_COUNT }, (_, i) => ({ index: i, dataUri: '' })) : drafts
|
||||
|
||||
return (
|
||||
<div className="flex flex-col gap-2 p-2">
|
||||
{error && <p className="px-1 text-[0.6875rem] text-(--ui-red)">{error}</p>}
|
||||
|
||||
{!hasDrafts && !generating && (
|
||||
<p className="px-1 py-1 text-xs text-muted-foreground">{prompt ? copy.readyHint : copy.promptHint}</p>
|
||||
)}
|
||||
|
||||
{(hasDrafts || generating) && (
|
||||
<div className="grid grid-cols-2 gap-2">
|
||||
{cells.map((draft, i) => {
|
||||
const isSelected = !generating && selected === draft.index
|
||||
|
||||
return (
|
||||
<button
|
||||
className={cn(
|
||||
'relative flex aspect-square items-center justify-center overflow-hidden rounded-lg border bg-(--ui-bg-quinary) transition-colors',
|
||||
isSelected
|
||||
? 'border-(--ui-accent) ring-2 ring-(--ui-accent)/40'
|
||||
: 'border-(--ui-stroke-tertiary) hover:border-foreground/40'
|
||||
)}
|
||||
disabled={generating || busy}
|
||||
key={generating ? i : draft.index}
|
||||
onClick={() => $petGenSelected.set(draft.index)}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
{generating ? (
|
||||
<Loader2 className="size-5 animate-spin text-muted-foreground" />
|
||||
) : (
|
||||
<img alt="" className="size-full object-contain" draggable={false} src={draft.dataUri} />
|
||||
)}
|
||||
{isSelected && (
|
||||
<span className="absolute right-1 top-1 rounded-full bg-(--ui-accent) p-0.5 text-(--ui-base)">
|
||||
<Check className="size-3" />
|
||||
</span>
|
||||
)}
|
||||
</button>
|
||||
)
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{hasDrafts ? (
|
||||
<div className="flex flex-col gap-2">
|
||||
<input
|
||||
className="w-full rounded-md border border-(--ui-stroke-tertiary) bg-transparent px-2 py-1.5 text-xs outline-none placeholder:text-muted-foreground focus:border-foreground/40"
|
||||
onChange={event => setName(event.target.value)}
|
||||
onKeyDown={event => {
|
||||
if (event.key === 'Enter') {
|
||||
event.preventDefault()
|
||||
hatch()
|
||||
}
|
||||
}}
|
||||
placeholder={copy.namePlaceholder}
|
||||
value={name}
|
||||
/>
|
||||
<div className="flex gap-2">
|
||||
<button
|
||||
className="flex flex-1 items-center justify-center gap-1.5 rounded-md border border-border px-2 py-1.5 text-xs font-medium transition-colors hover:bg-(--chrome-action-hover) disabled:opacity-50"
|
||||
disabled={busy || !prompt}
|
||||
onClick={generate}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
<RefreshCw className="size-3.5" />
|
||||
{copy.retry}
|
||||
</button>
|
||||
<button
|
||||
className="flex flex-1 items-center justify-center gap-1.5 rounded-md bg-primary px-2 py-1.5 text-xs font-medium text-primary-foreground transition-opacity hover:opacity-90 disabled:opacity-50"
|
||||
disabled={busy || selected === null}
|
||||
onClick={hatch}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
<PawPrint className="size-3.5" />
|
||||
{copy.hatch}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<button
|
||||
className="flex items-center justify-center gap-1.5 rounded-md bg-primary px-2 py-2 text-xs font-medium text-primary-foreground transition-opacity hover:opacity-90 disabled:opacity-50"
|
||||
disabled={busy || !prompt}
|
||||
onClick={generate}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
{generating ? <Loader2 className="size-3.5 animate-spin" /> : <Egg className="size-3.5" />}
|
||||
{generating ? copy.generating : copy.generate}
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
interface HatchPreviewProps {
|
||||
pet: PetInfo
|
||||
adopting: boolean
|
||||
error: string | null
|
||||
onAdopt: () => void
|
||||
onDiscard: () => void
|
||||
}
|
||||
|
||||
function HatchPreview({ pet, adopting, error, onAdopt, onDiscard }: HatchPreviewProps) {
|
||||
const { t } = useI18n()
|
||||
const copy = t.commandCenter.generatePet
|
||||
const [stateIndex, setStateIndex] = useState(0)
|
||||
|
||||
// Cycle through the animation rows so the preview showcases all frames.
|
||||
useEffect(() => {
|
||||
const id = setInterval(() => {
|
||||
setStateIndex(i => (i + 1) % PREVIEW_STATES.length)
|
||||
}, PREVIEW_STATE_MS)
|
||||
|
||||
return () => clearInterval(id)
|
||||
}, [])
|
||||
|
||||
const previewInfo: PetInfo = { ...pet, scale: PREVIEW_SCALE }
|
||||
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-2 p-2">
|
||||
<div className="flex min-h-[9rem] w-full items-center justify-center rounded-lg border border-(--ui-stroke-tertiary) bg-(--ui-bg-quinary) py-2">
|
||||
<PetSprite info={previewInfo} stateOverride={PREVIEW_STATES[stateIndex]} />
|
||||
</div>
|
||||
|
||||
{pet.displayName && <p className="text-xs font-medium text-foreground">{pet.displayName}</p>}
|
||||
|
||||
{error && <p className="px-1 text-[0.6875rem] text-(--ui-red)">{error}</p>}
|
||||
|
||||
<div className="flex w-full gap-2">
|
||||
<button
|
||||
className="flex flex-1 items-center justify-center gap-1.5 rounded-md border border-border px-2 py-1.5 text-xs font-medium transition-colors hover:bg-(--chrome-action-hover) disabled:opacity-50"
|
||||
disabled={adopting}
|
||||
onClick={onDiscard}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
<RefreshCw className="size-3.5" />
|
||||
{copy.startOver}
|
||||
</button>
|
||||
<button
|
||||
className="flex flex-1 items-center justify-center gap-1.5 rounded-md bg-primary px-2 py-1.5 text-xs font-medium text-primary-foreground transition-opacity hover:opacity-90 disabled:opacity-50"
|
||||
disabled={adopting}
|
||||
onClick={onAdopt}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
{adopting ? <Loader2 className="size-3.5 animate-spin" /> : <PawPrint className="size-3.5" />}
|
||||
{copy.adopt}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
function Status({ icon, text, tone }: { icon?: React.ReactNode; text: string; tone?: 'error' }) {
|
||||
return (
|
||||
<div
|
||||
className={cn(
|
||||
'flex items-center justify-center gap-2 px-2 py-6 text-xs',
|
||||
tone === 'error' ? 'text-(--ui-red)' : 'text-muted-foreground'
|
||||
)}
|
||||
>
|
||||
{icon}
|
||||
{text}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@ -15,7 +15,7 @@ import { useGatewayRequest } from '@/app/gateway/hooks/use-gateway-request'
|
||||
import { PetThumb } from '@/components/pet/pet-thumb'
|
||||
import { useI18n } from '@/i18n'
|
||||
import { triggerHaptic } from '@/lib/haptics'
|
||||
import { Check, Loader2, PawPrint } from '@/lib/icons'
|
||||
import { Check, Egg, Loader2, PawPrint } from '@/lib/icons'
|
||||
import { cn } from '@/lib/utils'
|
||||
import {
|
||||
$petBusy,
|
||||
@ -31,9 +31,11 @@ import {
|
||||
|
||||
interface PetPalettePageProps {
|
||||
search: string
|
||||
/** Navigate to the "generate a pet" page (rendered as a header action). */
|
||||
onGenerate?: () => void
|
||||
}
|
||||
|
||||
export function PetPalettePage({ search }: PetPalettePageProps) {
|
||||
export function PetPalettePage({ search, onGenerate }: PetPalettePageProps) {
|
||||
const { t } = useI18n()
|
||||
const copy = t.commandCenter.pets
|
||||
const { requestGateway } = useGatewayRequest()
|
||||
@ -72,6 +74,24 @@ export function PetPalettePage({ search }: PetPalettePageProps) {
|
||||
|
||||
return (
|
||||
<div role="listbox">
|
||||
{onGenerate && (
|
||||
<button
|
||||
className={cn(
|
||||
'flex w-full items-center gap-2 rounded-md text-left text-foreground transition-colors hover:bg-(--chrome-action-hover)',
|
||||
HUD_ITEM,
|
||||
HUD_TEXT
|
||||
)}
|
||||
onClick={onGenerate}
|
||||
onMouseDown={event => event.preventDefault()}
|
||||
type="button"
|
||||
>
|
||||
<span className="flex size-8 shrink-0 items-center justify-center rounded-md bg-(--chrome-action-hover)">
|
||||
<Egg className="size-4" />
|
||||
</span>
|
||||
<span className="font-medium">{t.commandCenter.generatePet.title}</span>
|
||||
</button>
|
||||
)}
|
||||
|
||||
{error && <p className="px-2 pb-1 pt-1.5 text-[0.6875rem] text-(--ui-red)">{error}</p>}
|
||||
|
||||
{shown.length === 0 ? (
|
||||
|
||||
@ -15,6 +15,12 @@ interface PetSpriteProps {
|
||||
info: PetInfo
|
||||
/** On-screen scale multiplier applied on top of the pet's native scale. */
|
||||
zoom?: number
|
||||
/**
|
||||
* Force a specific animation state instead of reading the live `$petState`.
|
||||
* Used by the generate-flow preview to showcase every row without driving (or
|
||||
* being driven by) the real agent activity that moves the floating mascot.
|
||||
*/
|
||||
stateOverride?: PetState
|
||||
}
|
||||
|
||||
/**
|
||||
@ -28,9 +34,15 @@ interface PetSpriteProps {
|
||||
* with `memo`, this component effectively never re-renders after mount until
|
||||
* the pet itself changes.
|
||||
*/
|
||||
function PetSpriteImpl({ info, zoom = 1 }: PetSpriteProps) {
|
||||
function PetSpriteImpl({ info, zoom = 1, stateOverride }: PetSpriteProps) {
|
||||
const canvasRef = useRef<HTMLCanvasElement | null>(null)
|
||||
const stateRef = useRef<PetState>($petState.get())
|
||||
const overrideRef = useRef<PetState | undefined>(stateOverride)
|
||||
|
||||
// Keep the override current without re-running the RAF setup effect.
|
||||
useEffect(() => {
|
||||
overrideRef.current = stateOverride
|
||||
}, [stateOverride])
|
||||
|
||||
const frameW = info.frameW ?? DEFAULT_FRAME_W
|
||||
const frameH = info.frameH ?? DEFAULT_FRAME_H
|
||||
@ -91,6 +103,7 @@ function PetSpriteImpl({ info, zoom = 1 }: PetSpriteProps) {
|
||||
// than flashing blank padding.
|
||||
const resolve = (s: PetState): { row: number; count: number } => {
|
||||
const real = framesByState?.[s] ?? frames
|
||||
|
||||
if (real > 0) {
|
||||
return { row: rowIndex(s), count: real }
|
||||
}
|
||||
@ -99,7 +112,7 @@ function PetSpriteImpl({ info, zoom = 1 }: PetSpriteProps) {
|
||||
}
|
||||
|
||||
const render = (now: number) => {
|
||||
const { row, count } = resolve(stateRef.current)
|
||||
const { row, count } = resolve(overrideRef.current ?? stateRef.current)
|
||||
// Per-state step keeps every state's loop ~loopMs even when frame counts
|
||||
// differ; counts vary per row so derive the cadence here, not once.
|
||||
const stepMs = loopMs / count
|
||||
|
||||
@ -751,6 +751,21 @@ export const en: Translations = {
|
||||
toggleFailed: 'Could not toggle the pet.',
|
||||
noneAvailable: 'No pets available — pick one below to install.'
|
||||
},
|
||||
generatePet: {
|
||||
title: 'Generate a pet',
|
||||
placeholder: 'Describe a pet to generate…',
|
||||
promptHint: 'Type a description, then press Enter to draft four looks.',
|
||||
readyHint: 'Press Enter to draft four looks from your description.',
|
||||
generate: 'Generate',
|
||||
generating: 'Generating…',
|
||||
retry: 'Retry',
|
||||
hatch: 'Hatch',
|
||||
hatching: 'Hatching your pet…',
|
||||
namePlaceholder: 'Name your pet',
|
||||
staleBackend: 'Update Hermes to generate pets.',
|
||||
adopt: 'Adopt',
|
||||
startOver: 'Start over'
|
||||
},
|
||||
installTheme: {
|
||||
title: 'Install theme...',
|
||||
placeholder: 'Search the VS Code Marketplace...',
|
||||
|
||||
@ -875,6 +875,21 @@ export const ja = defineLocale({
|
||||
toggleFailed: 'ペットを切り替えできませんでした。',
|
||||
noneAvailable: '利用可能なペットがありません。'
|
||||
},
|
||||
generatePet: {
|
||||
title: 'ペットを生成',
|
||||
placeholder: '生成するペットを説明…',
|
||||
promptHint: '説明を入力して Enter を押すと、4 つの見た目を生成します。',
|
||||
readyHint: 'Enter を押すと、説明から 4 つの見た目を生成します。',
|
||||
generate: '生成',
|
||||
generating: '生成中…',
|
||||
retry: '再試行',
|
||||
hatch: '孵化',
|
||||
hatching: 'ペットを孵化しています…',
|
||||
namePlaceholder: 'ペットに名前を付ける',
|
||||
staleBackend: 'ペットを生成するには Hermes を更新してください。',
|
||||
adopt: '迎え入れる',
|
||||
startOver: 'やり直す'
|
||||
},
|
||||
installTheme: {
|
||||
title: 'テーマをインストール...',
|
||||
placeholder: 'VS Code Marketplace を検索...',
|
||||
|
||||
@ -627,6 +627,21 @@ export interface Translations {
|
||||
toggleFailed: string
|
||||
noneAvailable: string
|
||||
}
|
||||
generatePet: {
|
||||
title: string
|
||||
placeholder: string
|
||||
promptHint: string
|
||||
readyHint: string
|
||||
generate: string
|
||||
generating: string
|
||||
retry: string
|
||||
hatch: string
|
||||
hatching: string
|
||||
namePlaceholder: string
|
||||
staleBackend: string
|
||||
adopt: string
|
||||
startOver: string
|
||||
}
|
||||
installTheme: {
|
||||
title: string
|
||||
placeholder: string
|
||||
|
||||
@ -846,6 +846,21 @@ export const zhHant = defineLocale({
|
||||
toggleFailed: '無法切換寵物顯示。',
|
||||
noneAvailable: '尚無可用寵物——請在下方選擇一個安裝。'
|
||||
},
|
||||
generatePet: {
|
||||
title: '生成寵物',
|
||||
placeholder: '描述要生成的寵物……',
|
||||
promptHint: '輸入描述，然後按 Enter 生成四種造型。',
|
||||
readyHint: '按 Enter 依描述生成四種造型。',
|
||||
generate: '生成',
|
||||
generating: '生成中……',
|
||||
retry: '重試',
|
||||
hatch: '孵化',
|
||||
hatching: '正在孵化你的寵物……',
|
||||
namePlaceholder: '為寵物命名',
|
||||
staleBackend: '請更新 Hermes 以生成寵物。',
|
||||
adopt: '領養',
|
||||
startOver: '重新開始'
|
||||
},
|
||||
installTheme: {
|
||||
title: '安裝主題...',
|
||||
placeholder: '搜尋 VS Code Marketplace...',
|
||||
|
||||
@ -939,6 +939,21 @@ export const zh: Translations = {
|
||||
toggleFailed: '无法切换宠物显示。',
|
||||
noneAvailable: '暂无可用宠物——请在下方选择一个安装。'
|
||||
},
|
||||
generatePet: {
|
||||
title: '生成宠物',
|
||||
placeholder: '描述要生成的宠物……',
|
||||
promptHint: '输入描述，然后按 Enter 生成四种造型。',
|
||||
readyHint: '按 Enter 根据描述生成四种造型。',
|
||||
generate: '生成',
|
||||
generating: '生成中……',
|
||||
retry: '重试',
|
||||
hatch: '孵化',
|
||||
hatching: '正在孵化你的宠物……',
|
||||
namePlaceholder: '给宠物起个名字',
|
||||
staleBackend: '请更新 Hermes 以生成宠物。',
|
||||
adopt: '领养',
|
||||
startOver: '重新开始'
|
||||
},
|
||||
installTheme: {
|
||||
title: '安装主题...',
|
||||
placeholder: '搜索 VS Code Marketplace...',
|
||||
|
||||
@ -29,6 +29,7 @@ import {
|
||||
IconCopy as CopyIcon,
|
||||
IconCpu as Cpu,
|
||||
IconDownload as Download,
|
||||
IconEgg as Egg,
|
||||
IconExternalLink as ExternalLink,
|
||||
IconEye as Eye,
|
||||
IconEyeOff as EyeOff,
|
||||
@ -132,6 +133,7 @@ export {
|
||||
CopyIcon,
|
||||
Cpu,
|
||||
Download,
|
||||
Egg,
|
||||
ExternalLink,
|
||||
Eye,
|
||||
EyeOff,
|
||||
|
||||
229
apps/desktop/src/store/pet-generate.ts
Normal file
229
apps/desktop/src/store/pet-generate.ts
Normal file
@ -0,0 +1,229 @@
|
||||
import { atom } from 'nanostores'
|
||||
|
||||
import { type PetInfo } from '@/store/pet'
|
||||
import { type GatewayRequest, loadPetGallery } from '@/store/pet-gallery'
|
||||
|
||||
/**
|
||||
* Feature store for the "generate a pet" flow (Cmd-K → Pets → Generate).
|
||||
*
|
||||
* Three backend steps, mirrored as state here:
|
||||
* - `pet.generate` produces N cheap base-look *drafts* keyed by a `token`.
|
||||
* - `pet.hatch` turns the chosen draft into a full animated pet — installed but
|
||||
* NOT active — and returns its renderer payload so we can preview all frames.
|
||||
* - the user then *adopts* (`pet.select`) or *discards* (`pet.remove`) it.
|
||||
*
|
||||
* The store owns the draft set, the selected variant, the hatched preview, and
|
||||
* the busy/error status so the page is a thin view. Retry == regenerate (new
|
||||
* token). Kept separate from `pet-gallery` because its lifecycle (ephemeral
|
||||
* drafts + an unadopted preview) is unrelated to the long-lived gallery cache.
|
||||
*/
|
||||
|
||||
export interface PetDraft {
|
||||
index: number
|
||||
/** Downscaled PNG data URI preview from the gateway. */
|
||||
dataUri: string
|
||||
}
|
||||
|
||||
export type PetGenStatus =
|
||||
| 'idle'
|
||||
| 'generating'
|
||||
| 'ready'
|
||||
| 'hatching'
|
||||
| 'preview'
|
||||
| 'adopting'
|
||||
| 'error'
|
||||
| 'stale'
|
||||
|
||||
export const $petGenStatus = atom<PetGenStatus>('idle')
|
||||
export const $petGenError = atom<string | null>(null)
|
||||
export const $petGenToken = atom<string | null>(null)
|
||||
export const $petGenDrafts = atom<PetDraft[]>([])
|
||||
export const $petGenSelected = atom<number | null>(null)
|
||||
/** The hatched-but-unadopted pet: its renderer payload, played in the preview. */
|
||||
export const $petGenPreview = atom<PetInfo | null>(null)
|
||||
|
||||
function isMissingMethod(error: unknown): boolean {
|
||||
const message = error instanceof Error ? error.message : String(error)
|
||||
|
||||
return /method not found|-32601|unknown method|no such method/i.test(message)
|
||||
}
|
||||
|
||||
/** Clear all generation state (on close, or before a fresh run). */
|
||||
export function resetPetGen(): void {
|
||||
$petGenStatus.set('idle')
|
||||
$petGenError.set(null)
|
||||
$petGenToken.set(null)
|
||||
$petGenDrafts.set([])
|
||||
$petGenSelected.set(null)
|
||||
$petGenPreview.set(null)
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset on palette close, deleting an unadopted preview pet first so a hatched-
|
||||
* but-never-adopted creature doesn't linger in the gallery. Fire-and-forget.
|
||||
*/
|
||||
export function cleanupPetGen(request: GatewayRequest): void {
|
||||
const preview = $petGenPreview.get()
|
||||
|
||||
if ($petGenStatus.get() === 'preview' && preview?.slug) {
|
||||
void request('pet.remove', { slug: preview.slug }).catch(() => {})
|
||||
}
|
||||
|
||||
resetPetGen()
|
||||
}
|
||||
|
||||
interface GenerateOptions {
|
||||
prompt: string
|
||||
style?: string
|
||||
count?: number
|
||||
}
|
||||
|
||||
/** Generate (or retry) a fresh set of base-look drafts for `prompt`. */
|
||||
export async function generateDrafts(request: GatewayRequest, options: GenerateOptions): Promise<boolean> {
|
||||
const prompt = options.prompt.trim()
|
||||
|
||||
if (!prompt) {
|
||||
return false
|
||||
}
|
||||
|
||||
$petGenStatus.set('generating')
|
||||
$petGenError.set(null)
|
||||
$petGenDrafts.set([])
|
||||
$petGenSelected.set(null)
|
||||
|
||||
try {
|
||||
const result = await request<{ ok: boolean; token: string; drafts: PetDraft[] }>('pet.generate', {
|
||||
prompt,
|
||||
style: options.style ?? 'auto',
|
||||
count: options.count ?? 4
|
||||
})
|
||||
|
||||
if (!result?.ok || !result.drafts?.length) {
|
||||
throw new Error('generation produced no drafts')
|
||||
}
|
||||
|
||||
$petGenToken.set(result.token)
|
||||
$petGenDrafts.set(result.drafts)
|
||||
$petGenSelected.set(result.drafts[0]?.index ?? 0)
|
||||
$petGenStatus.set('ready')
|
||||
|
||||
return true
|
||||
} catch (e) {
|
||||
if (isMissingMethod(e)) {
|
||||
$petGenStatus.set('stale')
|
||||
} else {
|
||||
$petGenStatus.set('error')
|
||||
$petGenError.set(e instanceof Error ? e.message : 'Could not generate pet drafts.')
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
interface HatchOptions {
|
||||
name: string
|
||||
description?: string
|
||||
prompt?: string
|
||||
style?: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Hatch the selected draft into a full pet (installed but NOT yet active) and
|
||||
* load its renderer payload into the preview. Adoption is a separate, explicit
|
||||
* step (`adoptHatched`) so the user sees every frame play before committing.
|
||||
* Returns true when the preview is ready.
|
||||
*/
|
||||
export async function hatchSelected(request: GatewayRequest, options: HatchOptions): Promise<boolean> {
|
||||
const token = $petGenToken.get()
|
||||
const index = $petGenSelected.get()
|
||||
const name = options.name.trim()
|
||||
|
||||
if (token === null || index === null || !name) {
|
||||
return false
|
||||
}
|
||||
|
||||
$petGenStatus.set('hatching')
|
||||
$petGenError.set(null)
|
||||
|
||||
try {
|
||||
const result = await request<{ ok: boolean; slug: string; displayName: string; pet?: PetInfo }>('pet.hatch', {
|
||||
token,
|
||||
index,
|
||||
name,
|
||||
description: options.description ?? '',
|
||||
prompt: options.prompt ?? name,
|
||||
style: options.style ?? 'auto'
|
||||
})
|
||||
|
||||
if (!result?.ok || !result.pet?.spritesheetBase64) {
|
||||
throw new Error('hatch produced no preview')
|
||||
}
|
||||
|
||||
$petGenPreview.set({ ...result.pet, enabled: true })
|
||||
$petGenStatus.set('preview')
|
||||
|
||||
return true
|
||||
} catch (e) {
|
||||
$petGenStatus.set('error')
|
||||
$petGenError.set(e instanceof Error ? e.message : 'Could not hatch the pet.')
|
||||
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
export interface AdoptOutcome {
|
||||
ok: boolean
|
||||
slug?: string
|
||||
displayName?: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Adopt the previewed pet: activate it (`pet.select`), refresh the gallery + live
|
||||
* mascot, and clear generation state. No-op unless a preview exists.
|
||||
*/
|
||||
export async function adoptHatched(request: GatewayRequest): Promise<AdoptOutcome> {
|
||||
const preview = $petGenPreview.get()
|
||||
|
||||
if (!preview?.slug) {
|
||||
return { ok: false }
|
||||
}
|
||||
|
||||
$petGenStatus.set('adopting')
|
||||
$petGenError.set(null)
|
||||
|
||||
try {
|
||||
const result = await request<{ ok: boolean; slug: string; displayName: string }>('pet.select', {
|
||||
slug: preview.slug
|
||||
})
|
||||
|
||||
if (!result?.ok) {
|
||||
throw new Error('adopt failed')
|
||||
}
|
||||
|
||||
await loadPetGallery(request, { force: true })
|
||||
resetPetGen()
|
||||
|
||||
return { ok: true, slug: result.slug, displayName: result.displayName }
|
||||
} catch (e) {
|
||||
$petGenStatus.set('preview')
|
||||
$petGenError.set(e instanceof Error ? e.message : 'Could not adopt the pet.')
|
||||
|
||||
return { ok: false }
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Throw away the previewed pet (`pet.remove`) and return to the draft picker so
|
||||
* the user can choose another base or regenerate. Best-effort on the delete.
|
||||
*/
|
||||
export async function discardHatched(request: GatewayRequest): Promise<void> {
|
||||
const preview = $petGenPreview.get()
|
||||
|
||||
if (preview?.slug) {
|
||||
await request('pet.remove', { slug: preview.slug }).catch(() => {})
|
||||
}
|
||||
|
||||
$petGenPreview.set(null)
|
||||
$petGenError.set(null)
|
||||
$petGenStatus.set($petGenDrafts.get().length > 0 ? 'ready' : 'idle')
|
||||
}
|
||||
@ -171,6 +171,29 @@ class OpenAIImageGenProvider(ImageGenProvider):
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _edit(client: Any, payload: Dict[str, Any], reference_images: List[str]) -> Any:
|
||||
"""Run gpt-image-2 image *edit* grounded on reference image(s).
|
||||
|
||||
``images.edit`` keeps the output anchored to the supplied character,
|
||||
which is what makes per-state pet rows stay the same creature. Opens the
|
||||
reference files as binary handles and forwards the same size/quality/
|
||||
background knobs as ``images.generate``.
|
||||
"""
|
||||
handles = [open(path, "rb") for path in reference_images]
|
||||
try:
|
||||
edit_payload = dict(payload)
|
||||
# The edit endpoint takes one image or a list; a single handle is
|
||||
# the broadly-supported shape.
|
||||
edit_payload["image"] = handles if len(handles) > 1 else handles[0]
|
||||
return client.images.edit(**edit_payload)
|
||||
finally:
|
||||
for handle in handles:
|
||||
try:
|
||||
handle.close()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
@ -223,9 +246,20 @@ class OpenAIImageGenProvider(ImageGenProvider):
|
||||
"quality": meta["quality"],
|
||||
}
|
||||
|
||||
# Optional sprite-oriented extras (used by pet generation; ignored by
|
||||
# the prompt-only ``image_generate`` tool). ``background=transparent``
|
||||
# asks gpt-image for a cutout; ``reference_images`` routes to the image
|
||||
# *edit* endpoint so the output stays grounded on a base character.
|
||||
if str(kwargs.get("background", "")).lower() == "transparent":
|
||||
payload["background"] = "transparent"
|
||||
reference_images = kwargs.get("reference_images") or []
|
||||
|
||||
try:
|
||||
client = openai.OpenAI()
|
||||
response = client.images.generate(**payload)
|
||||
if reference_images:
|
||||
response = self._edit(client, payload, reference_images)
|
||||
else:
|
||||
response = client.images.generate(**payload)
|
||||
except Exception as exc:
|
||||
logger.debug("OpenAI image generation failed", exc_info=True)
|
||||
return error_response(
|
||||
|
||||
309
tests/agent/test_pet_generate.py
Normal file
309
tests/agent/test_pet_generate.py
Normal file
@ -0,0 +1,309 @@
|
||||
"""Tests for pet generation: deterministic atlas ops, store register, orchestration.
|
||||
|
||||
No network/API calls — image generation is mocked with synthetic strips so the
|
||||
whole pipeline (segmentation → compose → validate → register → adopt) is
|
||||
exercised hermetically.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.pet.generate import atlas
|
||||
|
||||
PIL = pytest.importorskip("PIL")
|
||||
from PIL import Image, ImageDraw # noqa: E402
|
||||
|
||||
|
||||
def _strip(n_blobs: int, *, transparent: bool = True, bg=(0, 255, 0, 255), size=(208, 208)) -> Image.Image:
|
||||
"""A horizontal strip with *n_blobs* clearly-separated colored ellipses."""
|
||||
w = size[0] * n_blobs
|
||||
h = size[1]
|
||||
base = (0, 0, 0, 0) if transparent else bg
|
||||
img = Image.new("RGBA", (w, h), base)
|
||||
draw = ImageDraw.Draw(img)
|
||||
for i in range(n_blobs):
|
||||
cx = i * size[0] + size[0] // 2
|
||||
cy = h // 2
|
||||
r = size[0] // 3
|
||||
color = (40 + i * 30 % 200, 80, 200 - i * 20 % 180, 255)
|
||||
draw.ellipse((cx - r, cy - r, cx + r, cy + r), fill=color)
|
||||
return img
|
||||
|
||||
|
||||
# ───────────────────────── frame extraction ─────────────────────────
|
||||
|
||||
|
||||
def test_extract_strip_frames_transparent_returns_centered_cells():
|
||||
frames = atlas.extract_strip_frames(_strip(6), 6)
|
||||
assert len(frames) == 6
|
||||
for frame in frames:
|
||||
assert frame.size == (atlas.CELL_WIDTH, atlas.CELL_HEIGHT)
|
||||
# Background corners must be transparent.
|
||||
assert frame.getpixel((0, 0))[3] == 0
|
||||
# Something is drawn.
|
||||
assert frame.getchannel("A").getextrema()[1] > 0
|
||||
|
||||
|
||||
def test_extract_strip_frames_keys_out_solid_background():
|
||||
frames = atlas.extract_strip_frames(_strip(4, transparent=False), 4)
|
||||
assert len(frames) == 4
|
||||
# The green backdrop must be gone (corner transparent).
|
||||
assert frames[0].getpixel((0, 0))[3] == 0
|
||||
|
||||
|
||||
def test_extract_strip_frames_slot_fallback_when_unsegmentable():
|
||||
# A single connected smear can't be split into 5 components → slot fallback.
|
||||
img = Image.new("RGBA", (200 * 5, 208), (0, 0, 0, 0))
|
||||
ImageDraw.Draw(img).rectangle((0, 80, 200 * 5 - 1, 120), fill=(200, 50, 50, 255))
|
||||
frames = atlas.extract_strip_frames(img, 5, method="auto")
|
||||
assert len(frames) == 5
|
||||
|
||||
|
||||
def test_extract_components_method_raises_when_too_few():
|
||||
img = Image.new("RGBA", (400, 208), (0, 0, 0, 0))
|
||||
ImageDraw.Draw(img).ellipse((10, 10, 100, 100), fill=(255, 0, 0, 255))
|
||||
with pytest.raises(ValueError):
|
||||
atlas.extract_strip_frames(img, 6, method="components")
|
||||
|
||||
|
||||
# ───────────────────────── atlas compose / validate ─────────────────────────
|
||||
|
||||
|
||||
def _frames_for_all_states() -> dict[str, list]:
|
||||
out: dict[str, list] = {}
|
||||
for state, _row, count in atlas.ROW_SPECS:
|
||||
out[state] = atlas.extract_strip_frames(_strip(count), count)
|
||||
return out
|
||||
|
||||
|
||||
def test_compose_atlas_geometry_and_validation():
|
||||
sheet = atlas.compose_atlas(_frames_for_all_states())
|
||||
assert sheet.size == (atlas.ATLAS_WIDTH, atlas.ATLAS_HEIGHT)
|
||||
result = atlas.validate_atlas(sheet)
|
||||
assert result["ok"], result["errors"]
|
||||
assert set(result["filled_states"]) == {s for s, _, _ in atlas.ROW_SPECS}
|
||||
|
||||
|
||||
def test_compose_atlas_leaves_unused_tail_transparent():
|
||||
# wave has 4 frames; columns 4 and 5 of its row must be transparent.
|
||||
sheet = atlas.compose_atlas(_frames_for_all_states())
|
||||
wave_row = next(r for s, r, _ in atlas.ROW_SPECS if s == "wave")
|
||||
top = wave_row * atlas.CELL_HEIGHT
|
||||
for col in (4, 5):
|
||||
left = col * atlas.CELL_WIDTH
|
||||
cell = sheet.crop((left, top, left + atlas.CELL_WIDTH, top + atlas.CELL_HEIGHT))
|
||||
assert cell.getchannel("A").getextrema()[1] == 0
|
||||
|
||||
|
||||
def test_validate_atlas_rejects_wrong_size():
|
||||
bad = Image.new("RGBA", (100, 100), (0, 0, 0, 0))
|
||||
result = atlas.validate_atlas(bad)
|
||||
assert not result["ok"]
|
||||
assert any("expected" in e for e in result["errors"])
|
||||
|
||||
|
||||
def test_validate_atlas_rejects_rgb_residue():
|
||||
sheet = atlas.compose_atlas(_frames_for_all_states())
|
||||
# Poke a fully-transparent pixel with non-zero RGB.
|
||||
sheet.putpixel((0, 0), (120, 0, 0, 0))
|
||||
result = atlas.validate_atlas(sheet)
|
||||
assert not result["ok"]
|
||||
assert any("residue" in e for e in result["errors"])
|
||||
|
||||
|
||||
def test_validate_atlas_warns_on_empty_state():
|
||||
frames = _frames_for_all_states()
|
||||
frames["jump"] = []
|
||||
sheet = atlas.compose_atlas(frames)
|
||||
result = atlas.validate_atlas(sheet)
|
||||
assert result["ok"] # one empty row is a warning, not an error
|
||||
assert any("jump" in w for w in result["warnings"])
|
||||
|
||||
|
||||
def test_single_frame_fits_cell():
|
||||
frame = atlas.single_frame(_strip(1))
|
||||
assert frame.size == (atlas.CELL_WIDTH, atlas.CELL_HEIGHT)
|
||||
assert frame.getchannel("A").getextrema()[1] > 0
|
||||
|
||||
|
||||
# ───────────────────────── store register / adopt ─────────────────────────
|
||||
|
||||
|
||||
def test_slugify_and_unique_slug():
|
||||
from agent.pet import store
|
||||
|
||||
assert store.slugify("My Cool Pet!") == "my-cool-pet"
|
||||
assert store.slugify(" ") == "pet"
|
||||
first = store.unique_slug("Robo")
|
||||
(store.pets_dir() / first).mkdir(parents=True)
|
||||
assert store.unique_slug("Robo") == "robo-2"
|
||||
|
||||
|
||||
def test_register_local_pet_appears_and_is_adoptable():
|
||||
from agent.pet import store
|
||||
|
||||
sheet = atlas.compose_atlas(_frames_for_all_states())
|
||||
pet = store.register_local_pet(sheet, slug="Sparky", display_name="Sparky", description="zappy")
|
||||
assert pet.slug == "sparky"
|
||||
assert pet.exists
|
||||
assert any(p.slug == "sparky" for p in store.installed_pets())
|
||||
|
||||
# install_pet returns the on-disk pet without ever hitting the manifest.
|
||||
adopted = store.install_pet("sparky")
|
||||
assert adopted.slug == "sparky"
|
||||
assert adopted.display_name == "Sparky"
|
||||
|
||||
|
||||
def test_register_local_pet_accepts_bytes():
|
||||
from agent.pet import store
|
||||
|
||||
sheet = atlas.compose_atlas(_frames_for_all_states())
|
||||
data = atlas.atlas_to_webp_bytes(sheet)
|
||||
pet = store.register_local_pet(data, slug="bytey")
|
||||
assert pet.exists
|
||||
|
||||
|
||||
# ───────────────────────── orchestration (mocked imagegen) ─────────────────────────
|
||||
|
||||
|
||||
def test_generate_base_drafts_returns_n(monkeypatch, tmp_path):
|
||||
from agent.pet.generate import imagegen, orchestrate
|
||||
|
||||
calls = {"n": 0}
|
||||
|
||||
def fake_generate(prompt, *, n=1, reference_images=None, provider=None, prefix="pet"):
|
||||
paths = []
|
||||
for i in range(n):
|
||||
calls["n"] += 1
|
||||
p = tmp_path / f"{prefix}_{calls['n']}.png"
|
||||
_strip(1).save(p)
|
||||
paths.append(p)
|
||||
return paths
|
||||
|
||||
monkeypatch.setattr(imagegen, "resolve_provider", lambda **_: object())
|
||||
monkeypatch.setattr(imagegen, "generate", fake_generate)
|
||||
|
||||
drafts = orchestrate.generate_base_drafts("a fox", n=4)
|
||||
assert len(drafts) == 4
|
||||
|
||||
|
||||
def test_generate_base_drafts_hardens_opaque_background(monkeypatch, tmp_path):
|
||||
"""A provider that ignores background=transparent still yields a cutout."""
|
||||
from agent.pet.generate import imagegen, orchestrate
|
||||
|
||||
def fake_generate(prompt, *, n=1, reference_images=None, provider=None, prefix="pet"):
|
||||
# Solid-green backdrop with a blob — i.e. the provider painted a backdrop.
|
||||
p = tmp_path / f"{prefix}_opaque.png"
|
||||
_strip(1, transparent=False, bg=(0, 255, 0, 255)).save(p)
|
||||
return [p]
|
||||
|
||||
monkeypatch.setattr(imagegen, "resolve_provider", lambda **_: object())
|
||||
monkeypatch.setattr(imagegen, "generate", fake_generate)
|
||||
|
||||
drafts = orchestrate.generate_base_drafts("a fox", n=1)
|
||||
assert len(drafts) == 1
|
||||
|
||||
with Image.open(drafts[0]) as out:
|
||||
rgba = out.convert("RGBA")
|
||||
# The keyed backdrop is now transparent (corner pixel fully see-through).
|
||||
assert rgba.getpixel((0, 0))[3] == 0
|
||||
# The pet blob in the center is still opaque.
|
||||
assert rgba.getpixel((rgba.width // 2, rgba.height // 2))[3] > 0
|
||||
|
||||
|
||||
def test_hatch_pet_end_to_end(monkeypatch, tmp_path):
|
||||
from agent.pet import store
|
||||
from agent.pet.generate import atlas as atlas_mod
|
||||
from agent.pet.generate import imagegen, orchestrate
|
||||
|
||||
base = tmp_path / "base.png"
|
||||
_strip(1).save(base)
|
||||
|
||||
def fake_generate(prompt, *, n=1, reference_images=None, provider=None, prefix="pet"):
|
||||
# Return a synthetic row strip; frame count is inferable from the spec.
|
||||
state = prefix.replace("pet_row_", "")
|
||||
count = atlas_mod.FRAME_COUNTS.get(state, 6)
|
||||
p = tmp_path / f"{prefix}.png"
|
||||
_strip(count).save(p)
|
||||
return [p]
|
||||
|
||||
monkeypatch.setattr(imagegen, "resolve_provider", lambda **_: object())
|
||||
monkeypatch.setattr(imagegen, "generate", fake_generate)
|
||||
|
||||
events: list[tuple[str, str]] = []
|
||||
result = orchestrate.hatch_pet(
|
||||
base_image=base,
|
||||
slug="mocky",
|
||||
display_name="Mocky",
|
||||
description="a test pet",
|
||||
concept="a fox",
|
||||
on_progress=lambda ev, detail: events.append((ev, detail)),
|
||||
)
|
||||
|
||||
assert result.slug == "mocky"
|
||||
assert result.validation["ok"]
|
||||
assert set(result.states) == {s for s, _, _ in atlas_mod.ROW_SPECS}
|
||||
assert ("compose", "") in events
|
||||
# The pet is on disk and adoptable.
|
||||
assert store.load_pet("mocky").exists
|
||||
|
||||
|
||||
def test_hatch_pet_idle_fallback_when_row_fails(monkeypatch, tmp_path):
|
||||
from agent.pet.generate import atlas as atlas_mod
|
||||
from agent.pet.generate import imagegen, orchestrate
|
||||
from agent.pet.generate.imagegen import GenerationError
|
||||
|
||||
base = tmp_path / "base.png"
|
||||
_strip(1).save(base)
|
||||
|
||||
def fake_generate(prompt, *, n=1, reference_images=None, provider=None, prefix="pet"):
|
||||
if prefix == "pet_row_idle":
|
||||
raise GenerationError("boom")
|
||||
state = prefix.replace("pet_row_", "")
|
||||
count = atlas_mod.FRAME_COUNTS.get(state, 6)
|
||||
p = tmp_path / f"{prefix}.png"
|
||||
_strip(count).save(p)
|
||||
return [p]
|
||||
|
||||
monkeypatch.setattr(imagegen, "resolve_provider", lambda **_: object())
|
||||
monkeypatch.setattr(imagegen, "generate", fake_generate)
|
||||
|
||||
result = orchestrate.hatch_pet(base_image=base, slug="fallbacky", concept="a fox")
|
||||
assert "idle" in result.states # filled by the base-image fallback
|
||||
|
||||
|
||||
def test_resolve_provider_errors_without_backend(monkeypatch):
|
||||
from agent.pet.generate import imagegen
|
||||
|
||||
monkeypatch.setattr(imagegen, "_discover", lambda: None)
|
||||
monkeypatch.setattr("agent.image_gen_registry.get_active_provider", lambda: None)
|
||||
monkeypatch.setattr("agent.image_gen_registry.get_provider", lambda name: None)
|
||||
|
||||
with pytest.raises(imagegen.GenerationError):
|
||||
imagegen.resolve_provider(require_references=True)
|
||||
|
||||
|
||||
def test_generate_retries_without_transparent_background(monkeypatch, tmp_path):
|
||||
"""A model that rejects background=transparent still produces images."""
|
||||
from agent.pet.generate import imagegen
|
||||
|
||||
saved = tmp_path / "img.png"
|
||||
_strip(1).save(saved)
|
||||
calls: list[dict] = []
|
||||
|
||||
class FakeProvider:
|
||||
def generate(self, prompt, **kwargs):
|
||||
calls.append(kwargs)
|
||||
if kwargs.get("background") == "transparent":
|
||||
return {"success": False, "error": "Transparent background is not supported for this model."}
|
||||
return {"success": True, "image": str(saved)}
|
||||
|
||||
sprite = imagegen.SpriteProvider(name="openai", provider=FakeProvider(), supports_references=False)
|
||||
|
||||
out = imagegen.generate("a fox", n=2, provider=sprite)
|
||||
assert len(out) == 2
|
||||
# First variant probes transparent (rejected) then retries opaque; the second
|
||||
# variant skips the transparent probe entirely.
|
||||
backgrounds = [c.get("background") for c in calls]
|
||||
assert backgrounds == ["transparent", None, None]
|
||||
144
tests/tui_gateway/test_pet_generate_rpc.py
Normal file
144
tests/tui_gateway/test_pet_generate_rpc.py
Normal file
@ -0,0 +1,144 @@
|
||||
"""Gateway RPC tests for pet generation (pet.generate / pet.hatch).
|
||||
|
||||
Image generation is mocked, so these assert the RPC contract + staging behavior
|
||||
(draft tokens, data-URI previews, expiry, activation) without any API calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("PIL")
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
from tui_gateway import server # noqa: E402
|
||||
|
||||
|
||||
def _png(path):
|
||||
Image.new("RGBA", (64, 64), (200, 80, 80, 255)).save(path)
|
||||
|
||||
|
||||
def test_pet_generate_requires_prompt():
|
||||
resp = server._methods["pet.generate"]("r1", {"prompt": " "})
|
||||
assert "error" in resp
|
||||
|
||||
|
||||
def test_pet_generate_returns_token_and_previews(monkeypatch, tmp_path):
|
||||
import agent.pet.generate as gen
|
||||
|
||||
def fake_drafts(prompt, *, n=4, style="auto"):
|
||||
paths = []
|
||||
for i in range(n):
|
||||
p = tmp_path / f"d{i}.png"
|
||||
_png(p)
|
||||
paths.append(p)
|
||||
return paths
|
||||
|
||||
monkeypatch.setattr(gen, "generate_base_drafts", fake_drafts)
|
||||
|
||||
resp = server._methods["pet.generate"]("r2", {"prompt": "a robot fox", "count": 4})
|
||||
result = resp["result"]
|
||||
assert result["ok"]
|
||||
assert len(result["drafts"]) == 4
|
||||
assert all(d["dataUri"].startswith("data:image/png;base64,") for d in result["drafts"])
|
||||
|
||||
# Drafts are staged on disk under the returned token.
|
||||
staged = server._pet_gen_root() / result["token"] / "draft-0.png"
|
||||
assert staged.is_file()
|
||||
|
||||
|
||||
def test_pet_hatch_validates_params():
|
||||
assert "error" in server._methods["pet.hatch"]("r1", {"name": "x"}) # missing token
|
||||
assert "error" in server._methods["pet.hatch"]("r2", {"token": "abc"}) # missing name
|
||||
|
||||
|
||||
def test_pet_hatch_expired_draft():
|
||||
resp = server._methods["pet.hatch"]("r3", {"token": "nope", "index": 0, "name": "Ghost"})
|
||||
assert "error" in resp
|
||||
assert "expired" in resp["error"]["message"]
|
||||
|
||||
|
||||
def _fake_drafts_factory(tmp_path):
|
||||
def fake_drafts(prompt, *, n=4, style="auto"):
|
||||
paths = []
|
||||
for i in range(n):
|
||||
p = tmp_path / f"d{i}.png"
|
||||
_png(p)
|
||||
paths.append(p)
|
||||
return paths
|
||||
|
||||
return fake_drafts
|
||||
|
||||
|
||||
def _fake_hatch_factory(captured):
|
||||
"""A hatch that registers a real local pet (so the preview payload populates)."""
|
||||
import agent.pet.generate as gen
|
||||
from agent.pet import store
|
||||
|
||||
def fake_hatch(*, base_image, slug, display_name="", description="", concept="", style="auto", on_progress=None, provider=None):
|
||||
captured["base_image"] = str(base_image)
|
||||
captured["slug"] = slug
|
||||
pet = store.register_local_pet(
|
||||
Image.new("RGBA", (192, 208), (10, 20, 30, 255)),
|
||||
slug=slug,
|
||||
display_name=display_name,
|
||||
description=description,
|
||||
)
|
||||
return gen.HatchResult(
|
||||
slug=pet.slug,
|
||||
display_name=display_name or pet.display_name,
|
||||
spritesheet=pet.spritesheet,
|
||||
states=["idle", "wave"],
|
||||
validation={"ok": True, "warnings": ["state 'jump' has no frames"]},
|
||||
)
|
||||
|
||||
return fake_hatch
|
||||
|
||||
|
||||
def test_pet_generate_then_hatch_previews_without_activating(monkeypatch, tmp_path):
|
||||
import agent.pet.generate as gen
|
||||
from agent.pet import store
|
||||
|
||||
captured = {}
|
||||
monkeypatch.setattr(gen, "generate_base_drafts", _fake_drafts_factory(tmp_path))
|
||||
monkeypatch.setattr(gen, "hatch_pet", _fake_hatch_factory(captured))
|
||||
|
||||
token = server._methods["pet.generate"]("r1", {"prompt": "a fox"})["result"]["token"]
|
||||
|
||||
resp = server._methods["pet.hatch"](
|
||||
"r2",
|
||||
{"token": token, "index": 1, "name": "My Fox", "description": "vulpine"},
|
||||
)
|
||||
result = resp["result"]
|
||||
assert result["ok"]
|
||||
assert result["slug"] == "my-fox"
|
||||
assert result["displayName"] == "My Fox"
|
||||
assert result["warnings"] == ["state 'jump' has no frames"]
|
||||
# Hatched from the chosen draft index.
|
||||
assert captured["base_image"].endswith("draft-1.png")
|
||||
|
||||
# The pet is installed on disk and the preview payload carries the sheet,
|
||||
# but hatch must NOT activate it — adoption is a separate step.
|
||||
assert store.load_pet("my-fox") is not None
|
||||
assert result["pet"]["slug"] == "my-fox"
|
||||
assert result["pet"]["spritesheetBase64"]
|
||||
assert server._methods["pet.info"]("r3", {}).get("result", {}).get("enabled") in (False, None)
|
||||
|
||||
|
||||
def test_pet_hatch_then_adopt_activates(monkeypatch, tmp_path):
|
||||
import agent.pet.generate as gen
|
||||
|
||||
captured = {}
|
||||
monkeypatch.setattr(gen, "generate_base_drafts", _fake_drafts_factory(tmp_path))
|
||||
monkeypatch.setattr(gen, "hatch_pet", _fake_hatch_factory(captured))
|
||||
|
||||
activated = {}
|
||||
monkeypatch.setattr("hermes_cli.pets._set_active", lambda slug: activated.setdefault("slug", slug))
|
||||
|
||||
token = server._methods["pet.generate"]("r1", {"prompt": "a fox"})["result"]["token"]
|
||||
hatched = server._methods["pet.hatch"]("r2", {"token": token, "index": 0, "name": "My Fox"})["result"]
|
||||
|
||||
# Adoption is the existing pet.select path, against the now-installed slug.
|
||||
adopt = server._methods["pet.select"]("r3", {"slug": hatched["slug"]})["result"]
|
||||
assert adopt["ok"]
|
||||
assert activated["slug"] == "my-fox"
|
||||
@ -182,6 +182,10 @@ _LONG_HANDLERS = frozenset(
|
||||
# animation poll stutters. On the pool they run concurrently.
|
||||
"pet.cells",
|
||||
"pet.gallery",
|
||||
# Generation is the heaviest pet path by far — multiple image-model
|
||||
# round-trips per call — so it must never block the reader thread.
|
||||
"pet.generate",
|
||||
"pet.hatch",
|
||||
"pet.select",
|
||||
"pet.thumb",
|
||||
"plugins.manage",
|
||||
@ -5007,6 +5011,49 @@ def _pet_frame_counts(spritesheet) -> dict:
|
||||
return {}
|
||||
|
||||
|
||||
def _pet_config_scale() -> float:
|
||||
"""Configured ``display.pet.scale`` (or the engine default), never raises."""
|
||||
from agent.pet import constants
|
||||
|
||||
try:
|
||||
from hermes_cli.config import load_config
|
||||
|
||||
cfg = load_config()
|
||||
display = cfg.get("display", {}) if isinstance(cfg.get("display"), dict) else {}
|
||||
pet_cfg = display.get("pet", {}) if isinstance(display.get("pet"), dict) else {}
|
||||
return float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE)
|
||||
except Exception: # noqa: BLE001
|
||||
return constants.DEFAULT_SCALE
|
||||
|
||||
|
||||
def _pet_sprite_payload(pet, *, scale: float) -> dict:
|
||||
"""Build the renderer payload (spritesheet bytes + geometry) for *pet*.
|
||||
|
||||
Shared by ``pet.info`` (the active mascot) and ``pet.hatch`` (the unadopted
|
||||
preview) so both feed the desktop canvas / TUI from one shape.
|
||||
"""
|
||||
import base64
|
||||
|
||||
from agent.pet import constants
|
||||
|
||||
raw = pet.spritesheet.read_bytes()
|
||||
suffix = pet.spritesheet.suffix.lower()
|
||||
mime = "image/png" if suffix == ".png" else "image/webp"
|
||||
return {
|
||||
"slug": pet.slug,
|
||||
"displayName": pet.display_name,
|
||||
"mime": mime,
|
||||
"spritesheetBase64": base64.standard_b64encode(raw).decode("ascii"),
|
||||
"frameW": constants.FRAME_W,
|
||||
"frameH": constants.FRAME_H,
|
||||
"framesPerState": constants.FRAMES_PER_STATE,
|
||||
"framesByState": _pet_frame_counts(pet.spritesheet),
|
||||
"loopMs": constants.LOOP_MS,
|
||||
"scale": scale,
|
||||
"stateRows": list(constants.STATE_ROWS),
|
||||
}
|
||||
|
||||
|
||||
@method("pet.info")
|
||||
def _(rid, params: dict) -> dict:
|
||||
"""Return the active petdex pet for surfaces that render sprites.
|
||||
@ -5020,8 +5067,6 @@ def _(rid, params: dict) -> dict:
|
||||
before the agent finishes building. Fail-open: returns ``enabled=False``
|
||||
on any error rather than erroring the surface.
|
||||
"""
|
||||
import base64
|
||||
|
||||
try:
|
||||
from agent.pet import constants, store
|
||||
|
||||
@ -5041,26 +5086,8 @@ def _(rid, params: dict) -> dict:
|
||||
if not enabled or pet is None or not pet.exists:
|
||||
return _ok(rid, {"enabled": False})
|
||||
|
||||
raw = pet.spritesheet.read_bytes()
|
||||
suffix = pet.spritesheet.suffix.lower()
|
||||
mime = "image/png" if suffix == ".png" else "image/webp"
|
||||
return _ok(
|
||||
rid,
|
||||
{
|
||||
"enabled": True,
|
||||
"slug": pet.slug,
|
||||
"displayName": pet.display_name,
|
||||
"mime": mime,
|
||||
"spritesheetBase64": base64.standard_b64encode(raw).decode("ascii"),
|
||||
"frameW": constants.FRAME_W,
|
||||
"frameH": constants.FRAME_H,
|
||||
"framesPerState": constants.FRAMES_PER_STATE,
|
||||
"framesByState": _pet_frame_counts(pet.spritesheet),
|
||||
"loopMs": constants.LOOP_MS,
|
||||
"scale": float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE),
|
||||
"stateRows": list(constants.STATE_ROWS),
|
||||
},
|
||||
)
|
||||
scale = float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE)
|
||||
return _ok(rid, {"enabled": True, **_pet_sprite_payload(pet, scale=scale)})
|
||||
except Exception as exc: # noqa: BLE001 - cosmetic, never break the surface
|
||||
logger.debug("pet.info failed: %s", exc)
|
||||
return _ok(rid, {"enabled": False})
|
||||
@ -5357,6 +5384,162 @@ def _(rid, params: dict) -> dict:
|
||||
return _err(rid, 5031, f"pet.scale failed: {exc}")
|
||||
|
||||
|
||||
def _pet_gen_root():
|
||||
"""Profile-scoped staging dir for in-progress generation drafts."""
|
||||
from hermes_constants import get_hermes_home
|
||||
|
||||
root = get_hermes_home() / "cache" / "pet-gen"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root
|
||||
|
||||
|
||||
def _pet_gen_sweep(root, *, max_age_s: float = 3600.0) -> None:
|
||||
"""Drop stale draft staging dirs so cache never grows unbounded."""
|
||||
import shutil
|
||||
import time
|
||||
|
||||
try:
|
||||
now = time.time()
|
||||
for child in root.iterdir():
|
||||
if child.is_dir() and now - child.stat().st_mtime > max_age_s:
|
||||
shutil.rmtree(child, ignore_errors=True)
|
||||
except Exception as exc: # noqa: BLE001 - cleanup is best-effort
|
||||
logger.debug("pet-gen sweep failed: %s", exc)
|
||||
|
||||
|
||||
def _pet_png_data_uri(path, *, max_px: int = 160) -> str:
|
||||
"""Downscaled PNG data URI for a draft image (small preview payload)."""
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(path) as opened:
|
||||
img = opened.convert("RGBA")
|
||||
img.thumbnail((max_px, max_px), Image.LANCZOS)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return "data:image/png;base64," + base64.standard_b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
@method("pet.generate")
|
||||
def _(rid, params: dict) -> dict:
|
||||
"""Generate candidate base looks for a new pet (the draft/variant step).
|
||||
|
||||
Params: ``prompt`` (required), ``count`` (default 4), ``style`` (default
|
||||
``auto``). Returns ``{ok, token, drafts:[{index, dataUri}]}`` — the token
|
||||
keys the staged base images for a later ``pet.hatch``. Retry == call again
|
||||
(fresh token). Heavy (network): runs on the worker pool.
|
||||
"""
|
||||
prompt = str(params.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
return _err(rid, 4004, "missing prompt")
|
||||
try:
|
||||
count = max(1, min(4, int(params.get("count") or 4)))
|
||||
except (TypeError, ValueError):
|
||||
count = 4
|
||||
style = str(params.get("style") or "auto").strip() or "auto"
|
||||
|
||||
try:
|
||||
import shutil
|
||||
import uuid
|
||||
|
||||
from agent.pet.generate import generate_base_drafts
|
||||
from agent.pet.generate.imagegen import GenerationError
|
||||
|
||||
root = _pet_gen_root()
|
||||
_pet_gen_sweep(root)
|
||||
|
||||
try:
|
||||
drafts = generate_base_drafts(prompt, n=count, style=style)
|
||||
except GenerationError as exc:
|
||||
return _err(rid, 5031, str(exc))
|
||||
|
||||
token = uuid.uuid4().hex[:12]
|
||||
stage = root / token
|
||||
stage.mkdir(parents=True, exist_ok=True)
|
||||
out = []
|
||||
for i, src in enumerate(drafts):
|
||||
dest = stage / f"draft-{i}.png"
|
||||
try:
|
||||
shutil.copyfile(src, dest)
|
||||
out.append({"index": i, "dataUri": _pet_png_data_uri(dest)})
|
||||
except Exception as exc: # noqa: BLE001 - skip a bad draft, keep the rest
|
||||
logger.debug("pet.generate draft %d failed: %s", i, exc)
|
||||
|
||||
if not out:
|
||||
return _err(rid, 5031, "generation produced no usable drafts")
|
||||
return _ok(rid, {"ok": True, "token": token, "drafts": out})
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("pet.generate failed: %s", exc)
|
||||
return _err(rid, 5031, f"pet.generate failed: {exc}")
|
||||
|
||||
|
||||
@method("pet.hatch")
|
||||
def _(rid, params: dict) -> dict:
|
||||
"""Turn a chosen base draft into a full pet — installed but NOT yet active.
|
||||
|
||||
Generation is expensive and the result varies, so hatch produces a *preview*
|
||||
the surface plays (all frames) before the user commits: the pet is written to
|
||||
the store (so it can be rendered + later activated) but the active pet is left
|
||||
untouched. Adopt with ``pet.select`` or throw it away with ``pet.remove``.
|
||||
|
||||
Params: ``token`` + ``index`` (from ``pet.generate``), ``name`` (required),
|
||||
``description`` (optional), ``prompt`` (optional concept for row prompts),
|
||||
``style`` (optional). Returns ``{ok, slug, displayName, warnings, pet}`` where
|
||||
``pet`` is the renderer payload. Heavy (network + raster): worker pool.
|
||||
"""
|
||||
token = str(params.get("token") or "").strip()
|
||||
index = params.get("index", 0)
|
||||
name = str(params.get("name") or "").strip()
|
||||
if not token:
|
||||
return _err(rid, 4004, "missing token")
|
||||
if not name:
|
||||
return _err(rid, 4004, "missing name")
|
||||
try:
|
||||
index = int(index)
|
||||
except (TypeError, ValueError):
|
||||
index = 0
|
||||
|
||||
try:
|
||||
from agent.pet import store
|
||||
from agent.pet.generate import hatch_pet
|
||||
from agent.pet.generate.imagegen import GenerationError
|
||||
|
||||
base = _pet_gen_root() / token / f"draft-{index}.png"
|
||||
if not base.is_file():
|
||||
return _err(rid, 4004, "draft expired — generate again")
|
||||
|
||||
slug = store.unique_slug(name)
|
||||
try:
|
||||
result = hatch_pet(
|
||||
base_image=base,
|
||||
slug=slug,
|
||||
display_name=name,
|
||||
description=str(params.get("description") or ""),
|
||||
concept=str(params.get("prompt") or name),
|
||||
style=str(params.get("style") or "auto").strip() or "auto",
|
||||
)
|
||||
except GenerationError as exc:
|
||||
return _err(rid, 5031, str(exc))
|
||||
|
||||
pet = store.load_pet(result.slug)
|
||||
payload = _pet_sprite_payload(pet, scale=_pet_config_scale()) if pet else {}
|
||||
return _ok(
|
||||
rid,
|
||||
{
|
||||
"ok": True,
|
||||
"slug": result.slug,
|
||||
"displayName": result.display_name,
|
||||
"warnings": result.validation.get("warnings", []),
|
||||
"pet": payload,
|
||||
},
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("pet.hatch failed: %s", exc)
|
||||
return _err(rid, 5031, f"pet.hatch failed: {exc}")
|
||||
|
||||
|
||||
@method("credits.view")
|
||||
def _(rid, params: dict) -> dict:
|
||||
"""Structured Nous credit view for the TUI /credits command.
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user