adding RAG nodes to ComfyUI

__init__.py

@@ -22,6 +22,7 @@ from .nodes.agent.RunAgent import RunAgent
 
 # COMBINE
 from .nodes.combine.gtUIMergeInputs import gtUIMergeInputs
+from .nodes.combine.gtUIModuleList import gtUIModuleList
 from .nodes.combine.MergeTexts import MergeTexts
 from .nodes.combine.RulesList import RulesList
 from .nodes.combine.ToolList import ToolList
@@ -98,6 +99,7 @@ from .nodes.drivers.gtUIAzureOpenAiImageGenerationDriver import (
 # - Cohere
 from .nodes.drivers.gtUICohereEmbeddingDriver import gtUICohereEmbeddingDriver
 from .nodes.drivers.gtUICoherePromptDriver import gtUICoherePromptDriver
+from .nodes.drivers.gtUICohereRerankDriver import gtUICohereRerankDriver
 
 # - DuckDuckGo
 from .nodes.drivers.gtUIDuckDuckGoWebSearchDriver import gtUIDuckDuckGoWebSearchDriver
@@ -182,6 +184,16 @@ from .nodes.drivers.gtUIVoyageAiEmbeddingDriver import gtUIVoyageAiEmbeddingDriv
 from .nodes.loaders.gtUIFetchImage import gtUIFetchImage
 from .nodes.loaders.gtUILoadAudio import gtUILoadAudio
 from .nodes.loaders.gtUILoadText import gtUILoadText
+from .nodes.rag.gtUIPromptResponseRagModule import gtUIPromptResponseRagModule
+
+# RAG
+from .nodes.rag.gtUIRagEngine import gtUIRagEngine
+from .nodes.rag.gtUITextChunksRerankRagModule import gtUITextChunksRerankRagModule
+from .nodes.rag.gtUITextLoaderRetrievalRagModule import gtUITextLoaderRetrievalRagModule
+from .nodes.rag.gtUITranslateQueryRagModule import gtUITranslateQueryRagModule
+from .nodes.rag.gtUIVectorStoreRetrievalRagModule import (
+    gtUIVectorStoreRetrievalRagModule,
+)
 
 # RULES
 from .nodes.rules.gtUIRule import gtUIRule
@@ -224,6 +236,7 @@ from .nodes.tools.gtUIFileManager import gtUIFileManager
 from .nodes.tools.gtUIKnowledgeBaseTool import gtUIKnowledgeBaseTool
 from .nodes.tools.gtUIPromptSummaryTool import gtUIPromptSummaryTool
 from .nodes.tools.gtUIQueryTool import gtUIQueryTool
+from .nodes.tools.gtUIRagTool import gtUIRagTool
 from .nodes.tools.gtUITextToSpeechClient import gtUITextToSpeechClient
 from .nodes.tools.gtUIVectorStoreClient import gtUIVectorStoreClient
 from .nodes.tools.gtUIWebScraper import gtUIWebScraper
@@ -309,6 +322,8 @@ NODE_CLASS_MAPPINGS = {
     "Griptape Embedding Driver: OpenAI": gtUIOpenAiEmbeddingDriver,
     "Griptape Embedding Driver: OpenAI Compatible": gtUIOpenAiCompatibleEmbeddingDriver,
     "Griptape Embedding Driver: Voyage AI": gtUIVoyageAiEmbeddingDriver,
+    # RERANK DRIVER
+    "Griptape Rerank Driver: Cohere": gtUICohereRerankDriver,
     # VECTOR STORE DRIVERS
     "Griptape Vector Store Driver: Amazon OpenSearch": gtUIAmazonOpenSearchVectorStoreDriver,
     "Griptape Vector Store Driver: Azure MongoDB": gtUIAzureMongoDbVectorStoreDriver,
@@ -379,6 +394,15 @@ NODE_CLASS_MAPPINGS = {
     "Griptape Save: Text": gtUISaveText,
     "Griptape Vector Store: Add Text": gtUIVectorStoreUpsertTextTask,
     "Griptape Vector Store: Query": gtUIVectorStoreQueryTask,
+    # RAG
+    "Griptape RAG: Engine": gtUIRagEngine,
+    "Griptape RAG: Tool": gtUIRagTool,
+    "Griptape Combine: RAG Module List": gtUIModuleList,
+    "Griptape RAG Query: Translate Module": gtUITranslateQueryRagModule,
+    "Griptape RAG Retrieve: Text Loader Module": gtUITextLoaderRetrievalRagModule,
+    "Griptape RAG Retrieve: Vector Store Module": gtUIVectorStoreRetrievalRagModule,
+    "Griptape RAG Rerank: Text Chunks Module": gtUITextChunksRerankRagModule,
+    "Griptape RAG Response: Prompt Module": gtUIPromptResponseRagModule,
     # "Griptape Display: Artifact": gtUIOutputArtifactNode,
     # "Griptape Config: Environment Variables": gtUIEnv,
 }

js/CombineNodes.js

@@ -3,6 +3,7 @@ export function setupCombineNodes(nodeType, nodeData, app) {
         nodeData.name === "Griptape Combine: Merge Inputs" || 
         nodeData.name === "Griptape Combine: Rules List" ||
         nodeData.name === "Griptape Combine: Tool List" ||
+        nodeData.name === "Griptape Combine: RAG Module List" ||
         nodeData.name === "Griptape Create: Pipeline") {
         
         setupCombineNode(nodeType, nodeData, app);
@@ -42,6 +43,7 @@ function getInputName(nodeName) {
         case 'Griptape Combine: Rules List': return "rules_";
         case 'Griptape Combine: Tool List': return "tool_";
         case 'Griptape Create: Pipeline': return "task_";
+        case 'Griptape Combine: RAG Module List': return "module_";
         default: return "input_";
     }
 }

js/TextLoaderModuleNodes.js

@@ -0,0 +1,33 @@
+import { hideWidget, showWidget } from "./utils.js";
+export function setupTextLoaderModuleNodes(nodeType, nodeData, app) {
+    if (nodeData.name === "Griptape RAG Retrieve: Text Loader Module") {
+      setupLoaderAttr(nodeType, nodeData, app);
+    }
+  }
+  
+  function setupLoaderAttr(nodeType, nodeData, app) {
+    const onNodeCreated = nodeType.prototype.onNodeCreated
+    nodeType.prototype.onNodeCreated = function() {
+      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 widget_url = this.widgets.find(w=> w.name === 'url');
+
+      // Hide both widgets
+      widget_loader_type.callback = async() => {
+        hideWidget(this, widget_url);
+
+        switch (widget_loader_type.value) {
+          case "WebLoader":
+            showWidget(widget_url);
+            break;
+        }
+      }
+
+      setTimeout(() => { widget_loader_type.callback() }, 5);
+      return me;
+      // setupMessageStyle(this.message);
+    };
+  
+  }
+  

js/gtUINodes.js

@@ -6,6 +6,7 @@ import { setupNodeColors } from "./ColorNodes.js";
 import { setupDisplayNodes } from "./DisplayNodes.js";
 import { setupCombineNodes } from "./CombineNodes.js";
 import { setupExtractionNodes } from "./ExtractionNodes.js";
+import { setupTextLoaderModuleNodes } from "./TextLoaderModuleNodes.js";
 import { gtUIAddUploadWidget } from "./gtUIUtils.js";
 import {  setupMenuSeparator } from "./gtUIMenuSeparator.js";
 // app.extensionManager.registerSidebarTab({
@@ -59,6 +60,7 @@ app.registerExtension({
     setupDisplayNodes(nodeType, nodeData, app);
     setupCombineNodes(nodeType, nodeData, app);
     setupExtractionNodes(nodeType, nodeData, app);
+    setupTextLoaderModuleNodes(nodeType, nodeData, app);
 
     
     // Create Audio Node

nodes/combine/gtUIModuleList.py

@@ -0,0 +1,38 @@
+class gtUIModuleList:
+    """
+    Griptape Module List
+    """
+
+    DESCRIPTION = "Combine rag modules into a list to be used with RAG Engines."
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "optional": {
+                "module_1": (
+                    "MODULE_LIST",
+                    {
+                        "tooltip": "A module to add to the list. Connect an input to dynamically create more inputs."
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("MODULE_LIST",)
+    RETURN_NAMES = ("MODULE_LIST",)
+    FUNCTION = "create"
+    # OUTPUT_IS_LIST = (True,)
+
+    CATEGORY = "Griptape/RAG"
+
+    def create(self, **kwargs):
+        module_list = []
+
+        modules = [value for value in kwargs.values()]
+        if len(modules) > 0:
+            for module in modules:
+                module_list.append(module[0])
+        return (module_list,)

nodes/drivers/gtUIBaseRerankDriver.py

@@ -0,0 +1,23 @@
+from griptape.drivers import BaseRerankDriver
+
+from .gtUIBaseDriver import gtUIBaseDriver
+
+
+class gtUIBaseRerankDriver(gtUIBaseDriver):
+    DESCRIPTION = "Griptape Rerank Driver"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update()
+        inputs["optional"].update()
+
+        return inputs
+
+    CATEGORY = "Griptape/Agent Drivers/Rerank"
+
+    RETURN_TYPES = ("RERANK_DRIVER",)
+
+    def create(self, **kwargs):
+        driver = BaseRerankDriver()
+        return (driver,)

nodes/drivers/gtUICohereRerankDriver.py

@@ -0,0 +1,50 @@
+import os
+
+from griptape.drivers import CohereRerankDriver
+
+from .gtUIBaseRerankDriver import gtUIBaseRerankDriver
+
+DEFAULT_API_KEY_ENV_VAR = "COHERE_API_KEY"
+models = ["rerank-english-v3.0", "rerank-multilingual-v3.0"]
+
+
+class gtUICohereRerankDriver(gtUIBaseRerankDriver):
+    DESCRIPTION = "Cohere Rerank Driver"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update()
+        inputs["optional"].update(
+            {
+                "models": (
+                    models,
+                    {"default": models[0]},
+                ),
+                "top_n": (
+                    "INT",
+                    {"default": 5},
+                ),
+                "api_key_env_var": (
+                    "STRING",
+                    {"default": DEFAULT_API_KEY_ENV_VAR},
+                ),
+            }
+        )
+
+        return inputs
+
+    def create(self, **kwargs):
+        api_key_env_var = kwargs.get("api_key_env_var", DEFAULT_API_KEY_ENV_VAR)
+        model = kwargs.get("models", models[0])
+        top_n = kwargs.get("top_n", 5)
+        params = {}
+
+        if model:
+            params["model"] = model
+        if api_key_env_var:
+            params["api_key"] = os.getenv(api_key_env_var)
+        params["top_n"] = top_n
+
+        driver = CohereRerankDriver(**params)
+        return (driver,)

nodes/drivers/gtUIPgVectorVectorStoreDriver.py

@@ -65,7 +65,6 @@ class gtUIPgVectorVectorStoreDriver(gtUIBaseVectorStoreDriver):
         else:
             params["embedding_driver"] = self.get_default_embedding_driver()
 
-        print(f"{params=}")
         driver = PgVectorVectorStoreDriver(**params)
         # driver.setup()
         return (driver,)

nodes/rag/gtUIBaseQueryRagModule.py

@@ -0,0 +1,21 @@
+from griptape.engines.rag.modules import BaseQueryRagModule
+
+from .gtUIBaseRagModule import gtUIBaseRagModule
+
+
+class gtUIBaseQueryRagModule(gtUIBaseRagModule):
+    """
+    Griptape Base Query Rag Module
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        return inputs
+
+    def create(self):
+        return ([BaseQueryRagModule()],)

nodes/rag/gtUIBaseRagModule.py

@@ -0,0 +1,52 @@
+from griptape.engines.rag.modules import BaseRagModule
+
+
+class gtUIBaseRagModule:
+    """
+    Griptape Base Rag Module
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {},
+            "optional": {},
+            "hidden": {},
+        }
+
+    RETURN_TYPES = ("MODULE_LIST",)
+    RETURN_NAMES = ("MODULE",)
+    FUNCTION = "create"
+
+    CATEGORY = "Griptape/RAG"
+
+    def ensure_dict(params):
+        if isinstance(params, dict):
+            return params
+        elif isinstance(params, str):
+            # Try to evaluate the string as a dictionary
+            try:
+                return eval(f"dict({params})")
+            except Exception:
+                # If evaluation fails, try to parse it as JSON
+                import json
+
+                try:
+                    return json.loads(params)
+                except json.JSONDecodeError:
+                    # If JSON parsing fails, split by comma and colon
+                    return dict(item.split(":") for item in params.split(","))
+        elif params is None:
+            return {}
+        else:
+            # If it's any other type, try to convert it to a dict
+            try:
+                return dict(params)
+            except Exception:
+                raise ValueError(f"Could not convert {type(params)} to dictionary")
+
+    def create(self):
+        return ([BaseRagModule()],)

nodes/rag/gtUIBaseReponseRagModule.py

@@ -0,0 +1,12 @@
+from griptape.engines.rag.modules import BaseResponseRagModule
+
+from .gtUIBaseRagModule import gtUIBaseRagModule
+
+
+class gtUIBaseResponseRagModule(gtUIBaseRagModule):
+    """
+    Griptape Base Response Rag Module
+    """
+
+    def create(self, **kwargs):
+        return ([BaseResponseRagModule()],)

nodes/rag/gtUIBaseRerankRagModule.py

@@ -0,0 +1,29 @@
+from griptape.engines.rag.modules import BaseRerankRagModule
+
+from .gtUIBaseRagModule import gtUIBaseRagModule
+
+
+class gtUIBaseRerankRagModule(gtUIBaseRagModule):
+    """
+    Griptape Base Rerank Rag Module
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        inputs["optional"].update(
+            {
+                "rerank_driver": ("RERANK_DRIVER", {"default": None}),
+            }
+        )
+        return inputs
+
+    RETURN_TYPES = ("MODULE",)
+    RETURN_NAMES = ("MODULE",)
+
+    def create(self, **kwargs):
+        return (BaseRerankRagModule(),)

nodes/rag/gtUIBaseRetrievalRagModule.py

@@ -0,0 +1,55 @@
+from griptape.drivers import DummyVectorStoreDriver, LocalVectorStoreDriver
+from griptape.engines.rag.modules import BaseRetrievalRagModule
+
+from .gtUIBaseRagModule import gtUIBaseRagModule
+
+
+class gtUIBaseRetrievalRagModule(gtUIBaseRagModule):
+    """
+    Griptape Base Query Rag Module
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        inputs["optional"] = {
+            "vector_store_driver": ("VECTOR_STORE_DRIVER", {"default": None}),
+            "namespace": (
+                "STRING",
+                {"default": "default", "tooltip": "Namespace of the vector store."},
+            ),
+            # "top_n": (
+            #     "INT",
+            #     {"default": 20, "tooltip": "Return the # most similar items."},
+            # ),
+            "count": (
+                "INT",
+                {"default": 5, "tooltip": "Total number of results to return."},
+            ),
+        }
+        return inputs
+
+    def get_query_params(self, kwargs):
+        query_params = {}
+        query_params["namespace"] = kwargs.get("namespace", "default")
+        # query_params["top_n"] = kwargs.get("top_n", 20)
+
+        if kwargs.get("count") != -1:
+            query_params["count"] = kwargs.get("count")
+        return query_params
+
+    def get_vector_store_driver(self, vector_store_driver=None):
+        if not vector_store_driver or isinstance(
+            vector_store_driver, DummyVectorStoreDriver
+        ):
+            return LocalVectorStoreDriver()
+        else:
+            return vector_store_driver
+        return None
+
+    def create(self, **kwargs):
+        return ([BaseRetrievalRagModule()],)

nodes/rag/gtUIPromptResponseRagModule.py

@@ -0,0 +1,43 @@
+from griptape.engines.rag.modules import PromptResponseRagModule
+
+from .gtUIBaseReponseRagModule import gtUIBaseResponseRagModule
+
+
+class gtUIPromptResponseRagModule(gtUIBaseResponseRagModule):
+    """
+    Griptape Prompt Response Rag Module.
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        inputs["optional"] = {
+            "prompt_driver": (
+                "PROMPT_DRIVER",
+                {"tooltip": "Prompt driver to use for the module."},
+            ),
+            "rulesets": (
+                "RULESET",
+                {
+                    "forceInput": True,
+                    "tooltip": "One or more rules to use with RAG.\nUse these to control the RAG response.",
+                },
+            ),
+        }
+        return inputs
+
+    def create(self, **kwargs):
+        params = {}
+        prompt_driver = kwargs.get("prompt_driver", None)
+        rulesets = kwargs.get("rulesets", [])
+        if prompt_driver:
+            params["prompt_driver"] = prompt_driver
+        if len(rulesets) > 0:
+            params["rulesets"] = rulesets
+
+        module = PromptResponseRagModule(**params)
+        return ([module],)

nodes/rag/gtUIRagEngine.py

@@ -0,0 +1,81 @@
+from griptape.engines.rag import RagEngine
+from griptape.engines.rag.stages import (
+    QueryRagStage,
+    ResponseRagStage,
+    RetrievalRagStage,
+)
+
+
+class gtUIRagEngine:
+    DESCRIPTION = "Griptape Rag Tool"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {},
+            "optional": {
+                "query_stage_modules": (
+                    "MODULE_LIST",
+                    {
+                        "forceInput:": True,
+                        "tooltip": "(Optional) Used for modifying user queries. For example, translating the query to another language.",
+                    },
+                ),
+                "retrieval_stage_modules": (
+                    "MODULE_LIST",
+                    {
+                        "forceInput": True,
+                        "tooltip": "Modules Used for retrieving relevant documents and re-ranking results.",
+                    },
+                ),
+                "rerank_module": (
+                    "MODULE",
+                    {
+                        "forceInput": True,
+                        "tooltip": "(Optional) Module used for re-ranking results. ",
+                    },
+                ),
+                "response_stage_modules": (
+                    "MODULE_LIST",
+                    {
+                        "forceInput": True,
+                        "tooltip": "Modules used to append metadata, rulesets, generate responses, or adding footnotes.",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("RAG_ENGINE",)
+    RETURN_NAMES = ("RAG_ENGINE",)
+
+    FUNCTION = "create"
+
+    CATEGORY = "Griptape/RAG"
+
+    def create(self, **kwargs):
+        query_stage_modules = kwargs.get("query_stage_modules", [])
+        retrieval_stage_modules = kwargs.get("retrieval_stage_modules", [])
+        rerank_module = kwargs.get("rerank_module", None)
+        response_stage_modules = kwargs.get("response_stage_modules", [])
+
+        engine_params = {}
+
+        if len(query_stage_modules) > 0:
+            engine_params["query_stage"] = QueryRagStage(
+                query_modules=query_stage_modules
+            )
+        if len(retrieval_stage_modules) > 0:
+            retrieval_stage_params = {}
+            if rerank_module:
+                retrieval_stage_params["rerank_module"] = rerank_module
+            retrieval_stage_params["retrieval_modules"] = retrieval_stage_modules
+            engine_params["retrieval_stage"] = RetrievalRagStage(
+                **retrieval_stage_params
+            )
+        if len(response_stage_modules) > 0:
+            engine_params["response_stage"] = ResponseRagStage(
+                response_modules=response_stage_modules
+            )
+
+        engine = RagEngine(**engine_params)
+        return (engine,)

nodes/rag/gtUITextChunksRerankRagModule.py

@@ -0,0 +1,21 @@
+from griptape.engines.rag.modules import TextChunksRerankRagModule
+
+from .gtUIBaseRerankRagModule import gtUIBaseRerankRagModule
+
+
+class gtUITextChunksRerankRagModule(gtUIBaseRerankRagModule):
+    """
+    Griptape Text Loader Retrieval Rag Module. Used for the Retrieval Stage of the RAG Engine.
+    """
+
+    def __init__(self):
+        pass
+
+    def create(self, **kwargs):
+        params = {}
+        rerank_driver = kwargs.get("rerank_driver", None)
+        if rerank_driver:
+            params["rerank_driver"] = rerank_driver
+
+        module = TextChunksRerankRagModule(**params)
+        return (module,)

nodes/rag/gtUITextLoaderRetrievalRagModule.py

@@ -0,0 +1,66 @@
+from griptape.engines.rag.modules import TextLoaderRetrievalRagModule
+from griptape.loaders import CsvLoader, TextLoader, WebLoader
+
+from .gtUIBaseRetrievalRagModule import gtUIBaseRetrievalRagModule
+
+loaders = ["TextLoader", "WebLoader", "CsvLoader"]
+
+
+class gtUITextLoaderRetrievalRagModule(gtUIBaseRetrievalRagModule):
+    """
+    Griptape Text Loader Retrieval Rag Module. Used for the Retrieval Stage of the RAG Engine.
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        inputs["required"].update(
+            {
+                "loader": (loaders, {"default": "TextLoader"}),
+            }
+        )
+        inputs["optional"].update(
+            {
+                "text": (
+                    "STRING",
+                    {
+                        "forceInput": True,
+                        "dynamicPrompts": True,
+                        "tooltip": "Text to be loaded.",
+                    },
+                ),
+                "url": (
+                    "STRING",
+                    {"tooltip": "URL to be loaded.", "default": "https://griptape.ai"},
+                ),
+            }
+        )
+        return inputs
+
+    def create(self, **kwargs):
+        vector_store_driver = self.get_vector_store_driver(
+            kwargs.get("vector_store_driver", None)
+        )
+        text = kwargs.get("text", None)
+        url = kwargs.get("url", "https://griptape.ai")
+        loader = kwargs.get("loader", "TextLoader")
+
+        params = {}
+        params["query_params"] = self.get_query_params(kwargs)
+        params["vector_store_driver"] = vector_store_driver
+        if loader == "TextLoader":
+            params["loader"] = TextLoader()
+            params["source"] = text
+        if loader == "WebLoader":
+            params["loader"] = WebLoader()
+            params["source"] = url
+        if loader == "CsvLoader":
+            params["loader"] = CsvLoader()
+            params["source"] = text
+
+        module = TextLoaderRetrievalRagModule(**params)
+        return ([module],)

nodes/rag/gtUITranslateQueryRagModule.py

@@ -0,0 +1,43 @@
+from griptape.drivers import OpenAiChatPromptDriver
+from griptape.engines.rag.modules import TranslateQueryRagModule
+
+from .gtUIBaseQueryRagModule import gtUIBaseQueryRagModule
+
+
+class gtUITranslateQueryRagModule(gtUIBaseQueryRagModule):
+    """
+    Griptape Translate Query Rag Module. Used for the Query Stage of the RAG Engine.
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        inputs["optional"] = {
+            "language": (
+                "STRING",
+                {
+                    "default": "english",
+                    "tooltip": "Language to translate the query to.",
+                },
+            ),
+            "prompt_driver": (
+                "PROMPT_DRIVER",
+                {"tooltip": "Prompt driver to use for the module."},
+            ),
+        }
+        return inputs
+
+    def create(self, **kwargs):
+        prompt_driver = kwargs.get("prompt_driver", None)
+        if not prompt_driver:
+            prompt_driver = OpenAiChatPromptDriver(model="gpt-4o")
+
+        params = {}
+        params["language"] = kwargs.get("language", "english")
+        params["prompt_driver"] = prompt_driver
+        module = TranslateQueryRagModule(**params)
+        return ([module],)

nodes/rag/gtUIVectorStoreRetrievalRagModule.py

@@ -0,0 +1,53 @@
+from griptape.engines.rag.modules import VectorStoreRetrievalRagModule
+
+from .gtUIBaseRetrievalRagModule import gtUIBaseRetrievalRagModule
+
+
+class gtUIVectorStoreRetrievalRagModule(gtUIBaseRetrievalRagModule):
+    """
+    Griptape Text Loader Retrieval Rag Module. Used for the Retrieval Stage of the RAG Engine.
+    """
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+
+        inputs["optional"].update(
+            {
+                # "namespace": (
+                #     "STRING",
+                #     {"default": "default", "tooltip": "Namespace of the vector store."},
+                # ),
+                # "query_params": (
+                #     "STRING",
+                #     {
+                #         "dynamicPrompts": True,
+                #         "tooltip": "Dict of any query parameters.",
+                #         "multiline": True,
+                #         "default": '{"top_n": 20}',
+                #     },
+                # ),
+            }
+        )
+        return inputs
+
+    def create(self, **kwargs):
+        vector_store_driver = self.get_vector_store_driver(
+            kwargs.get("vector_store_driver", None)
+        )
+        # namespace = kwargs.get("namespace", "default")
+        # query_params = kwargs.get("query_params", {})
+
+        # query_params_dict = self.ensure_dict(query_params)
+        # query_params_dict["namespace"] = namespace
+
+        params = {}
+        params["query_params"] = self.get_query_params(kwargs)
+
+        params["vector_store_driver"] = vector_store_driver
+        # params["query_params"] = query_params_dict
+        module = VectorStoreRetrievalRagModule(**params)
+        return ([module],)

nodes/tools/gtUIRagTool.py

@@ -0,0 +1,41 @@
+from griptape.tools import RagTool
+
+from .gtUIBaseTool import gtUIBaseTool
+
+
+class gtUIRagTool(gtUIBaseTool):
+    DESCRIPTION = "Griptape Rag Tool"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "description": (
+                    "STRING",
+                    {
+                        "default": "Contains information about...",
+                        "tooltip": "Description of the type of information you're querying for.",
+                    },
+                ),
+                "rag_engine": (
+                    "RAG_ENGINE",
+                    {"tooltip": "Rag Engine used for the tool."},
+                ),
+            },
+        }
+
+    CATEGORY = "Griptape/RAG"
+
+    def create(self, **kwargs):
+        description = kwargs.get("description", "Contains information.")
+        off_prompt = kwargs.get("off_prompt", False)
+        rag_engine = kwargs.get("rag_engine", None)
+
+        tool_params = {}
+
+        tool_params["description"] = description
+        tool_params["off_prompt"] = off_prompt
+        tool_params["rag_engine"] = rag_engine
+
+        tool = RagTool(**tool_params)
+        return ([tool],)