From 06ebbc2c2936ad58b1674f69409afc162ba81168 Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Sat, 16 Nov 2024 14:55:41 +1300 Subject: [PATCH] added ability to load a file from disk, or choose text --- js/TextLoaderModuleNodes.js | 34 +++++++- js/gtUINodes.js | 7 +- nodes/rag/gtUITextLoaderRetrievalRagModule.py | 85 +++++++++++++++---- 3 files changed, 106 insertions(+), 20 deletions(-) diff --git a/js/TextLoaderModuleNodes.js b/js/TextLoaderModuleNodes.js index caf23e4..26ccfb4 100644 --- a/js/TextLoaderModuleNodes.js +++ b/js/TextLoaderModuleNodes.js @@ -11,20 +11,48 @@ export function setupTextLoaderModuleNodes(nodeType, nodeData, app) { const me = onNodeCreated?.apply(this); const widget_loader_type = this.widgets.find(w => w.name === 'loader'); const widget_text = this.widgets.find(w=> w.name === 'text'); + const file_path = this.widgets.find(w=> w.name === 'file_path'); const widget_url = this.widgets.find(w=> w.name === 'url'); + const upload_button = this.widgets.find(w=> w.name === 'choose text to upload'); + const input_source = this.widgets.find(w=> w.name === 'input_source'); - // Hide both widgets + // Hide widgets widget_loader_type.callback = async() => { hideWidget(this, widget_url); - + hideWidget(this, input_source); switch (widget_loader_type.value) { case "WebLoader": showWidget(widget_url); break; + case "TextLoader": + showWidget(input_source); + break; + case "CsvLoader": + showWidget(input_source); + break; } } + input_source.callback = async() => { + hideWidget(this, file_path); + hideWidget(this, upload_button); + console.log(input_source.value); + if (input_source.value == "File Path") { + showWidget(file_path); + showWidget(upload_button); + } + else { + hideWidget(this, file_path); + hideWidget(this, upload_button); + } - setTimeout(() => { widget_loader_type.callback() }, 5); + } + + setTimeout(() => { + input_source.callback(); + }, 5); + setTimeout(() => { + widget_loader_type.callback(); + }, 5); return me; // setupMessageStyle(this.message); }; diff --git a/js/gtUINodes.js b/js/gtUINodes.js index dfb363b..4a055a1 100644 --- a/js/gtUINodes.js +++ b/js/gtUINodes.js @@ -55,8 +55,7 @@ app.registerExtension({ setupDisplayNodes(nodeType, nodeData, app); setupCombineNodes(nodeType, nodeData, app); setupExtractionNodes(nodeType, nodeData, app); - setupTextLoaderModuleNodes(nodeType, nodeData, app); - + // Create Audio Node if (nodeData.name === "Griptape Load: Audio") { @@ -66,5 +65,9 @@ app.registerExtension({ if (nodeData.name === "Griptape Load: Text") { gtUIAddUploadWidget(nodeType, nodeData, "text", "text") } + if (nodeData.name === "Griptape RAG Retrieve: Text Loader Module") { + gtUIAddUploadWidget(nodeType, nodeData, "file_path", "text") + } + setupTextLoaderModuleNodes(nodeType, nodeData, app); } }); diff --git a/nodes/rag/gtUITextLoaderRetrievalRagModule.py b/nodes/rag/gtUITextLoaderRetrievalRagModule.py index 9351d2e..1fccb61 100644 --- a/nodes/rag/gtUITextLoaderRetrievalRagModule.py +++ b/nodes/rag/gtUITextLoaderRetrievalRagModule.py @@ -1,11 +1,14 @@ +import os import tempfile +import folder_paths from griptape.engines.rag.modules import TextLoaderRetrievalRagModule from griptape.loaders import CsvLoader, TextLoader, WebLoader from .gtUIBaseRetrievalRagModule import gtUIBaseRetrievalRagModule -loaders = ["TextLoader", "WebLoader", "CsvLoader"] +loaders = ["TextLoader", "CsvLoader", "WebLoader"] +input_source_options = ["Text Input", "File Path"] class gtUITextLoaderRetrievalRagModule(gtUIBaseRetrievalRagModule): @@ -13,16 +16,46 @@ class gtUITextLoaderRetrievalRagModule(gtUIBaseRetrievalRagModule): Griptape Text Loader Retrieval Rag Module. Used for the Retrieval Stage of the RAG Engine. """ + SUPPORTED_FORMATS = ( + ".data", + ".env", + ".info", + ".json", + ".log", + ".text", + ".txt", + ".yaml", + ".yml", + ".csv", + ".tsv", + ".md", + ) + def __init__(self): pass @classmethod def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [ + f + for f in os.listdir(input_dir) + if ( + os.path.isfile(os.path.join(input_dir, f)) + and f.endswith(gtUITextLoaderRetrievalRagModule.SUPPORTED_FORMATS) + ) + ] inputs = super().INPUT_TYPES() inputs["required"].update( { - "loader": (loaders, {"default": "TextLoader"}), + "loader": ( + loaders, + { + "default": "TextInput", + "tooltip": "The type of text to load. If TextLoader or CsvLoader, it won't use the Text Input Port.", + }, + ), } ) inputs["optional"].update( @@ -35,6 +68,22 @@ class gtUITextLoaderRetrievalRagModule(gtUIBaseRetrievalRagModule): "tooltip": "Text to be loaded.", }, ), + "input_source": ( + input_source_options, + { + "default": input_source_options[0], + "tooltip": "Use a specified file path or the text input port.", + "label_on": "Use file path", + "label_off": "Use incoming text", + }, + ), + "file_path": ( + sorted(files), + { + "text_upload": True, + "tooltip": "File to be loaded.", + }, + ), "url": ( "STRING", {"tooltip": "URL to be loaded.", "default": "https://griptape.ai"}, @@ -50,25 +99,31 @@ class gtUITextLoaderRetrievalRagModule(gtUIBaseRetrievalRagModule): text = kwargs.get("text", None) url = kwargs.get("url", "https://griptape.ai") loader = kwargs.get("loader", "TextLoader") - + input_source = kwargs.get("input_source", input_source_options[0]) params = {} params["query_params"] = self.get_query_params(kwargs) params["vector_store_driver"] = vector_store_driver - if loader == "TextLoader": - # write the text to a temnporary file - with tempfile.NamedTemporaryFile(delete=False) as temp_file: - temp_file.write(text.encode()) - temp_filename = temp_file.name - temp_file.close() - params["loader"] = TextLoader() - params["source"] = temp_filename - params["loader"] = TextLoader() if loader == "WebLoader": params["loader"] = WebLoader() params["source"] = url - if loader == "CsvLoader": - params["loader"] = CsvLoader() - params["source"] = text + else: + if loader == "CsvLoader": + params["loader"] = CsvLoader() + elif loader == "TextLoader": + params["loader"] = TextLoader() + + # Get the input source + if input_source == input_source_options[1]: + params["source"] = folder_paths.get_annotated_filepath( + kwargs.get("file_path") + ) + else: + # write the text to a temporary file + with tempfile.NamedTemporaryFile(delete=False) as temp_file: + temp_file.write(text.encode()) + temp_filename = temp_file.name + temp_file.close() + params["source"] = temp_filename module = TextLoaderRetrievalRagModule(**params) return ([module],)