From 19ebd9d2bfc91a82f57719ad7bb9e005135f8aad Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Sun, 1 Sep 2024 05:46:01 +1200 Subject: [PATCH] added ruleset toggle to RAGTool to make it behave better --- __init__.py | 4 +++ nodes/agent/BaseAgent.py | 19 ++++++++++-- .../gtUIFootnotePromptResponseRagModule.py | 30 +++++++++++++++++++ nodes/tools/gtUIRagTool.py | 19 ++++++++++++ 4 files changed, 70 insertions(+), 2 deletions(-) create mode 100644 nodes/rag/gtUIFootnotePromptResponseRagModule.py diff --git a/__init__.py b/__init__.py index 507bda8..b99584d 100644 --- a/__init__.py +++ b/__init__.py @@ -184,6 +184,9 @@ 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.gtUIFootnotePromptResponseRagModule import ( + gtUIFootnotePromptResponseRagModule, +) from .nodes.rag.gtUIPromptResponseRagModule import gtUIPromptResponseRagModule # RAG @@ -403,6 +406,7 @@ NODE_CLASS_MAPPINGS = { "Griptape RAG Retrieve: Vector Store Module": gtUIVectorStoreRetrievalRagModule, "Griptape RAG Rerank: Text Chunks Module": gtUITextChunksRerankRagModule, "Griptape RAG Response: Prompt Module": gtUIPromptResponseRagModule, + "Griptape RAG Response: Footnote Module": gtUIFootnotePromptResponseRagModule, # "Griptape Display: Artifact": gtUIOutputArtifactNode, # "Griptape Config: Environment Variables": gtUIEnv, } diff --git a/nodes/agent/BaseAgent.py b/nodes/agent/BaseAgent.py index c24fc23..9835036 100644 --- a/nodes/agent/BaseAgent.py +++ b/nodes/agent/BaseAgent.py @@ -1,5 +1,5 @@ from griptape.drivers import DummyVectorStoreDriver -from griptape.tools import QueryTool, VectorStoreTool +from griptape.tools import QueryTool, RagTool, VectorStoreTool from openai import OpenAIError # from server import PromptServer @@ -94,11 +94,23 @@ class BaseAgent: CATEGORY = "Griptape/Agent" + def rag_tool_ruleset(self, tools): + for tool in tools: + # Check and see if any of the tools are RAGTools + if isinstance(tool, RagTool): + # See if there's an additional attribute on the RagTool + if hasattr(tool, "use_rules"): + # if it's true, get the ruleset + if tool.use_rules: + return tool.ruleset + return None + def tool_check(self, config, tools): tool_list = [] if len(tools) > 0: - # Check and see if any of the tools are VectorStoreClients + # Logic per tool for tool in tools: + # Check and see if any of the tools are VectorStoreTools if isinstance(tool, VectorStoreTool): # Check and see if the driver is a DummyVectorStoreDriver # If it is, replace it with the agent's vector store driver @@ -109,6 +121,7 @@ class BaseAgent: tool.vector_store_driver = vector_store_driver except Exception as e: print(f"Error: {str(e)}") + # Check and see if any of the tools have been set to off_prompt off_prompt = False for tool in tools: @@ -145,6 +158,8 @@ class BaseAgent: # Tools create_dict["tools"] = self.tool_check(config, tools) + # add a ruleset to handle RAGTools if it's on. + rulesets.append(self.rag_tool_ruleset(tools)) # Rulesets if len(rulesets) > 0: create_dict["rulesets"] = rulesets diff --git a/nodes/rag/gtUIFootnotePromptResponseRagModule.py b/nodes/rag/gtUIFootnotePromptResponseRagModule.py new file mode 100644 index 0000000..919d0a2 --- /dev/null +++ b/nodes/rag/gtUIFootnotePromptResponseRagModule.py @@ -0,0 +1,30 @@ +from griptape.engines.rag.modules import FootnotePromptResponseRagModule + +from .gtUIPromptResponseRagModule import gtUIPromptResponseRagModule + + +class gtUIFootnotePromptResponseRagModule(gtUIPromptResponseRagModule): + """ + Griptape Footnote Prompt Response Rag Module. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + inputs = super().INPUT_TYPES() + + 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 = FootnotePromptResponseRagModule(**params) + return ([module],) diff --git a/nodes/tools/gtUIRagTool.py b/nodes/tools/gtUIRagTool.py index 78059dc..6b17fda 100644 --- a/nodes/tools/gtUIRagTool.py +++ b/nodes/tools/gtUIRagTool.py @@ -1,3 +1,4 @@ +from griptape.rules import Rule, Ruleset from griptape.tools import RagTool from .gtUIBaseTool import gtUIBaseTool @@ -21,6 +22,13 @@ class gtUIRagTool(gtUIBaseTool): "RAG_ENGINE", {"tooltip": "Rag Engine used for the tool."}, ), + "use_rules": ( + "BOOLEAN", + { + "default": True, + "tooltip": "If enabled, will include helpful rules to sculpt the response of the agent to suit RAG.", + }, + ), }, } @@ -30,6 +38,7 @@ class gtUIRagTool(gtUIBaseTool): description = kwargs.get("description", "Contains information.") off_prompt = kwargs.get("off_prompt", False) rag_engine = kwargs.get("rag_engine", None) + use_rules = kwargs.get("use_rules", True) tool_params = {} @@ -38,4 +47,14 @@ class gtUIRagTool(gtUIBaseTool): tool_params["rag_engine"] = rag_engine tool = RagTool(**tool_params) + if use_rules: + tool.use_rules = True + tool.ruleset = Ruleset( + name="GriptapeRagToolWithRules", + rules=[ + Rule("Always include footnotes in your responses if specified."), + Rule("Use the RAG Tool to get answers to your questions."), + ], + ) + return ([tool],)