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

Merge pull request #46 from griptape-ai/dev

Browse Source 
Dev
This commit is contained in:
Jason Schleifer
2024-07-01 09:54:05 -07:00
committed by GitHub
parent 1da5037ce8 07c179226c
commit 8a29ebcdb6
17 changed files with 639 additions and 404 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+5 -1
View File
@@ -11,7 +11,10 @@ import os
from dotenv import load_dotenv
# Load the griptape_config.json data
from .nodes.agent import CreateAgent, ExpandAgent, RunAgent
from .nodes.agent.create_agent import CreateAgent
from .nodes.agent.expand_agent import ExpandAgent
from .nodes.agent.run_agent import RunAgent
from .nodes.agent.set_default_agent import gtUISetDefaultAgent
from .nodes.audio_drivers import gtUIOpenAiAudioTranscriptionDriver
from .nodes.audio_nodes import gtUILoadAudio
from .nodes.combine_nodes import MergeTexts, RulesList, ToolList, gtUIMergeInputs
@@ -97,6 +100,7 @@ NODE_CLASS_MAPPINGS = {
"Griptape Create: Agent": CreateAgent,
"Griptape Run: Agent": RunAgent,
"Griptape Expand: Agent Nodes": ExpandAgent,
"Griptape Set: Default Agent": gtUISetDefaultAgent,
# AGENT CONFIG
"Griptape Agent Config: OpenAI": gtUIOpenAiStructureConfig,
"Griptape Agent Config: Amazon Bedrock": gtUIAmazonBedrockStructureConfig,
js/gtUINodes.js
+96 -30
View File
@@ -162,7 +162,7 @@ class GriptapeNodes extends EventTarget {
getGriptapeContextMenuItems() {
const that = this;
return [
// TODO: Figure out how to pull and set the environment variables
// // TODO: Figure out how to pull and set the environment variables
// {
// content: "⚙️ Environment Variables",
// callback: (...args) => {
@@ -225,6 +225,68 @@ class GriptapeNodes extends EventTarget {
document.head.appendChild(link);
}
}
async function getOllamaModels(baseUrl, port) {
const url = `${baseUrl}:${port}/api/tags`;
try {
const response = await fetch(url);
if (!response.ok) {
throw new Error(`Failed to fetch models: ${response.status}`);
}
const modelsInfo = await response.json();
// Extract the model names
const models = modelsInfo.models.map(model => model.name);
return models;
} catch (error) {
console.error("Error fetching Ollama models:", error);
return [];
}
}
async function getLMStudioModels(baseUrl, port) {
const url = `${baseUrl}:${port}/v1/models`;
try {
const response = await fetch(url);
if (!response.ok) {
throw new Error(`Failed to fetch models: ${response.status}`);
}
const modelsInfo = await response.json();
// Extract the model names
const models = modelsInfo.data.map(model => model.id);
return models;
} catch (error) {
console.error("Error fetching LM Studio models:", error);
return [];
}
}
async function updatePromptModelList(node, models) {
const modelWidget = node.widgets.find((w) => w.name === "prompt_model");
const selectedItem = modelWidget.value;
modelWidget.options.values = models;
console.log("Models" + models);
// if (models.length > 0) {
// modelWidget.value = models[0];
// }
const warning = "No models returned from server"
if (models.length == 0){
modelWidget.value = warning
}
if ((selectedItem === "" || selectedItem === null || selectedItem == warning) && models.length > 0 && !models.includes(selectedItem)) {
modelWidget.value = models[0];
}
}
function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") {
chainCallback(nodeType.prototype, "onNodeCreated", function() {
const pathWidget = this.widgets.find((w) => w.name === widgetName);
@@ -300,24 +362,6 @@ function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") {
},
});
} else if (type == "audio") {
// Create an audio element if not already created
// let audioWidget = document.getElementById('audio-widget');
// let audioWidget = document.createElement('audio');
// audioWidget.controls = true; // Add play controls
// document.body.append(audioWidget); // Append the audio element to the body or a specific container
// let audioButton = this.addWidget("button", "Play Audio ▶️", "audio", () => {
// audioWidget.play();
// });
// document.body.append(audioWidget);
// if (!audioWidget) {
// audioWidget = document.createElement('audio');
// audioWidget.id = 'audio-widget';
// audioWidget.controls = true; // Add play controls
// document.body.append(audioWidget); // Append the audio element to the body or a specific container
// let audioButton = this.addWidget("button", "Play Audio ▶️", "audio", () => {
// audioWidget.play();
// });
// }
Object.assign(fileInput, {
type: "file",
accept: "audio/mpeg,audio/wav,audio/x-wav,audio/ogg",
@@ -357,7 +401,6 @@ function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") {
for (const file of e.dataTransfer.files) {
if (file.type.startsWith("audio/")) {
handled = true;
console.log(file);
uploadFile(file, !handled);
handled = true;
const filename = file.name;
@@ -366,13 +409,6 @@ function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") {
if (pathWidget.callback) {
pathWidget.callback(filename)
}
// Create a URL for the audio file and set it as the source of the audio element
// const audioURL = URL.createObjectURL(file);
// audioWidget.src = audioURL;
// audioWidget.play(); // Optionally play the audio automatically
}
}
return handled;
@@ -406,18 +442,48 @@ app.registerExtension({
// this.color = getColor(nodeData.category);
// };
// Set Config node randomization to Fixed
// Configuration Nodes
if (nodeData.name.includes("Griptape Agent Config")) {
const onNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = async function () {
// Ollama Config Node
if (nodeData.name.includes("Ollama")) {
// get the base_url
const base_url = this.widgets.find((w) => w.name === "base_url");
const port = this.widgets.find((w) => w.name === "port");
getOllamaModels(base_url.value, port.value).then((models) => {
const model = this.widgets.find((w) => w.name === "prompt_model");
model.options.values = models;
updatePromptModelList(this, models);
})
}
// LMStudio Config Node
if (nodeData.name.includes("LM Studio")) {
chainCallback(nodeType.prototype, "onNodeUpdated", function() {
});
// get the base_url
const base_url = this.widgets.find((w) => w.name === "base_url");
const port = this.widgets.find((w) => w.name === "port");
getLMStudioModels(base_url.value, port.value).then((models) => {
const model = this.widgets.find((w) => w.name === "prompt_model");
model.options.values = models;
updatePromptModelList(this, models);
})
}
// Set Config node randomization to Fixed
for (const widget of this.widgets) {
if (widget.name === "control_after_generate") {
widget.value = "fixed";
}
}
}
}
}
// Create Audio Node
if (nodeData.name === "Griptape Load: Audio") {
gtUIAddUploadWidget(nodeType, nodeData, "audio", "audio")
nodes/agent.py
-224
View File
@@ -1,224 +0,0 @@
from textwrap import dedent
from griptape.config import (
OpenAiStructureConfig,
)
from griptape.structures import Agent as gtAgent
from griptape.tools import TaskMemoryClient
from .base_agent import BaseAgent
default_prompt = "{{ input_string }}"
def model_check(agent):
# There are certain models that can't handle Tools well.
# If this agent is using one of those models AND they have tools supplied, we'll
# warn the user.
simple_models = ["llama3", "mistral", "LLama-3"]
drivers = ["OllamaPromptDriver", "LMStudioPromptDriver"]
agent_prompt_driver_name = agent.config.prompt_driver.__class__.__name__
model = agent.config.prompt_driver.model
print(f"Model: {model}")
if agent_prompt_driver_name in drivers:
if model == "":
return (model, True)
for simple in simple_models:
if simple in model:
if len(agent.tools) > 0:
return (model, True)
return (model, False)
class RunAgent(BaseAgent):
DESCRIPTION = "Run a simple Griptape Agent"
def run(
self,
STRING,
# input_artifact=None,
agent=None,
input_string=None,
):
if not agent:
agent = gtAgent()
# Get the prompt text
if not input_string:
prompt_text = STRING
else:
prompt_text = STRING + "\n\n" + input_string
# There are certain models that can't handle Tools well.
model, simple_model = model_check(agent)
if simple_model:
if model == "":
return (
"You have provided a blank model for the Agent Configuration.\n\nPlease specify a model configuration, or disconnect it from the agent.",
agent,
)
else:
return (
dedent(
f"""This Agent Configuration Model: **{ agent.config.prompt_driver.model }** may run into issues using tools.\n\nPlease consider using a different configuration, a different model, or removing tools from the agent and use the **Griptape Run: Tool Task** node for specific tool use."""
),
agent,
)
result = agent.run(prompt_text)
artifact = result.output_task.output
output_string = result.output_task.output.value
return (
output_string,
# artifact,
agent,
)
class CreateAgent(BaseAgent):
"""
Create a Griptape Agent
"""
DESCRIPTION = (
"Create a Griptape Agent with inputs for configuration, tools, and rulesets."
)
@classmethod
def INPUT_TYPES(s):
inputs = super().INPUT_TYPES()
# Update optional inputs to include 'tool' and adjust others as necessary
inputs["optional"].update(
{
"config": (
"CONFIG",
{
"forceInput": True,
},
),
"tools": ("TOOL_LIST", {"forceInput": True, "INPUT_IS_LIST": True}),
"rulesets": ("RULESET", {"forceInput": True}),
},
)
return inputs
RETURN_TYPES = ("STRING", "AGENT")
RETURN_NAMES = ("OUTPUT", "AGENT")
FUNCTION = "run"
CATEGORY = "Griptape/Agent"
def run(
self,
STRING,
agent=None,
config=None,
input_string=None,
tools=[],
rulesets=[],
):
if not config:
config = OpenAiStructureConfig()
task_memory_client = [TaskMemoryClient(off_prompt=False)]
agent_tools = []
if len(tools) > 0:
agent_tools += tools
agent_tools += task_memory_client if len(agent_tools) > 0 else []
agent_rulesets = []
for ruleset in rulesets:
agent_rulesets.append(ruleset)
if not agent:
agent = gtAgent(config=config, tools=agent_tools, rulesets=agent_rulesets)
else:
if config:
agent.config = config
if len(tools) > 0:
agent.tools = agent_tools
if len(rulesets) > 0:
agent.rulesets = agent_rulesets
# There are certain models that can't handle Tools well.
# If this agent is using one of those models AND they have tools supplied, we'll
# warn the user.
model, simple_model = model_check(agent)
if simple_model:
if model == "":
return (
"You have provided a blank model for the Agent Configuration.\n\nPlease specify a model configuration, or disconnect it from the agent.",
agent,
)
else:
return (
dedent(
f"""This Agent Configuration Model: **{ agent.config.prompt_driver.model }** may run into issues using tools.\n\nPlease consider using a different configuration, a different model, or removing tools from the agent and use the **Griptape Run: Tool Task** node for specific tool use."""
),
agent,
)
# Run the agent if there's a prompt
if input_string or STRING not in [default_prompt, ""]:
if not input_string:
prompt_text = STRING
else:
prompt_text = STRING + "\n" + input_string
result = agent.run(prompt_text)
output_string = result.output_task.output.value
else:
output_string = "Agent Created"
return (output_string, agent)
class ExpandAgent:
DESCRIPTION = "Expand the components of a Griptape Agent."
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"agent": (
"AGENT",
{
"forceInput": True,
},
),
},
}
RETURN_TYPES = (
"AGENT",
"CONFIG",
"RULESET",
"TOOL_LIST",
"MEMORY",
)
RETURN_NAMES = (
"AGENT",
"CONFIG",
"RULESETS",
"TOOLS",
"MEMORY",
)
FUNCTION = "expand"
CATEGORY = "Griptape/Agent"
OUTPUT_NODE = True
def expand(self, agent):
rulesets = agent.rulesets
tools = agent.tools
conversation_memory = agent.conversation_memory
config = agent.config
# Run the agent
return (
agent,
config,
rulesets,
tools,
conversation_memory,
)
nodes/agent/agent.py
+64
View File
@@ -0,0 +1,64 @@
from griptape.config import BaseStructureConfig
from griptape.structures import Agent
from ...py.griptape_config import get_config
default_prompt = "{{ input_string }}"
class gtComfyAgent(Agent):
def __init__(self, *args, **kwargs):
# Initialize the parent class
super().__init__(*args, **kwargs)
# Add any additional initialization here
self.set_default_config()
def set_default_config(self):
agent_config = get_config("agent_config")
if agent_config:
self.config = BaseStructureConfig.from_dict(agent_config)
def model_check(self):
# There are certain models that can't handle Tools well.
# If this agent is using one of those models AND they have tools supplied, we'll
# warn the user.
simple_models = ["llama3", "mistral", "LLama-3"]
drivers = ["OllamaPromptDriver", "LMStudioPromptDriver"]
agent_prompt_driver_name = self.config.prompt_driver.__class__.__name__
print(agent_prompt_driver_name)
model = self.config.prompt_driver.model
if agent_prompt_driver_name in drivers:
if model == "":
return (model, True)
for simple in simple_models:
if simple in model:
if len(self.tools) > 0:
return (model, True)
return (model, False)
def model_response(self, model):
if model == "":
return "You have provided a blank model for the Agent Configuration.\n\nPlease specify a model configuration, or disconnect it from the agent."
else:
return f"This Agent Configuration Model: **{ self.config.prompt_driver.model }** may run into issues using tools.\n\nPlease consider using a different configuration, a different model, or removing tools from the agent and use the **Griptape Run: Tool Task** node for specific tool use."
def model_check(agent):
# There are certain models that can't handle Tools well.
# If this agent is using one of those models AND they have tools supplied, we'll
# warn the user.
simple_models = ["llama3", "mistral", "LLama-3"]
drivers = ["OllamaPromptDriver", "LMStudioPromptDriver"]
agent_prompt_driver_name = agent.config.prompt_driver.__class__.__name__
model = agent.config.prompt_driver.model
print(f"Model: {model}")
if agent_prompt_driver_name in drivers:
if model == "":
return (model, True)
for simple in simple_models:
if simple in model:
if len(agent.tools) > 0:
return (model, True)
return (model, False)
nodes/agent/base_agent.py
+140
View File
@@ -0,0 +1,140 @@
from ...py.griptape_config import get_config
from .agent import gtComfyAgent
default_prompt = "{{ input_string }}"
def get_default_config():
return get_config("agent_config")
class BaseAgent:
"""
Create a Griptape Agent
"""
def __init__(self):
self.default_prompt = default_prompt
self.agent = gtComfyAgent()
self.agent.set_default_config()
@classmethod
def INPUT_TYPES(s):
return {
"required": {},
"optional": {
"agent": (
"AGENT",
{
"forceInput": True,
},
),
"config": (
"CONFIG",
{
"forceInput": True,
},
),
"tools": ("TOOL_LIST", {"forceInput": True, "INPUT_IS_LIST": True}),
"rulesets": ("RULESET", {"forceInput": True}),
"input_string": (
"STRING",
{
"forceInput": True,
},
),
"STRING": (
"STRING",
{
"multiline": True,
},
),
},
}
RETURN_TYPES = (
"STRING",
"AGENT",
)
RETURN_NAMES = (
"OUTPUT",
"AGENT",
)
FUNCTION = "run"
OUTPUT_NODE = True
CATEGORY = "Griptape/Agent"
# def set_default_config(self):
# agent_config = get_config("agent_config")
# if agent_config:
# self.agent.config = BaseStructureConfig.from_dict(agent_config)
# def model_check(self):
# # There are certain models that can't handle Tools well.
# # If this agent is using one of those models AND they have tools supplied, we'll
# # warn the user.
# simple_models = ["llama3", "mistral", "LLama-3"]
# drivers = ["OllamaPromptDriver", "LMStudioPromptDriver"]
# agent_prompt_driver_name = self.agent.config.prompt_driver.__class__.__name__
# model = self.agent.config.prompt_driver.model
# if agent_prompt_driver_name in drivers:
# if model == "":
# return (model, True)
# for simple in simple_models:
# if simple in model:
# if len(self.agent.tools) > 0:
# return (model, True)
# return (model, False)
# def model_response(self, model):
# if model == "":
# return "You have provided a blank model for the Agent Configuration.\n\nPlease specify a model configuration, or disconnect it from the agent."
# else:
# return f"This Agent Configuration Model: **{ self.agent.config.prompt_driver.model }** may run into issues using tools.\n\nPlease consider using a different configuration, a different model, or removing tools from the agent and use the **Griptape Run: Tool Task** node for specific tool use."
def run(
self,
STRING,
config=None,
tools=[],
rulesets=[],
agent=None,
input_string=None,
):
if agent:
self.agent = agent
# Replace bits of the agent based off the inputs
if config:
self.agent.config = config
else:
self.agent.set_default_config()
if len(tools) > 0:
self.agent.tools = tools
if len(rulesets) > 0:
self.agent.rulesets = rulesets
# Warn for models
model, simple_model = self.agent.model_check()
if simple_model:
return (self.agent.model_response(model), self.agent)
# Check for inputs. If none, then just create the agent
if not input_string and STRING == "":
output_string = "Agent created."
else:
# Get the prompt text
if not input_string:
prompt_text = STRING
else:
prompt_text = STRING + "\n\n" + input_string
result = self.agent.run(prompt_text)
output_string = result.output_task.output.value
return (
output_string,
self.agent,
)
nodes/agent/create_agent.py
+6
View File
@@ -0,0 +1,6 @@
from .base_agent import BaseAgent
default_prompt = "{{ input_string }}"
class CreateAgent(BaseAgent): ...
nodes/agent/expand_agent.py
+51
View File
@@ -0,0 +1,51 @@
class ExpandAgent:
DESCRIPTION = "Expand the components of a Griptape Agent."
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"agent": (
"AGENT",
{
"forceInput": True,
},
),
},
}
RETURN_TYPES = (
"AGENT",
"CONFIG",
"RULESET",
"TOOL_LIST",
"MEMORY",
)
RETURN_NAMES = (
"AGENT",
"CONFIG",
"RULESETS",
"TOOLS",
"MEMORY",
)
FUNCTION = "expand"
CATEGORY = "Griptape/Agent"
OUTPUT_NODE = True
def expand(self, agent):
rulesets = agent.rulesets
tools = agent.tools
conversation_memory = agent.conversation_memory
config = agent.config
return (
agent,
config,
rulesets,
tools,
conversation_memory,
)
nodes/agent/run_agent.py
+41
View File
@@ -0,0 +1,41 @@
from .base_agent import BaseAgent
class RunAgent(BaseAgent):
DESCRIPTION = "Run a simple Griptape Agent"
@classmethod
def INPUT_TYPES(s):
inputs = super().INPUT_TYPES()
del inputs["optional"]["config"]
del inputs["optional"]["tools"]
del inputs["optional"]["rulesets"]
return inputs
def run(
self,
STRING,
agent=None,
input_string=None,
):
if agent:
self.agent = agent
# Warn for models
model, simple_model = self.agent.model_check()
if simple_model:
return (self.agent.model_response(model), self.agent)
# Get the prompt text
if not input_string:
prompt_text = STRING
else:
prompt_text = STRING + "\n\n" + input_string
result = self.agent.run(prompt_text)
output_string = result.output_task.output.value
return (
output_string,
self.agent,
)
nodes/agent/set_default_agent.py
+28
View File
@@ -0,0 +1,28 @@
from ...py.griptape_config import update_config_with_dict
from .base_agent import BaseAgent
default_prompt = "{{ input_string }}"
class gtUISetDefaultAgent(BaseAgent):
DESCRIPTION = "Set the default agent."
@classmethod
def INPUT_TYPES(s):
inputs = super().INPUT_TYPES()
del inputs["optional"]["tools"]
del inputs["optional"]["rulesets"]
del inputs["optional"]["agent"]
del inputs["optional"]["input_string"]
del inputs["optional"]["STRING"]
return inputs
RETURN_TYPES = ("CONFIG",)
OUTPUT_NODE = True
def run(self, config=None):
if config:
self.agent.config = config
update_config_with_dict(self.agent.config.to_dict())
return (config,)
nodes/audio_nodes.py
+18 -11
View File
@@ -2,6 +2,7 @@ import mimetypes
import os
import folder_paths
import torchaudio
def is_audio_file(filepath):
@@ -11,6 +12,7 @@ def is_audio_file(filepath):
class gtUILoadAudio:
DESCRIPTION = "Load an audio file for processing."
SUPPORTED_FORMATS = (".wav", ".mp3", ".ogg", ".flac", ".aiff", ".aif")
@classmethod
def INPUT_TYPES(s):
@@ -18,20 +20,25 @@ class gtUILoadAudio:
files = [
f
for f in os.listdir(input_dir)
if os.path.isfile(os.path.join(input_dir, f))
and is_audio_file(os.path.join(input_dir, f))
if (
os.path.isfile(os.path.join(input_dir, f))
and f.endswith(gtUILoadAudio.SUPPORTED_FORMATS)
)
]
return {
"required": {"audio": (sorted(files), {"audio_upload": True})},
}
return {"required": {"audio": (sorted(files), {"audio_upload": False})}}
CATEGORY = "Griptape/Audio"
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("AUDIO PATH",)
FUNCTION = "load_audio"
RETURN_TYPES = ("STRING", "AUDIO")
RETURN_NAMES = ("AUDIO_PATH", "AUDIO")
FUNCTION = "gt_load_audio"
def load_audio(self, audio):
def gt_load_audio(self, audio):
audio_path = folder_paths.get_annotated_filepath(audio)
return (audio_path,)
waveform, sample_rate = torchaudio.load(audio_path)
multiplier = 1.0
audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
return (
audio_path,
audio,
)
nodes/base_agent.py
-85
View File
@@ -1,85 +0,0 @@
from griptape.structures import Agent as gtAgent
default_prompt = "{{ input_string }}"
class BaseAgent:
"""
Create a Griptape Agent
"""
def __init__(self):
self.default_prompt = default_prompt
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {},
"optional": {
"agent": (
"AGENT",
{
"forceInput": True,
},
),
"input_string": (
"STRING",
{
"forceInput": True,
},
),
# "inputs": (
# "ARTIFACT",
# {
# "forceInput": True,
# },
# ),
"STRING": (
"STRING",
{
"multiline": True,
},
),
},
}
RETURN_TYPES = (
"STRING",
# "ARTIFACT",
"AGENT",
)
RETURN_NAMES = (
"OUTPUT",
# "ARTIFACT",
"AGENT",
)
FUNCTION = "run"
OUTPUT_NODE = True
CATEGORY = "Griptape/Agent"
def run(
self,
STRING,
# inputs=None,
agent=None,
input_string=None,
):
if not agent:
agent = gtAgent()
# Get the prompt text
if not input_string:
prompt_text = STRING
else:
prompt_text = STRING + "\n\n" + input_string
result = agent.run(prompt_text)
output_string = result.output_task.output.value
return (
output_string,
# result.output_task.output,
agent,
)
nodes/base_audio_task.py
+30 -3
View File
@@ -1,3 +1,7 @@
import io
import tempfile
import torchaudio
from griptape.drivers import OpenAiAudioTranscriptionDriver
from .base_task import gtUIBaseTask
@@ -10,15 +14,17 @@ class gtUIBaseAudioTask(gtUIBaseTask):
inputs["optional"].update(
{
"audio": (
"audio_filepath": (
"STRING",
{
"forceInput": True,
},
),
"audio": ("AUDIO", {}),
"driver": ("DRIVER", {}),
},
)
inputs["optional"].update({"driver": ("DRIVER",)})
# inputs["optional"].update({"driver": ("DRIVER",)})
del inputs["optional"]["input_string"]
del inputs["optional"]["agent"]
del inputs["required"]["STRING"]
@@ -28,7 +34,28 @@ class gtUIBaseAudioTask(gtUIBaseTask):
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("OUTPUT",)
def run(self, audio, driver=None):
def save_audio_tempfile(self, audio_data):
temp_files = []
for waveform in audio_data["waveform"]:
with tempfile.NamedTemporaryFile(suffix=".flac", delete=False) as temp_file:
temp_filename = temp_file.name
# Save the audio to the buffer
buff = io.BytesIO()
torchaudio.save(
buff, waveform, audio_data["sample_rate"], format="FLAC"
)
# Write the buffer's contents to the temporary file
with open(temp_filename, "wb") as f:
f.write(buff.getbuffer())
temp_files.append(temp_filename)
return temp_files
def run(self, audio=None, audio_filepath=None, driver=None):
if not driver:
driver = OpenAiAudioTranscriptionDriver(model="whisper-1")
output = "Output"
nodes/base_config.py
+5 -1
View File
@@ -21,6 +21,10 @@ class gtUIBaseConfig:
{"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01},
),
"seed": ("INT", {"default": 10342349342}),
"image_generation_driver": (
"DRIVER",
{},
),
},
}
@@ -32,5 +36,5 @@ class gtUIBaseConfig:
CATEGORY = "Griptape/Agent Configs"
def create(self, temperature, seed):
def create(self, temperature, seed, image_generation_driver=None):
return (OpenAiStructureConfig(),)
nodes/base_task.py
+2 -1
View File
@@ -1,6 +1,7 @@
from griptape.structures import Agent
from griptape.tasks import PromptTask
from .agent.agent import gtComfyAgent as Agent
class gtUIBaseTask:
def __init__(self):
nodes/config_nodes.py
+94 -31
View File
@@ -4,6 +4,7 @@ from griptape.config import (
AmazonBedrockStructureConfig,
AnthropicStructureConfig,
GoogleStructureConfig,
OpenAiStructureConfig,
StructureConfig,
)
@@ -14,6 +15,7 @@ from griptape.drivers import (
AnthropicImageQueryDriver,
AnthropicPromptDriver,
BedrockClaudeImageQueryModelDriver,
DummyImageGenerationDriver,
GooglePromptDriver,
OllamaPromptDriver,
OpenAiChatPromptDriver,
@@ -65,9 +67,13 @@ class gtUIEnv:
ollama_models = get_ollama_models()
ollama_models.append("")
ollama_port = "11434"
ollama_base_url = "http://localhost"
lmstudio_models = get_lmstudio_models(port="1234")
lmstudio_models.append("")
lmstudio_port = "1234"
lmstudio_base_url = "http://localhost"
class LMStudioPromptDriver(OpenAiChatPromptDriver):
@@ -100,24 +106,33 @@ class gtUILMStudioStructureConfig(gtUIBaseConfig):
inputs = super().INPUT_TYPES()
inputs["required"].update(
{
# "prompt_model": (
# lmstudio_models,
# {"default": lmstudio_models[0]},
# ),
"prompt_model": ("STRING", {"default": ""}),
"prompt_model": (
[],
{"default": ""},
),
# "prompt_model": ("STRING", {"default": ""}),
"base_url": ("STRING", {"default": lmstudio_base_url}),
"port": (
"STRING",
{"default": "1234"},
{"default": lmstudio_port},
),
},
)
return inputs
def create(self, prompt_model, port, temperature, seed):
def create(
self,
prompt_model,
base_url,
port,
temperature,
seed,
image_generation_driver=DummyImageGenerationDriver(),
):
custom_config = StructureConfig(
prompt_driver=LMStudioPromptDriver(
model=prompt_model,
base_url=f"http://localhost:{port}/v1",
base_url=f"{base_url}:{port}/v1",
api_key="lm_studio",
temperature=temperature,
tokenizer=SimpleTokenizer(
@@ -126,6 +141,7 @@ class gtUILMStudioStructureConfig(gtUIBaseConfig):
max_output_tokens=1024,
),
),
image_generation_driver=image_generation_driver,
)
return (custom_config,)
@@ -144,18 +160,31 @@ class gtUIOllamaStructureConfig(gtUIBaseConfig):
inputs["required"].update(
{
"prompt_model": (
ollama_models,
{"default": ollama_models[0]},
[""],
{"default": ""},
),
"base_url": ("STRING", {"default": ollama_base_url}),
"port": ("STRING", {"default": ollama_port}),
},
)
return inputs
def create(self, prompt_model, temperature, seed):
def create(
self,
prompt_model,
temperature,
base_url,
port,
seed,
image_generation_driver=DummyImageGenerationDriver(),
):
custom_config = StructureConfig(
prompt_driver=OllamaPromptDriver(
model=prompt_model, options={"temperature": temperature}
model=prompt_model,
options={"temperature": temperature},
host=f"{base_url}:{port}",
),
image_generation_driver=image_generation_driver,
)
return (custom_config,)
@@ -202,7 +231,14 @@ class gtUIAmazonBedrockStructureConfig(gtUIBaseConfig):
)
return inputs
def create(self, prompt_model, image_query_model, temperature, seed):
def create(
self,
prompt_model,
image_query_model,
temperature,
seed,
image_generation_driver=None,
):
custom_config = AmazonBedrockStructureConfig()
custom_config.prompt_driver = AmazonBedrockPromptDriver(
model=prompt_model, temperature=temperature
@@ -211,7 +247,8 @@ class gtUIAmazonBedrockStructureConfig(gtUIBaseConfig):
image_query_model_driver=BedrockClaudeImageQueryModelDriver(),
model=image_query_model,
)
if image_generation_driver:
custom_config.image_generation_driver = image_generation_driver
return (custom_config,)
@@ -229,11 +266,14 @@ class gtUIGoogleStructureConfig(gtUIBaseConfig):
"Google Structure Config. Use Google's models for prompt and image query."
)
def create(self, temperature, seed):
def create(
self, temperature, seed, image_generation_driver=DummyImageGenerationDriver()
):
custom_config = GoogleStructureConfig(
prompt_driver=GooglePromptDriver(
model="gemini-pro", temperature=temperature
)
),
image_generation_driver=image_generation_driver,
)
return (custom_config,)
@@ -282,7 +322,14 @@ class gtUIAnthropicStructureConfig(gtUIBaseConfig):
)
return inputs
def create(self, prompt_model, image_query_model, temperature, seed):
def create(
self,
prompt_model,
image_query_model,
temperature,
seed,
image_generation_driver=DummyImageGenerationDriver(),
):
custom_config = AnthropicStructureConfig()
custom_config.prompt_driver = AnthropicPromptDriver(
model=prompt_model, temperature=temperature
@@ -290,6 +337,7 @@ class gtUIAnthropicStructureConfig(gtUIBaseConfig):
custom_config.image_query_driver = AnthropicImageQueryDriver(
model=image_query_model
)
custom_config.image_generation_driver = image_generation_driver
return (custom_config,)
@@ -318,23 +366,38 @@ class gtUIOpenAiStructureConfig(gtUIBaseConfig):
)
return inputs
def create(self, prompt_model, image_query_model, temperature, seed):
def create(
self,
prompt_model,
image_query_model,
temperature,
seed,
image_generation_driver=None,
):
OPENAI_API_KEY = get_config("env.OPENAI_API_KEY")
custom_config = StructureConfig(
prompt_driver=OpenAiChatPromptDriver(
model=prompt_model,
api_key=OPENAI_API_KEY,
temperature=temperature,
seed=seed,
),
embedding_driver=OpenAiEmbeddingDriver(api_key=OPENAI_API_KEY),
image_generation_driver=OpenAiImageGenerationDriver(
prompt_driver = OpenAiChatPromptDriver(
model=prompt_model,
api_key=OPENAI_API_KEY,
temperature=temperature,
seed=seed,
)
embedding_driver = OpenAiEmbeddingDriver(api_key=OPENAI_API_KEY)
if not image_generation_driver:
image_generation_driver = OpenAiImageGenerationDriver(
api_key=OPENAI_API_KEY,
model="dalle-e-3",
),
image_query_driver=OpenAiImageQueryDriver(
api_key=OPENAI_API_KEY, model=image_query_model
),
)
image_query_driver = OpenAiImageQueryDriver(
api_key=OPENAI_API_KEY, model=image_query_model
)
# OpenAiStructureConfig()
custom_config = OpenAiStructureConfig(
prompt_driver=prompt_driver,
embedding_driver=embedding_driver,
image_generation_driver=image_generation_driver,
image_query_driver=image_query_driver,
)
return (custom_config,)
nodes/tasks.py
+46 -17
View File
@@ -8,6 +8,7 @@ from griptape.drivers import (
AmazonBedrockImageQueryDriver,
AnthropicImageQueryDriver,
DummyAudioTranscriptionDriver,
DummyImageGenerationDriver,
DummyImageQueryDriver,
OpenAiAudioTranscriptionDriver,
OpenAiImageGenerationDriver,
@@ -19,7 +20,7 @@ from griptape.engines import (
VariationImageGenerationEngine,
)
from griptape.loaders import AudioLoader, ImageLoader
from griptape.structures import Agent, Pipeline, Workflow
from griptape.structures import Pipeline, Workflow
from griptape.tasks import (
AudioTranscriptionTask,
CodeExecutionTask,
@@ -37,7 +38,7 @@ from griptape.utils import load_file
from schema import Schema
from ..py.griptape_config import get_config
from .agent import model_check
from .agent.agent import gtComfyAgent as Agent
from .base_audio_task import gtUIBaseAudioTask
from .base_image_task import gtUIBaseImageTask
from .base_task import gtUIBaseTask
@@ -176,10 +177,24 @@ class gtUIPromptImageGenerationTask(gtUIBaseTask):
agent = Agent()
prompt_text = self.get_prompt_text(STRING, input_string)
if not driver:
driver = OpenAiImageGenerationDriver(
model="dall-e-3", quality="hd", style="natural", api_key=OPENAI_API_KEY
)
# Check and see if the agent.config has an image_generation_driver
if isinstance(
agent.config.image_generation_driver, DummyImageGenerationDriver
):
# create a default driver
driver = OpenAiImageGenerationDriver(
model="dall-e-3",
quality="hd",
style="natural",
api_key=OPENAI_API_KEY,
)
print(
"Current driver doesn't have an image_generation - using OpenAI by default."
)
else:
driver = agent.config.image_generation_driver
# Create an engine configured to use the driver.
engine = PromptImageGenerationEngine(
image_generation_driver=driver,
@@ -272,7 +287,15 @@ class gtUIAudioTranscriptionTask(gtUIBaseAudioTask):
DESCRIPTION = "Transcribe an audio file."
CATEGORY = "Griptape/Audio"
def run(self, audio, driver=None):
def run(self, audio=None, audio_filepath=None, driver=None):
audio_artifact = None
if audio:
audio = self.save_audio_tempfile(audio)[0]
elif audio_filepath:
audio = audio_filepath
else:
return ("There is no audio file.",)
try:
audio_artifact = AudioLoader().load(load_file(audio))
except Exception as e:
@@ -547,26 +570,32 @@ class gtUIToolkitTask(gtUIBaseTask):
if len(tools) == 0:
return super().run(STRING, input_string, agent)
if prompt_text.strip() == "":
return ("No prompt provided", agent)
# if the tool is provided, keep going
if not agent:
agent = Agent()
if model_check(agent):
return (
dedent(
"""
I'm sorry, this agent uses a simple model that can't handle ToolkitTasks.
You might want to try using a different Agent Configuration, or use a simple ToolTask instead.
model, simple_model = agent.model_check()
if simple_model:
response = agent.model_response(model)
return (response, agent)
# return (
# dedent(
# """
# I'm sorry, this agent uses a simple model that can't handle ToolkitTasks.
# You might want to try using a different Agent Configuration, or use a simple ToolTask instead.
Reach out for help on Discord (https://discord.gg/gnWRz88eym) if you would like some help.
"""
),
agent,
)
# Reach out for help on Discord (https://discord.gg/gnWRz88eym) if you would like some help.
# """
# ),
# agent,
# )
task = ToolkitTask(prompt_text, tools=tools)
try:
agent.add_task(task)
except Exception as e:
print(e)
result = agent.run()
return (result.output_task.output.value, agent)
py/griptape_config.py
+13
View File
@@ -103,3 +103,16 @@ def get_config(key, default=None):
else:
return default
return config
def update_config_with_dict(config_dict={}):
with open(USER_CONFIG_FILE, "r") as file:
existing_data = json.load(file)
agent_config = {"agent_config": config_dict}
existing_data.update(agent_config)
with open(USER_CONFIG_FILE, "w") as file:
json.dump(
existing_data, file, indent=4
) # Use indent=4 for pretty formatting, optional