From a4daa645ad94fb1cfdb2cdf7afb3cf29a93fa170 Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Sat, 13 Jul 2024 12:14:48 +1200 Subject: [PATCH] separated out config nodes --- __init__.py | 25 +- nodes/agent/base_agent.py | 7 + nodes/config/amazon_bedrock_config.py | 73 ++++ nodes/config/anthropic_config.py | 74 +++++ nodes/{ => config}/base_config.py | 0 nodes/config/env_config.py | 37 +++ nodes/config/google_config.py | 59 ++++ nodes/config/lmstudio_config.py | 67 ++++ nodes/config/ollama_config.py | 60 ++++ nodes/config/openai_compatible_config.py | 76 +++++ nodes/config/openai_config.py | 75 +++++ nodes/config_nodes.py | 402 ----------------------- 12 files changed, 544 insertions(+), 411 deletions(-) create mode 100644 nodes/config/amazon_bedrock_config.py create mode 100644 nodes/config/anthropic_config.py rename nodes/{ => config}/base_config.py (100%) create mode 100644 nodes/config/env_config.py create mode 100644 nodes/config/google_config.py create mode 100644 nodes/config/lmstudio_config.py create mode 100644 nodes/config/ollama_config.py create mode 100644 nodes/config/openai_compatible_config.py create mode 100644 nodes/config/openai_config.py delete mode 100644 nodes/config_nodes.py diff --git a/__init__.py b/__init__.py index 2fe0a3c..09e59c3 100644 --- a/__init__.py +++ b/__init__.py @@ -20,15 +20,21 @@ 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 -from .nodes.config_nodes import ( - gtUIAmazonBedrockStructureConfig, - gtUIAnthropicStructureConfig, - # gtUIEnv, - gtUIGoogleStructureConfig, - gtUILMStudioStructureConfig, - gtUIOllamaStructureConfig, - gtUIOpenAiStructureConfig, -) +from .nodes.config.amazon_bedrock_config import gtUIAmazonBedrockStructureConfig +from .nodes.config.anthropic_config import gtUIAnthropicStructureConfig +from .nodes.config.google_config import gtUIGoogleStructureConfig +from .nodes.config.lmstudio_config import gtUILMStudioStructureConfig +from .nodes.config.ollama_config import gtUIOllamaStructureConfig +from .nodes.config.openai_compatible_config import gtUIOpenAiCompatibleConfig +from .nodes.config.openai_config import gtUIOpenAiStructureConfig + +# from .nodes.config.config_nodes import ( +# # gtUIEnv, +# # gtUIGoogleStructureConfig, +# # gtUILMStudioStructureConfig, +# # gtUIOllamaStructureConfig, +# # gtUIOpenAiStructureConfig, +# ) from .nodes.display_nodes import ( # gtUIOutputArtifactNode, gtUIOutputDataNode, @@ -117,6 +123,7 @@ NODE_CLASS_MAPPINGS = { "Griptape Set: Default Agent": gtUISetDefaultAgent, # AGENT CONFIG "Griptape Agent Config: OpenAI": gtUIOpenAiStructureConfig, + "Griptape Agent Config: OpenAI Compatible": gtUIOpenAiCompatibleConfig, "Griptape Agent Config: Amazon Bedrock": gtUIAmazonBedrockStructureConfig, "Griptape Agent Config: Google": gtUIGoogleStructureConfig, "Griptape Agent Config: Anthropic": gtUIAnthropicStructureConfig, diff --git a/nodes/agent/base_agent.py b/nodes/agent/base_agent.py index 1eb44e4..74903f4 100644 --- a/nodes/agent/base_agent.py +++ b/nodes/agent/base_agent.py @@ -1,5 +1,6 @@ from griptape.tools import TaskMemoryClient +# from server import PromptServer from ...py.griptape_config import get_config from .agent import gtComfyAgent @@ -180,6 +181,12 @@ class BaseAgent: else: prompt_text = STRING + "\n\n" + input_string + # # Start to think about sending update messages + # PromptServer.instance.send_sync( + # "comfy.gtUI.textmessage", + # {"message": f"Created agent with prompt: {prompt_text}"}, + # ) + result = self.agent.run(prompt_text) output_string = result.output_task.output.value return ( diff --git a/nodes/config/amazon_bedrock_config.py b/nodes/config/amazon_bedrock_config.py new file mode 100644 index 0000000..fe88ff8 --- /dev/null +++ b/nodes/config/amazon_bedrock_config.py @@ -0,0 +1,73 @@ +from griptape.config import ( + AmazonBedrockStructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + AmazonBedrockImageQueryDriver, + AmazonBedrockPromptDriver, + BedrockClaudeImageQueryModelDriver, +) + +from .base_config import gtUIBaseConfig + +amazonBedrockPromptModels = [ + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-opus-20240229-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", + "anthropic.claude-3-haiku-20240307-v1:0", + "amazon.titan-text-premier-v1:0", + "amazon.titan-text-express-v1", + "amazon.titan-text-lite-v1", +] +amazonBedrockImageQueryModels = [ + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-opus-20240229-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", + "anthropic.claude-3-haiku-20240307-v1:0", +] + + +class gtUIAmazonBedrockStructureConfig(gtUIBaseConfig): + """ + The Griptape Amazon Bedrock Structure Config + """ + + DESCRIPTION = "Amazon Bedrock Prompt Driver." + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update( + { + "prompt_model": ( + amazonBedrockPromptModels, + {"default": amazonBedrockPromptModels[0]}, + ), + "image_query_model": ( + amazonBedrockImageQueryModels, + {"default": amazonBedrockImageQueryModels[0]}, + ), + }, + ) + return inputs + + 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 + ) + custom_config.image_query_driver = AmazonBedrockImageQueryDriver( + image_query_model_driver=BedrockClaudeImageQueryModelDriver(), + model=image_query_model, + ) + if image_generation_driver: + custom_config.image_generation_driver = image_generation_driver + return (custom_config,) diff --git a/nodes/config/anthropic_config.py b/nodes/config/anthropic_config.py new file mode 100644 index 0000000..231527a --- /dev/null +++ b/nodes/config/anthropic_config.py @@ -0,0 +1,74 @@ +from griptape.config import ( + AnthropicStructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + AnthropicImageQueryDriver, + AnthropicPromptDriver, + DummyImageGenerationDriver, +) + +from .base_config import gtUIBaseConfig + +anthropicPromptModels = [ + "claude-3-5-sonnet-20240620", + "claude-3-opus-20240229", + "claude-3-sonnet-20240229", + "claude-3-haiku-20240307", +] +anthropicImageQueryModels = [ + "claude-3-5-sonnet-20240620", + "claude-3-opus-20240229", + "claude-3-sonnet-20240229", + "claude-3-haiku-20240307", +] +voyageAiEmbeddingModels = [ + "voyage-large-2", +] + + +class gtUIAnthropicStructureConfig(gtUIBaseConfig): + """ + The Griptape Anthropic Structure Config + """ + + DESCRIPTION = ( + "Anthropic Structure Config. Use Anthropic's models for prompt and image query." + ) + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update( + { + "prompt_model": ( + anthropicPromptModels, + {"default": anthropicPromptModels[0]}, + ), + "image_query_model": ( + anthropicImageQueryModels, + {"default": anthropicImageQueryModels[0]}, + ), + }, + ) + return inputs + + 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 + ) + custom_config.image_query_driver = AnthropicImageQueryDriver( + model=image_query_model + ) + custom_config.image_generation_driver = image_generation_driver + + return (custom_config,) diff --git a/nodes/base_config.py b/nodes/config/base_config.py similarity index 100% rename from nodes/base_config.py rename to nodes/config/base_config.py diff --git a/nodes/config/env_config.py b/nodes/config/env_config.py new file mode 100644 index 0000000..4028858 --- /dev/null +++ b/nodes/config/env_config.py @@ -0,0 +1,37 @@ +import os + +# StructureGlobalDriversConfig, + + +class gtUIEnv: + """ + The Griptape Environment Config + Setting environment variables + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": {}, + "optional": { + "Environment Vars": ("STRING", {"default": "ENV=", "multiline": True}) + }, + } + + FUNCTION = "run" + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("ENVIRS",) + OUTPUT_NODE = True + + CATEGORY = "Griptape" + + def run(self, **kwargs): + envirs = kwargs.get("Environment Vars", "") + environment_vars = [] + for envir in envirs.split("\n"): + if "=" in envir: + key, value = envir.split("=", 1) + if key and value: + os.environ[key] = value + environment_vars.append(f"{key}={value}") + return (environment_vars,) diff --git a/nodes/config/google_config.py b/nodes/config/google_config.py new file mode 100644 index 0000000..b64bb7b --- /dev/null +++ b/nodes/config/google_config.py @@ -0,0 +1,59 @@ +from griptape.config import ( + GoogleStructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + DummyImageGenerationDriver, + GooglePromptDriver, +) + +from .base_config import gtUIBaseConfig + +google_models = [ + "gemini-1.5-pro", + "gemini-1.5-flash", + "gemini-1.0-pro", +] + + +class gtUIGoogleStructureConfig(gtUIBaseConfig): + """ + The Griptape Google Structure Config + """ + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + + inputs["required"].update( + { + "prompt_model": ( + google_models, + {"default": google_models[0]}, + ), + }, + ) + return inputs + + DESCRIPTION = ( + "Google Structure Config. Use Google's models for prompt and image query." + ) + + def create( + self, + temperature, + seed, + prompt_model, + image_generation_driver=DummyImageGenerationDriver(), + ): + # custom_config = GoogleStructureConfig() + + custom_config = GoogleStructureConfig( + prompt_driver=GooglePromptDriver( + model=prompt_model, temperature=temperature + ), + image_generation_driver=image_generation_driver, + ) + + return (custom_config,) diff --git a/nodes/config/lmstudio_config.py b/nodes/config/lmstudio_config.py new file mode 100644 index 0000000..dd3a6b6 --- /dev/null +++ b/nodes/config/lmstudio_config.py @@ -0,0 +1,67 @@ +from griptape.config import ( + StructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + DummyImageGenerationDriver, + OpenAiChatPromptDriver, +) + +from ..utilities import get_lmstudio_models +from .base_config import gtUIBaseConfig + +lmstudio_models = get_lmstudio_models(port="1234") +lmstudio_models.append("") +lmstudio_port = "1234" +lmstudio_base_url = "http://127.0.0.1" + + +class gtUILMStudioStructureConfig(gtUIBaseConfig): + """ + The Griptape LM Studio Structure Config + """ + + DESCRIPTION = ( + "LM Studio Prompt Driver. LMStudio is available at https://lmstudio.ai " + ) + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update( + { + "prompt_model": ( + [], + {"default": ""}, + ), + # "prompt_model": ("STRING", {"default": ""}), + "base_url": ("STRING", {"default": lmstudio_base_url}), + "port": ( + "STRING", + {"default": lmstudio_port}, + ), + }, + ) + return inputs + + def create( + self, + prompt_model, + base_url, + port, + temperature, + seed, + image_generation_driver=DummyImageGenerationDriver(), + ): + custom_config = StructureConfig( + prompt_driver=OpenAiChatPromptDriver( + model=prompt_model, + base_url=f"{base_url}:{port}/v1", + api_key="lm_studio", + temperature=temperature, + ), + image_generation_driver=image_generation_driver, + ) + + return (custom_config,) diff --git a/nodes/config/ollama_config.py b/nodes/config/ollama_config.py new file mode 100644 index 0000000..8ed3307 --- /dev/null +++ b/nodes/config/ollama_config.py @@ -0,0 +1,60 @@ +from griptape.config import ( + StructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + DummyImageGenerationDriver, + OllamaPromptDriver, +) + +from ..utilities import get_ollama_models +from .base_config import gtUIBaseConfig + +ollama_models = get_ollama_models() +ollama_models.append("") +ollama_port = "11434" +ollama_base_url = "http://127.0.0.1" + + +class gtUIOllamaStructureConfig(gtUIBaseConfig): + """ + The Griptape Ollama Structure Config + """ + + DESCRIPTION = "Ollama Prompt Driver. Use local models with Ollama. Available at https://ollama.com" + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update( + { + "prompt_model": ( + [""], + {"default": ""}, + ), + "base_url": ("STRING", {"default": ollama_base_url}), + "port": ("STRING", {"default": ollama_port}), + }, + ) + return inputs + + 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}, + host=f"{base_url}:{port}", + ), + image_generation_driver=image_generation_driver, + ) + + return (custom_config,) diff --git a/nodes/config/openai_compatible_config.py b/nodes/config/openai_compatible_config.py new file mode 100644 index 0000000..949d066 --- /dev/null +++ b/nodes/config/openai_compatible_config.py @@ -0,0 +1,76 @@ +from griptape.config import ( + StructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + OpenAiChatPromptDriver, + OpenAiImageGenerationDriver, + OpenAiImageQueryDriver, + OpenAiTextToSpeechDriver, +) + +from .base_config import gtUIBaseConfig + + +class gtUIOpenAiCompatibleConfig(gtUIBaseConfig): + """ + Create an OpenAI Compatible Structure Config + """ + + DESCRIPTION = "OpenAI Compatible Structure Config." + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + del inputs["optional"]["image_generation_driver"] + inputs["optional"].update( + { + "prompt_model": ("STRING", {"default": "gpt-4o"}), + "image_generation_model": ("STRING", {"default": "dall-e-3"}), + "image_query_model": ("STRING", {"default": "gpt-4o"}), + "text_to_speech_model": ("STRING", {"default": "tts-1"}), + "prompt_base_url": ("STRING", {"default": "https://url/v1"}), + "api_key": ("STRING", {"default": ""}), + } + ) + return inputs + + def create(self, **kwargs): + prompt_model = kwargs.get("prompt_model", None) + image_generation_model = kwargs.get("image_generation_model", None) + image_query_model = kwargs.get("image_query_model", None) + text_to_speech_model = kwargs.get("text_to_speech_model", None) + base_url = kwargs.get("prompt_base_url", None) + api_key = kwargs.get("api_key", None) + + configs = {} + if prompt_model and base_url and api_key: + configs["prompt_driver"] = OpenAiChatPromptDriver( + model=prompt_model, + base_url=base_url, + api_key=api_key, + ) + if image_generation_model and base_url and api_key: + configs["image_generation_driver"] = OpenAiImageGenerationDriver( + model=image_generation_model, + base_url=base_url, + api_key=api_key, + ) + + if image_query_model and base_url and api_key: + configs["image_query_driver"] = OpenAiImageQueryDriver( + model=image_query_model, + base_url=base_url, + api_key=api_key, + ) + + if text_to_speech_model and base_url and api_key: + configs["text_to_speech_driver"] = OpenAiTextToSpeechDriver( + model=text_to_speech_model, + base_url=base_url, + api_key=api_key, + ) + custom_config = StructureConfig(**configs) + + return (custom_config,) diff --git a/nodes/config/openai_config.py b/nodes/config/openai_config.py new file mode 100644 index 0000000..a3fb6c9 --- /dev/null +++ b/nodes/config/openai_config.py @@ -0,0 +1,75 @@ +from griptape.config import ( + OpenAiStructureConfig, +) + +# StructureGlobalDriversConfig, +from griptape.drivers import ( + OpenAiChatPromptDriver, + OpenAiEmbeddingDriver, + OpenAiImageGenerationDriver, + OpenAiImageQueryDriver, +) + +from ...py.griptape_config import get_config +from .base_config import gtUIBaseConfig + + +class gtUIOpenAiStructureConfig(gtUIBaseConfig): + """ + The Griptape OpenAI Structure Config + """ + + DESCRIPTION = "OpenAI Structure Config. Use OpenAI's models for prompt, embedding, image generation, and image query." + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + inputs["required"].update( + { + "prompt_model": ( + ["gpt-4o", "gpt-4", "gpt-3.5-turbo"], + {"default": "gpt-4o"}, + ), + "image_query_model": ( + ["gpt-4o", "gpt-4-vision-preview"], + {"default": "gpt-4o"}, + ), + } + ) + return inputs + + def create( + self, + prompt_model, + image_query_model, + temperature, + seed, + image_generation_driver=None, + ): + OPENAI_API_KEY = get_config("env.OPENAI_API_KEY") + 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="dall-e-3", + ) + + 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,) diff --git a/nodes/config_nodes.py b/nodes/config_nodes.py deleted file mode 100644 index 0be760c..0000000 --- a/nodes/config_nodes.py +++ /dev/null @@ -1,402 +0,0 @@ -import os - -from griptape.config import ( - AmazonBedrockStructureConfig, - AnthropicStructureConfig, - GoogleStructureConfig, - OpenAiStructureConfig, - StructureConfig, -) - -# StructureGlobalDriversConfig, -from griptape.drivers import ( - AmazonBedrockImageQueryDriver, - AmazonBedrockPromptDriver, - AnthropicImageQueryDriver, - AnthropicPromptDriver, - BedrockClaudeImageQueryModelDriver, - DummyImageGenerationDriver, - GooglePromptDriver, - OllamaPromptDriver, - OpenAiChatPromptDriver, - OpenAiEmbeddingDriver, - OpenAiImageGenerationDriver, - OpenAiImageQueryDriver, -) - -from ..py.griptape_config import get_config -from .base_config import gtUIBaseConfig -from .utilities import get_lmstudio_models, get_ollama_models - - -class gtUIEnv: - """ - The Griptape Environment Config - Setting environment variables - """ - - @classmethod - def INPUT_TYPES(s): - return { - "required": {}, - "optional": { - "Environment Vars": ("STRING", {"default": "ENV=", "multiline": True}) - }, - } - - FUNCTION = "run" - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("ENVIRS",) - OUTPUT_NODE = True - - CATEGORY = "Griptape" - - def run(self, **kwargs): - envirs = kwargs.get("Environment Vars", "") - environment_vars = [] - for envir in envirs.split("\n"): - if "=" in envir: - key, value = envir.split("=", 1) - if key and value: - os.environ[key] = value - environment_vars.append(f"{key}={value}") - return (environment_vars,) - - -ollama_models = get_ollama_models() -ollama_models.append("") -ollama_port = "11434" -ollama_base_url = "http://127.0.0.1" - -lmstudio_models = get_lmstudio_models(port="1234") -lmstudio_models.append("") -lmstudio_port = "1234" -lmstudio_base_url = "http://127.0.0.1" - - -class gtUILMStudioStructureConfig(gtUIBaseConfig): - """ - The Griptape LM Studio Structure Config - """ - - DESCRIPTION = ( - "LM Studio Prompt Driver. LMStudio is available at https://lmstudio.ai " - ) - - @classmethod - def INPUT_TYPES(s): - inputs = super().INPUT_TYPES() - inputs["required"].update( - { - "prompt_model": ( - [], - {"default": ""}, - ), - # "prompt_model": ("STRING", {"default": ""}), - "base_url": ("STRING", {"default": lmstudio_base_url}), - "port": ( - "STRING", - {"default": lmstudio_port}, - ), - }, - ) - return inputs - - def create( - self, - prompt_model, - base_url, - port, - temperature, - seed, - image_generation_driver=DummyImageGenerationDriver(), - ): - custom_config = StructureConfig( - prompt_driver=OpenAiChatPromptDriver( - model=prompt_model, - base_url=f"{base_url}:{port}/v1", - api_key="lm_studio", - temperature=temperature, - ), - image_generation_driver=image_generation_driver, - ) - - return (custom_config,) - - -class gtUIOllamaStructureConfig(gtUIBaseConfig): - """ - The Griptape Ollama Structure Config - """ - - DESCRIPTION = "Ollama Prompt Driver. Use local models with Ollama. Available at https://ollama.com" - - @classmethod - def INPUT_TYPES(s): - inputs = super().INPUT_TYPES() - inputs["required"].update( - { - "prompt_model": ( - [""], - {"default": ""}, - ), - "base_url": ("STRING", {"default": ollama_base_url}), - "port": ("STRING", {"default": ollama_port}), - }, - ) - return inputs - - 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}, - host=f"{base_url}:{port}", - ), - image_generation_driver=image_generation_driver, - ) - - return (custom_config,) - - -amazonBedrockPromptModels = [ - "anthropic.claude-3-5-sonnet-20240620-v1:0", - "anthropic.claude-3-opus-20240229-v1:0", - "anthropic.claude-3-sonnet-20240229-v1:0", - "anthropic.claude-3-haiku-20240307-v1:0", - "amazon.titan-text-premier-v1:0", - "amazon.titan-text-express-v1", - "amazon.titan-text-lite-v1", -] -amazonBedrockImageQueryModels = [ - "anthropic.claude-3-5-sonnet-20240620-v1:0", - "anthropic.claude-3-opus-20240229-v1:0", - "anthropic.claude-3-sonnet-20240229-v1:0", - "anthropic.claude-3-haiku-20240307-v1:0", -] - - -class gtUIAmazonBedrockStructureConfig(gtUIBaseConfig): - """ - The Griptape Amazon Bedrock Structure Config - """ - - DESCRIPTION = "Amazon Bedrock Prompt Driver." - - @classmethod - def INPUT_TYPES(s): - inputs = super().INPUT_TYPES() - inputs["required"].update( - { - "prompt_model": ( - amazonBedrockPromptModels, - {"default": amazonBedrockPromptModels[0]}, - ), - "image_query_model": ( - amazonBedrockImageQueryModels, - {"default": amazonBedrockImageQueryModels[0]}, - ), - }, - ) - return inputs - - 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 - ) - custom_config.image_query_driver = AmazonBedrockImageQueryDriver( - image_query_model_driver=BedrockClaudeImageQueryModelDriver(), - model=image_query_model, - ) - if image_generation_driver: - custom_config.image_generation_driver = image_generation_driver - return (custom_config,) - - -google_models = [ - "gemini-1.5-pro", - "gemini-1.5-flash", - "gemini-1.0-pro", -] - - -class gtUIGoogleStructureConfig(gtUIBaseConfig): - """ - The Griptape Google Structure Config - """ - - @classmethod - def INPUT_TYPES(s): - inputs = super().INPUT_TYPES() - - inputs["required"].update( - { - "prompt_model": ( - google_models, - {"default": google_models[0]}, - ), - }, - ) - return inputs - - DESCRIPTION = ( - "Google Structure Config. Use Google's models for prompt and image query." - ) - - def create( - self, - temperature, - seed, - prompt_model, - image_generation_driver=DummyImageGenerationDriver(), - ): - # custom_config = GoogleStructureConfig() - - custom_config = GoogleStructureConfig( - prompt_driver=GooglePromptDriver( - model=prompt_model, temperature=temperature - ), - image_generation_driver=image_generation_driver, - ) - - return (custom_config,) - - -anthropicPromptModels = [ - "claude-3-5-sonnet-20240620", - "claude-3-opus-20240229", - "claude-3-sonnet-20240229", - "claude-3-haiku-20240307", -] -anthropicImageQueryModels = [ - "claude-3-5-sonnet-20240620", - "claude-3-opus-20240229", - "claude-3-sonnet-20240229", - "claude-3-haiku-20240307", -] -voyageAiEmbeddingModels = [ - "voyage-large-2", -] - - -class gtUIAnthropicStructureConfig(gtUIBaseConfig): - """ - The Griptape Anthropic Structure Config - """ - - DESCRIPTION = ( - "Anthropic Structure Config. Use Anthropic's models for prompt and image query." - ) - - @classmethod - def INPUT_TYPES(s): - inputs = super().INPUT_TYPES() - inputs["required"].update( - { - "prompt_model": ( - anthropicPromptModels, - {"default": anthropicPromptModels[0]}, - ), - "image_query_model": ( - anthropicImageQueryModels, - {"default": anthropicImageQueryModels[0]}, - ), - }, - ) - return inputs - - 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 - ) - custom_config.image_query_driver = AnthropicImageQueryDriver( - model=image_query_model - ) - custom_config.image_generation_driver = image_generation_driver - - return (custom_config,) - - -class gtUIOpenAiStructureConfig(gtUIBaseConfig): - """ - The Griptape OpenAI Structure Config - """ - - DESCRIPTION = "OpenAI Structure Config. Use OpenAI's models for prompt, embedding, image generation, and image query." - - @classmethod - def INPUT_TYPES(s): - inputs = super().INPUT_TYPES() - inputs["required"].update( - { - "prompt_model": ( - ["gpt-4o", "gpt-4", "gpt-3.5-turbo"], - {"default": "gpt-4o"}, - ), - "image_query_model": ( - ["gpt-4o", "gpt-4-vision-preview"], - {"default": "gpt-4o"}, - ), - } - ) - return inputs - - def create( - self, - prompt_model, - image_query_model, - temperature, - seed, - image_generation_driver=None, - ): - OPENAI_API_KEY = get_config("env.OPENAI_API_KEY") - 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="dall-e-3", - ) - - 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,)