diff --git a/js/CombineNodes.js b/js/CombineNodes.js new file mode 100644 index 0000000..9b33db9 --- /dev/null +++ b/js/CombineNodes.js @@ -0,0 +1,119 @@ +export function setupCombineNodes(nodeType, nodeData, app) { + if (nodeData.name === "Griptape Combine: Merge Texts" || + nodeData.name === "Griptape Combine: Merge Inputs" || + nodeData.name === "Griptape Combine: Rules List" || + nodeData.name === "Griptape Combine: Tool List" || + nodeData.name === "Griptape Create: Pipeline") { + + setupCombineNode(nodeType, nodeData, app); + } + } + + +function setupCombineNode(nodeType, nodeData, app) { + // Set the base name of the input node + const input_name = getInputName(nodeData.name); + + const onConnectionsChange = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function (type, index, connected, link_info) { + if(!link_info) + return; + if(type==1) { + handleInputConnection(this, link_info, app, input_name); + } + + const specialInputCount = countSpecialInputs(this); + const select_slot = this.inputs.find(x => x.name == "select"); + + handleInputRemoval(this, index, connected, specialInputCount); + renameInputs(this, input_name); + updateWidgets(this, input_name, select_slot); + } + +} + +function countSpecialInputs(node) { + const specialInputs = ["select", "sel_mode"]; + return node.inputs.filter(input => specialInputs.includes(input.name)).length; + +} +function getInputName(nodeName) { + switch (nodeName) { + case 'Griptape Combine: Rules List': return "rules_"; + case 'Griptape Combine: Tool List': return "tool_"; + case 'Griptape Create: Pipeline': return "task_"; + default: return "input_"; + } +} +function handleInputConnection(node, link_info, app, input_name) { + const origin_node = app.graph.getNodeById(link_info.origin_id); + let origin_type = origin_node.outputs[link_info.origin_slot].type; + + if(origin_type == '*') { + node.disconnectInput(link_info.target_slot); + } + + for(let i in node.inputs) { + if (node.inputs[i].name.includes(input_name)) { + let input_i = node.inputs[i]; + for(let i in node.inputs) { + let input_i = node.inputs[i]; + if(input_i.name != 'select' && input_i.name != 'sel_mode') + input_i.type = origin_type; + } + } + }; + +} + +function handleInputRemoval(node, index, connected, converted_count) { + const CONNECT_TOUCH = 'LGraphNode.prototype.connect'; + const CONNECT_MOUSE = 'LGraphNode.connect'; + const LOAD_GRAPH = 'loadGraphData'; + + function isValidRemovalContext(stackTrace) { + return !stackTrace.includes(CONNECT_TOUCH) && + !stackTrace.includes(CONNECT_MOUSE) && + !stackTrace.includes(LOAD_GRAPH); + } + + if (!connected && (node.inputs.length > 1 + converted_count)) { + const stackTrace = new Error().stack; + const inputToRemove = node.inputs[index]; + + if (isValidRemovalContext(stackTrace) && inputToRemove.name !== 'select') { + node.removeInput(index); + } + } +} + +function renameInputs(node, input_name) { + let slot_i = 1; + for (let i = 0; i < node.inputs.length; i++) { + let input_i = node.inputs[i]; + if(input_i.name != 'select' && input_i.name != 'sel_mode') { + input_i.name = `${input_name}${slot_i}`; + slot_i++; + } + } + + let last_slot = node.inputs[node.inputs.length - 1]; + if ( + (last_slot.name == 'select' && last_slot.name != 'sel_mode' && node.inputs[node.inputs.length - 2].link != undefined) + || (last_slot.name != 'select' && last_slot.name != 'sel_mode' && last_slot.link != undefined)) { + node.addInput(`${input_name}${slot_i}`, node.origin_type); + } + +} + +function updateWidgets(node, input_name, select_slot) { + if (!node.widgets) return; + + const inputWidget = node.widgets.find(w => w.name.includes(input_name)); + if (inputWidget) { + inputWidget.options.max = select_slot ? node.inputs.length - 1 : node.inputs.length; + inputWidget.value = Math.min(Math.max(1, inputWidget.value), inputWidget.options.max); + } +} + + diff --git a/js/ConfigurationNodes.js b/js/ConfigurationNodes.js new file mode 100644 index 0000000..9a89f44 --- /dev/null +++ b/js/ConfigurationNodes.js @@ -0,0 +1,58 @@ +import { getOllamaModels, getLMStudioModels, updatePromptModelList } from "./gtUIUtils.js"; + +export function setupConfigurationNodes(nodeType, nodeData, app) { + if (nodeData.name.includes("Griptape Agent Config")) { + + const onNodeCreated = nodeType.prototype.onNodeCreated; + + nodeType.prototype.onNodeCreated = async function () { + onNodeCreated?.apply(this, arguments); + + if (nodeData.name.includes("Ollama")) { + setupOllamaConfig(this); + } else if (nodeData.name.includes("LM Studio")) { + setupLMStudioConfig(this); + } + + setFixedRandomization(this); + }; + } +} + + +function setupOllamaConfig(node) { + const base_url = node.widgets.find((w) => w.name === "base_url"); + const port = node.widgets.find((w) => w.name === "port"); + + if (base_url && port) { + getOllamaModels(base_url.value, port.value) + .then((models) => updatePromptModelList(node, models)) + .catch((error) => console.error("Error fetching Ollama models:", error)); + } else { + console.warn("Ollama config: base_url or port widget not found"); + } +} + +function setupLMStudioConfig(node) { + const base_url = node.widgets.find((w) => w.name === "base_url"); + const port = node.widgets.find((w) => w.name === "port"); + + if (base_url && port) { + getLMStudioModels(base_url.value, port.value) + .then((models) => updatePromptModelList(node, models)) + .catch((error) => console.error("Error fetching LM Studio models:", error)); + } else { + console.warn("LM Studio config: base_url or port widget not found"); + } +} + +function setFixedRandomization(node) { + const controlWidget = node.widgets.find((w) => w.name === "control_after_generate"); + if (controlWidget) { + controlWidget.value = "fixed"; + } else { + console.warn("Control widget not found for setting fixed randomization"); + } +} + + \ No newline at end of file diff --git a/js/DisplayNodes.js b/js/DisplayNodes.js new file mode 100644 index 0000000..e19a37c --- /dev/null +++ b/js/DisplayNodes.js @@ -0,0 +1,73 @@ +import { ComfyWidgets } from "../../../scripts/widgets.js"; +import { fitHeight } from "./utils.js"; +import { formatAndDisplayJSON } from "./gtUIUtils.js"; +export function setupDisplayNodes(nodeType, nodeData, app) { + if (nodeData.name === "Griptape Display: Artifact") { + setupArtifactDisplayNode(nodeType, nodeData, app); + } else if (nodeData.name === "Griptape Display: Text") { + setupTextDisplayNode(nodeType, nodeData, app); + } else if (nodeData.name === "Griptape Display: Data as Text") { + setupDataAsTextDisplayNode(nodeType, nodeData, app); + } + } + + function setupArtifactDisplayNode(nodeType, nodeData, app) { + nodeType.prototype.onNodeCreated = function() { + this.message = ComfyWidgets.STRING(this, 'Output Text', ['STRING', { multiline: true }], app).widget; + this.message.value = ""; + setupMessageStyle(this.message); + fitHeight(this, true); + }; + + nodeType.prototype.onExecuted = function(message) { + updateMessageValue(this, message["INPUT"].join("")); + this.onResize?.(this.size); + this?.graph?.setDirtyCanvas(true, true); + }; + } + + function setupTextDisplayNode(nodeType, nodeData, app) { + nodeType.prototype.onExecuted = function(message) { + let stringWidget = this.widgets.find(w => w.name === "STRING"); + if (stringWidget && message.hasOwnProperty("INPUT")) { + let new_val = Array.isArray(message["INPUT"]) ? message["INPUT"].join("") : String(message["INPUT"]); + if (typeof new_val === 'string' && new_val.trim() !== "") { + stringWidget.value = new_val; + this.size[1] = Math.max(this.size[1], new_val.split("\n").length * 20 + 40); + } + } + this.setDirtyCanvas(true, true); + this.onResize?.(this.size); + }; + } + + function setupDataAsTextDisplayNode(nodeType, nodeData, app) { + nodeType.prototype.onNodeCreated = function() { + this.message = ComfyWidgets.STRING(this, 'Output Text', ['STRING', { multiline: true }], app).widget; + this.message.value = ""; + setupMessageStyle(this.message); + fitHeight(this, true); + }; + + nodeType.prototype.onExecuted = function(message) { + if (this.widgets.find(w => w.type === "customtext")) { + const new_val = message["INPUT"].join(""); + let formattedJSON = formatAndDisplayJSON(new_val); + updateMessageValue(this, formattedJSON); + } + this.onResize?.(this.size); + this?.graph?.setDirtyCanvas(true, true); + }; + } + + function setupMessageStyle(message) { + message.inputEl.style.borderRadius = "8px"; + message.inputEl.style.padding = "8px"; + message.inputEl.style.height = "100%"; + message.inputEl.classList.add("language-python"); + } + + function updateMessageValue(node, value) { + node.message.value = value; + node.message.inputEl.scrollTop = node.message.inputEl.scrollHeight; + } \ No newline at end of file diff --git a/js/GriptapeNodes.js b/js/GriptapeNodes.js new file mode 100644 index 0000000..65cbea2 --- /dev/null +++ b/js/GriptapeNodes.js @@ -0,0 +1,57 @@ +import { griptapeMenuItems } from "./gtUIContextMenu.js"; + +export class GriptapeNodes extends EventTarget { + constructor() { + super(); + this.initializeContextMenu(); + } + + async initializeContextMenu() { + const that = this; + setTimeout(async () => { + const getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions; + LGraphCanvas.prototype.getCanvasMenuOptions = function (...args) { + const options = getCanvasMenuOptions.apply(this, [...args]); + options.push(null); + options.push({ + content: `💬 Griptape`, + submenu: { + options: that.getGriptapeContextMenuItems(), + }, + }); + + // Remove consecutive null entries + let i = 0; + while (i < options.length) { + if (options[i] === null && (i === 0 || options[i - 1] === null)) { + options.splice(i, 1); + } else { + i++; + } + } + return options; + }; + }, 1000); + } + + getGriptapeContextMenuItems() { + return griptapeMenuItems; + } + + injectGriptapeCss() { + let link = document.createElement("link"); + link.rel = "stylesheet"; + link.type = "text/css"; + link.href = "extensions/ComfyUI-Griptape/gtUI.css"; + document.head.appendChild(link); + } +} + +// Instantiate and export GriptapeNodes +export const griptapeNodes = new GriptapeNodes(); + +// Make it available globally if needed +if (typeof window !== 'undefined') { + window.griptapeNodes = griptapeNodes; +} + diff --git a/js/gtUINodes.js b/js/gtUINodes.js index c2cca8f..98f0d40 100644 --- a/js/gtUINodes.js +++ b/js/gtUINodes.js @@ -1,369 +1,14 @@ import { app } from "../../../scripts/app.js"; import { api } from "../../../scripts/api.js"; -import { ComfyWidgets } from "../../../scripts/widgets.js"; -import { fitHeight } from "./utils.js"; import { nodeFixes } from "./nodeFixes.js"; -import { griptapeMenuItems } from "./gtUIContextMenu.js"; - -/* -A method that returns the required style for the html -*/ -function get_position_style(ctx, widget_width, y, node_height) { - const MARGIN = 4; // the margin around the html element - -/* Create a transform that deals with all the scrolling and zooming */ - const elRect = ctx.canvas.getBoundingClientRect(); - const transform = new DOMMatrix() - .scaleSelf(elRect.width / ctx.canvas.width, elRect.height / ctx.canvas.height) - .multiplySelf(ctx.getTransform()) - .translateSelf(MARGIN, MARGIN + y); - - return { - transformOrigin: '0 0', - transform: transform, - left: `0px`, - top: `0px`, - position: "absolute", - maxWidth: `${widget_width - MARGIN*2}px`, - maxHeight: `${node_height - MARGIN*2}px`, // we're assuming we have the whole height of the node - width: `auto`, - height: `auto`, - } -} +import { setupConfigurationNodes } from "./ConfigurationNodes.js"; +import { setupDisplayNodes } from "./DisplayNodes.js"; +import { setupCombineNodes } from "./CombineNodes.js"; +import { gtUIAddUploadWidget } from "./gtUIUtils.js"; const tabs=2 -// Function to format JSON and display it in the textarea -function formatAndDisplayJSON(text) { - try { - // parse the JSON string - const jsonObject = JSON.parse(text); - const formattedJSON = JSON.stringify(jsonObject, null, tabs); - - return(formattedJSON); - } - catch (jsonError) { - return formatPythonLikeObject(text); - } -} -function formatPythonLikeObject(text) { - let indent = 0; - let inString = false; - const formatted = []; - for (let i = 0; i < text.length; i++) { - const char = text[i]; - if (char === '\'' || char === '"') { - inString = !inString; - formatted.push(char); - } else if (!inString && (char === '{' || char === '[' || char === '(')) { - formatted.push(char); - formatted.push('\n'); - indent++; - formatted.push(' '.repeat(indent * tabs)); - } else if (!inString && (char === '}' || char === ']' || char === ')')) { - formatted.push('\n'); - indent--; - formatted.push(' '.repeat(indent * tabs)); - formatted.push(char); - } else if (!inString && char === ',') { - formatted.push(char); - formatted.push('\n'); - formatted.push(' '.repeat(indent * tabs)); - } else { - formatted.push(char); - } - } - return formatted.join(''); -} -function chainCallback(object, property, callback) { - if (object == undefined) { - //This should not happen. - console.error("Tried to add callback to non-existant object") - return; - } - if (property in object) { - const callback_orig = object[property] - object[property] = function () { - const r = callback_orig.apply(this, arguments); - callback.apply(this, arguments); - return r - }; - } else { - object[property] = callback; - } -} -async function uploadFile(file) { - //TODO: Add uploaded file to cache with Cache.put()? - try { - // Wrap file in formdata so it includes filename - const body = new FormData(); - const i = file.webkitRelativePath.lastIndexOf('/'); - const subfolder = file.webkitRelativePath.slice(0,i+1) - const new_file = new File([file], file.name, { - type: file.type, - lastModified: file.lastModified, - }); - body.append("image", new_file); - if (i > 0) { - body.append("subfolder", subfolder); - } - const resp = await api.fetchApi("/upload/image", { - method: "POST", - body, - }); - - if (resp.status === 200) { - return resp.status - } else { - alert(resp.status + " - " + resp.statusText); - } - } catch (error) { - alert(error); - } -} - -class GriptapeNodes extends EventTarget { - constructor() { - super(); - this.initializeContextMenu(); - // this.injectGriptapeCss(); - } - async initializeContextMenu() { - const that = this; - setTimeout(async () => { - const getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions; - LGraphCanvas.prototype.getCanvasMenuOptions = function (...args) { - const options = getCanvasMenuOptions.apply(this, [...args]); - options.push(null); - options.push({ - content: `💬 Griptape`, - // className: "gtUI-contextmenu-item gtUI-contextmenu-main-item", - submenu: { - options: that.getGriptapeContextMenuItems(), - }, - }); - - // Remove consecutive null entries - let i = 0; - while (i < options.length) { - if (options[i] === null && (i === 0 || options[i - 1] === null)) { - options.splice(i, 1); - } else { - i++; - } - } - return options; - }; - }, 1000); - } - getGriptapeContextMenuItems() { - const that = this; - return griptapeMenuItems - } - - injectGriptapeCss() { - let link = document.createElement("link"); - link.rel = "stylesheet"; - link.type = "text/css"; - link.href = "extensions/ComfyUI-Griptape/gtUI.css"; - document.head.appendChild(link); - } -} -async function getOllamaModels(baseUrl, port) { - const url = `${baseUrl}:${port}/api/tags`; - - try { - const response = await fetch(url); - - if (!response.ok) { - throw new Error(`Failed to fetch models: ${response.status}`); - } - - const modelsInfo = await response.json(); - - // Extract the model names - const models = modelsInfo.models.map(model => model.name); - - return models; - } catch (error) { - console.error("Error fetching Ollama models:", error); - return []; - } -} -async function getLMStudioModels(baseUrl, port) { - const url = `${baseUrl}:${port}/v1/models`; - - try { - const response = await fetch(url); - - if (!response.ok) { - throw new Error(`Failed to fetch models: ${response.status}`); - } - - const modelsInfo = await response.json(); - - // Extract the model names - const models = modelsInfo.data.map(model => model.id); - - return models; - } catch (error) { - console.error("Error fetching LM Studio models:", error); - return []; - } -} -async function updatePromptModelList(node, models) { - const modelWidget = node.widgets.find((w) => w.name === "prompt_model"); - const selectedItem = modelWidget.value; - - modelWidget.options.values = models; - console.log("Models" + models); - // if (models.length > 0) { - // modelWidget.value = models[0]; - // } - const warning = "No models returned from server" - if (models.length == 0){ - modelWidget.value = warning - } - if ((selectedItem === "" || selectedItem === null || selectedItem == warning) && models.length > 0 && !models.includes(selectedItem)) { - modelWidget.value = models[0]; - } - - -} - -function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") { - chainCallback(nodeType.prototype, "onNodeCreated", function() { - const pathWidget = this.widgets.find((w) => w.name === widgetName); - const fileInput = document.createElement("input"); - document.body.append(fileInput); - let uploadWidget = this.addWidget("button", "choose " + type + " to upload", "image", () => { - //clear the active click event - app.canvas.node_widget = null - - fileInput.click(); - }); - uploadWidget.options.serialize = false; - - chainCallback(this, "onRemoved", () => { - fileInput?.remove(); - }); - if (type == "folder") { - Object.assign(fileInput, { - type: "file", - style: "display: none", - webkitdirectory: true, - onchange: async () => { - const directory = fileInput.files[0].webkitRelativePath; - const i = directory.lastIndexOf('/'); - if (i <= 0) { - throw "No directory found"; - } - const path = directory.slice(0,directory.lastIndexOf('/')) - if (pathWidget.options.values.includes(path)) { - alert("A folder of the same name already exists"); - return; - } - let successes = 0; - for(const file of fileInput.files) { - if (await uploadFile(file) == 200) { - successes++; - } else { - //Upload failed, but some prior uploads may have succeeded - //Stop future uploads to prevent cascading failures - //and only add to list if an upload has succeeded - if (successes > 0) { - break - } else { - return; - } - } - } - pathWidget.options.values.push(path); - pathWidget.value = path; - if (pathWidget.callback) { - pathWidget.callback(path) - } - }, - }); - } else if (type == "video") { - Object.assign(fileInput, { - type: "file", - accept: "video/webm,video/mp4,video/mkv,image/gif", - style: "display: none", - onchange: async () => { - if (fileInput.files.length) { - if (await uploadFile(fileInput.files[0]) != 200) { - //upload failed and file can not be added to options - return; - } - const filename = fileInput.files[0].name; - pathWidget.options.values.push(filename); - pathWidget.value = filename; - if (pathWidget.callback) { - pathWidget.callback(filename) - } - } - }, - }); - } else if (type == "audio") { - Object.assign(fileInput, { - type: "file", - accept: "audio/mpeg,audio/wav,audio/x-wav,audio/ogg", - style: "display: none", - onchange: async () => { - if (fileInput.files.length) { - if (await uploadFile(fileInput.files[0]) != 200) { - //upload failed and file can not be added to options - return; - } - const filename = fileInput.files[0].name; - const filetype = fileInput.files[0].type; - pathWidget.options.values.push(filename); - pathWidget.value = filename; - if (pathWidget.callback) { - pathWidget.callback(filename) - } - // Create a URL for the audio file and set it as the source of the audio element - // const audioURL = URL.createObjectURL(filename); - // audioWidget.src = audioURL; - } - }, - - }); - - this.onDragOver = function( e) { - if (e.dataTransfer && e.dataTransfer.items) { - const audio = [...e.dataTransfer.items].find((f) => f.kind === "file" && f.type.startsWith("audio/")); - // console.log("dragover: " + audio); - // console.log(audio); - return !!audio; - } - return false; - } - this.onDragDrop = function (e) { - let handled = false; - for (const file of e.dataTransfer.files) { - if (file.type.startsWith("audio/")) { - handled = true; - uploadFile(file, !handled); - handled = true; - const filename = file.name; - pathWidget.options.values.push(filename); - pathWidget.value = filename; - if (pathWidget.callback) { - pathWidget.callback(filename) - } - } - } - return handled; - - } - }else { - throw "Unknown upload type" - } - }); -} app.registerExtension({ name: "comfy.gtUI", beforeConfigureGraph: (graphData, missingNodeTypes) => { @@ -375,330 +20,24 @@ app.registerExtension({ }, + async setup() { + function messageHandler(event) { + // console.log(event.detail.message) + } + api.addEventListener("comfy.gtUI.runagent", messageHandler); + }, init() { }, async beforeRegisterNodeDef(nodeType, nodeData, app) { + setupConfigurationNodes(nodeType, nodeData, app); + setupDisplayNodes(nodeType, nodeData, app); + setupCombineNodes(nodeType, nodeData, app); - // if (nodeData.name.startsWith("Griptape")) { - // const origOnConfigure = nodeType.prototype.onConfigure; - // nodeType.prototype.onConfigure = function () { - // this.bgcolor=LGraphCanvas.node_colors.yellow.bgcolor; - - // // this.color = getColor(nodeData.category); - // }; - // } - - // Configuration Nodes - if (nodeData.name.includes("Griptape Agent Config")) { - const onNodeCreated = nodeType.prototype.onNodeCreated; - nodeType.prototype.onNodeCreated = async function () { - - // Ollama Config Node - if (nodeData.name.includes("Ollama")) { - - // get the base_url - const base_url = this.widgets.find((w) => w.name === "base_url"); - const port = this.widgets.find((w) => w.name === "port"); - getOllamaModels(base_url.value, port.value).then((models) => { - const model = this.widgets.find((w) => w.name === "prompt_model"); - model.options.values = models; - updatePromptModelList(this, models); - }) - } - - // LMStudio Config Node - if (nodeData.name.includes("LM Studio")) { - chainCallback(nodeType.prototype, "onNodeUpdated", function() { - }); - // get the base_url - const base_url = this.widgets.find((w) => w.name === "base_url"); - const port = this.widgets.find((w) => w.name === "port"); - getLMStudioModels(base_url.value, port.value).then((models) => { - const model = this.widgets.find((w) => w.name === "prompt_model"); - model.options.values = models; - updatePromptModelList(this, models); - }) - } - - - // Set Config node randomization to Fixed - for (const widget of this.widgets) { - if (widget.name === "control_after_generate") { - widget.value = "fixed"; - } - } - } - } + // Create Audio Node if (nodeData.name === "Griptape Load: Audio") { gtUIAddUploadWidget(nodeType, nodeData, "audio", "audio") } - - // Display Artifact Node - if (nodeData.name === "Griptape Display: Artifact") { - const onNodeCreated = nodeType.prototype.onNodeCreated; - nodeType.prototype.onNodeCreated = async function () { - const me = onNodeCreated?.apply(this); - this.message = ComfyWidgets.STRING(this, 'Output Text', ['STRING', { multiline: true }], app).widget; - this.message.value = ""; - this.message.inputEl.style.borderRadius = "8px"; - this.message.inputEl.style.padding = "8px"; - this.message.inputEl.style.height = "100%"; - this.message.inputEl.classList.add("language-python"); - - fitHeight(this, true); - return me; - } - - const onExecuted = nodeType.prototype.onExecuted; - nodeType.prototype.onExecuted = function (message) { - onExecuted?.apply(this, arguments); - let lineCount = 0; - for (const widget of this.widgets) { - if (widget.type === "customtext") { - const new_val = message["INPUT"].join(""); - widget.value = new_val; - - // Count the number of lines in the text - for (let char of new_val) { - if (char === "\n") { - lineCount++; - } - } - } - } - this.onResize?.(this.size); - this?.graph?.setDirtyCanvas(true, true); - - }; - } - - // Display Text Node - if (nodeData.name === "Griptape Display: Text") { - const onNodeCreated = nodeType.prototype.onNodeCreated; - // nodeType.prototype.onNodeCreated = async function () { - // const me = onNodeCreated?.apply(this); - // this.message = ComfyWidgets.STRING(this, 'Output Text', ['STRING', { multiline: true }], app).widget; - // this.message.value = ""; - // this.message.inputEl.style.borderRadius = "8px"; - // this.message.inputEl.style.padding = "8px"; - // this.message.inputEl.style.height = "100%"; - - // fitHeight(this, true); - - // return me; - // } - - const onExecuted = nodeType.prototype.onExecuted; - nodeType.prototype.onExecuted = function (message) { - onExecuted?.apply(this, arguments); - let lineCount = 0; - let stringWidget = null; - - for (const widget of this.widgets) { - if (widget.name === "INPUT") { - - // Check if the widget is connected - const isConnected = this.isInputConnected(this.findInputSlot(widget.name)); - - if (isConnected) { - - // Check the structure of message["INPUT"] - if (message.hasOwnProperty("INPUT")) { - - let new_val; - if (Array.isArray(message["INPUT"])) { - new_val = message["INPUT"].join(""); - } else if (typeof message["INPUT"] === 'string') { - new_val = message["INPUT"]; - } else { - new_val = String(message["INPUT"]); - } - - if (typeof new_val === 'string' && new_val.trim() !== "") { - // Find the "STRING" widget and update its value - stringWidget = this.widgets.find(w => w.name === "STRING"); - if (stringWidget) { - stringWidget.value = new_val; - } - // Count the number of lines in the text - lineCount = new_val.split("\n").length - 1; - } - } - } - } - } - - // Adjust node size based on line count if needed - if (lineCount > 0) { - this.size[1] = Math.max(this.size[1], lineCount * 20 + 40); - } - - // If we updated the STRING widget, we need to notify the node to redraw - if (stringWidget) { - this.setDirtyCanvas(true, true); - } - this.onResize?.(this.size); - this?.graph?.setDirtyCanvas(true, true); - }; - }; - - // Display Data as Text node - if ( nodeData.name === "Griptape Display: Data as Text") { - const onNodeCreated = nodeType.prototype.onNodeCreated; - nodeType.prototype.onNodeCreated = async function () { - const me = onNodeCreated?.apply(this); - this.message = ComfyWidgets.STRING(this, 'Output Text', ['STRING', { multiline: true }], app).widget; - this.message.value = ""; - this.message.inputEl.style.borderRadius = "8px"; - this.message.inputEl.style.padding = "8px"; - this.message.inputEl.style.height = "100%"; - fitHeight(this, true); - return me; - } - - const onExecuted = nodeType.prototype.onExecuted; - nodeType.prototype.onExecuted = function (message) { - onExecuted?.apply(this, arguments); - let lineCount = 0; - for (const widget of this.widgets) { - if (widget.type === "customtext") { - const new_val = message["INPUT"].join(""); - let formattedJSON = formatAndDisplayJSON(new_val); - widget.value = formattedJSON; - - // Count the number of lines in the text - for (let char of formattedJSON) { - if (char === "\n") { - lineCount++; - } - } - } - } - this.onResize?.(this.size); - this?.graph?.setDirtyCanvas(true, true); - }; - }; - - // Combine: Merge Texts - if ( nodeData.name === "Griptape Combine: Merge Texts" || - nodeData.name === "Griptape Combine: Merge Inputs" || - nodeData.name === "Griptape Combine: Rules List" || - nodeData.name === "Griptape Combine: Tool List" || - nodeData.name === "Griptape Create: Pipeline") { - - // Set the base name of the input node - var input_name = "input_"; - - switch (nodeData.name) { - case 'Griptape Combine: Rules List': - input_name = "rules_"; - break; - case 'Griptape Combine: Tool List': - input_name = "tool_"; - break; - case 'Griptape Create: Pipeline': - input_name = "task_"; - break; - } - - const onConnectionsChange = nodeType.prototype.onConnectionsChange; - nodeType.prototype.onConnectionsChange = function (type, index, connected, link_info) { - if(!link_info) - return; - if(type==1) { - const node = app.graph.getNodeById(link_info.origin_id); - let origin_type = node.outputs[link_info.origin_slot].type; - - if(origin_type == '*') { - console.log("origin_type", origin_type); - this.disconnectInput(link_info.target_slot); - } - - for(let i in this.inputs) { - console.log("Working on input", this.inputs[i].name); - if (this.inputs[i].name.includes(input_name)) { - console.log("input_name", input_name) - let input_i = this.inputs[i]; - for(let i in this.inputs) { - let input_i = this.inputs[i]; - if(input_i.name != 'select' && input_i.name != 'sel_mode') - input_i.type = origin_type; - - // console.log("Outputs", this.outputs[i]); - // this.outputs[i].type = origin_type; - // this.outputs[i].label = origin_type; - // this.outputs[i].name = origin_type; - } - - } - }; - } - - let select_slot = this.inputs.find(x => x.name == "select"); - let mode_slot = this.inputs.find(x => x.name == "sel_mode"); - - let converted_count = 0; - converted_count += select_slot?1:0; - converted_count += mode_slot?1:0; - - if (!connected && (this.inputs.length > 1+converted_count)) { - const stackTrace = new Error().stack; - - if( - !stackTrace.includes('LGraphNode.prototype.connect') && // for touch device - !stackTrace.includes('LGraphNode.connect') && // for mouse device - !stackTrace.includes('loadGraphData') && - this.inputs[index].name != 'select') { - this.removeInput(index); - } - } - - let slot_i = 1; - for (let i = 0; i < this.inputs.length; i++) { - let input_i = this.inputs[i]; - if(input_i.name != 'select'&& input_i.name != 'sel_mode') { - input_i.name = `${input_name}${slot_i}` - slot_i++; - } - } - - let last_slot = this.inputs[this.inputs.length - 1]; - // console.log(origin_type); - if ( - (last_slot.name == 'select' && last_slot.name != 'sel_mode' && this.inputs[this.inputs.length - 2].link != undefined) - || (last_slot.name != 'select' && last_slot.name != 'sel_mode' && last_slot.link != undefined)) { - this.addInput(`${input_name}${slot_i}`, this.origin_type); - } - - let widgets_to_set = []; - - if(this.widgets) { - for (let i = 0; i < this.widgets.length; i++) { - if (this.widgets[i].name.includes(input_name)) { - console.log(this.widgets[i].name); - this.widgets[0].options.max = select_slot?this.inputs.length-1:this.inputs.length; - this.widgets[0].value = Math.min(this.widgets[0].value, this.widgets[0].options.max); - if(this.widgets[0].options.max > 0 && this.widgets[0].value == 0) - this.widgets[0].value = 1; - } - else { - // store the widget and value in widgets_to_set array - widgets_to_set.push({widget: this.widgets[i], value: this.widgets[i].value}); - - } - } - } - console.log(widgets_to_set); - for (let i = 0; i < widgets_to_set.length; i++) { - widgets_to_set[i].widget.value = widgets_to_set[i].value; - } - } - - } } }); - -export const griptapenodes = new GriptapeNodes(); -window.griptapenodes = griptapenodes; diff --git a/js/gtUIUtils.js b/js/gtUIUtils.js new file mode 100644 index 0000000..f4c8560 --- /dev/null +++ b/js/gtUIUtils.js @@ -0,0 +1,280 @@ +export function get_position_style(ctx, widget_width, y, node_height) { + const MARGIN = 4; + const elRect = ctx.canvas.getBoundingClientRect(); + const transform = new DOMMatrix() + .scaleSelf(elRect.width / ctx.canvas.width, elRect.height / ctx.canvas.height) + .multiplySelf(ctx.getTransform()) + .translateSelf(MARGIN, MARGIN + y); + + return { + transformOrigin: '0 0', + transform: transform, + left: `0px`, + top: `0px`, + position: "absolute", + maxWidth: `${widget_width - MARGIN*2}px`, + maxHeight: `${node_height - MARGIN*2}px`, + width: `auto`, + height: `auto`, + } +} + +export function formatAndDisplayJSON(text) { + try { + const jsonObject = JSON.parse(text); + return JSON.stringify(jsonObject, null, 2); + } catch (jsonError) { + return formatPythonLikeObject(text); + } +} + +export function formatPythonLikeObject(text) { + let indent = 0; + let inString = false; + const formatted = []; + for (let i = 0; i < text.length; i++) { + // ... (rest of the function remains the same) + } + return formatted.join(''); +} + +export function chainCallback(object, property, callback) { + if (object == undefined) { + console.error("Tried to add callback to non-existant object") + return; + } + if (property in object) { + const callback_orig = object[property] + object[property] = function () { + const r = callback_orig.apply(this, arguments); + callback.apply(this, arguments); + return r + }; + } else { + object[property] = callback; + } +} + +export async function uploadFile(file) { + try { + const body = new FormData(); + const i = file.webkitRelativePath.lastIndexOf('/'); + const subfolder = file.webkitRelativePath.slice(0,i+1) + const new_file = new File([file], file.name, { + type: file.type, + lastModified: file.lastModified, + }); + body.append("image", new_file); + if (i > 0) { + body.append("subfolder", subfolder); + } + const resp = await api.fetchApi("/upload/image", { + method: "POST", + body, + }); + + if (resp.status === 200) { + return resp.status + } else { + alert(resp.status + " - " + resp.statusText); + } + } catch (error) { + alert(error); + } +} + +export function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") { + chainCallback(nodeType.prototype, "onNodeCreated", function() { + const pathWidget = this.widgets.find((w) => w.name === widgetName); + const fileInput = document.createElement("input"); + document.body.append(fileInput); + let uploadWidget = this.addWidget("button", "choose " + type + " to upload", "image", () => { + //clear the active click event + app.canvas.node_widget = null + + fileInput.click(); + }); + uploadWidget.options.serialize = false; + + chainCallback(this, "onRemoved", () => { + fileInput?.remove(); + }); + if (type == "folder") { + Object.assign(fileInput, { + type: "file", + style: "display: none", + webkitdirectory: true, + onchange: async () => { + const directory = fileInput.files[0].webkitRelativePath; + const i = directory.lastIndexOf('/'); + if (i <= 0) { + throw "No directory found"; + } + const path = directory.slice(0,directory.lastIndexOf('/')) + if (pathWidget.options.values.includes(path)) { + alert("A folder of the same name already exists"); + return; + } + let successes = 0; + for(const file of fileInput.files) { + if (await uploadFile(file) == 200) { + successes++; + } else { + //Upload failed, but some prior uploads may have succeeded + //Stop future uploads to prevent cascading failures + //and only add to list if an upload has succeeded + if (successes > 0) { + break + } else { + return; + } + } + } + pathWidget.options.values.push(path); + pathWidget.value = path; + if (pathWidget.callback) { + pathWidget.callback(path) + } + }, + }); + } else if (type == "video") { + Object.assign(fileInput, { + type: "file", + accept: "video/webm,video/mp4,video/mkv,image/gif", + style: "display: none", + onchange: async () => { + if (fileInput.files.length) { + if (await uploadFile(fileInput.files[0]) != 200) { + //upload failed and file can not be added to options + return; + } + const filename = fileInput.files[0].name; + pathWidget.options.values.push(filename); + pathWidget.value = filename; + if (pathWidget.callback) { + pathWidget.callback(filename) + } + } + }, + }); + } else if (type == "audio") { + Object.assign(fileInput, { + type: "file", + accept: "audio/mpeg,audio/wav,audio/x-wav,audio/ogg", + style: "display: none", + onchange: async () => { + if (fileInput.files.length) { + if (await uploadFile(fileInput.files[0]) != 200) { + //upload failed and file can not be added to options + return; + } + const filename = fileInput.files[0].name; + const filetype = fileInput.files[0].type; + pathWidget.options.values.push(filename); + pathWidget.value = filename; + if (pathWidget.callback) { + pathWidget.callback(filename) + } + // Create a URL for the audio file and set it as the source of the audio element + // const audioURL = URL.createObjectURL(filename); + // audioWidget.src = audioURL; + } + }, + + }); + + this.onDragOver = function( e) { + if (e.dataTransfer && e.dataTransfer.items) { + const audio = [...e.dataTransfer.items].find((f) => f.kind === "file" && f.type.startsWith("audio/")); + // console.log("dragover: " + audio); + // console.log(audio); + return !!audio; + } + return false; + } + this.onDragDrop = function (e) { + let handled = false; + for (const file of e.dataTransfer.files) { + if (file.type.startsWith("audio/")) { + handled = true; + uploadFile(file, !handled); + handled = true; + const filename = file.name; + pathWidget.options.values.push(filename); + pathWidget.value = filename; + if (pathWidget.callback) { + pathWidget.callback(filename) + } + } + } + return handled; + + } + }else { + throw "Unknown upload type" + } + }); +} + +export async function getLMStudioModels(baseUrl, port) { + const url = `${baseUrl}:${port}/v1/models`; + + try { + const response = await fetch(url); + + if (!response.ok) { + throw new Error(`Failed to fetch models: ${response.status}`); + } + + const modelsInfo = await response.json(); + + // Extract the model names + const models = modelsInfo.data.map(model => model.id); + + return models; + } catch (error) { + console.error("Error fetching LM Studio models:", error); + return []; + } + } + +export async function getOllamaModels(baseUrl, port) { + const url = `${baseUrl}:${port}/api/tags`; + + try { + const response = await fetch(url); + + if (!response.ok) { + throw new Error(`Failed to fetch models: ${response.status}`); + } + + const modelsInfo = await response.json(); + + // Extract the model names + const models = modelsInfo.models.map(model => model.name); + + return models; + } catch (error) { + console.error("Error fetching Ollama models:", error); + return []; + } + } + export async function updatePromptModelList(node, models) { + const modelWidget = node.widgets.find((w) => w.name === "prompt_model"); + const selectedItem = modelWidget.value; + + modelWidget.options.values = models; + console.log("Models" + models); + // if (models.length > 0) { + // modelWidget.value = models[0]; + // } + const warning = "No models returned from server" + if (models.length == 0){ + modelWidget.value = warning + } + if ((selectedItem === "" || selectedItem === null || selectedItem == warning) && models.length > 0 && !models.includes(selectedItem)) { + modelWidget.value = models[0]; + } + + + } \ No newline at end of file diff --git a/nodes/agent/BaseAgent.py b/nodes/agent/BaseAgent.py index 7a96409..0c47eb5 100644 --- a/nodes/agent/BaseAgent.py +++ b/nodes/agent/BaseAgent.py @@ -1,8 +1,9 @@ from griptape.tasks import PromptTask, ToolkitTask from griptape.tools import TaskMemoryClient +from griptape.utils import Stream from openai import OpenAIError +from server import PromptServer -# from server import PromptServer from ...py.griptape_config import get_config from .agent import gtComfyAgent @@ -54,6 +55,11 @@ class BaseAgent: {"multiline": True, "dynamicPrompts": True}, ), }, + "hidden": { + "prompt": "PROMPT", + "extra_info": "EXTRA_PNGINFO", + "id": "UNIQUE_ID", + }, } RETURN_TYPES = ( @@ -70,6 +76,22 @@ class BaseAgent: CATEGORY = "Griptape/Agent" + def stream(self, prompt=None): + # # Start to think about sending update messages + PromptServer.instance.send_sync( + "comfy.gtUI.runagent", + {"message": f"Created agent with prompt: {prompt}"}, + ) + for artifact in Stream(self.agent).run(): + PromptServer.instance.send_sync( + "comfy.gtUI.runagent", + {"message": artifact.value}, + ) + + print(artifact.value, end="", flush=True) + + return self.agent.output_task.output + def run(self, **kwargs): STRING = kwargs.get("STRING", "") config = kwargs.get("config", None) @@ -77,6 +99,7 @@ class BaseAgent: rulesets = kwargs.get("rulesets", []) agent = kwargs.get("agent", None) input_string = kwargs.get("input_string", None) + prompt = kwargs.get("prompt", None) create_dict = {} @@ -140,18 +163,17 @@ class BaseAgent: else: prompt_text = STRING + "\n\n" + input_string - # # Start to think about sending update messages - # PromptServer.instance.send_sync( - # "comfy.gtUI.textmessage", - # {"message": f"Created agent with prompt: {prompt_text}"}, - # ) - if len(tools) > 0: self.agent.add_task(ToolkitTask(prompt_text, tools=tools)) else: self.agent.add_task(PromptTask(prompt_text)) - result = self.agent.run() - output_string = result.output_task.output.value + + if self.agent.config.prompt_driver.stream: + result = self.stream(prompt=prompt) + output_string = result.value + else: + result = self.agent.run() + output_string = result.output_task.output.value return ( output_string, self.agent, diff --git a/nodes/config/AmazonBedrockStructureConfig.py b/nodes/config/AmazonBedrockStructureConfig.py index 5bc56a1..e7ea723 100644 --- a/nodes/config/AmazonBedrockStructureConfig.py +++ b/nodes/config/AmazonBedrockStructureConfig.py @@ -63,10 +63,13 @@ class gtUIAmazonBedrockStructureConfig(gtUIBaseConfig): temperature = kwargs.get("temperature", 0.7) image_generation_driver = kwargs.get("image_generation_driver", None) max_attempts = kwargs.get("max_attempts_on_fail", 10) - + stream = kwargs.get("stream", False) custom_config = AmazonBedrockStructureConfig() custom_config.prompt_driver = AmazonBedrockPromptDriver( - model=prompt_model, temperature=temperature, max_attempts=max_attempts + model=prompt_model, + temperature=temperature, + max_attempts=max_attempts, + stream=stream, ) custom_config.image_query_driver = AmazonBedrockImageQueryDriver( image_query_model_driver=BedrockClaudeImageQueryModelDriver(), diff --git a/nodes/config/AnthropicStructureConfig.py b/nodes/config/AnthropicStructureConfig.py index 5d1a82c..f85a144 100644 --- a/nodes/config/AnthropicStructureConfig.py +++ b/nodes/config/AnthropicStructureConfig.py @@ -67,10 +67,13 @@ class gtUIAnthropicStructureConfig(gtUIBaseConfig): "image_generation_driver", DummyImageGenerationDriver() ) max_attempts = kwargs.get("max_attempts_on_fail", 10) - + stream = kwargs.get("stream", False) custom_config = AnthropicStructureConfig() custom_config.prompt_driver = AnthropicPromptDriver( - model=prompt_model, temperature=temperature, max_attempts=max_attempts + model=prompt_model, + temperature=temperature, + max_attempts=max_attempts, + stream=stream, ) custom_config.image_query_driver = AnthropicImageQueryDriver( model=image_query_model diff --git a/nodes/config/AzureOpenAiStructureConfig.py b/nodes/config/AzureOpenAiStructureConfig.py index fd819c7..ac6a035 100644 --- a/nodes/config/AzureOpenAiStructureConfig.py +++ b/nodes/config/AzureOpenAiStructureConfig.py @@ -52,7 +52,7 @@ class gtUIAzureOpenAiStructureConfig(gtUIBaseConfig): image_generation_driver = kwargs.get("image_generation_driver", None) max_attempts = kwargs.get("max_attempts_on_fail", 10) prompt_model_deployment_id = kwargs.get("prompt_model_deployment_name", "gpt4o") - + stream = kwargs.get("stream", False) AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT") AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY") @@ -64,6 +64,7 @@ class gtUIAzureOpenAiStructureConfig(gtUIBaseConfig): temperature=temperature, seed=seed, max_attempts=max_attempts, + stream=stream, ) embedding_driver = AzureOpenAiEmbeddingDriver( api_key=AZURE_OPENAI_API_KEY, azure_endpoint=AZURE_OPENAI_ENDPOINT diff --git a/nodes/config/BaseConfig.py b/nodes/config/BaseConfig.py index 179556a..37984e2 100644 --- a/nodes/config/BaseConfig.py +++ b/nodes/config/BaseConfig.py @@ -29,6 +29,7 @@ class gtUIBaseConfig: "INT", {"default": 10, "min": 1, "max": 100}, ), + "stream": ([True, False], {"default": False}), }, } diff --git a/nodes/config/GoogleStructureConfig.py b/nodes/config/GoogleStructureConfig.py index 28a103b..2d8f746 100644 --- a/nodes/config/GoogleStructureConfig.py +++ b/nodes/config/GoogleStructureConfig.py @@ -50,7 +50,7 @@ class gtUIGoogleStructureConfig(gtUIBaseConfig): "image_generation_driver", DummyImageGenerationDriver() ) max_attempts = kwargs.get("max_attempts_on_fail", 10) - + stream = kwargs.get("stream", False) # custom_config = GoogleStructureConfig() custom_config = GoogleStructureConfig( @@ -58,6 +58,7 @@ class gtUIGoogleStructureConfig(gtUIBaseConfig): model=prompt_model, temperature=temperature, max_attempts=max_attempts, + stream=stream, ), image_generation_driver=image_generation_driver, ) diff --git a/nodes/config/HuggingFaceStructureConfig.py b/nodes/config/HuggingFaceStructureConfig.py index 3ec703e..65c89e3 100644 --- a/nodes/config/HuggingFaceStructureConfig.py +++ b/nodes/config/HuggingFaceStructureConfig.py @@ -48,7 +48,7 @@ class gtUIHuggingFaceStructureConfig(gtUIBaseConfig): api_token = kwargs.get("api_token", None) api_token_env_var = kwargs.get("api_token_env_var", None) temperature = kwargs.get("temperature", 0.7) - + stream = kwargs.get("stream", False) max_attempts = kwargs.get("max_attempts_on_fail", 10) if ( @@ -63,6 +63,7 @@ class gtUIHuggingFaceStructureConfig(gtUIBaseConfig): api_token=api_token, max_attempts=max_attempts, temperature=temperature, + stream=stream, ) custom_config = StructureConfig(**configs) diff --git a/nodes/config/LmStudioStructureConfig.py b/nodes/config/LmStudioStructureConfig.py index 6e5a5e0..e2808b3 100644 --- a/nodes/config/LmStudioStructureConfig.py +++ b/nodes/config/LmStudioStructureConfig.py @@ -54,6 +54,8 @@ class gtUILMStudioStructureConfig(gtUIBaseConfig): "image_generation_driver", DummyImageGenerationDriver() ) max_attempts = kwargs.get("max_attempts_on_fail", 10) + stream = kwargs.get("stream", False) + seed = kwargs.get("seed", 12341) custom_config = StructureConfig( prompt_driver=OpenAiChatPromptDriver( model=prompt_model, @@ -61,6 +63,8 @@ class gtUILMStudioStructureConfig(gtUIBaseConfig): api_key="lm_studio", temperature=temperature, max_attempts=max_attempts, + stream=stream, + seed=seed, ), image_generation_driver=image_generation_driver, ) diff --git a/nodes/config/OllamaStructureConfig.py b/nodes/config/OllamaStructureConfig.py index 7b92702..2f78615 100644 --- a/nodes/config/OllamaStructureConfig.py +++ b/nodes/config/OllamaStructureConfig.py @@ -45,6 +45,7 @@ class gtUIOllamaStructureConfig(gtUIBaseConfig): temperature = kwargs.get("temperature", 0.7) base_url = kwargs.get("base_url", ollama_base_url) port = kwargs.get("port", ollama_port) + stream = kwargs.get("stream", False) image_generation_driver = kwargs.get( "image_generation_driver", DummyImageGenerationDriver() @@ -57,6 +58,7 @@ class gtUIOllamaStructureConfig(gtUIBaseConfig): temperature=temperature, host=f"{base_url}:{port}", max_attempts=max_attempts, + stream=stream, ), embedding_driver=DummyEmbeddingDriver(), image_generation_driver=image_generation_driver, diff --git a/nodes/config/OpenAiCompatableConfig.py b/nodes/config/OpenAiCompatableConfig.py index ae9c614..5704c34 100644 --- a/nodes/config/OpenAiCompatableConfig.py +++ b/nodes/config/OpenAiCompatableConfig.py @@ -59,6 +59,7 @@ class gtUIOpenAiCompatableConfig(gtUIBaseConfig): api_key = kwargs.get("api_key", None) api_key_env_var = kwargs.get("api_key_env_var", None) max_attempts = kwargs.get("max_attempts_on_fail", 10) + stream = kwargs.get("stream", False) if ( not api_key or api_key.strip() == "" or api_key == default_string @@ -72,6 +73,7 @@ class gtUIOpenAiCompatableConfig(gtUIBaseConfig): base_url=base_url, api_key=api_key, max_attempts=max_attempts, + stream=stream, ) if image_generation_model and base_url and api_key: configs["image_generation_driver"] = OpenAiImageGenerationDriver( diff --git a/nodes/config/OpenAiConfig.py b/nodes/config/OpenAiConfig.py index ba3642c..b31853f 100644 --- a/nodes/config/OpenAiConfig.py +++ b/nodes/config/OpenAiConfig.py @@ -52,14 +52,16 @@ class gtUIOpenAiStructureConfig(gtUIBaseConfig): seed = kwargs.get("seed", 12341) image_generation_driver = kwargs.get("image_generation_driver", None) max_attempts = kwargs.get("max_attempts_on_fail", 10) - + stream = kwargs.get("stream", False) OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") + prompt_driver = OpenAiChatPromptDriver( model=prompt_model, api_key=OPENAI_API_KEY, temperature=temperature, seed=seed, max_attempts=max_attempts, + stream=stream, ) embedding_driver = OpenAiEmbeddingDriver(api_key=OPENAI_API_KEY) if not image_generation_driver: