Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 120 additions & 1 deletion hermes_cli/tools_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,14 +249,16 @@ def _get_plugin_toolset_keys() -> set:
"requires_nous_auth": True,
"managed_nous_feature": "image_gen",
"override_env_vars": ["FAL_KEY"],
"imagegen_backend": "fal",
},
{
"name": "FAL.ai",
"badge": "paid",
"tag": "FLUX 2 Pro with auto-upscaling",
"tag": "Pick from flux-2-klein, flux-2-pro, gpt-image, nano-banana, etc.",
"env_vars": [
{"key": "FAL_KEY", "prompt": "FAL API key", "url": "https://fal.ai/dashboard/keys"},
],
"imagegen_backend": "fal",
},
],
},
Expand Down Expand Up @@ -941,6 +943,106 @@ def _detect_active_provider_index(providers: list, config: dict) -> int:
return 0


# ─── Image Generation Model Pickers ───────────────────────────────────────────
#
# IMAGEGEN_BACKENDS is a per-backend catalog. Each entry exposes:
# - config_key: top-level config.yaml key for this backend's settings
# - model_catalog_fn: returns an OrderedDict-like {model_id: metadata}
# - default_model: fallback when nothing is configured
#
# This prepares for future imagegen backends (Replicate, Stability, etc.):
# each new backend registers its own entry; the FAL provider entry in
# TOOL_CATEGORIES tags itself with `imagegen_backend: "fal"` to select the
# right catalog at picker time.


def _fal_model_catalog():
"""Lazy-load the FAL model catalog from the tool module."""
from tools.image_generation_tool import FAL_MODELS, DEFAULT_MODEL
return FAL_MODELS, DEFAULT_MODEL


IMAGEGEN_BACKENDS = {
"fal": {
"display": "FAL.ai",
"config_key": "image_gen",
"catalog_fn": _fal_model_catalog,
},
}


def _format_imagegen_model_row(model_id: str, meta: dict, widths: dict) -> str:
"""Format a single picker row with column-aligned speed / strengths / price."""
return (
f"{model_id:<{widths['model']}} "
f"{meta.get('speed', ''):<{widths['speed']}} "
f"{meta.get('strengths', ''):<{widths['strengths']}} "
f"{meta.get('price', '')}"
)


def _configure_imagegen_model(backend_name: str, config: dict) -> None:
"""Prompt the user to pick a model for the given imagegen backend.

Writes selection to ``config[backend_config_key]["model"]``. Safe to
call even when stdin is not a TTY — curses_radiolist falls back to
keeping the current selection.
"""
backend = IMAGEGEN_BACKENDS.get(backend_name)
if not backend:
return

catalog, default_model = backend["catalog_fn"]()
if not catalog:
return

cfg_key = backend["config_key"]
cur_cfg = config.setdefault(cfg_key, {})
if not isinstance(cur_cfg, dict):
cur_cfg = {}
config[cfg_key] = cur_cfg
current_model = cur_cfg.get("model") or default_model
if current_model not in catalog:
current_model = default_model

model_ids = list(catalog.keys())
# Put current model at the top so the cursor lands on it by default.
ordered = [current_model] + [m for m in model_ids if m != current_model]

# Column widths
widths = {
"model": max(len(m) for m in model_ids),
"speed": max((len(catalog[m].get("speed", "")) for m in model_ids), default=6),
"strengths": max((len(catalog[m].get("strengths", "")) for m in model_ids), default=0),
}

print()
header = (
f" {'Model':<{widths['model']}} "
f"{'Speed':<{widths['speed']}} "
f"{'Strengths':<{widths['strengths']}} "
f"Price"
)
print(color(header, Colors.CYAN))

rows = []
for mid in ordered:
row = _format_imagegen_model_row(mid, catalog[mid], widths)
if mid == current_model:
row += " ← currently in use"
rows.append(row)

idx = _prompt_choice(
f" Choose {backend['display']} model:",
rows,
default=0,
)

chosen = ordered[idx]
cur_cfg["model"] = chosen
_print_success(f" Model set to: {chosen}")


def _configure_provider(provider: dict, config: dict):
"""Configure a single provider - prompt for API keys and set config."""
env_vars = provider.get("env_vars", [])
Expand Down Expand Up @@ -997,6 +1099,10 @@ def _configure_provider(provider: dict, config: dict):
_print_success(f" {provider['name']} - no configuration needed!")
if managed_feature:
_print_info(" Requests for this tool will be billed to your Nous subscription.")
# Imagegen backends prompt for model selection after backend pick.
backend = provider.get("imagegen_backend")
if backend:
_configure_imagegen_model(backend, config)
return

# Prompt for each required env var
Expand Down Expand Up @@ -1031,6 +1137,10 @@ def _configure_provider(provider: dict, config: dict):

if all_configured:
_print_success(f" {provider['name']} configured!")
# Imagegen backends prompt for model selection after env vars are in.
backend = provider.get("imagegen_backend")
if backend:
_configure_imagegen_model(backend, config)


