99 lines
2.8 KiB
JavaScript
99 lines
2.8 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
|
|
const NODE_NAME = "PM_JsonPrompt";
|
|
|
|
function extractPlaceholderKeys(template) {
|
|
if (!template) return [];
|
|
const re = /\{\{\s*\$json\[['\"]([^'\"]+)['\"]\]\.prompt\s*\}\}/g;
|
|
const keys = new Set();
|
|
let m;
|
|
while ((m = re.exec(template)) !== null) {
|
|
keys.add(m[1]);
|
|
}
|
|
return Array.from(keys);
|
|
}
|
|
|
|
function syncInputsFromTemplate(node) {
|
|
const templateWidget = node.widgets?.find((w) => w.name === "template");
|
|
const template = templateWidget?.value ?? "";
|
|
const desiredKeys = extractPlaceholderKeys(template);
|
|
|
|
// Remove internal backend placeholder input if it exists
|
|
for (let i = (node.inputs?.length || 0) - 1; i >= 0; i--) {
|
|
const inp = node.inputs[i];
|
|
if (inp?.name === "_dyn") {
|
|
node.removeInput(i);
|
|
}
|
|
}
|
|
|
|
const existing = new Map();
|
|
(node.inputs || []).forEach((inp, idx) => {
|
|
existing.set(inp.name, { inp, idx });
|
|
});
|
|
|
|
let changed = false;
|
|
|
|
for (const key of desiredKeys) {
|
|
if (!existing.has(key)) {
|
|
node.addInput(key, "STRING");
|
|
changed = true;
|
|
}
|
|
}
|
|
|
|
// Remove inputs that are no longer referenced only if they are not connected
|
|
for (let i = (node.inputs?.length || 0) - 1; i >= 0; i--) {
|
|
const inp = node.inputs[i];
|
|
if (!desiredKeys.includes(inp.name) && inp.link == null) {
|
|
node.removeInput(i);
|
|
changed = true;
|
|
}
|
|
}
|
|
|
|
if (changed) {
|
|
node.setDirtyCanvas(true, true);
|
|
}
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "Polymath.JsonPromptTemplate",
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
if (nodeData.name !== NODE_NAME) return;
|
|
|
|
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
const r = onNodeCreated?.apply(this, arguments);
|
|
|
|
const attach = () => {
|
|
const templateWidget = this.widgets?.find((w) => w.name === "template");
|
|
if (!templateWidget) return;
|
|
|
|
const oldCb = templateWidget.callback;
|
|
templateWidget.callback = (...args) => {
|
|
const rr = oldCb?.apply(templateWidget, args);
|
|
syncInputsFromTemplate(this);
|
|
return rr;
|
|
};
|
|
|
|
syncInputsFromTemplate(this);
|
|
};
|
|
|
|
setTimeout(attach, 0);
|
|
return r;
|
|
};
|
|
|
|
const onConfigure = nodeType.prototype.onConfigure;
|
|
nodeType.prototype.onConfigure = function () {
|
|
const r = onConfigure?.apply(this, arguments);
|
|
setTimeout(() => syncInputsFromTemplate(this), 0);
|
|
return r;
|
|
};
|
|
|
|
const onConnectionsChange = nodeType.prototype.onConnectionsChange;
|
|
nodeType.prototype.onConnectionsChange = function () {
|
|
const r = onConnectionsChange?.apply(this, arguments);
|
|
syncInputsFromTemplate(this);
|
|
return r;
|
|
};
|
|
},
|
|
});
|