griptape-ai-ComfyUI-Griptape/js/ConfigurationNodes.js

import {
  getOllamaModels,
  getLMStudioModels,
  updatePromptModelList,
} from "./gtUIUtils.js";
import { commentWidget } from "./utils.js";

export function setupConfigurationNodes(nodeType, nodeData, app) {
  if (
    nodeData.name.includes("Griptape Agent Config") ||
    nodeData.name.includes("Griptape Util: Remove Ollama Model") ||
    nodeData.name.includes("Griptape Util: Create Agent Modelfile") ||
    nodeData.name.includes("Griptape Embedding Driver") ||
    nodeData.name.includes("Griptape Prompt Driver")
  ) {
    const onNodeCreated = nodeType.prototype.onNodeCreated;

    nodeType.prototype.onNodeCreated = async function () {
      onNodeCreated?.apply(this, arguments);

      // Ensure this.widgets exists before trying to filter it
      if (this.widgets && Array.isArray(this.widgets)) {
        // Find widgets you want to turn into comments
        const commentWidgets = this.widgets.filter((w) =>
          w.name.toLowerCase().includes("comment")
        );
        commentWidgets.forEach((widget) => commentWidget(this, widget));
      } else {
        // If widgets aren't available yet, set up a timeout to check again
        setTimeout(() => {
          if (this.widgets && Array.isArray(this.widgets)) {
            const commentWidgets = this.widgets.filter((w) =>
              w.name.toLowerCase().includes("comment")
            );
            commentWidgets.forEach((widget) => commentWidget(this, widget));
          } else {
            // console.warn('Widgets not available for node:', this.title);
          }
        }, 100); // Wait 100ms and try again
      }

      let engine = null;
      if (
        nodeData.name.includes("Ollama") ||
        nodeData.name.includes("Griptape Util: Create Agent Modelfile")
      ) {
        engine = "ollama";
      }
      if (nodeData.name.includes("LM Studio")) {
        // engine="lmstudio"
      }
      if (engine) {
        if (this.widgets) {
          const modelWidgets = this.widgets.filter((w) =>
            ["model", "prompt_model", "embedding_model", "base_model"].includes(
              w.name
            )
          );
          const baseIpWidget = this.widgets.find((w) => w.name === "base_url");
          const portWidget = this.widgets.find((w) => w.name === "port");
          const fetchModels = async (engine, baseIp, port) => {
            try {
              const response = await fetch("/Griptape/get_models", {
                method: "POST",
                headers: {
                  "Content-Type": "application/json",
                },
                body: JSON.stringify({
                  engine: engine,
                  base_ip: baseIp,
                  port: port,
                }),
              });

              if (response.ok) {
                const models = await response.json();
                return models;
              } else {
                console.error(`Failed to fetch models: ${response.status}`);
                return [];
              }
            } catch (error) {
              console.error(
                `Error fetching models for engine ${engine}:`,
                error
              );
              return [];
            }
          };
          const updateModels = async () => {
            let engine = "ollama";
            if (nodeData.name.includes("LM Studio")) {
              engine = "lmstudio";
            }

            const baseIp = baseIpWidget.value;
            const port = portWidget.value;

            const allModels = await fetchModels(engine, baseIp, port);

            // Update each modelWidget's options and value
            modelWidgets.forEach((modelWidget) => {
              let filteredModels;
              if (modelWidget.name.includes("embed")) {
                filteredModels = allModels.filter((model) =>
                  model.toLowerCase().includes("embed")
                );
              } else {
                filteredModels = allModels.filter(
                  (model) => !model.toLowerCase().includes("embed")
                );
              }

              modelWidget.options.values = filteredModels;

              if (filteredModels.includes(modelWidget.value)) {
                modelWidget.value = modelWidget.value;
              } else if (filteredModels.length > 0) {
                modelWidget.value = filteredModels[0];
              } else {
                modelWidget.value = "";
              }
            });

            this.triggerSlot(0);
          };

          baseIpWidget.callback = updateModels;
          portWidget.callback = updateModels;

          // Initial update
          await updateModels();
          fetchModels(engine, baseIpWidget.value, portWidget.value);
        }
      }
      setFixedRandomization(this);
      // console.log("original size: ", this.size);
      // setTimeout(() =>{
      //   this.onResize?.(this.size);
      //   this.graph.setDirtyCanvas(true, true);

      // },100)
    };
  }
}

function setFixedRandomization(node) {
  if (node.widgets) {
    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");
    }
  }
}