added ruleset toggle to RAGTool to make it behave better

__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,
 }

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

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],)

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],)