From b619f7bc953b2dadb4b09fb47c66e5cab23701e6 Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Wed, 15 Jan 2025 05:05:36 +1300 Subject: [PATCH] added patch for QueryTool for Google Gemini --- CHANGELOG.md | 4 +++ __init__.py | 7 ++-- nodes/agent/gtComfyAgent.py | 23 ++++++++++++ nodes/patches/gemini_query_tool.py | 57 ++++++++++++++++++++++++++++++ nodes/tasks/gtUITask.py | 8 +++-- nodes/tools/gtUIQueryTool.py | 7 +++- pyproject.toml | 4 +-- 7 files changed, 103 insertions(+), 7 deletions(-) create mode 100644 nodes/patches/gemini_query_tool.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b77f8d1..f5a668f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,10 +8,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [2.1.09] - 2025-10-01 diff --git a/__init__.py b/__init__.py index 08649a6..973fccc 100644 --- a/__init__.py +++ b/__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 diff --git a/nodes/agent/gtComfyAgent.py b/nodes/agent/gtComfyAgent.py index 5ad20f6..06b23c7 100644 --- a/nodes/agent/gtComfyAgent.py +++ b/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() diff --git a/nodes/patches/gemini_query_tool.py b/nodes/patches/gemini_query_tool.py new file mode 100644 index 0000000..ab23a89 --- /dev/null +++ b/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") diff --git a/nodes/tasks/gtUITask.py b/nodes/tasks/gtUITask.py index 9a6c76b..1cf413a 100644 --- a/nodes/tasks/gtUITask.py +++ b/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 diff --git a/nodes/tools/gtUIQueryTool.py b/nodes/tools/gtUIQueryTool.py index aa5fd1a..0ba5e77 100644 --- a/nodes/tools/gtUIQueryTool.py +++ b/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],) diff --git a/pyproject.toml b/pyproject.toml index 7ff922f..98f21dc 100644 --- a/pyproject.toml +++ b/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 "] 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"]