malkuthro
1732030175
* fix(setups): resolve switch-selected run results * fix(setups): harden switch result reporting * fix(setups): reject stale execution map writers
1011 lines
36 KiB
Python
1011 lines
36 KiB
Python
"""Execution boundary for published Koolook workflow setups."""
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
try:
|
|
from .koolook_setups import PublishedSetupRegistry
|
|
except ImportError: # pragma: no cover - standalone test/import context
|
|
from koolook_setups import PublishedSetupRegistry
|
|
|
|
|
|
class SetupRunError(Exception):
|
|
"""API-facing setup runner error with stable HTTP status semantics."""
|
|
|
|
def __init__(self, status_code: int, errors: list[str]) -> None:
|
|
super().__init__("; ".join(errors))
|
|
self.status_code = status_code
|
|
self.errors = errors
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RunRecord:
|
|
run_id: str
|
|
setup_id: str
|
|
prompt_id: str
|
|
status: str
|
|
inputs: dict[str, Any]
|
|
|
|
|
|
class InMemorySetupRunStore:
|
|
"""Small run-id mapping store for a running ComfyUI process."""
|
|
|
|
def __init__(self) -> None:
|
|
self._next_id = 1
|
|
self._records: dict[str, RunRecord] = {}
|
|
|
|
def create(self, *, setup_id: str, prompt_id: str, inputs: dict[str, Any] | None = None) -> RunRecord:
|
|
run_id = f"run-{self._next_id:06d}"
|
|
self._next_id += 1
|
|
record = RunRecord(
|
|
run_id=run_id,
|
|
setup_id=setup_id,
|
|
prompt_id=prompt_id,
|
|
status="queued",
|
|
inputs=deepcopy(inputs) if isinstance(inputs, dict) else {},
|
|
)
|
|
self._records[run_id] = record
|
|
return record
|
|
|
|
def get(self, run_id: str) -> RunRecord | None:
|
|
return self._records.get(run_id)
|
|
|
|
|
|
class AiohttpComfyClient:
|
|
"""HTTP adapter for the managed ComfyUI server."""
|
|
|
|
def __init__(self, base_url: str) -> None:
|
|
self._base_url = base_url.rstrip("/")
|
|
|
|
async def queue_prompt(self, prompt: dict) -> dict:
|
|
import aiohttp
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(f"{self._base_url}/prompt", json={"prompt": prompt}) as response:
|
|
try:
|
|
payload = await response.json()
|
|
except Exception:
|
|
payload = {"error": await response.text()}
|
|
if response.status >= 400:
|
|
detail = payload.get("error") if isinstance(payload, dict) else payload
|
|
raise SetupRunError(
|
|
502,
|
|
[f"ComfyUI queue request failed with HTTP {response.status}: {detail}"],
|
|
)
|
|
return payload if isinstance(payload, dict) else {}
|
|
|
|
async def get_history(self, prompt_id: str) -> dict:
|
|
import aiohttp
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(f"{self._base_url}/history/{prompt_id}") as response:
|
|
try:
|
|
payload = await response.json()
|
|
except Exception:
|
|
payload = {"error": await response.text()}
|
|
if response.status >= 400:
|
|
detail = payload.get("error") if isinstance(payload, dict) else payload
|
|
raise SetupRunError(
|
|
502,
|
|
[f"ComfyUI history request failed with HTTP {response.status}: {detail}"],
|
|
)
|
|
return payload if isinstance(payload, dict) else {}
|
|
|
|
async def get_queue(self) -> dict:
|
|
import aiohttp
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(f"{self._base_url}/queue") as response:
|
|
try:
|
|
payload = await response.json()
|
|
except Exception:
|
|
payload = {"error": await response.text()}
|
|
if response.status >= 400:
|
|
detail = payload.get("error") if isinstance(payload, dict) else payload
|
|
raise SetupRunError(
|
|
502,
|
|
[f"ComfyUI queue status request failed with HTTP {response.status}: {detail}"],
|
|
)
|
|
return payload if isinstance(payload, dict) else {}
|
|
|
|
|
|
class PublishedSetupRunner:
|
|
"""Runs callable published setups without exposing ComfyUI internals."""
|
|
|
|
def __init__(
|
|
self,
|
|
registry: PublishedSetupRegistry,
|
|
comfy_client: Any,
|
|
run_store: InMemorySetupRunStore | None = None,
|
|
) -> None:
|
|
self._registry = registry
|
|
self._comfy_client = comfy_client
|
|
self._run_store = run_store or InMemorySetupRunStore()
|
|
|
|
async def runSetup(self, setup_id: str, inputs: dict[str, Any]) -> dict[str, Any]:
|
|
setup = self._registry.getSetup(setup_id)
|
|
if setup is None:
|
|
raise SetupRunError(404, [f"Published setup '{setup_id}' not found."])
|
|
if not isinstance(inputs, dict):
|
|
raise SetupRunError(400, ["inputs must be a JSON object"])
|
|
|
|
api_prompt = setup.get("apiPrompt")
|
|
if not isinstance(api_prompt, dict) or setup.get("validation", {}).get("status") != "valid":
|
|
raise SetupRunError(400, [f"Published setup '{setup_id}' is not callable."])
|
|
|
|
prompt = deepcopy(api_prompt)
|
|
input_fields = _declared_input_fields(setup)
|
|
errors = _validate_run_inputs(input_fields, inputs)
|
|
errors.extend(_validate_execution_map_inputs(setup, inputs))
|
|
if errors:
|
|
raise SetupRunError(400, errors)
|
|
|
|
for key, value in inputs.items():
|
|
field = input_fields[key]
|
|
target = field["target"]
|
|
prompt[str(target["node"])]["inputs"][str(target["input"])] = _prompt_value_for_field(field, value)
|
|
|
|
mapped_prompt = _prune_prompt_for_execution_map(setup, prompt, inputs)
|
|
prompt = (
|
|
mapped_prompt
|
|
if mapped_prompt is not None
|
|
else _prune_prompt_for_selected_app_results(setup, prompt, inputs)
|
|
)
|
|
|
|
try:
|
|
queued = await self._comfy_client.queue_prompt(prompt)
|
|
except SetupRunError:
|
|
raise
|
|
except Exception as exc:
|
|
raise SetupRunError(502, [f"ComfyUI queue request failed: {exc}"]) from exc
|
|
|
|
prompt_id = queued.get("prompt_id") if isinstance(queued, dict) else None
|
|
if not isinstance(prompt_id, str) or not prompt_id:
|
|
raise SetupRunError(502, ["ComfyUI queue response did not include prompt_id"])
|
|
|
|
record = self._run_store.create(setup_id=setup_id, prompt_id=prompt_id, inputs=inputs)
|
|
return {"runId": record.run_id, "promptId": record.prompt_id, "status": record.status}
|
|
|
|
async def getRun(self, run_id: str) -> dict[str, Any]:
|
|
record = self._run_store.get(run_id)
|
|
if record is None:
|
|
raise SetupRunError(404, [f"Koolook run '{run_id}' not found."])
|
|
|
|
setup = self._registry.getSetup(record.setup_id)
|
|
if setup is None:
|
|
raise SetupRunError(404, [f"Published setup '{record.setup_id}' not found."])
|
|
|
|
try:
|
|
history = await self._comfy_client.get_history(record.prompt_id)
|
|
except SetupRunError:
|
|
raise
|
|
except Exception as exc:
|
|
raise SetupRunError(502, [f"ComfyUI history request failed: {exc}"]) from exc
|
|
|
|
history_entry = history.get(record.prompt_id) if isinstance(history, dict) else None
|
|
if isinstance(history_entry, dict):
|
|
status = _status_from_history(history_entry)
|
|
return {
|
|
"runId": record.run_id,
|
|
"setupId": record.setup_id,
|
|
"promptId": record.prompt_id,
|
|
"status": status,
|
|
"comfyStatus": history_entry.get("status", {}),
|
|
"outputs": _summarize_outputs(
|
|
setup,
|
|
history_entry.get("outputs", {}),
|
|
record.inputs,
|
|
_history_prompt(history_entry),
|
|
),
|
|
}
|
|
|
|
try:
|
|
queue = await self._comfy_client.get_queue()
|
|
except SetupRunError:
|
|
raise
|
|
except Exception as exc:
|
|
raise SetupRunError(502, [f"ComfyUI queue status request failed: {exc}"]) from exc
|
|
return {
|
|
"runId": record.run_id,
|
|
"setupId": record.setup_id,
|
|
"promptId": record.prompt_id,
|
|
"status": _status_from_queue(record.prompt_id, queue),
|
|
"outputs": [],
|
|
}
|
|
|
|
|
|
def _declared_input_fields(setup: dict[str, Any]) -> dict[str, dict[str, Any]]:
|
|
fields: dict[str, dict[str, Any]] = {}
|
|
for field in setup.get("inputContract", {}).get("inputs", []):
|
|
if not isinstance(field, dict):
|
|
continue
|
|
key = field.get("key")
|
|
target = field.get("target")
|
|
if isinstance(key, str) and key and isinstance(target, dict):
|
|
fields[key] = field
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
if isinstance(app, dict):
|
|
for field in app.get("inputs", []):
|
|
if not isinstance(field, dict):
|
|
continue
|
|
key = field.get("key")
|
|
target = field.get("target")
|
|
if isinstance(key, str) and key and isinstance(target, dict):
|
|
fields.setdefault(key, field)
|
|
for field in app.get("outputs", []):
|
|
if not isinstance(field, dict):
|
|
continue
|
|
key = field.get("key")
|
|
target = field.get("target")
|
|
if isinstance(key, str) and key and isinstance(target, dict):
|
|
fields.setdefault(key, field)
|
|
switch = app.get("switch")
|
|
if isinstance(switch, dict):
|
|
key = switch.get("key")
|
|
target = switch.get("target")
|
|
if isinstance(key, str) and key and isinstance(target, dict):
|
|
fields.setdefault(key, switch)
|
|
return fields
|
|
|
|
|
|
def _validate_run_inputs(input_fields: dict[str, dict[str, Any]], inputs: dict[str, Any]) -> list[str]:
|
|
errors: list[str] = []
|
|
for key in inputs:
|
|
if key not in input_fields:
|
|
errors.append(f"input '{key}' is not declared by this setup")
|
|
continue
|
|
field = input_fields[key]
|
|
if field.get("key") == "switch" and _switch_option_value(field, inputs[key]) is None:
|
|
errors.append(_switch_value_error(key, field))
|
|
for key, field in input_fields.items():
|
|
if field.get("required") is True and key not in inputs:
|
|
errors.append(f"required input '{key}' is missing")
|
|
return errors
|
|
|
|
|
|
def _prompt_value_for_field(field: dict[str, Any], value: Any) -> Any:
|
|
if field.get("key") != "switch":
|
|
return value
|
|
selected_value = _switch_option_value(field, value)
|
|
for option in field.get("options", []):
|
|
if not isinstance(option, dict):
|
|
continue
|
|
if option.get("value") == selected_value and isinstance(option.get("label"), str):
|
|
return option["label"]
|
|
return value
|
|
|
|
|
|
def _switch_option_value(switch: dict[str, Any], selected: Any) -> int | None:
|
|
if isinstance(selected, bool):
|
|
return None
|
|
for option in switch.get("options", []):
|
|
if not isinstance(option, dict):
|
|
continue
|
|
value = option.get("value")
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
continue
|
|
if selected == value:
|
|
return value
|
|
if isinstance(selected, str) and (selected == str(value) or selected == option.get("label")):
|
|
return value
|
|
return None
|
|
|
|
|
|
def _switch_value_error(key: str, switch: dict[str, Any]) -> str:
|
|
choices: list[str] = []
|
|
for option in switch.get("options", []):
|
|
if not isinstance(option, dict):
|
|
continue
|
|
value = option.get("value")
|
|
label = option.get("label")
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
continue
|
|
choices.append(f"{value} ({label})" if isinstance(label, str) and label else str(value))
|
|
if choices:
|
|
return f"input '{key}' must be one of: {', '.join(choices)}"
|
|
return f"input '{key}' has no valid switch options"
|
|
|
|
|
|
def _status_from_history(history_entry: dict[str, Any]) -> str:
|
|
status = history_entry.get("status")
|
|
if not isinstance(status, dict):
|
|
return "succeeded"
|
|
status_text = str(status.get("status_str", "")).lower()
|
|
if status.get("completed") is True and status_text not in {"error", "failed", "failure"}:
|
|
return "succeeded"
|
|
return "failed"
|
|
|
|
|
|
def _status_from_queue(prompt_id: str, queue: Any) -> str:
|
|
if not isinstance(queue, dict):
|
|
return "lost"
|
|
if _queue_contains_prompt(queue.get("queue_running", []), prompt_id):
|
|
return "running"
|
|
if _queue_contains_prompt(queue.get("queue_pending", []), prompt_id):
|
|
return "queued"
|
|
return "lost"
|
|
|
|
|
|
def _queue_contains_prompt(entries: Any, prompt_id: str) -> bool:
|
|
if not isinstance(entries, list):
|
|
return False
|
|
for entry in entries:
|
|
if isinstance(entry, dict) and entry.get("prompt_id") == prompt_id:
|
|
return True
|
|
if isinstance(entry, list) and prompt_id in entry:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _summarize_outputs(
|
|
setup: dict[str, Any],
|
|
raw_outputs: Any,
|
|
run_inputs: dict[str, Any] | None = None,
|
|
history_prompt: dict[str, Any] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
output_items = _flatten_history_outputs(raw_outputs)
|
|
summaries = []
|
|
for output in setup.get("outputContract", {}).get("outputs", []):
|
|
if not isinstance(output, dict):
|
|
continue
|
|
summaries.append(
|
|
{
|
|
"key": output.get("key", ""),
|
|
"label": output.get("label", output.get("key", "")),
|
|
"type": output.get("type", ""),
|
|
"items": [
|
|
item
|
|
for item in output_items
|
|
if _history_output_matches_contract(item, str(output.get("type", "")))
|
|
],
|
|
}
|
|
)
|
|
summaries.extend(_summarize_app_surface_outputs(setup, output_items, run_inputs or {}, history_prompt))
|
|
return summaries
|
|
|
|
|
|
def _summarize_app_surface_outputs(
|
|
setup: dict[str, Any],
|
|
output_items: list[dict[str, Any]],
|
|
run_inputs: dict[str, Any],
|
|
history_prompt: dict[str, Any] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
if not isinstance(app, dict):
|
|
return []
|
|
summaries: list[dict[str, Any]] = []
|
|
for field_type, fields in (("output", app.get("outputs")), ("result", app.get("results"))):
|
|
if not isinstance(fields, list):
|
|
continue
|
|
for field in fields:
|
|
if not isinstance(field, dict):
|
|
continue
|
|
target = field.get("target")
|
|
target_node = str(target.get("node")) if isinstance(target, dict) else ""
|
|
items = [
|
|
item
|
|
for item in output_items
|
|
if target_node and item.get("nodeId") == target_node
|
|
]
|
|
if field_type == "result" and not items:
|
|
items = _selected_switch_result_items(setup, field, output_items, run_inputs)
|
|
if field_type == "result" and not items and history_prompt is not None:
|
|
items = _execution_map_writer_result_items(setup, history_prompt, run_inputs)
|
|
summary = {
|
|
"key": field.get("key", ""),
|
|
"label": field.get("label", field.get("key", "")),
|
|
"type": field_type,
|
|
"visible": field.get("visible", True),
|
|
"target": target if isinstance(target, dict) else {},
|
|
"default": field.get("default"),
|
|
"items": items,
|
|
}
|
|
summaries.append(summary)
|
|
return summaries
|
|
|
|
|
|
def _history_prompt(history_entry: dict[str, Any]) -> dict[str, Any] | None:
|
|
prompt_record = history_entry.get("prompt")
|
|
if (
|
|
isinstance(prompt_record, list)
|
|
and len(prompt_record) >= 3
|
|
and isinstance(prompt_record[2], dict)
|
|
):
|
|
return prompt_record[2]
|
|
return None
|
|
|
|
|
|
def _execution_map_writer_result_items(
|
|
setup: dict[str, Any],
|
|
prompt: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> list[dict[str, Any]]:
|
|
writer_nodes = _selected_execution_map_writer_nodes(setup, run_inputs)
|
|
items: list[dict[str, Any]] = []
|
|
for writer_node in writer_nodes:
|
|
node = prompt.get(writer_node)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
value = _resolve_prompt_value(prompt, inputs.get("filepath"))
|
|
if isinstance(value, str) and value:
|
|
items.append({"nodeId": writer_node, "kind": "text", "value": value})
|
|
return items
|
|
|
|
|
|
def _selected_execution_map_writer_nodes(
|
|
setup: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> list[str]:
|
|
execution_map = setup.get("executionMap")
|
|
if not isinstance(execution_map, dict) or execution_map.get("version") != 1:
|
|
return []
|
|
routers = execution_map.get("routers")
|
|
if not isinstance(routers, list):
|
|
return []
|
|
out: list[str] = []
|
|
for router in routers:
|
|
if not isinstance(router, dict):
|
|
continue
|
|
selected_value = _execution_map_selected_value(setup, router, run_inputs)
|
|
branches = router.get("branches")
|
|
branch = branches.get(str(selected_value)) if isinstance(branches, dict) else None
|
|
writer_nodes = branch.get("writerNodes") if isinstance(branch, dict) else None
|
|
if isinstance(writer_nodes, list):
|
|
out.extend(str(node_id) for node_id in writer_nodes)
|
|
return list(dict.fromkeys(out))
|
|
|
|
|
|
def _selected_switch_result_items(
|
|
setup: dict[str, Any],
|
|
result_field: dict[str, Any],
|
|
output_items: list[dict[str, Any]],
|
|
run_inputs: dict[str, Any],
|
|
) -> list[dict[str, Any]]:
|
|
selected_branch = _selected_result_branch_ref(setup, result_field, run_inputs)
|
|
if selected_branch is None:
|
|
return []
|
|
selected_node = str(selected_branch[0])
|
|
return [item for item in output_items if item.get("nodeId") == selected_node]
|
|
|
|
|
|
def _prune_prompt_for_execution_map(
|
|
setup: dict[str, Any],
|
|
prompt: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> dict[str, Any] | None:
|
|
execution_map = setup.get("executionMap")
|
|
if not isinstance(execution_map, dict) or execution_map.get("version") != 1:
|
|
return None
|
|
routers = execution_map.get("routers")
|
|
if not isinstance(routers, list):
|
|
return None
|
|
|
|
roots: set[str] = set()
|
|
for router in routers:
|
|
if not isinstance(router, dict):
|
|
continue
|
|
router_node = str(router.get("node", ""))
|
|
selected_value = _execution_map_selected_value(setup, router, run_inputs)
|
|
if selected_value is None:
|
|
continue
|
|
branches = router.get("branches")
|
|
branch = branches.get(str(selected_value)) if isinstance(branches, dict) else None
|
|
if not isinstance(branch, dict):
|
|
continue
|
|
writer_nodes = branch.get("writerNodes")
|
|
if isinstance(writer_nodes, list):
|
|
roots.update(str(node_id) for node_id in writer_nodes if str(node_id) in prompt)
|
|
if not writer_nodes and router_node in prompt:
|
|
roots.add(router_node)
|
|
if not roots:
|
|
return None
|
|
|
|
return _prune_prompt_from_roots(
|
|
prompt,
|
|
roots,
|
|
_selected_app_switches(setup, prompt, run_inputs),
|
|
)
|
|
|
|
|
|
def _execution_map_selected_value(
|
|
setup: dict[str, Any],
|
|
router: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> int | None:
|
|
key = router.get("switchKey")
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
switch = app.get("switch") if isinstance(app, dict) else None
|
|
if isinstance(switch, dict) and switch.get("key") == key:
|
|
return _selected_switch_value(switch, run_inputs)
|
|
selected = run_inputs.get(key) if isinstance(key, str) else None
|
|
if isinstance(selected, bool):
|
|
return None
|
|
if isinstance(selected, int):
|
|
return selected
|
|
if isinstance(selected, str) and selected.isdigit():
|
|
return int(selected)
|
|
return None
|
|
|
|
|
|
def _validate_execution_map_inputs(setup: dict[str, Any], run_inputs: dict[str, Any]) -> list[str]:
|
|
execution_map = setup.get("executionMap")
|
|
if not isinstance(execution_map, dict) or execution_map.get("version") != 1:
|
|
return []
|
|
routers = execution_map.get("routers")
|
|
if not isinstance(routers, list):
|
|
return []
|
|
api_prompt = setup.get("apiPrompt")
|
|
prompt_node_ids = set(api_prompt) if isinstance(api_prompt, dict) else set()
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
switch = app.get("switch") if isinstance(app, dict) else None
|
|
errors: list[str] = []
|
|
for router in routers:
|
|
if not isinstance(router, dict):
|
|
continue
|
|
selected_value = _execution_map_selected_value(setup, router, run_inputs)
|
|
key = router.get("switchKey")
|
|
if selected_value is None:
|
|
if isinstance(switch, dict) and switch.get("key") == key and isinstance(key, str):
|
|
if key not in run_inputs:
|
|
errors.append(_switch_value_error(key, switch))
|
|
elif isinstance(key, str):
|
|
errors.append(f"input '{key}' must select a valid execution branch")
|
|
continue
|
|
branches = router.get("branches")
|
|
if not isinstance(branches, dict) or str(selected_value) not in branches:
|
|
errors.append(f"input '{key}' selects branch {selected_value}, but this setup has no execution branch for it")
|
|
continue
|
|
branch = branches.get(str(selected_value))
|
|
writer_nodes = branch.get("writerNodes") if isinstance(branch, dict) else None
|
|
if isinstance(writer_nodes, list) and writer_nodes:
|
|
missing_writer_nodes = [
|
|
str(node_id)
|
|
for node_id in writer_nodes
|
|
if str(node_id) not in prompt_node_ids
|
|
]
|
|
if missing_writer_nodes:
|
|
errors.append(
|
|
f"execution map branch {selected_value} for switch '{key}' references "
|
|
f"writer node(s) not present in the prompt: {', '.join(missing_writer_nodes)}"
|
|
)
|
|
elif not writer_nodes and str(router.get("node", "")) not in prompt_node_ids:
|
|
errors.append(
|
|
f"execution map branch {selected_value} for switch '{key}' has no writer nodes "
|
|
"and its router node is not present in the prompt"
|
|
)
|
|
return errors
|
|
|
|
|
|
def _prune_prompt_for_selected_app_results(
|
|
setup: dict[str, Any],
|
|
prompt: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
result_switches = _selected_result_switches(setup, prompt, run_inputs)
|
|
if not result_switches:
|
|
return prompt
|
|
selected_switches = _selected_app_switches(setup, prompt, run_inputs)
|
|
|
|
roots = {
|
|
node_id
|
|
for node_id, _selected_value in result_switches.values()
|
|
if node_id in prompt
|
|
}
|
|
for field in _app_result_fields(setup):
|
|
target = field.get("target")
|
|
if isinstance(target, dict) and str(target.get("node")) in prompt:
|
|
roots.add(str(target["node"]))
|
|
|
|
selected_branch_nodes = {
|
|
str(prompt[switch_node_id]["inputs"][f"value{selected_value}"][0])
|
|
for switch_node_id, (_node_id, selected_value) in result_switches.items()
|
|
if (
|
|
switch_node_id in prompt
|
|
and isinstance(prompt[switch_node_id], dict)
|
|
and isinstance(prompt[switch_node_id].get("inputs"), dict)
|
|
and _is_api_ref(prompt[switch_node_id]["inputs"].get(f"value{selected_value}"))
|
|
)
|
|
}
|
|
roots.update(
|
|
_selected_output_descendant_roots(
|
|
prompt,
|
|
selected_branch_nodes,
|
|
_setup_output_surface_node_ids(setup),
|
|
)
|
|
)
|
|
if not roots:
|
|
return prompt
|
|
|
|
return _prune_prompt_from_roots(prompt, roots, selected_switches)
|
|
|
|
|
|
def _prune_prompt_from_roots(
|
|
prompt: dict[str, Any],
|
|
roots: set[str],
|
|
selected_switches: dict[str, tuple[str, int]],
|
|
) -> dict[str, Any]:
|
|
keep: set[str] = set()
|
|
|
|
def visit(node_id: str) -> None:
|
|
if node_id in keep or node_id not in prompt:
|
|
return
|
|
keep.add(node_id)
|
|
node = prompt.get(node_id)
|
|
if not isinstance(node, dict):
|
|
return
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
return
|
|
selected_value = selected_switches.get(node_id, (None, None))[1]
|
|
for input_name, value in inputs.items():
|
|
if (
|
|
selected_value is not None
|
|
and input_name.startswith("value")
|
|
and input_name != f"value{selected_value}"
|
|
):
|
|
continue
|
|
for ref in _api_refs(value):
|
|
visit(str(ref[0]))
|
|
|
|
for root in roots:
|
|
visit(root)
|
|
if not keep:
|
|
return prompt
|
|
pruned = {node_id: deepcopy(node) for node_id, node in prompt.items() if node_id in keep}
|
|
for node_id, (_switch_node_id, selected_value) in selected_switches.items():
|
|
node = pruned.get(node_id)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
node["inputs"] = {
|
|
key: value
|
|
for key, value in inputs.items()
|
|
if not key.startswith("value") or key == f"value{selected_value}"
|
|
}
|
|
return pruned
|
|
|
|
|
|
def _selected_result_switches(
|
|
setup: dict[str, Any],
|
|
api_prompt: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> dict[str, tuple[str, int]]:
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
if not isinstance(app, dict):
|
|
return {}
|
|
switch = app.get("switch")
|
|
if not isinstance(switch, dict):
|
|
return {}
|
|
selected_value = _selected_switch_value(switch, run_inputs)
|
|
if selected_value is None:
|
|
return {}
|
|
selected: dict[str, tuple[str, int]] = {}
|
|
for field in _app_result_fields(setup):
|
|
result_node_id = _result_switch_node_id(setup, field, run_inputs, api_prompt)
|
|
if result_node_id is not None:
|
|
selected[result_node_id] = (result_node_id, selected_value)
|
|
return selected
|
|
|
|
|
|
def _selected_app_switches(
|
|
setup: dict[str, Any],
|
|
api_prompt: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> dict[str, tuple[str, int]]:
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
if not isinstance(app, dict):
|
|
return {}
|
|
switch = app.get("switch")
|
|
if not isinstance(switch, dict):
|
|
return {}
|
|
switch_target = switch.get("target")
|
|
if not isinstance(switch_target, dict):
|
|
return {}
|
|
selected_value = _selected_switch_value(switch, run_inputs)
|
|
if selected_value is None:
|
|
return {}
|
|
|
|
selected: dict[str, tuple[str, int]] = {}
|
|
for node_id, node in api_prompt.items():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
selector_ref = _switch_selector_ref(inputs)
|
|
if not _selector_matches_switch_target(selector_ref, switch_target, api_prompt):
|
|
continue
|
|
if _is_api_ref(inputs.get(f"value{selected_value}")):
|
|
selected[str(node_id)] = (str(node_id), selected_value)
|
|
return selected
|
|
|
|
|
|
def _selected_output_descendant_roots(
|
|
api_prompt: dict[str, Any],
|
|
selected_branch_nodes: set[str],
|
|
output_surface_node_ids: set[str],
|
|
) -> set[str]:
|
|
if not selected_branch_nodes or not output_surface_node_ids:
|
|
return set()
|
|
children_by_node = _prompt_children_by_node(api_prompt)
|
|
roots: set[str] = set()
|
|
seen: set[str] = set()
|
|
pending = list(selected_branch_nodes)
|
|
while pending:
|
|
node_id = pending.pop()
|
|
if node_id in seen:
|
|
continue
|
|
seen.add(node_id)
|
|
for child_id in children_by_node.get(node_id, set()):
|
|
if child_id in output_surface_node_ids and child_id not in selected_branch_nodes:
|
|
roots.add(child_id)
|
|
pending.append(child_id)
|
|
return roots
|
|
|
|
|
|
def _prompt_children_by_node(api_prompt: dict[str, Any]) -> dict[str, set[str]]:
|
|
children: dict[str, set[str]] = {}
|
|
for node_id, node in api_prompt.items():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
for value in inputs.values():
|
|
for ref in _api_refs(value):
|
|
children.setdefault(str(ref[0]), set()).add(str(node_id))
|
|
return children
|
|
|
|
|
|
def _setup_output_surface_node_ids(setup: dict[str, Any]) -> set[str]:
|
|
setup_surface = setup.get("setupSurface", {})
|
|
if not isinstance(setup_surface, dict):
|
|
return set()
|
|
output_ids: set[str] = set()
|
|
outputs = setup_surface.get("outputs")
|
|
if not isinstance(outputs, list):
|
|
return output_ids
|
|
for group in outputs:
|
|
if not isinstance(group, dict) or not isinstance(group.get("nodes"), list):
|
|
continue
|
|
for node in group["nodes"]:
|
|
if isinstance(node, dict) and node.get("id") is not None:
|
|
output_ids.add(str(node["id"]))
|
|
return output_ids
|
|
|
|
|
|
def _app_result_fields(setup: dict[str, Any]) -> list[dict[str, Any]]:
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
if not isinstance(app, dict) or not isinstance(app.get("results"), list):
|
|
return []
|
|
return [field for field in app["results"] if isinstance(field, dict)]
|
|
|
|
|
|
def _result_switch_node_id(
|
|
setup: dict[str, Any],
|
|
result_field: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
api_prompt: dict[str, Any] | None = None,
|
|
) -> str | None:
|
|
app = setup.get("setupSurface", {}).get("app", {})
|
|
if not isinstance(app, dict):
|
|
return None
|
|
switch = app.get("switch")
|
|
if not isinstance(switch, dict):
|
|
return None
|
|
selected_value = _selected_switch_value(switch, run_inputs)
|
|
if selected_value is None:
|
|
return None
|
|
|
|
api_prompt = api_prompt if isinstance(api_prompt, dict) else setup.get("apiPrompt")
|
|
if not isinstance(api_prompt, dict):
|
|
return None
|
|
target = result_field.get("target")
|
|
if not isinstance(target, dict):
|
|
return None
|
|
result_node = api_prompt.get(str(target.get("node")))
|
|
if not isinstance(result_node, dict):
|
|
return None
|
|
result_inputs = result_node.get("inputs")
|
|
if not isinstance(result_inputs, dict):
|
|
return None
|
|
result_ref = result_inputs.get(str(target.get("input")))
|
|
if not _is_api_ref(result_ref):
|
|
return None
|
|
|
|
switch_node_id = str(result_ref[0])
|
|
switch_node = api_prompt.get(switch_node_id)
|
|
if not isinstance(switch_node, dict):
|
|
return None
|
|
switch_inputs = switch_node.get("inputs")
|
|
if not isinstance(switch_inputs, dict):
|
|
return None
|
|
selector_ref = _switch_selector_ref(switch_inputs)
|
|
switch_target = switch.get("target")
|
|
if isinstance(switch_target, dict) and not _selector_matches_switch_target(
|
|
selector_ref,
|
|
switch_target,
|
|
api_prompt,
|
|
):
|
|
return None
|
|
if not _is_api_ref(switch_inputs.get(f"value{selected_value}")):
|
|
return None
|
|
return switch_node_id
|
|
|
|
|
|
def _selected_result_branch_ref(
|
|
setup: dict[str, Any],
|
|
result_field: dict[str, Any],
|
|
run_inputs: dict[str, Any],
|
|
) -> list[Any] | None:
|
|
api_prompt = setup.get("apiPrompt")
|
|
if not isinstance(api_prompt, dict):
|
|
return None
|
|
switch_node_id = _result_switch_node_id(setup, result_field, run_inputs, api_prompt)
|
|
if switch_node_id is None:
|
|
return None
|
|
switch_node = api_prompt.get(switch_node_id)
|
|
if not isinstance(switch_node, dict):
|
|
return None
|
|
switch_inputs = switch_node.get("inputs")
|
|
if not isinstance(switch_inputs, dict):
|
|
return None
|
|
switch = setup.get("setupSurface", {}).get("app", {}).get("switch")
|
|
selected_value = _selected_switch_value(switch, run_inputs) if isinstance(switch, dict) else None
|
|
if selected_value is None:
|
|
return None
|
|
branch_ref = switch_inputs.get(f"value{selected_value}")
|
|
return branch_ref if _is_api_ref(branch_ref) else None
|
|
|
|
|
|
def _selected_switch_value(switch: dict[str, Any], run_inputs: dict[str, Any]) -> int | None:
|
|
key = switch.get("key")
|
|
selected = run_inputs.get(key) if isinstance(key, str) and key in run_inputs else switch.get("default")
|
|
return _switch_option_value(switch, selected)
|
|
|
|
|
|
def _selector_matches_switch_target(selector_ref: Any, switch_target: dict[str, Any], api_prompt: dict[str, Any]) -> bool:
|
|
if not _is_api_ref(selector_ref):
|
|
return False
|
|
source_node_id = str(selector_ref[0])
|
|
target_node_id = str(switch_target.get("node"))
|
|
if source_node_id != target_node_id:
|
|
return False
|
|
source_node = api_prompt.get(source_node_id)
|
|
if not isinstance(source_node, dict):
|
|
return False
|
|
outputs = _publish_input_output_slots(source_node)
|
|
return outputs.get("switch") == int(selector_ref[1])
|
|
|
|
|
|
def _switch_selector_ref(switch_inputs: dict[str, Any]) -> Any:
|
|
if "select" in switch_inputs:
|
|
return switch_inputs.get("select")
|
|
return switch_inputs.get("index")
|
|
|
|
|
|
def _publish_input_output_slots(api_node: dict[str, Any]) -> dict[str, int]:
|
|
if api_node.get("class_type") != "Koolook_PublishInput":
|
|
return {}
|
|
return {
|
|
"sequence_folder": 0,
|
|
"qt_file": 1,
|
|
"single_file": 2,
|
|
"prompt": 3,
|
|
"switch": 4,
|
|
}
|
|
|
|
|
|
def _resolve_prompt_value(prompt: dict[str, Any], value: Any) -> Any:
|
|
if not _is_api_ref(value):
|
|
return value
|
|
node_id = str(value[0])
|
|
output_index = int(value[1])
|
|
node = prompt.get(node_id)
|
|
if not isinstance(node, dict):
|
|
return None
|
|
class_type = node.get("class_type")
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
return None
|
|
if class_type == "Koolook_PublishOutput":
|
|
values = (
|
|
inputs.get("folder"),
|
|
inputs.get("name"),
|
|
inputs.get("version"),
|
|
)
|
|
return values[output_index] if output_index < len(values) else None
|
|
if class_type == "EasyAIPipeline":
|
|
if "WRITE_file_path" in inputs:
|
|
values = (
|
|
inputs.get("WRITE_file_path"),
|
|
inputs.get("output_name", ""),
|
|
inputs.get("version_string", ""),
|
|
inputs.get("output_directory", ""),
|
|
inputs.get("shot_duration", 0),
|
|
inputs.get("seed_value", 0),
|
|
inputs.get("shot_name", ""),
|
|
)
|
|
return values[output_index] if output_index < len(values) else None
|
|
resolved_inputs = {
|
|
key: _resolve_prompt_value(prompt, input_value)
|
|
for key, input_value in inputs.items()
|
|
}
|
|
try:
|
|
try:
|
|
from .k_ai_pipeline import build_pipeline_outputs
|
|
except ImportError: # pragma: no cover - standalone test/import context
|
|
from k_ai_pipeline import build_pipeline_outputs
|
|
|
|
values = build_pipeline_outputs(
|
|
**resolved_inputs,
|
|
create_directory=False,
|
|
check_overwrite=False,
|
|
)
|
|
except Exception:
|
|
return None
|
|
return values[output_index] if output_index < len(values) else None
|
|
return None
|
|
|
|
|
|
def _is_api_ref(value: Any) -> bool:
|
|
return (
|
|
isinstance(value, list)
|
|
and len(value) == 2
|
|
and isinstance(value[0], (str, int))
|
|
and not isinstance(value[1], bool)
|
|
and isinstance(value[1], int)
|
|
)
|
|
|
|
|
|
def _api_refs(value: Any) -> list[list[Any]]:
|
|
if _is_api_ref(value):
|
|
return [value]
|
|
if isinstance(value, dict):
|
|
refs: list[list[Any]] = []
|
|
for child in value.values():
|
|
refs.extend(_api_refs(child))
|
|
return refs
|
|
if isinstance(value, list):
|
|
refs = []
|
|
for child in value:
|
|
refs.extend(_api_refs(child))
|
|
return refs
|
|
return []
|
|
|
|
|
|
def _flatten_history_outputs(raw_outputs: Any) -> list[dict[str, Any]]:
|
|
if not isinstance(raw_outputs, dict):
|
|
return []
|
|
items: list[dict[str, Any]] = []
|
|
for node_id, node_outputs in raw_outputs.items():
|
|
if not isinstance(node_outputs, dict):
|
|
continue
|
|
for kind, values in node_outputs.items():
|
|
if not isinstance(values, list):
|
|
continue
|
|
for value in values:
|
|
item = {"nodeId": str(node_id), "kind": str(kind)}
|
|
if isinstance(value, dict):
|
|
item.update(value)
|
|
else:
|
|
item["value"] = value
|
|
items.append(item)
|
|
return items
|
|
|
|
|
|
def _history_output_matches_contract(item: dict[str, Any], contract_type: str) -> bool:
|
|
kind = str(item.get("kind", "")).lower()
|
|
if contract_type == "video":
|
|
return kind in {"video", "videos", "animated"}
|
|
if contract_type == "image":
|
|
return kind in {"image", "images"}
|
|
return True
|