diff --git a/__init__.py b/__init__.py index d8f3fc9..507bda8 100644 --- a/__init__.py +++ b/__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, } diff --git a/js/CombineNodes.js b/js/CombineNodes.js index 9b33db9..84786b5 100644 --- a/js/CombineNodes.js +++ b/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_"; } } diff --git a/js/TextLoaderModuleNodes.js b/js/TextLoaderModuleNodes.js new file mode 100644 index 0000000..caf23e4 --- /dev/null +++ b/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); + }; + + } + diff --git a/js/gtUINodes.js b/js/gtUINodes.js index ad3b9c5..34ab515 100644 --- a/js/gtUINodes.js +++ b/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 diff --git a/nodes/combine/gtUIModuleList.py b/nodes/combine/gtUIModuleList.py new file mode 100644 index 0000000..54253cd --- /dev/null +++ b/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,) diff --git a/nodes/drivers/gtUIBaseRerankDriver.py b/nodes/drivers/gtUIBaseRerankDriver.py new file mode 100644 index 0000000..a2d0bf2 --- /dev/null +++ b/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,) diff --git a/nodes/drivers/gtUICohereRerankDriver.py b/nodes/drivers/gtUICohereRerankDriver.py new file mode 100644 index 0000000..0ac1854 --- /dev/null +++ b/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,) diff --git a/nodes/drivers/gtUIPgVectorVectorStoreDriver.py b/nodes/drivers/gtUIPgVectorVectorStoreDriver.py index f5f44f2..00b44a3 100644 --- a/nodes/drivers/gtUIPgVectorVectorStoreDriver.py +++ b/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,) diff --git a/nodes/rag/__init__.py b/nodes/rag/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/rag/gtUIBaseQueryRagModule.py b/nodes/rag/gtUIBaseQueryRagModule.py new file mode 100644 index 0000000..5789b51 --- /dev/null +++ b/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()],) diff --git a/nodes/rag/gtUIBaseRagModule.py b/nodes/rag/gtUIBaseRagModule.py new file mode 100644 index 0000000..75c0bcb --- /dev/null +++ b/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()],) diff --git a/nodes/rag/gtUIBaseReponseRagModule.py b/nodes/rag/gtUIBaseReponseRagModule.py new file mode 100644 index 0000000..e12f722 --- /dev/null +++ b/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()],) diff --git a/nodes/rag/gtUIBaseRerankRagModule.py b/nodes/rag/gtUIBaseRerankRagModule.py new file mode 100644 index 0000000..b5bc374 --- /dev/null +++ b/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(),) diff --git a/nodes/rag/gtUIBaseRetrievalRagModule.py b/nodes/rag/gtUIBaseRetrievalRagModule.py new file mode 100644 index 0000000..9de0b7a --- /dev/null +++ b/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()],) diff --git a/nodes/rag/gtUIPromptResponseRagModule.py b/nodes/rag/gtUIPromptResponseRagModule.py new file mode 100644 index 0000000..bb542da --- /dev/null +++ b/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],) diff --git a/nodes/rag/gtUIRagEngine.py b/nodes/rag/gtUIRagEngine.py new file mode 100644 index 0000000..a038660 --- /dev/null +++ b/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,) diff --git a/nodes/rag/gtUITextChunksRerankRagModule.py b/nodes/rag/gtUITextChunksRerankRagModule.py new file mode 100644 index 0000000..bae2851 --- /dev/null +++ b/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,) diff --git a/nodes/rag/gtUITextLoaderRetrievalRagModule.py b/nodes/rag/gtUITextLoaderRetrievalRagModule.py new file mode 100644 index 0000000..cb49745 --- /dev/null +++ b/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],) diff --git a/nodes/rag/gtUITranslateQueryRagModule.py b/nodes/rag/gtUITranslateQueryRagModule.py new file mode 100644 index 0000000..63f4acd --- /dev/null +++ b/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],) diff --git a/nodes/rag/gtUIVectorStoreRetrievalRagModule.py b/nodes/rag/gtUIVectorStoreRetrievalRagModule.py new file mode 100644 index 0000000..f54281d --- /dev/null +++ b/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],) diff --git a/nodes/tools/gtUIRagTool.py b/nodes/tools/gtUIRagTool.py new file mode 100644 index 0000000..78059dc --- /dev/null +++ b/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],)