nodes/griptape-ai-ComfyUI-Griptape
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #101 from griptape-ai:dev

Browse Source 
Dev
This commit is contained in:
Jason Schleifer
2024-08-04 09:04:02 -06:00
committed by GitHub
parent a15abc0367 83960f47ea
commit bf985e4d8a
10 changed files with 156 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4
View File
@@ -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
__init__.py
+6 -1
View File
@@ -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)
js/ConfigurationNodes.js
+67 -42
View File
@@ -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) {
nodes/config/gtUILMStudioStructureConfig.py
+9 -7
View File
@@ -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": (
nodes/config/gtUIOllamaStructureConfig.py
+4 -8
View File
@@ -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}),
},
nodes/custom_routes.py
+22
View File
@@ -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")
nodes/drivers/gtUILMStudioChatPromptDriver.py
+6 -4
View File
@@ -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}),
}
nodes/drivers/gtUIOllamaPromptDriver.py
+5 -4
View File
@@ -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}),
}
nodes/tasks/gtUIVectorStoreQueryTask.py
+1 -1
View File
@@ -18,7 +18,7 @@ class gtUIVectorStoreQueryTask(gtUIBaseVectorStoreTask):
"STRING": (
"STRING",
{
"default": None,
"default": "",
"description": "The prompt to query the vector store with.",
"multiline": True,
},
nodes/utilities.py
+32 -72
View File
@@ -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 []