Merge pull request #228 from griptape-ai/fix/GoogleQueryTool

added patch for QueryTool for Google Gemini

CHANGELOG.md

@@ -8,10 +8,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 <!-- ## [Unreleased]
 
 ### Added 
+- `GeminiQueryTool` - a patch for the QueryTool: https://github.com/griptape-ai/griptape/issues/1572
 ### Changed 
+- `gtComfyAgent` checks to see if the `prompt_driver` is `GooglePromptDriver`. if so, it uses the `GeminiQueryTool` isntead of `QueryTool`.
+- `QueryTool` now checks to see if the `prompt_driver` is `GooglePromptDriver`. If so, it uses the `GeminiQueryTool` instead of `QueryTool`.
 ### Deprecated 
 ### Removed 
 ### Fixed 
+- `GooglePromptDriver` can now handle tools that are `off_prompt`.
 ### Security   -->
 
 ## [2.1.09] - 2025-10-01

__init__.py

@@ -220,12 +220,15 @@ from .nodes.expand.gtUIExpandConfig import gtUIExpandConfig
 from .nodes.loaders.gtUIFetchImage import gtUIFetchImage
 from .nodes.loaders.gtUILoadAudio import gtUILoadAudio
 from .nodes.loaders.gtUILoadText import gtUILoadText
+
+# PATCHES
+from .nodes.patches.gemini_query_tool import GeminiQueryTool
+
+# RAG
 from .nodes.rag.gtUIFootnotePromptResponseRagModule import (
     gtUIFootnotePromptResponseRagModule,
 )
 from .nodes.rag.gtUIPromptResponseRagModule import gtUIPromptResponseRagModule
-
-# RAG
 from .nodes.rag.gtUIRagEngine import gtUIRagEngine
 from .nodes.rag.gtUITextChunksRerankRagModule import gtUITextChunksRerankRagModule
 from .nodes.rag.gtUITextChunksResponseRagModule import gtUITextChunksResponseRagModule

nodes/agent/gtComfyAgent.py

@@ -1,9 +1,12 @@
 from dotenv import load_dotenv
 from griptape.configs import Defaults
 from griptape.configs.drivers import DriversConfig, OpenAiDriversConfig
+from griptape.drivers import GooglePromptDriver
 from griptape.structures import Agent
+from griptape.tools import QueryTool
 
 from ...py.griptape_settings import GriptapeSettings
+from ..patches.gemini_query_tool import GeminiQueryTool
 
 default_prompt = "{{ input_string }}"
 
@@ -47,12 +50,32 @@ class gtComfyAgent(Agent):
         if "prompt_driver" in kwargs:
             Defaults.drivers_config.prompt_driver = kwargs["prompt_driver"]
 
+        if "tools" in kwargs:
+            self._fix_google_query_tool(
+                Defaults.drivers_config.prompt_driver, kwargs["tools"]
+            )
         # Initialize the parent class
         super().__init__(*args, **kwargs)
 
         # Add any additional initialization here
         self.drivers_config = Defaults.drivers_config
 
+    def _fix_google_query_tool(self, prompt_driver, tools):
+        # Check and see if any of the tools have been set to off_prompt
+        off_prompt = False
+        return_tools = tools
+        for tool in tools:
+            if tool.off_prompt and not off_prompt:
+                off_prompt = True
+        if off_prompt:
+            for x, tool in enumerate(tools):
+                if isinstance(tool, QueryTool):
+                    # Check and see if the prompt driver is a GooglePromptDriver
+                    print(f"Prompt Driver: {prompt_driver}")
+                    if isinstance(prompt_driver, GooglePromptDriver):
+                        tools[x] = GeminiQueryTool()
+        return return_tools
+
     def set_default_config(self):
         # agent_config = get_config("agent_config")
         settings = GriptapeSettings()

nodes/patches/gemini_query_tool.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from attrs import define
+from griptape.artifacts import ErrorArtifact, ListArtifact, TextArtifact
+from griptape.engines.rag.rag_context import RagContext
+from griptape.tools import QueryTool
+from griptape.utils.decorators import activity
+from schema import Literal, Schema
+
+
+@define(kw_only=True)
+class GeminiQueryTool(QueryTool):
+    @activity(
+        config={
+            "description": "Can be used to search through textual content.",
+            "schema": Schema(
+                {
+                    Literal(
+                        "query", description="A natural language search query"
+                    ): str,
+                    Literal("content"): Schema(
+                        {
+                            "memory_name": str,
+                            "artifact_namespace": str,
+                        }
+                    ),
+                }
+            ),
+        },
+    )
+    def query(self, params: dict) -> ListArtifact | ErrorArtifact:
+        query = params["values"]["query"]
+        content = params["values"]["content"]
+
+        if isinstance(content, str):
+            text_artifacts = [TextArtifact(content)]
+        else:
+            memory = self.find_input_memory(content["memory_name"])
+            artifact_namespace = content["artifact_namespace"]
+
+            if memory is not None:
+                artifacts = memory.load_artifacts(artifact_namespace)
+            else:
+                return ErrorArtifact("memory not found")
+
+            text_artifacts = [
+                artifact for artifact in artifacts if isinstance(artifact, TextArtifact)
+            ]
+
+        outputs = self._rag_engine.process(
+            RagContext(query=query, text_chunks=text_artifacts)
+        ).outputs
+
+        if len(outputs) > 0:
+            return ListArtifact(outputs)
+        else:
+            return ErrorArtifact("query output is empty")

nodes/tasks/gtUITask.py

@@ -1,9 +1,10 @@
 from griptape.artifacts import TextArtifact
-from griptape.drivers import DummyVectorStoreDriver
+from griptape.drivers import DummyVectorStoreDriver, GooglePromptDriver
 from griptape.tasks import PromptTask, ToolkitTask, ToolTask
 from griptape.tools import QueryTool, VectorStoreTool
 
 from ..agent.gtComfyAgent import gtComfyAgent as Agent
+from ..patches.gemini_query_tool import GeminiQueryTool
 from .gtUIBaseTask import gtUIBaseTask
 
 
@@ -52,7 +53,10 @@ class gtUITask(gtUIBaseTask):
                     if isinstance(tool, QueryTool):
                         taskMemoryClient = True
                 if not taskMemoryClient:
-                    tools.append(QueryTool(off_prompt=False))
+                    if isinstance(config.prompt_driver, GooglePromptDriver):
+                        tools.append(GeminiQueryTool())
+                    else:
+                        tools.append(QueryTool(off_prompt=False))
             tool_list = tools
         return tool_list
 

nodes/tools/gtUIQueryTool.py

@@ -1,5 +1,7 @@
+from griptape.drivers import GooglePromptDriver
 from griptape.tools import QueryTool
 
+from ..patches.gemini_query_tool import GeminiQueryTool
 from .gtUIBaseTool import gtUIBaseTool
 
 
@@ -29,5 +31,8 @@ class gtUIQueryTool(gtUIBaseTool):
         params = {}
         if prompt_driver:
             params["prompt_driver"] = prompt_driver
-        tool = QueryTool(**params)
+        if isinstance(prompt_driver, GooglePromptDriver):
+            tool = GeminiQueryTool()
+        else:
+            tool = QueryTool(**params)
         return ([tool],)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "comfyui-griptape"
-version = "2.1.09"
+version = "2.1.10"
 description = "Griptape LLM(Large Language Model) Nodes for ComfyUI."
 authors = ["Jason Schleifer <jason.schleifer@gmail.com>"]
 readme = "README.md"
@@ -9,7 +9,7 @@ readme = "README.md"
 [project]
 name = "comfyui-griptape"
 description = "Griptape LLM(Large Language Model) Nodes for ComfyUI."
-version = "2.1.09" 
+version = "2.1.10" 
 license = {file = "LICENSE"}
 dependencies = ["attrs>=24.3.0,<25.0.0", "openai>=1.58.1,<2.0.0", "griptape[all]>=1.1.0", "python-dotenv", "poetry==1.8.5", "griptape-black-forest @ git+https://github.com/griptape-ai/griptape-black-forest.git"]