diff --git a/README.md b/README.md index 30c286a..9926c87 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,9 @@ You can previous and download more examples [here](examples/README.md). ## Recent Changelog +### Aug 3, 2024 +* Reverted ollama and lmstudio configuration nodes to a list of installed models using new method for grabbing them. + ### July 29, 2024 * Temporarily replaceing the ollama config nodes with a string input for specifying the model instead of a list of installed models. @@ -393,3 +396,4 @@ Massive thank you for help and inspiration from the following people and repos! * Jovieux from https://github.com/Amorano/Jovimetrix * rgthree https://github.com/rgthree/rgthree-comfy +* IF_AI_tools https://github.com/if-ai/ComfyUI-IF_AI_tools/tree/main \ No newline at end of file diff --git a/__init__.py b/__init__.py index dd6fc5b..6757578 100644 --- a/__init__.py +++ b/__init__.py @@ -10,7 +10,6 @@ import sys from dotenv import load_dotenv -# Load the griptape_config.json data # AGENT from .nodes.agent.CreateAgent import CreateAgent from .nodes.agent.ExpandAgent import ExpandAgent @@ -45,6 +44,9 @@ from .nodes.config.gtUIStructureConfig import gtUIStructureConfig from .nodes.convert.gtUITextToClipEncode import gtUITextToClipEncode from .nodes.convert.gtUITextToCombo import gtUITextToCombo +# Load the routes +from .nodes.custom_routes import init_routes + # DISPLAY from .nodes.display.gtUIOutputDataNode import gtUIOutputDataNode from .nodes.display.gtUIOutputImageNode import gtUIOutputImageNode @@ -242,6 +244,9 @@ load_dotenv() print("\n\033[34m[Griptape Custom Nodes]:\033[0m") +# Initialize the routes +init_routes() + # Now load and prepare the configuration config = load_and_prepare_config(DEFAULT_CONFIG_FILE, USER_CONFIG_FILE) diff --git a/js/ConfigurationNodes.js b/js/ConfigurationNodes.js index 9411d9d..7861b75 100644 --- a/js/ConfigurationNodes.js +++ b/js/ConfigurationNodes.js @@ -7,55 +7,80 @@ export function setupConfigurationNodes(nodeType, nodeData, app) { nodeType.prototype.onNodeCreated = async function () { onNodeCreated?.apply(this, arguments); - - // TODO: Find a better way to identify the config nodes - // - // if (nodeData.name.includes("Ollama")) { - // console.log("Ollama config node"); - // console.log(nodeData.name); - // setupOllamaConfig(this); - // } else if (nodeData.name.includes("LM Studio")) { - // setupLMStudioConfig(this); - // } + let engine = "ollama" + if (nodeData.name.includes("LM Studio")) { + engine="lmstudio" + } + const modelWidget = this.widgets.find((w) => w.name === "model" || w.name === "prompt_model"); + 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 models = await fetchModels(engine, baseIp, port); + + // Update modelWidget options and value + modelWidget.options.values = models; + + if (models.includes(modelWidget.value)) { + modelWidget.value = modelWidget.value; + } else if (models.length > 0) { + modelWidget.value = models[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); }; + } } -function setupOllamaConfig(node) { - const base_url = node.widgets.find((w) => w.name === "base_url"); - const port = node.widgets.find((w) => w.name === "port"); - console.log(node.widgets); - console.log("base_url widget:", base_url); - console.log("port widget:", port); - - if (base_url && port) { - console.log("base_url value:", base_url.value); - console.log("port value:", port.value); - console.log("port value type:", typeof port.value); - - 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) { if(node.widgets) { diff --git a/nodes/config/gtUILMStudioStructureConfig.py b/nodes/config/gtUILMStudioStructureConfig.py index 89cf8fa..e605c0c 100644 --- a/nodes/config/gtUILMStudioStructureConfig.py +++ b/nodes/config/gtUILMStudioStructureConfig.py @@ -7,13 +7,14 @@ from griptape.drivers import ( OpenAiChatPromptDriver, ) -from ..utilities import get_lmstudio_models +from ..utilities import get_models from .gtUIBaseConfig import gtUIBaseConfig -lmstudio_models = get_lmstudio_models(port="1234") -lmstudio_models.append("") lmstudio_port = "1234" lmstudio_base_url = "http://127.0.0.1" +lmstudio_models = get_models("lmstudio", lmstudio_base_url, lmstudio_port) +if len(lmstudio_models) == 0: + lmstudio_models.append("") class gtUILMStudioStructureConfig(gtUIBaseConfig): @@ -30,10 +31,11 @@ class gtUILMStudioStructureConfig(gtUIBaseConfig): inputs = super().INPUT_TYPES() inputs["required"].update( { - "model": ( - "STRING", - {"default": lmstudio_models[0]}, - ), + "model": ((), {}), + # "model": ( + # "STRING", + # {"default": lmstudio_models[0]}, + # ), # "prompt_model": ("STRING", {"default": ""}), "base_url": ("STRING", {"default": lmstudio_base_url}), "port": ( diff --git a/nodes/config/gtUIOllamaStructureConfig.py b/nodes/config/gtUIOllamaStructureConfig.py index b6d5da2..e17652b 100644 --- a/nodes/config/gtUIOllamaStructureConfig.py +++ b/nodes/config/gtUIOllamaStructureConfig.py @@ -7,13 +7,13 @@ from griptape.drivers import ( OllamaPromptDriver, ) -from ..utilities import get_ollama_models +from ..utilities import get_models from .gtUIBaseConfig import gtUIBaseConfig -ollama_models = get_ollama_models() -ollama_models.append("") ollama_port = "11434" ollama_base_url = "http://127.0.0.1" +ollama_models = get_models("ollama", ollama_base_url, ollama_port) +ollama_models.append("") class gtUIOllamaStructureConfig(gtUIBaseConfig): @@ -28,11 +28,7 @@ class gtUIOllamaStructureConfig(gtUIBaseConfig): inputs = super().INPUT_TYPES() inputs["required"].update( { - "prompt_model": ( - # [""], - "STRING", - {"default": ollama_models[0]}, - ), + "prompt_model": ((), {}), "base_url": ("STRING", {"default": ollama_base_url}), "port": ("STRING", {"default": ollama_port}), }, diff --git a/nodes/custom_routes.py b/nodes/custom_routes.py new file mode 100644 index 0000000..24e2523 --- /dev/null +++ b/nodes/custom_routes.py @@ -0,0 +1,22 @@ +from aiohttp import web +from server import PromptServer + +# Import your route handlers here +from .utilities import get_models + + +def setup_routes(): + @PromptServer.instance.routes.post("/Griptape/get_models") + async def get_models_endpoint(request): + data = await request.json() + engine = data.get("engine") + base_ip = data.get("base_ip") + port = data.get("port") + models = get_models(engine, base_ip, port) + return web.json_response(models) + + +# Call this function to set up all routes +def init_routes(): + setup_routes() + print(" \033[34m- Custom routes initialized.\033[0m") diff --git a/nodes/drivers/gtUILMStudioChatPromptDriver.py b/nodes/drivers/gtUILMStudioChatPromptDriver.py index 5af7251..6325c33 100644 --- a/nodes/drivers/gtUILMStudioChatPromptDriver.py +++ b/nodes/drivers/gtUILMStudioChatPromptDriver.py @@ -1,13 +1,14 @@ from griptape.drivers import OpenAiChatPromptDriver -from ..utilities import get_lmstudio_models +from ..utilities import get_models from .gtUIBasePromptDriver import gtUIBasePromptDriver -models = get_lmstudio_models() -# models.append("") default_port = "1234" default_base_url = "http://127.0.0.1" DEFAULT_API_KEY = "lm_studio" +models = get_models("lmstudio", default_base_url, default_port) +if len(models) == 0: + models.append("") class gtUILMStudioChatPromptDriver(gtUIBasePromptDriver): @@ -17,7 +18,8 @@ class gtUILMStudioChatPromptDriver(gtUIBasePromptDriver): inputs["required"].update( { - "model": ("STRING", {"default": models[0]}), + "model": ((), {}), + # "model": ("STRING", {"default": models[0]}), "base_url": ("STRING", {"default": default_base_url}), "port": ("STRING", {"default": default_port}), } diff --git a/nodes/drivers/gtUIOllamaPromptDriver.py b/nodes/drivers/gtUIOllamaPromptDriver.py index 6ce1057..1ca493d 100644 --- a/nodes/drivers/gtUIOllamaPromptDriver.py +++ b/nodes/drivers/gtUIOllamaPromptDriver.py @@ -1,12 +1,13 @@ from griptape.drivers import OllamaPromptDriver -from ..utilities import get_ollama_models +from ..utilities import get_models from .gtUIBasePromptDriver import gtUIBasePromptDriver -models = get_ollama_models() -models.append("") default_port = "11434" default_base_url = "http://127.0.0.1" +models = get_models("ollama", default_base_url, default_port) +if len(models) == 0: + models.append("") class gtUIOllamaPromptDriver(gtUIBasePromptDriver): @@ -16,7 +17,7 @@ class gtUIOllamaPromptDriver(gtUIBasePromptDriver): inputs["required"].update( { - "model": ("STRING", {"default": models[0]}), + "model": ((), {}), "base_url": ("STRING", {"default": default_base_url}), "port": ("STRING", {"default": default_port}), } diff --git a/nodes/tasks/gtUIVectorStoreQueryTask.py b/nodes/tasks/gtUIVectorStoreQueryTask.py index 67eda13..3614ae4 100644 --- a/nodes/tasks/gtUIVectorStoreQueryTask.py +++ b/nodes/tasks/gtUIVectorStoreQueryTask.py @@ -18,7 +18,7 @@ class gtUIVectorStoreQueryTask(gtUIBaseVectorStoreTask): "STRING": ( "STRING", { - "default": None, + "default": "", "description": "The prompt to query the vector store with.", "multiline": True, }, diff --git a/nodes/utilities.py b/nodes/utilities.py index b106b4d..9be3dd9 100644 --- a/nodes/utilities.py +++ b/nodes/utilities.py @@ -1,5 +1,6 @@ import base64 from io import BytesIO +from urllib.parse import urlparse, urlunparse import numpy as np import requests @@ -9,84 +10,43 @@ from jinja2 import Template from PIL import Image, ImageOps, ImageSequence +def construct_base_url(base_url: str, port: int) -> str: + # Parse the base_url + parsed_url = urlparse(base_url) + + # If there's no scheme (http:// or https://), add http:// + if not parsed_url.scheme: + parsed_url = parsed_url._replace(scheme="http") + + # Replace the port + parsed_url = parsed_url._replace(netloc=f"{parsed_url.hostname}:{port}") + + # Reconstruct the URL without any path + return urlunparse(parsed_url._replace(path="")) + + def get_models(engine, base_url, port) -> list[str]: + reconstructed_base_url = construct_base_url(base_url, port) + if engine == "ollama": - api_url = f"http://{base_url}:{port}/api/tags" - try: - response = requests.get(api_url) - response.raise_for_status() - models = [model["name"] for model in response.json().get("models", [])] - return models - except Exception as e: - print(f"Failed to fetch models from Ollama: {e}") - return [] + api_url = f"{reconstructed_base_url}/api/tags" elif engine == "lmstudio": - api_url = f"http://{base_url}:{port}/v1/models" - try: - response = requests.get(api_url) - response.raise_for_status() + api_url = f"{reconstructed_base_url}/v1/models" + else: + raise ValueError(f"Unsupported engine: {engine}") + + try: + response = requests.get(api_url) + response.raise_for_status() + + if engine == "ollama": + models = [model["name"] for model in response.json().get("models", [])] + else: # lmstudio models = [model["id"] for model in response.json().get("data", [])] - return models - except Exception as e: - print(f"Failed to fetch models from LM Studio: {e}") - return [] - - -def get_lmstudio_models(port="1234") -> list[str]: - url = f"http://127.0.0.1:{port}/v1/models" - - try: - # Make the GET request - response = requests.get(url) - - # Check if the request was successful - if response.status_code != 200: - raise Exception(f"Failed to fetch models: {response.status_code}") - - # Parse the JSON response - models_info = response.json() - - # Extract the model names - models = [model["id"] for model in models_info["data"]] return models - except ( - requests.exceptions.ConnectionError, - requests.exceptions.HTTPError, - requests.exceptions.RequestException, - KeyError, - ): - # Return an empty list if there is any error - return [] - - -def get_ollama_models() -> list[str]: - # URL to fetch the local models - url = "http://127.0.0.1:11434/api/tags" - - try: - # Make the GET request - response = requests.get(url) - - # Check if the request was successful - if response.status_code != 200: - raise Exception(f"Failed to fetch models: {response.status_code}") - - # Parse the JSON response - models_info = response.json() - - # Extract the model names - # models = [model["name"].split(":")[0] for model in models_info["models"]] - models = [model["name"] for model in models_info["models"]] - - return models - except ( - requests.exceptions.ConnectionError, - requests.exceptions.HTTPError, - requests.exceptions.RequestException, - KeyError, - ): - # Return an empty list if there is any error + except Exception as e: + print(f"Failed to fetch models from {engine.capitalize()}: {e}") return []