fn_registry/python/functions/ml/comfyui_generate_styles_llm.py

"""comfyui_generate_styles_llm — genera estilos WAS para ComfyUI con el LLM (grupo claude-direct).

Pide al modelo (vía `ask_llm`, API directa de Anthropic, arranque 0) que produzca N estilos de
una CATEGORÍA temática en el formato exacto del selector WAS:

    { "NombreEstilo": {"prompt": "modificadores de estilo", "negative_prompt": "..."}, ... }

Cada `prompt` son MODIFICADORES de estilo (cámara, lente, iluminación, render, medio artístico,
paleta, mood) — NO una descripción de sujeto y SIN el placeholder `{prompt}` (el selector
múltiple concatena los prompts elegidos).

Robusta por diseño: extrae el primer bloque JSON de la respuesta, lo valida entrada por entrada
(descarta las que no tengan `prompt`, rellena `negative_prompt` por defecto si falta) y, ante
CUALQUIER fallo (rate-limit 429, JSON corrupto, respuesta vacía), devuelve `{}` en vez de lanzar.
Así el caller puede iterar varias categorías y seguir aunque una falle (error-path del DoD).

Impura: llama a la API del LLM (red). No escribe disco (eso es trabajo de
`comfyui_append_styles`). Pensada para componer: generar -> validar -> append.
"""
from __future__ import annotations

import json
import os
import re
import sys

# Importa ask_llm del registry (grupo claude-direct). Path al paquete de funciones Python.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

try:
    from core.ask_llm import ask_llm  # type: ignore
except Exception:  # pragma: no cover - sólo si el registry no está en el path
    ask_llm = None  # type: ignore

DEFAULT_NEGATIVE = (
    "ugly, deformed, noisy, blurry, low quality, distorted, disfigured, "
    "bad anatomy, watermark, signature, text, NSFW"
)

_SYSTEM = (
    "You are an expert prompt engineer for Stable Diffusion. You output ONLY valid JSON, "
    "no prose, no markdown fences. Styles must be SFW."
)

_PROMPT_TMPL = """Generate exactly {n} distinct image STYLE presets for the category: "{category}".

Output a single JSON object mapping a short unique style name to an object with keys
"prompt" and "negative_prompt". Rules:
- "prompt" = a comma-separated list of powerful STYLE modifiers ONLY (camera, lens, lighting,
  render engine, art medium, palette, texture, mood). NO subject description. Do NOT include
  the literal token {{prompt}}.
- "negative_prompt" = comma-separated terms to avoid for this style. Keep it SFW.
- Style names must be short, kebab-case, prefixed with "{prefix}", and must NOT repeat these
  existing names: {avoid}.
- Output ONLY the JSON object, nothing else.

Example of ONE entry:
"{prefix}example": {{"prompt": "cinematic film still, anamorphic lens, teal and orange grade, dramatic key light", "negative_prompt": "lowres, blurry, deformed, watermark, text"}}
"""


def _extract_json_object(text: str) -> dict | None:
    """Extrae el primer objeto JSON {...} de un texto, tolerando fences de markdown."""
    if not text:
        return None
    # Quitar fences ```json ... ``` si los hay.
    fenced = re.search(r"```(?:json)?\s*(\{.*\})\s*```", text, re.DOTALL)
    candidate = fenced.group(1) if fenced else None
    if candidate is None:
        # Buscar desde la primera { hasta la última } (greedy) — robusto a prosa alrededor.
        start = text.find("{")
        end = text.rfind("}")
        if start == -1 or end == -1 or end <= start:
            return None
        candidate = text[start : end + 1]
    try:
        obj = json.loads(candidate)
    except (json.JSONDecodeError, ValueError):
        return None
    return obj if isinstance(obj, dict) else None


def _validate(obj: dict, avoid: set[str]) -> dict:
    """Filtra el dict del LLM a entradas válidas WAS, descartando duplicados y vacíos."""
    clean: dict[str, dict[str, str]] = {}
    for name, v in obj.items():
        if not isinstance(name, str) or not name.strip():
            continue
        if name in avoid:
            continue
        if not isinstance(v, dict):
            continue
        prompt = v.get("prompt")
        if not isinstance(prompt, str) or not prompt.strip():
            continue
        if "{prompt}" in prompt:
            prompt = prompt.replace("{prompt}", "").strip().strip(",").strip()
            if not prompt:
                continue
        neg = v.get("negative_prompt")
        if not isinstance(neg, str) or not neg.strip():
            neg = DEFAULT_NEGATIVE
        clean[name.strip()] = {"prompt": prompt.strip(), "negative_prompt": neg.strip()}
    return clean


def comfyui_generate_styles_llm(
    category: str,
    n: int = 8,
    prefix: str = "",
    avoid: list[str] | None = None,
    model: str = "claude-haiku-4-5-20251001",
) -> dict:
    """Genera estilos WAS de una categoría con el LLM. Devuelve {} ante cualquier fallo.

    Args:
        category: tema de los estilos (ej. "vaporwave aesthetics", "baroque painting",
            "macro insect photography"). Texto libre, dirige el contenido.
        n: número de estilos a pedir (el modelo puede devolver menos; se validan todos).
        prefix: prefijo para los nombres generados (ej. "vapor-"), para agrupar y reducir
            colisiones con estilos existentes. Si vacío, los nombres van sin prefijo.
        avoid: lista de nombres ya existentes que el modelo NO debe repetir (dedup previo).
        model: id del modelo Anthropic. Default haiku (más cuota, rápido). Otros:
            claude-opus-4-8, claude-sonnet-4-6.

    Returns:
        dict {nombre: {"prompt", "negative_prompt"}} ya validado (prompt no vacío, sin
        `{prompt}`, negative rellenado). Vacío `{}` si el LLM falla, la respuesta no trae JSON
        válido, o no hay entradas utilizables. NUNCA lanza: el error-path es devolver {}.
    """
    if ask_llm is None:
        return {}
    avoid_set = set(avoid or [])
    avoid_str = ", ".join(sorted(avoid_set)[:40]) if avoid_set else "(none)"
    prompt = _PROMPT_TMPL.format(
        n=int(n), category=category, prefix=prefix, avoid=avoid_str
    )
    try:
        raw = ask_llm(prompt, model=model, system=_SYSTEM, max_tokens=4096, echo=False)
    except Exception:
        return {}
    obj = _extract_json_object(raw or "")
    if obj is None:
        return {}
    return _validate(obj, avoid_set)


if __name__ == "__main__":
    import json as _json

    cat = sys.argv[1] if len(sys.argv) > 1 else "synthwave retro aesthetics"
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 6
    pref = sys.argv[3] if len(sys.argv) > 3 else "synth-"
    out = comfyui_generate_styles_llm(cat, n=n, prefix=pref)
    print(_json.dumps(out, ensure_ascii=False, indent=2))
    print("GENERADOS:", len(out), file=sys.stderr)