From 342d7d3289a5a9dede238a147cb45e09e99050b2 Mon Sep 17 00:00:00 2001 From: ImpactFrames Date: Mon, 21 Apr 2025 16:25:31 +0800 Subject: [PATCH] Add files via upload --- comfy-nodes/llmtoolkit_providers.py | 2 +- comfy-nodes/openai_provider.py | 158 +++++++++++++++++++++++ web/js/LLMToolkitProviderSelectorNode.js | 87 ++++++++++--- web/js/OpenAIProviderNode.js | 63 +++++++++ 4 files changed, 288 insertions(+), 22 deletions(-) create mode 100644 comfy-nodes/openai_provider.py create mode 100644 web/js/OpenAIProviderNode.js diff --git a/comfy-nodes/llmtoolkit_providers.py b/comfy-nodes/llmtoolkit_providers.py index e39889b..f8cf867 100644 --- a/comfy-nodes/llmtoolkit_providers.py +++ b/comfy-nodes/llmtoolkit_providers.py @@ -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"}), diff --git a/comfy-nodes/openai_provider.py b/comfy-nodes/openai_provider.py new file mode 100644 index 0000000..4a2a5d8 --- /dev/null +++ b/comfy-nodes/openai_provider.py @@ -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)", +} \ No newline at end of file diff --git a/web/js/LLMToolkitProviderSelectorNode.js b/web/js/LLMToolkitProviderSelectorNode.js index db1a2f5..0054ce5 100644 --- a/web/js/LLMToolkitProviderSelectorNode.js +++ b/web/js/LLMToolkitProviderSelectorNode.js @@ -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); }; } } diff --git a/web/js/OpenAIProviderNode.js b/web/js/OpenAIProviderNode.js new file mode 100644 index 0000000..1147097 --- /dev/null +++ b/web/js/OpenAIProviderNode.js @@ -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); + }; + }, +}); \ No newline at end of file