diff --git a/nodes/config_nodes.py b/nodes/config_nodes.py index 26116e2..3d8d6c7 100644 --- a/nodes/config_nodes.py +++ b/nodes/config_nodes.py @@ -10,6 +10,7 @@ from griptape.config import ( # StructureGlobalDriversConfig, from griptape.drivers import ( AnthropicImageQueryDriver, + OllamaPromptDriver, OpenAiChatPromptDriver, OpenAiEmbeddingDriver, OpenAiImageGenerationDriver, @@ -54,6 +55,34 @@ class gtUIEnv: return (environment_vars,) +class gtUIOllamaStructureConfig(gtUIBaseConfig): + """ + The Griptape Ollama Structure Config + """ + + @classmethod + def INPUT_TYPES(s): + return { + "optional": {}, + "required": { + "prompt_model": ( + "STRING", + {"default": "llama3"}, + ), + }, + } + + def create( + self, + prompt_model, + ): + custom_config = StructureConfig( + prompt_driver=OllamaPromptDriver(model=prompt_model), + ) + + return (custom_config,) + + class gtUIAmazonBedrockStructureConfig(gtUIBaseConfig): """ The Griptape Amazon Bedrock Structure Config