From d6f27e9534bae447fbbf80a5dde8cc28f9669cd9 Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Fri, 19 Jul 2024 08:32:53 +1200 Subject: [PATCH] added vectorStoreDrivers Fixed tasks replacing agents in PromptTask, Create Image from text task, and text to speech --- __init__.py | 41 ++++++++- js/nodeFixes.js | 4 +- nodes/agent/gtUICreateAgentFromConfig.py | 35 ++++++++ nodes/config/gtUIOpenAiStructureConfig.py | 4 +- nodes/config/gtUIStructureConfig.py | 35 ++++++-- .../gtUIAmazonOpenSearchVectorStoreDriver.py | 48 +++++++++++ .../gtUIAzureMongoDbVectorStoreDriver.py | 85 +++++++++++++++++++ nodes/drivers/gtUIBaseTextToSpeechDriver.py | 2 + .../gtUIElevenLabsTextToSpeechDriver.py | 2 + .../gtUIMongoDbAtlasVectorStoreDriver.py | 85 +++++++++++++++++++ nodes/drivers/gtUIOpenAiChatPromptDriver.py | 2 +- nodes/drivers/gtUIOpenAiImageQueryDriver.py | 5 +- nodes/drivers/gtUIOpenAiTextToSpeechDriver.py | 44 ++++++++++ .../drivers/gtUIPgVectorVectorStoreDriver.py | 72 ++++++++++++++++ nodes/drivers/gtUIQdrantVectorStoreDriver.py | 61 +++++++++++++ nodes/drivers/gtUIRedisVectorStoreDriver.py | 62 ++++++++++++++ nodes/tasks/gtUIAudioTranscriptionTask.py | 21 ++++- nodes/tasks/gtUIBaseTask.py | 4 +- nodes/tasks/gtUIPromptImageGenerationTask.py | 8 +- nodes/tasks/gtUITextToSpeechTask.py | 21 +++-- 20 files changed, 610 insertions(+), 31 deletions(-) create mode 100644 nodes/agent/gtUICreateAgentFromConfig.py create mode 100644 nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py create mode 100644 nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py create mode 100644 nodes/drivers/gtUIMongoDbAtlasVectorStoreDriver.py create mode 100644 nodes/drivers/gtUIOpenAiTextToSpeechDriver.py create mode 100644 nodes/drivers/gtUIPgVectorVectorStoreDriver.py create mode 100644 nodes/drivers/gtUIQdrantVectorStoreDriver.py create mode 100644 nodes/drivers/gtUIRedisVectorStoreDriver.py diff --git a/__init__.py b/__init__.py index 5c0e816..11544a0 100644 --- a/__init__.py +++ b/__init__.py @@ -14,6 +14,7 @@ from dotenv import load_dotenv # AGENT from .nodes.agent.CreateAgent import CreateAgent from .nodes.agent.ExpandAgent import ExpandAgent +from .nodes.agent.gtUICreateAgentFromConfig import gtUICreateAgentFromConfig from .nodes.agent.gtUISetDefaultAgent import gtUISetDefaultAgent from .nodes.agent.RunAgent import RunAgent @@ -62,6 +63,11 @@ from .nodes.drivers.gtUIAmazonBedrockTitanImageGenerationDriver import ( gtUIAmazonBedrockTitanImageGenerationDriver, ) +# - Amazon OpenSearch +from .nodes.drivers.gtUIAmazonOpenSearchVectorStoreDriver import ( + gtUIAmazonOpenSearchVectorStoreDriver, +) + # - Amazon SageMaker Jumpstart from .nodes.drivers.gtUIAmazonSageMakerJumpstartEmbeddingDriver import ( gtUIAmazonSageMakerJumpstartEmbeddingDriver, @@ -73,8 +79,11 @@ from .nodes.drivers.gtUIAmazonSageMakerJumpstartPromptDriver import ( # - Anthropic from .nodes.drivers.gtUIAnthropicImageQueryDriver import gtUIAnthropicImageQueryDriver from .nodes.drivers.gtUIAnthropicPromptDriver import gtUIAnthropicPromptDriver +from .nodes.drivers.gtUIAzureMongoDbVectorStoreDriver import ( + gtUIAzureMongoDbVectorStoreDriver, +) -# - Azure OpenAI +# - Azure from .nodes.drivers.gtUIAzureOpenAiChatPromptDriver import ( gtUIAzureOpenAiChatPromptDriver, ) @@ -123,6 +132,11 @@ from .nodes.drivers.gtUILocalVectorStoreDriver import gtUILocalVectorStoreDriver # - Marqo from .nodes.drivers.gtUIMarqoVectorStoreDriver import gtUIMarqoVectorStoreDriver +# - Mongodb +from .nodes.drivers.gtUIMongoDbAtlasVectorStoreDriver import ( + gtUIMongoDbAtlasVectorStoreDriver, +) + # - Ollama from .nodes.drivers.gtUIOllamaPromptDriver import gtUIOllamaPromptDriver @@ -142,10 +156,20 @@ from .nodes.drivers.gtUIOpenAiImageGenerationDriver import ( gtUIOpenAiImageGenerationDriver, ) from .nodes.drivers.gtUIOpenAiImageQueryDriver import gtUIOpenAiImageQueryDriver +from .nodes.drivers.gtUIOpenAiTextToSpeechDriver import gtUIOpenAiTextToSpeechDriver + +# - PGVector +from .nodes.drivers.gtUIPgVectorVectorStoreDriver import gtUIPgVectorVectorStoreDriver # - Pinecone from .nodes.drivers.gtUIPineconeVectorStoreDriver import gtUIPineconeVectorStoreDriver +# - Qdrant +from .nodes.drivers.gtUIQdrantVectorStoreDriver import gtUIQdrantVectorStoreDriver + +# - Redis +from .nodes.drivers.gtUIRedisVectorStoreDriver import gtUIRedisVectorStoreDriver + # - Voyage AI from .nodes.drivers.gtUIVoyageAiEmbeddingDriver import gtUIVoyageAiEmbeddingDriver @@ -209,6 +233,7 @@ WEB_DIRECTORY = "./js" NODE_CLASS_MAPPINGS = { # AGENT "Griptape Create: Agent": CreateAgent, + "Griptape Create: Agent from Config": gtUICreateAgentFromConfig, "Griptape Run: Agent": RunAgent, "Griptape Expand: Agent Nodes": ExpandAgent, "Griptape Set: Default Agent": gtUISetDefaultAgent, @@ -281,12 +306,20 @@ NODE_CLASS_MAPPINGS = { "Griptape Embedding Driver: OpenAI Compatable": gtUIOpenAiCompatableEmbeddingDriver, "Griptape Embedding Driver: Voyage AI": gtUIVoyageAiEmbeddingDriver, # VECTOR STORE DRIVERS + "Griptape Driver: Amazon OpenSearch Vector Store": gtUIAmazonOpenSearchVectorStoreDriver, + "Griptape Driver: Azure MongoDB Vector Store": gtUIAzureMongoDbVectorStoreDriver, "Griptape Driver: Marqo Vector Store": gtUIMarqoVectorStoreDriver, + "Griptape Driver: MongoDB Atlas Vector Store": gtUIMongoDbAtlasVectorStoreDriver, "Griptape Driver: Local Vector Store": gtUILocalVectorStoreDriver, + "Griptape Driver: PGVector Vector Store": gtUIPgVectorVectorStoreDriver, "Griptape Driver: Pinecone Vector Store": gtUIPineconeVectorStoreDriver, - # AUDIO DRIVER - "Griptape Audio Driver: ElevenLabs": gtUIElevenLabsTextToSpeechDriver, - "Griptape Audio Driver: OpenAI": gtUIOpenAiAudioTranscriptionDriver, + "Griptape Driver: Redis Vector Store": gtUIRedisVectorStoreDriver, + "Griptape Driver: Qdrant Vector Store": gtUIQdrantVectorStoreDriver, + # TEXT TO SPEECH DRIVERS + "Griptape Text To Speech Driver: ElevenLabs": gtUIElevenLabsTextToSpeechDriver, + "Griptape Text To Speech Driver: OpenAI": gtUIOpenAiTextToSpeechDriver, + # AUDIO DRIVERS + "Griptape Audio Transcription Driver: OpenAI": gtUIOpenAiAudioTranscriptionDriver, # WEBSEARCH DRIVERS "Griptape Driver: DuckDuckGo WebSearch": gtUIDuckDuckGoWebSearchDriver, "Griptape Driver: Google WebSearch": gtUIGoogleWebSearchDriver, diff --git a/js/nodeFixes.js b/js/nodeFixes.js index bada9a4..cb83d1a 100644 --- a/js/nodeFixes.js +++ b/js/nodeFixes.js @@ -42,6 +42,8 @@ export const nodeFixes = { "MergeTexts": "Griptape Combine: Merge Texts", "EnvironmentConfig": "Griptape Config: Environment Variables", "gtUIOpenAiImageGenerationDriver": "Griptape Audio Driver: OpenAI", - "gtUILoadAudio": "Griptape Load: Audio" + "gtUILoadAudio": "Griptape Load: Audio", + "Griptape Audio Driver: ElevenLabs": "Griptape Text To Speech Driver: ElevenLabs", + "Griptape Audio Driver: OpenAI": "Griptape Audio Transcription Driver: OpenAI" } } \ No newline at end of file diff --git a/nodes/agent/gtUICreateAgentFromConfig.py b/nodes/agent/gtUICreateAgentFromConfig.py new file mode 100644 index 0000000..4495081 --- /dev/null +++ b/nodes/agent/gtUICreateAgentFromConfig.py @@ -0,0 +1,35 @@ +from .BaseAgent import BaseAgent +from .gtComfyAgent import gtComfyAgent + + +class gtUICreateAgentFromConfig(BaseAgent): + DESCRIPTION = "Create an Agent from a Config" + + def __init__(self): + self.agent = None + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + del inputs["optional"]["tools"] + del inputs["optional"]["rulesets"] + del inputs["optional"]["STRING"] + del inputs["optional"]["agent"] + del inputs["optional"]["input_string"] + return inputs + + RETURN_TYPES = ("AGENT",) + RETURN_NAMES = ("AGENT",) + + def run(self, **kwargs): + config = kwargs.get("config", None) + + create_dict = {} + + if config: + create_dict["config"] = config + + # Now create the agent + self.agent = gtComfyAgent(**create_dict) + + return (self.agent,) diff --git a/nodes/config/gtUIOpenAiStructureConfig.py b/nodes/config/gtUIOpenAiStructureConfig.py index 568762a..babe983 100644 --- a/nodes/config/gtUIOpenAiStructureConfig.py +++ b/nodes/config/gtUIOpenAiStructureConfig.py @@ -31,11 +31,11 @@ class gtUIOpenAiStructureConfig(gtUIBaseConfig): inputs["required"].update( { "prompt_model": ( - ["gpt-4o", "gpt-4", "gpt-3.5-turbo"], + ["gpt-4o", "gpt-4", "gpt-4o-mini", "gpt-3.5-turbo"], {"default": default_prompt_model}, ), "image_query_model": ( - ["gpt-4o", "gpt-4-vision-preview"], + ["gpt-4o", "gpt-4-vision-preview", "gpt-4o-mini"], {"default": default_image_query_model}, ), } diff --git a/nodes/config/gtUIStructureConfig.py b/nodes/config/gtUIStructureConfig.py index 45a0dbb..5f347e2 100644 --- a/nodes/config/gtUIStructureConfig.py +++ b/nodes/config/gtUIStructureConfig.py @@ -1,6 +1,21 @@ from griptape.config import ( StructureConfig, ) +from griptape.drivers import ( + OpenAiAudioTranscriptionDriver, + OpenAiChatPromptDriver, + OpenAiEmbeddingDriver, + OpenAiImageGenerationDriver, + OpenAiImageQueryDriver, + OpenAiTextToSpeechDriver, +) + +default_chat_prompt_driver = OpenAiChatPromptDriver(model="gpt-4o") +default_image_generation_driver = OpenAiImageGenerationDriver(model="dall-e-3") +default_image_query_driver = OpenAiImageQueryDriver(model="gpt-4o") +default_embedding_driver = OpenAiEmbeddingDriver() +default_text_to_speech_driver = OpenAiTextToSpeechDriver(model="tts-1", voice="alloy") +default_audio_transcription_driver = OpenAiAudioTranscriptionDriver(model="whisper-1") class gtUIStructureConfig: @@ -35,13 +50,21 @@ class gtUIStructureConfig: CATEGORY = "Griptape/Agent Configs" def create(self, **kwargs): - prompt_driver = kwargs.get("prompt_driver", None) - image_generation_driver = kwargs.get("image_generation_driver", None) - image_query_driver = kwargs.get("image_query_driver", None) - embedding_driver = kwargs.get("embedding_driver", None) + prompt_driver = kwargs.get("prompt_driver", default_chat_prompt_driver) + image_generation_driver = kwargs.get( + "image_generation_driver", default_image_generation_driver + ) + image_query_driver = kwargs.get( + "image_query_driver", default_image_query_driver + ) + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) vector_store_driver = kwargs.get("vector_store_driver", None) - text_to_speech_driver = kwargs.get("text_to_speech_driver", None) - audio_transcription_driver = kwargs.get("audio_transcription_driver", None) + text_to_speech_driver = kwargs.get( + "text_to_speech_driver", default_text_to_speech_driver + ) + audio_transcription_driver = kwargs.get( + "audio_transcription_driver", default_audio_transcription_driver + ) drivers = {} diff --git a/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py b/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py new file mode 100644 index 0000000..ab604b0 --- /dev/null +++ b/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py @@ -0,0 +1,48 @@ +import os + +from griptape.drivers import AmazonOpenSearchVectorStoreDriver, OpenAiEmbeddingDriver + +from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver + +default_embedding_driver = OpenAiEmbeddingDriver() + +default_host_env = "AMAZON_OPENSEARCH_HOST" +default_index_env = "AMAZON_OPENSEARCH_INDEX_NAME" + + +class gtUIAmazonOpenSearchVectorStoreDriver(gtUIBaseVectorStoreDriver): + DESCRIPTION = "Griptape Open Search Vector Store Driver: https://aws.amazon.com/opensearch-service/" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update() + inputs["optional"].update( + { + "host_env": ("STRING", {"default": default_host_env}), + "index_env": ("STRING", {"default": default_index_env}), + } + ) + + return inputs + + def create(self, **kwargs): + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + host_env = kwargs.get("host_env", default_host_env) + index_env_var = kwargs.get("index_name_env", default_index_env) + + if host_env: + host = os.getenv(host_env) + if index_env_var: + index = os.getenv(index_env_var) + params = {} + if host: + params["host"] = host + if index: + params["index_name"] = index + if embedding_driver: + params["embedding_driver"] = embedding_driver + else: + params["embedding_driver"] = default_embedding_driver + driver = AmazonOpenSearchVectorStoreDriver(**params) + return (driver,) diff --git a/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py b/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py new file mode 100644 index 0000000..d21aa92 --- /dev/null +++ b/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py @@ -0,0 +1,85 @@ +import os + +from griptape.drivers import AzureMongoDbVectorStoreDriver, OpenAiEmbeddingDriver + +from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver + +default_embedding_driver = OpenAiEmbeddingDriver() + +default_host_env = "AZURE_MONGODB_HOST" +default_username_env = "AZURE_MONGODB_USERNAME" +default_password_env = "AZURE_MONGODB_PASSWORD" +default_database_name_env = "AZURE_MONGODB_DATABASE_NAME" +default_collection_name_env = "AZURE_MONGODB_COLLECTION_NAME" +default_index_name_env = "AZURE_MONGODB_INDEX_NAME" +default_vector_path_env = "AZURE_MONGODB_VECTOR_PATH" + + +class gtUIAzureMongoDbVectorStoreDriver(gtUIBaseVectorStoreDriver): + DESCRIPTION = "Griptape Azure Mongodb Atlas Vector Store Driver: https://www.mongodb.com/products/platform/atlas-database" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update() + inputs["optional"].update( + { + "host_env": ("STRING", {"default": default_host_env}), + "username_env": ("STRING", {"default": default_username_env}), + "password_env": ("STRING", {"default": default_password_env}), + "database_name_env": ("STRING", {"default": default_database_name_env}), + "collection_name_env": ( + "STRING", + {"default": default_collection_name_env}, + ), + "index_name_env": ("STRING", {"default": default_index_name_env}), + "vector_path_env": ("STRING", {"default": default_vector_path_env}), + } + ) + + return inputs + + def create(self, **kwargs): + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + host_env = kwargs.get("host_env", default_host_env) + username_env = kwargs.get("username_env", default_username_env) + password_env = kwargs.get("password_env", default_password_env) + database_name_env = kwargs.get("database_name_env", default_database_name_env) + collection_name_env = kwargs.get( + "collection_name_env", default_collection_name_env + ) + index_name_env_var = kwargs.get("index_name_env", default_index_name_env) + vector_path_env = kwargs.get("vector_path_env", default_vector_path_env) + + if username_env: + username = os.getenv(username_env) + if password_env: + password = os.getenv(password_env) + if host_env: + host = os.getenv(host_env) + if database_name_env: + database_name = os.getenv(database_name_env) + if collection_name_env: + collection_name = os.getenv(collection_name_env) + if vector_path_env: + vector_path = os.getenv(vector_path_env) + + params = {} + if username and password and host and database_name: + params["connection_string"] = ( + f"mongodb+srv://{username}:{password}@{host}/{database_name}?tls=true&authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000", + ) + if database_name: + params["database_name"] = database_name + if collection_name: + params["collection_name"] = collection_name + if vector_path: + params["vector_path"] = vector_path + if index_name_env_var: + params["index_name"] = os.getenv(index_name_env_var) + if embedding_driver: + params["embedding_driver"] = embedding_driver + else: + params["embedding_driver"] = default_embedding_driver + driver = AzureMongoDbVectorStoreDriver(**params) + return (driver,) diff --git a/nodes/drivers/gtUIBaseTextToSpeechDriver.py b/nodes/drivers/gtUIBaseTextToSpeechDriver.py index 6531276..9e15b22 100644 --- a/nodes/drivers/gtUIBaseTextToSpeechDriver.py +++ b/nodes/drivers/gtUIBaseTextToSpeechDriver.py @@ -2,6 +2,8 @@ from griptape.drivers import BaseTextToSpeechDriver class gtUIBaseTextToSpeechDriver: + DESCRIPTION = "Griptape Text to Speech Driver" + @classmethod def INPUT_TYPES(s): return { diff --git a/nodes/drivers/gtUIElevenLabsTextToSpeechDriver.py b/nodes/drivers/gtUIElevenLabsTextToSpeechDriver.py index a013921..5a83e7e 100644 --- a/nodes/drivers/gtUIElevenLabsTextToSpeechDriver.py +++ b/nodes/drivers/gtUIElevenLabsTextToSpeechDriver.py @@ -8,6 +8,8 @@ premade_voices = [] class gtUIElevenLabsTextToSpeechDriver(gtUIBaseTextToSpeechDriver): + DESCRIPTION = "Griptape ElevenLabs to Speech Driver: https://elevenlabs.io/apps/" + @classmethod def INPUT_TYPES(s): inputs = super().INPUT_TYPES() diff --git a/nodes/drivers/gtUIMongoDbAtlasVectorStoreDriver.py b/nodes/drivers/gtUIMongoDbAtlasVectorStoreDriver.py new file mode 100644 index 0000000..6d9a460 --- /dev/null +++ b/nodes/drivers/gtUIMongoDbAtlasVectorStoreDriver.py @@ -0,0 +1,85 @@ +import os + +from griptape.drivers import MongoDbAtlasVectorStoreDriver, OpenAiEmbeddingDriver + +from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver + +default_embedding_driver = OpenAiEmbeddingDriver() + +default_host_env = "MONGODB_HOST" +default_username_env = "MONGODB_USERNAME" +default_password_env = "MONGODB_PASSWORD" +default_database_name_env = "MONGODB_DATABASE_NAME" +default_collection_name_env = "MONGODB_COLLECTION_NAME" +default_index_name_env = "MONGODB_INDEX_NAME" +default_vector_path_env = "MONGODB_VECTOR_PATH" + + +class gtUIMongoDbAtlasVectorStoreDriver(gtUIBaseVectorStoreDriver): + DESCRIPTION = "Griptape Mongodb Atlas Vector Store Driver: https://www.mongodb.com/products/platform/atlas-database" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update() + inputs["optional"].update( + { + "host_env": ("STRING", {"default": default_host_env}), + "username_env": ("STRING", {"default": default_username_env}), + "password_env": ("STRING", {"default": default_password_env}), + "database_name_env": ("STRING", {"default": default_database_name_env}), + "collection_name_env": ( + "STRING", + {"default": default_collection_name_env}, + ), + "index_name_env_var": ("STRING", {"default": default_index_name_env}), + "vector_path_env": ("STRING", {"default": default_vector_path_env}), + } + ) + + return inputs + + def create(self, **kwargs): + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + host_env = kwargs.get("host_env", default_host_env) + username_env = kwargs.get("username_env", default_username_env) + password_env = kwargs.get("password_env", default_password_env) + database_name_env = kwargs.get("database_name_env", default_database_name_env) + collection_name_env = kwargs.get( + "collection_name_env", default_collection_name_env + ) + index_name_env_var = kwargs.get("index_name_env_var", default_index_name_env) + vector_path_env = kwargs.get("vector_path_env", default_vector_path_env) + + if username_env: + username = os.getenv(username_env) + if password_env: + password = os.getenv(password_env) + if host_env: + host = os.getenv(host_env) + if database_name_env: + database_name = os.getenv(database_name_env) + if collection_name_env: + collection_name = os.getenv(collection_name_env) + if vector_path_env: + vector_path = os.getenv(vector_path_env) + + params = {} + if username and password and host and database_name: + params["connection_string"] = ( + f"mongodb+srv://{username}:{password}@{host}/{database_name}", + ) + if database_name: + params["database_name"] = database_name + if collection_name: + params["collection_name"] = collection_name + if vector_path: + params["vector_path"] = vector_path + if index_name_env_var: + params["index_name"] = os.getenv(index_name_env_var) + if embedding_driver: + params["embedding_driver"] = embedding_driver + else: + params["embedding_driver"] = default_embedding_driver + driver = MongoDbAtlasVectorStoreDriver(**params) + return (driver,) diff --git a/nodes/drivers/gtUIOpenAiChatPromptDriver.py b/nodes/drivers/gtUIOpenAiChatPromptDriver.py index a8a360c..6453155 100644 --- a/nodes/drivers/gtUIOpenAiChatPromptDriver.py +++ b/nodes/drivers/gtUIOpenAiChatPromptDriver.py @@ -4,7 +4,7 @@ from griptape.drivers import OpenAiChatPromptDriver from .gtUIBasePromptDriver import gtUIBasePromptDriver -models = ["gpt-4o", "gpt-4", "gpt-3.5-turbo"] +models = ["gpt-4o", "gpt-4", "gpt-4o-mini", "gpt-3.5-turbo"] class gtUIOpenAiChatPromptDriver(gtUIBasePromptDriver): diff --git a/nodes/drivers/gtUIOpenAiImageQueryDriver.py b/nodes/drivers/gtUIOpenAiImageQueryDriver.py index 0e421ba..8578c58 100644 --- a/nodes/drivers/gtUIOpenAiImageQueryDriver.py +++ b/nodes/drivers/gtUIOpenAiImageQueryDriver.py @@ -4,10 +4,7 @@ from griptape.drivers import OpenAiImageQueryDriver from .gtUIBaseImageQueryDriver import gtUIBaseImageQueryDriver -models = [ - "gpt-4o", - "gpt-4-vision-preview", -] +models = ["gpt-4o", "gpt-4-vision-preview", "gpt-4o-mini"] default_api_key_env = "OPENAI_API_KEY" diff --git a/nodes/drivers/gtUIOpenAiTextToSpeechDriver.py b/nodes/drivers/gtUIOpenAiTextToSpeechDriver.py new file mode 100644 index 0000000..e819330 --- /dev/null +++ b/nodes/drivers/gtUIOpenAiTextToSpeechDriver.py @@ -0,0 +1,44 @@ +import os + +from griptape.drivers import OpenAiTextToSpeechDriver + +from .gtUIBaseTextToSpeechDriver import gtUIBaseTextToSpeechDriver + +voices = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"] +default_model = "tts-1" + + +class gtUIOpenAiTextToSpeechDriver(gtUIBaseTextToSpeechDriver): + DESCRIPTION = "Griptape OpenAi Text to Speech Driver" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + + inputs["optional"].update( + { + "model": ("STRING", {"default": default_model}), + "voice": (voices, {"default": voices[0]}), + } + ) + return inputs + + RETURN_TYPES = ("DRIVER",) + RETURN_NAMES = ("DRIVER",) + + FUNCTION = "create" + + def create(self, **kwargs): + api_key = os.getenv("OPENAI_API_KEY") + model = kwargs.get("model", default_model) + voice = kwargs.get("voice", voices[0]) + + params = {} + if api_key: + params["api_key"] = api_key + if voice: + params["voice"] = voice + if model: + params["model"] = model + driver = OpenAiTextToSpeechDriver(**params) + return (driver,) diff --git a/nodes/drivers/gtUIPgVectorVectorStoreDriver.py b/nodes/drivers/gtUIPgVectorVectorStoreDriver.py new file mode 100644 index 0000000..c81bc73 --- /dev/null +++ b/nodes/drivers/gtUIPgVectorVectorStoreDriver.py @@ -0,0 +1,72 @@ +import os + +from griptape.drivers import OpenAiEmbeddingDriver, PgVectorVectorStoreDriver + +from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver + +default_embedding_driver = OpenAiEmbeddingDriver() + +default_user_env = "POSTGRES_USER" +default_pass_env = "POSTGRES_PASSWORD" +default_host_env = "POSTGRES_HOST" +default_port_env = "POSTGRES_PORT" +default_name_env = "POSTGRES_DB" +default_table_name = "griptape_vectors" + + +class gtUIPgVectorVectorStoreDriver(gtUIBaseVectorStoreDriver): + DESCRIPTION = ( + "Griptape PGVector Vector Store Driver: https://github.com/pgvector/pgvector" + ) + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update() + inputs["optional"].update( + { + "host_env": ("STRING", {"default": default_host_env}), + "user_env": ("STRING", {"default": default_user_env}), + "pass_env": ("STRING", {"default": default_pass_env}), + "port_env": ("STRING", {"default": default_port_env}), + "name_env": ("STRING", {"default": default_name_env}), + "table_name": ("STRING", {"default": default_table_name}), + } + ) + + return inputs + + def create(self, **kwargs): + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + host_env = kwargs.get("host_env", default_host_env) + user_env = kwargs.get("user_env", default_user_env) + pass_env = kwargs.get("pass_env", default_pass_env) + port_env = kwargs.get("port_env", default_port_env) + name_env = kwargs.get("name_env", default_name_env) + table_name = kwargs.get("table_name", default_table_name) + + if host_env: + host = os.getenv(host_env) + if user_env: + user = os.getenv(user_env) + if pass_env: + password = os.getenv(pass_env) + if port_env: + port = os.getenv(port_env) + if name_env: + name = os.getenv(name_env) + + params = {} + + if user and password and host and port and name: + params["connection_string"] = ( + f"postgresql://{user}:{password}@{host}:{port}/{name}" + ) + if table_name: + params["table_name"] = table_name + if embedding_driver: + params["embedding_driver"] = embedding_driver + else: + params["embedding_driver"] = default_embedding_driver + driver = PgVectorVectorStoreDriver(**params) + return (driver,) diff --git a/nodes/drivers/gtUIQdrantVectorStoreDriver.py b/nodes/drivers/gtUIQdrantVectorStoreDriver.py new file mode 100644 index 0000000..e76df8d --- /dev/null +++ b/nodes/drivers/gtUIQdrantVectorStoreDriver.py @@ -0,0 +1,61 @@ +import os + +from griptape.drivers import OpenAiEmbeddingDriver, QdrantVectorStoreDriver + +from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver + +default_embedding_driver = OpenAiEmbeddingDriver() + + +default_api_key_env = "QDRANT_CLUSTER_API_KEY" +default_url_env = "QDRANT_CLUSTER_ENDPOINT" +default_collection_name = "griptape" +default_payload_key = "content" + + +class gtUIQdrantVectorStoreDriver(gtUIBaseVectorStoreDriver): + DESCRIPTION = "Griptape Qdrant Vector Store Driver: https://qdrant.tech/" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update() + inputs["optional"].update( + { + "api_key_env": ("STRING", {"default": default_api_key_env}), + "url_env": ("STRING", {"default": default_url_env}), + "collection_name": ("STRING", {"default": default_collection_name}), + "content_payload_key": ("STRING", {"default": default_payload_key}), + } + ) + + return inputs + + def create(self, **kwargs): + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + url_env = kwargs.get("url_env", default_url_env) + api_key_env = kwargs.get("api_key_env", default_api_key_env) + collection_name = kwargs.get("collection_name", default_collection_name) + content_payload_key = kwargs.get("content_payload_key", default_payload_key) + + if url_env: + url = os.getenv(url_env) + if api_key_env: + api_key = os.getenv(api_key_env) + + params = {} + + if api_key: + params["api_key"] = api_key + if url: + params["url"] = url + if collection_name: + params["collection_name"] = collection_name + if content_payload_key: + params["content_payload_key"] = content_payload_key + if embedding_driver: + params["embedding_driver"] = embedding_driver + else: + params["embedding_driver"] = default_embedding_driver + driver = QdrantVectorStoreDriver(**params) + return (driver,) diff --git a/nodes/drivers/gtUIRedisVectorStoreDriver.py b/nodes/drivers/gtUIRedisVectorStoreDriver.py new file mode 100644 index 0000000..f3f8a81 --- /dev/null +++ b/nodes/drivers/gtUIRedisVectorStoreDriver.py @@ -0,0 +1,62 @@ +import os + +from griptape.drivers import OpenAiEmbeddingDriver, RedisVectorStoreDriver + +from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver + +default_embedding_driver = OpenAiEmbeddingDriver() + +default_host_env = "REDIS_HOST" +default_port_env = "REDIS_PORT" +default_password_env = "REDIS_PASSWORD" +default_index_env = "REDIS_INDEX" + + +class gtUIRedisVectorStoreDriver(gtUIBaseVectorStoreDriver): + DESCRIPTION = "Griptape Redis Atlas Vector Store Driver: https://redis.io/" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update() + inputs["optional"].update( + { + "host_env": ("STRING", {"default": default_host_env}), + "port_env": ("STRING", {"default": default_port_env}), + "password_env": ("STRING", {"default": default_password_env}), + "index_env": ("STRING", {"default": default_index_env}), + } + ) + + return inputs + + def create(self, **kwargs): + embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + host_env = kwargs.get("host_env", default_host_env) + port_env = kwargs.get("port_env", default_port_env) + password_env = kwargs.get("password_env", default_password_env) + index_env_var = kwargs.get("index_name_env", default_index_env) + + if password_env: + password = os.getenv(password_env) + if host_env: + host = os.getenv(host_env) + if port_env: + port = os.getenv(port_env) + if index_env_var: + index = os.getenv(index_env_var) + params = {} + if password: + params["password"] = password + if host: + params["host"] = host + if port: + params["port"] = port + if index: + params["index"] = index + if embedding_driver: + params["embedding_driver"] = embedding_driver + else: + params["embedding_driver"] = default_embedding_driver + driver = RedisVectorStoreDriver(**params) + return (driver,) diff --git a/nodes/tasks/gtUIAudioTranscriptionTask.py b/nodes/tasks/gtUIAudioTranscriptionTask.py index 65b948a..45f26ac 100644 --- a/nodes/tasks/gtUIAudioTranscriptionTask.py +++ b/nodes/tasks/gtUIAudioTranscriptionTask.py @@ -15,6 +15,7 @@ from griptape.tasks import ( ) from griptape.utils import load_file +from ..agent.gtComfyAgent import gtComfyAgent as Agent from .gtUIBaseAudioTask import gtUIBaseAudioTask default_prompt = "{{ input_string }}" @@ -24,6 +25,19 @@ OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") class gtUIAudioTranscriptionTask(gtUIBaseAudioTask): DESCRIPTION = "Transcribe an audio file." + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + + # Update optional inputs to include 'image' and adjust others as necessary + inputs["optional"].update( + { + "agent": ("AGENT",), + } + ) + return inputs + CATEGORY = "Griptape/Audio" RETURN_TYPES = ("STRING",) @@ -33,6 +47,9 @@ class gtUIAudioTranscriptionTask(gtUIBaseAudioTask): audio = kwargs.get("audio", None) audio_filepath = kwargs.get("audio_filepath", None) driver = kwargs.get("driver", None) + agent = kwargs.get("agent", None) + if not agent: + agent = Agent() audio_artifact = None if audio: @@ -47,10 +64,12 @@ class gtUIAudioTranscriptionTask(gtUIBaseAudioTask): except Exception as e: print(f"Error loading audio file: {e}") if audio_artifact: - if not driver: + if not driver and not agent: audio_transcription_driver = OpenAiAudioTranscriptionDriver( model="whisper-1" ) + elif not driver and agent: + audio_transcription_driver = agent.config.audio_transcription_driver else: audio_transcription_driver = driver diff --git a/nodes/tasks/gtUIBaseTask.py b/nodes/tasks/gtUIBaseTask.py index 37c8727..d729dbd 100644 --- a/nodes/tasks/gtUIBaseTask.py +++ b/nodes/tasks/gtUIBaseTask.py @@ -105,10 +105,10 @@ class gtUIBaseTask: print(e) else: try: - agent.add_task(task) + result = agent.run(prompt_text) except Exception as e: print(e) - result = agent.run() + # result = agent.run() return ( result.output_task.output.value, agent, diff --git a/nodes/tasks/gtUIPromptImageGenerationTask.py b/nodes/tasks/gtUIPromptImageGenerationTask.py index 6166171..9ae5c9b 100644 --- a/nodes/tasks/gtUIPromptImageGenerationTask.py +++ b/nodes/tasks/gtUIPromptImageGenerationTask.py @@ -8,6 +8,7 @@ from griptape.drivers import ( from griptape.engines import ( PromptImageGenerationEngine, ) +from griptape.structures import Pipeline from griptape.tasks import ( PromptImageGenerationTask, ) @@ -72,6 +73,7 @@ class gtUIPromptImageGenerationTask(gtUIBaseTask): else: driver = agent.config.image_generation_driver # Create an engine configured to use the driver. + engine = PromptImageGenerationEngine( image_generation_driver=driver, ) @@ -83,11 +85,13 @@ class gtUIPromptImageGenerationTask(gtUIBaseTask): output_dir=output_dir, ) try: - agent.add_task(prompt_task) + pipeline = Pipeline() + pipeline.add_task(prompt_task) + # agent.add_task(prompt_task) except Exception as e: print(e) - result = agent.run() + result = pipeline.run() filename = result.output_task.output.name image_path = os.path.join(output_dir, filename) diff --git a/nodes/tasks/gtUITextToSpeechTask.py b/nodes/tasks/gtUITextToSpeechTask.py index c890be6..7c9e001 100644 --- a/nodes/tasks/gtUITextToSpeechTask.py +++ b/nodes/tasks/gtUITextToSpeechTask.py @@ -1,10 +1,11 @@ import os from griptape.artifacts import AudioArtifact -from griptape.drivers import ElevenLabsTextToSpeechDriver +from griptape.drivers import DummyTextToSpeechDriver, ElevenLabsTextToSpeechDriver from griptape.engines import ( TextToSpeechEngine, ) +from griptape.structures import Pipeline # from ..agent.agent import gtComfyAgent as Agent # from griptape.structures import Agent @@ -38,20 +39,24 @@ class gtUITextToSpeechTask(gtUIBaseTask): input_string = kwargs.get("input_string", "") agent = kwargs.get("agent", Agent()) driver = kwargs.get("driver", None) + if not driver: - driver = ElevenLabsTextToSpeechDriver( - api_key=ELEVEN_LABS_API_KEY, - model="eleven_multilingual_v2", - voice="Matilda", - ) + driver = agent.config.text_to_speech_driver + if isinstance(driver, DummyTextToSpeechDriver): + driver = ElevenLabsTextToSpeechDriver( + api_key=ELEVEN_LABS_API_KEY, + model="eleven_multilingual_v2", + voice="Matilda", + ) prompt_text = self.get_prompt_text(STRING, input_string) task = TextToSpeechTask( prompt_text, text_to_speech_engine=TextToSpeechEngine(text_to_speech_driver=driver), ) - agent.add_task(task) - result = agent.run() + pipeline = Pipeline() + pipeline.add_task(task) + result = pipeline.run() # Check for output artifact artifact = result.output_task.output