Add files via upload
This commit is contained in:
ImpactFrames
@@ -161,7 +161,7 @@ class LLMToolkitProviderSelector:
|
||||
return {
|
||||
"required": {
|
||||
"llm_provider": (cls.SUPPORTED_PROVIDERS, {"default": "transformers"}),
|
||||
"llm_model": ("STRING", {"default": "Provider not selected or models not fetched", "tooltip": "Select the model. Updates when provider changes."}),
|
||||
"llm_model": ("STRING", {"default": "Provider not selected or models not fetched", "tooltip": "Select the model. Updates when provider/connection changes."}),
|
||||
},
|
||||
"optional": {
|
||||
"base_ip": ("STRING", {"default": "localhost", "tooltip": "IP address for local providers"}),
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
# openai_provider.py
|
||||
"""OpenAI Provider Node for ComfyUI‑LLM‑Toolkit
|
||||
|
||||
This node is a streamlined version of the generic LLMToolkitProviderSelector that is
|
||||
hard‑wired to the OpenAI provider. It only exposes the *model* selection to the UI
|
||||
while everything else (provider name and API key) is resolved automatically from the
|
||||
environment / .env file with the existing helper utilities.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
from typing import Any, Tuple
|
||||
|
||||
# Ensure parent directory (repository root) is on sys.path so we can import utils
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
parent_dir = os.path.dirname(current_dir)
|
||||
if parent_dir not in sys.path:
|
||||
sys.path.insert(0, parent_dir)
|
||||
|
||||
# Import utility helpers from the main toolkit
|
||||
try:
|
||||
from llmtoolkit_utils import get_api_key, get_models
|
||||
except ImportError as e:
|
||||
raise ImportError("OpenAIProviderNode: Failed to import llmtoolkit_utils. "
|
||||
"Ensure it exists and is error‑free.") from e
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# ComfyUI Server Endpoint for Model List (optional but convenient)
|
||||
# -------------------------------------------------------------------------
|
||||
try:
|
||||
from server import PromptServer
|
||||
from aiohttp import web
|
||||
|
||||
@PromptServer.instance.routes.get("/ComfyLLMToolkit/get_openai_models")
|
||||
async def get_openai_models_endpoint(request):
|
||||
"""Return the list of available OpenAI chat/completion models for the account."""
|
||||
try:
|
||||
api_key = ""
|
||||
try:
|
||||
api_key = get_api_key("OPENAI_API_KEY", OpenAIProviderNode.PROVIDER_NAME)
|
||||
except ValueError:
|
||||
logger.warning("get_openai_models_endpoint: OPENAI_API_KEY not set; returning empty list")
|
||||
|
||||
models = get_models(OpenAIProviderNode.PROVIDER_NAME, None, None, api_key)
|
||||
if not models:
|
||||
models = ["No models found"]
|
||||
|
||||
return web.json_response(models)
|
||||
except Exception as e:
|
||||
logger.error("get_openai_models_endpoint: error %s", e, exc_info=True)
|
||||
return web.json_response(["Error fetching models"], status=500)
|
||||
|
||||
logger.info("OpenAIProviderNode: /ComfyLLMToolkit/get_openai_models endpoint registered")
|
||||
|
||||
except (ImportError, AttributeError) as e:
|
||||
# Running outside ComfyUI server context; skip endpoint.
|
||||
logger.debug("OpenAIProviderNode: PromptServer not available (%s); endpoint not registered", e)
|
||||
|
||||
|
||||
class OpenAIProviderNode:
|
||||
"""A minimal provider‑configuration node dedicated to OpenAI."""
|
||||
|
||||
PROVIDER_NAME = "openai"
|
||||
|
||||
# Dummy placeholder value – JS will replace this with a dynamic combo‑box.
|
||||
PLACEHOLDER_MODEL = "Select OpenAI model…"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
# Backend sees a STRING; frontend JS converts to COMBO.
|
||||
"llm_model": (
|
||||
"STRING",
|
||||
{
|
||||
"default": cls.PLACEHOLDER_MODEL,
|
||||
"tooltip": "OpenAI model name. Populated dynamically from your account.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
# Pass through pipeline data
|
||||
"any": ("*", {}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*",)
|
||||
RETURN_NAMES = ("any",)
|
||||
FUNCTION = "configure_openai"
|
||||
CATEGORY = "llm_toolkit"
|
||||
|
||||
# This allows ComfyUI to invalidate the node if inputs that matter change.
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, llm_model: str, any: Any = None):
|
||||
# Model name is the only changing factor that should trigger a recompute.
|
||||
return llm_model
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Main execution
|
||||
# ---------------------------------------------------------------------
|
||||
def configure_openai(self, llm_model: str, any: Any = None) -> Tuple[Any]:
|
||||
"""Build a provider‑config dictionary for OpenAI and merge with `any`."""
|
||||
logger.info("OpenAIProviderNode: configuring model '%s'", llm_model)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Resolve API key (env / .env). If not present we return an error but
|
||||
# still allow the graph to proceed so the user gets feedback.
|
||||
# ------------------------------------------------------------------
|
||||
api_key = ""
|
||||
try:
|
||||
api_key = get_api_key("OPENAI_API_KEY", self.PROVIDER_NAME)
|
||||
logger.info("OpenAIProviderNode: API key fetched from environment / .env")
|
||||
except ValueError as e:
|
||||
logger.warning("OpenAIProviderNode: could not retrieve OPENAI_API_KEY – %s", e)
|
||||
# Leave api_key=""; downstream node will handle auth error.
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Prepare provider configuration dictionary
|
||||
# ------------------------------------------------------------------
|
||||
provider_config = {
|
||||
"provider_name": self.PROVIDER_NAME,
|
||||
"model_name": llm_model if llm_model != self.PLACEHOLDER_MODEL else "",
|
||||
"api_key": api_key,
|
||||
"base_ip": None,
|
||||
"port": None,
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Merge with incoming `any` object to preserve pipeline data
|
||||
# ------------------------------------------------------------------
|
||||
if any is not None:
|
||||
if isinstance(any, dict):
|
||||
any["provider_config"] = provider_config
|
||||
output = any
|
||||
else:
|
||||
output = {"provider_config": provider_config, "passthrough_data": any}
|
||||
else:
|
||||
output = provider_config
|
||||
|
||||
return (output,)
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# ComfyUI Node registration helpers
|
||||
# -------------------------------------------------------------------------
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OpenAIProviderNode": OpenAIProviderNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OpenAIProviderNode": "OpenAI Provider (LLMToolkit)",
|
||||
}
|
||||
@@ -16,20 +16,62 @@ app.registerExtension({
|
||||
const llmProviderWidget = this.widgets.find((w) => w.name === "llm_provider");
|
||||
const baseIpWidget = this.widgets.find((w) => w.name === "base_ip");
|
||||
const portWidget = this.widgets.find((w) => w.name === "port");
|
||||
const llmModelWidget = this.widgets.find((w) => w.name === "llm_model");
|
||||
const externalApiKeyWidget = this.widgets.find((w) => w.name === "external_api_key");
|
||||
|
||||
const targetModelIndex = 1; // Desired index (0=provider, 1=model)
|
||||
|
||||
// 1. Find and remove the original string widget
|
||||
const originalLlmModelWidget = this.widgets.find((w) => w.name === "llm_model");
|
||||
const originalIndex = this.widgets.findIndex((w) => w === originalLlmModelWidget);
|
||||
if (originalIndex !== -1) {
|
||||
this.widgets.splice(originalIndex, 1);
|
||||
console.log("LLM Toolkit Provider Node: Removed original llm_model string widget.");
|
||||
} else {
|
||||
console.warn("LLM Toolkit Provider Node: Original llm_model widget not found for removal.");
|
||||
}
|
||||
|
||||
// 2. Create the new COMBO widget using addWidget (this appends it)
|
||||
const llmModelCombo = this.addWidget(
|
||||
"COMBO",
|
||||
"llm_model",
|
||||
"Provider not selected or models not fetched",
|
||||
(value) => { this.properties['llm_model'] = value; },
|
||||
{ values: ["Provider not selected or models not fetched"] }
|
||||
);
|
||||
if (!this.properties) { this.properties = {}; }
|
||||
this.properties['llm_model'] = llmModelCombo.value; // Initialize property
|
||||
|
||||
// 3. Move the newly added widget (currently last) to the target index
|
||||
const currentIndex = this.widgets.indexOf(llmModelCombo); // Find where addWidget put it
|
||||
if (currentIndex !== -1 && currentIndex !== targetModelIndex) {
|
||||
// Remove from current position
|
||||
this.widgets.splice(currentIndex, 1);
|
||||
// Insert at target position (ensure index is valid)
|
||||
const insertionIndex = Math.min(targetModelIndex, this.widgets.length);
|
||||
this.widgets.splice(insertionIndex, 0, llmModelCombo);
|
||||
console.log(`LLM Toolkit Provider Node: Moved llm_model combo widget to index ${insertionIndex}.`);
|
||||
} else if (currentIndex === -1) {
|
||||
console.warn("LLM Toolkit Provider Node: Added llm_model combo widget not found for move.");
|
||||
}
|
||||
|
||||
// Ensure property link is updated after potential reordering/widget recreation
|
||||
const originalCallback = llmModelCombo.callback;
|
||||
llmModelCombo.callback = (value) => {
|
||||
if(originalCallback) originalCallback.call(this, value);
|
||||
this.properties['llm_model'] = value;
|
||||
}
|
||||
|
||||
const updateLLMModels = async () => {
|
||||
if (!llmProviderWidget || !baseIpWidget || !portWidget || !llmModelWidget) {
|
||||
console.warn("LLM Toolkit Provider Node: Required widgets not found.");
|
||||
// Use the llmModelCombo reference (which should now be correctly placed)
|
||||
if (!llmProviderWidget || !baseIpWidget || !portWidget || !llmModelCombo) {
|
||||
console.warn("LLM Toolkit Provider Node: Required widgets not found (using new combo).");
|
||||
return;
|
||||
}
|
||||
|
||||
const currentModelValue = llmModelWidget.value; // Store current value
|
||||
const currentModelValue = llmModelCombo.value;
|
||||
|
||||
// Show fetching status
|
||||
llmModelWidget.options.values = ["Fetching models..."];
|
||||
llmModelWidget.value = "Fetching models...";
|
||||
llmModelCombo.options.values = ["Fetching models..."];
|
||||
llmModelCombo.value = "Fetching models...";
|
||||
this.setDirtyCanvas(true, true);
|
||||
|
||||
try {
|
||||
@@ -53,42 +95,45 @@ app.registerExtension({
|
||||
console.log("Fetched models:", models);
|
||||
|
||||
if (Array.isArray(models) && models.length > 0 && models[0] !== "Error fetching models" && models[0] !== "No models found") {
|
||||
llmModelWidget.options.values = models;
|
||||
// Try to restore previous value if it exists in the new list, otherwise set to the first model
|
||||
llmModelCombo.options.values = models;
|
||||
if (models.includes(currentModelValue)) {
|
||||
llmModelWidget.value = currentModelValue;
|
||||
llmModelCombo.value = currentModelValue;
|
||||
} else {
|
||||
llmModelWidget.value = models[0];
|
||||
llmModelCombo.value = models[0];
|
||||
}
|
||||
} else {
|
||||
llmModelWidget.options.values = ["No models found or Error"];
|
||||
llmModelWidget.value = "No models found or Error";
|
||||
llmModelCombo.options.values = ["No models found or Error"];
|
||||
llmModelCombo.value = "No models found or Error";
|
||||
console.warn("No valid models received or error fetching.");
|
||||
}
|
||||
this.properties['llm_model'] = llmModelCombo.value;
|
||||
|
||||
} catch (error) {
|
||||
console.error("Error updating LLM models:", error);
|
||||
llmModelWidget.options.values = ["Error fetching models"];
|
||||
llmModelWidget.value = "Error fetching models";
|
||||
llmModelCombo.options.values = ["Error fetching models"];
|
||||
llmModelCombo.value = "Error fetching models";
|
||||
this.properties['llm_model'] = llmModelCombo.value;
|
||||
} finally {
|
||||
this.setDirtyCanvas(true, true); // Ensure UI update
|
||||
this.setDirtyCanvas(true, true);
|
||||
}
|
||||
};
|
||||
|
||||
// Add callbacks to update models when provider, IP, port, or key changes
|
||||
[llmProviderWidget, baseIpWidget, portWidget, externalApiKeyWidget].forEach(widget => {
|
||||
if (widget) {
|
||||
const originalCallback = widget.callback;
|
||||
widget.callback = async (value) => {
|
||||
if (originalCallback) {
|
||||
originalCallback.call(this, value); // Call original callback if exists
|
||||
originalCallback.call(this, value);
|
||||
}
|
||||
await updateLLMModels(); // Then update models
|
||||
if (widget.name && this.properties && widget.name in this.properties) {
|
||||
this.properties[widget.name] = value;
|
||||
}
|
||||
await updateLLMModels();
|
||||
};
|
||||
}
|
||||
});
|
||||
|
||||
// Initial model fetch when the node is created or loaded
|
||||
setTimeout(updateLLMModels, 100); // Delay slightly to ensure widgets are ready
|
||||
setTimeout(updateLLMModels, 100);
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "Comfy.OpenAIProviderNode",
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, appInstance) {
|
||||
if (nodeData.name !== "OpenAIProviderNode") return;
|
||||
|
||||
const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
if (originalOnNodeCreated) originalOnNodeCreated.apply(this, arguments);
|
||||
|
||||
// Locate original llm_model text widget and remove it
|
||||
const original = this.widgets.find((w) => w.name === "llm_model");
|
||||
const originalIndex = this.widgets.findIndex((w) => w === original);
|
||||
if (originalIndex !== -1) this.widgets.splice(originalIndex, 1);
|
||||
|
||||
// Create combo widget and insert at index 0 (top)
|
||||
const combo = this.addWidget(
|
||||
"COMBO",
|
||||
"llm_model",
|
||||
"Select OpenAI model…",
|
||||
(v) => {
|
||||
this.properties.llm_model = v;
|
||||
},
|
||||
{ values: ["Fetching models…"] }
|
||||
);
|
||||
|
||||
// Move to top if not already
|
||||
const idx = this.widgets.indexOf(combo);
|
||||
if (idx > 0) {
|
||||
this.widgets.splice(idx, 1);
|
||||
this.widgets.splice(0, 0, combo);
|
||||
}
|
||||
|
||||
// Fetch OpenAI models once
|
||||
const fetchModels = async () => {
|
||||
try {
|
||||
const res = await fetch("/ComfyLLMToolkit/get_openai_models");
|
||||
if (!res.ok) throw new Error(res.statusText);
|
||||
const models = await res.json();
|
||||
if (Array.isArray(models) && models.length) {
|
||||
combo.options.values = models;
|
||||
combo.value = models[0];
|
||||
this.properties.llm_model = combo.value;
|
||||
} else {
|
||||
combo.options.values = ["No models"];
|
||||
combo.value = "No models";
|
||||
}
|
||||
} catch (e) {
|
||||
console.error("OpenAIProviderNode: model fetch failed", e);
|
||||
combo.options.values = ["Error fetching models"];
|
||||
combo.value = "Error fetching models";
|
||||
} finally {
|
||||
this.setDirtyCanvas(true, true);
|
||||
}
|
||||
};
|
||||
|
||||
setTimeout(fetchModels, 100);
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user