def _configure_simple_requirements(ts_key: str):
Expand Down Expand Up @@ -1202,6 +1312,10 @@ def _reconfigure_provider(provider: dict, config: dict):
_print_success(f" {provider['name']} - no configuration needed!")
if managed_feature:
_print_info(" Requests for this tool will be billed to your Nous subscription.")
# Imagegen backends prompt for model selection on reconfig too.
backend = provider.get("imagegen_backend")
if backend:
_configure_imagegen_model(backend, config)
return

for var in env_vars:
Expand All @@ -1219,6 +1333,11 @@ def _reconfigure_provider(provider: dict, config: dict):
else:
_print_info(" Kept current")

# Imagegen backends prompt for model selection on reconfig too.
backend = provider.get("imagegen_backend")
if backend:
_configure_imagegen_model(backend, config)


def _reconfigure_simple_requirements(ts_key: str):
"""Reconfigure simple env var requirements."""
Expand Down
87 changes: 87 additions & 0 deletions tests/hermes_cli/test_tools_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -466,3 +466,90 @@ def test_numeric_mcp_server_name_does_not_crash_sorted():

# sorted() must not raise TypeError
sorted(enabled)


# ─── Imagegen Backend Picker Wiring ────────────────────────────────────────

class TestImagegenBackendRegistry:
"""IMAGEGEN_BACKENDS tags drive the model picker flow in tools_config."""

def test_fal_backend_registered(self):
from hermes_cli.tools_config import IMAGEGEN_BACKENDS
assert "fal" in IMAGEGEN_BACKENDS

def test_fal_catalog_loads_lazily(self):
"""catalog_fn should defer import to avoid import cycles."""
from hermes_cli.tools_config import IMAGEGEN_BACKENDS
catalog, default = IMAGEGEN_BACKENDS["fal"]["catalog_fn"]()
assert default == "fal-ai/flux-2/klein/9b"
assert "fal-ai/flux-2/klein/9b" in catalog
assert "fal-ai/flux-2-pro" in catalog

def test_image_gen_providers_tagged_with_fal_backend(self):
"""Both Nous Subscription and FAL.ai providers must carry the
imagegen_backend tag so _configure_provider fires the picker."""
from hermes_cli.tools_config import TOOL_CATEGORIES
providers = TOOL_CATEGORIES["image_gen"]["providers"]
for p in providers:
assert p.get("imagegen_backend") == "fal", (
f"{p['name']} missing imagegen_backend tag"
)


class TestImagegenModelPicker:
"""_configure_imagegen_model writes selection to config and respects
curses fallback semantics (returns default when stdin isn't a TTY)."""

def test_picker_writes_chosen_model_to_config(self):
from hermes_cli.tools_config import _configure_imagegen_model
config = {}
# Force _prompt_choice to pick index 1 (second-in-ordered-list).
with patch("hermes_cli.tools_config._prompt_choice", return_value=1):
_configure_imagegen_model("fal", config)
# ordered[0] == current (default klein), ordered[1] == first non-default
assert config["image_gen"]["model"] != "fal-ai/flux-2/klein/9b"
assert config["image_gen"]["model"].startswith("fal-ai/")

def test_picker_with_gpt_image_does_not_prompt_quality(self):
"""GPT-Image quality is pinned to medium in the tool's defaults —
no follow-up prompt, no config write for quality_setting."""
from hermes_cli.tools_config import (
_configure_imagegen_model,
IMAGEGEN_BACKENDS,
)
catalog, default_model = IMAGEGEN_BACKENDS["fal"]["catalog_fn"]()
model_ids = list(catalog.keys())
ordered = [default_model] + [m for m in model_ids if m != default_model]
gpt_idx = ordered.index("fal-ai/gpt-image-1.5")

# Only ONE picker call is expected (for model) — not two (model + quality).
call_count = {"n": 0}
def fake_prompt(*a, **kw):
call_count["n"] += 1
return gpt_idx

config = {}
with patch("hermes_cli.tools_config._prompt_choice", side_effect=fake_prompt):
_configure_imagegen_model("fal", config)

assert call_count["n"] == 1, (
f"Expected 1 picker call (model only), got {call_count['n']}"
)
assert config["image_gen"]["model"] == "fal-ai/gpt-image-1.5"
assert "quality_setting" not in config["image_gen"]

def test_picker_no_op_for_unknown_backend(self):
from hermes_cli.tools_config import _configure_imagegen_model
config = {}
_configure_imagegen_model("nonexistent-backend", config)
assert config == {} # untouched

def test_picker_repairs_corrupt_config_section(self):
"""When image_gen is a non-dict (user-edit YAML), the picker should
replace it with a fresh dict rather than crash."""
from hermes_cli.tools_config import _configure_imagegen_model
config = {"image_gen": "some-garbage-string"}
with patch("hermes_cli.tools_config._prompt_choice", return_value=0):
_configure_imagegen_model("fal", config)
assert isinstance(config["image_gen"], dict)
assert config["image_gen"]["model"] == "fal-ai/flux-2/klein/9b"
Loading
Loading