diff --git a/nodes/tasks/gtUIImageQueryTask.py b/nodes/tasks/gtUIImageQueryTask.py index e0ed065..60124a6 100644 --- a/nodes/tasks/gtUIImageQueryTask.py +++ b/nodes/tasks/gtUIImageQueryTask.py @@ -2,6 +2,8 @@ import base64 from griptape.drivers import AmazonBedrockPromptDriver, AnthropicPromptDriver from griptape.loaders import ImageLoader +from griptape.structures import Workflow +from griptape.tasks import PromptTask from ..agent.gtComfyAgent import gtComfyAgent as Agent from ..utilities import ( @@ -27,10 +29,8 @@ class gtUIImageQueryTask(gtUIBaseImageTask): if images: if not agent: agent = Agent() - - prompt_text = self.get_prompt_text(STRING, input_string) - prompt_driver = agent.config.prompt_driver + prompt_text = self.get_prompt_text(STRING, input_string) # If the driver is AmazonBedrock or Anthropic, the prompt_text cannot be empty if prompt_text.strip() == "": if isinstance( @@ -47,11 +47,28 @@ class gtUIImageQueryTask(gtUIBaseImageTask): ) except Exception as e: raise (f"Couldn't load image {e}") - - # if deferred_evaluation: - # task = PromptTask([prompt_text, *image_artifacts]) - # return ("Image Query Task Created", task) - result = agent.run([prompt_text, *image_artifacts]) + # Depending on the model, we might need to use a workflow instead of a simple prompt + rulesets = agent.rulesets + tasks = [] + if len(image_artifacts) > 2: + task_args = {} + if agent.config.prompt_driver: + task_args["prompt_driver"] = prompt_driver + if len(rulesets) > 0: + task_args["rulesets"] = rulesets + end_task = PromptTask( + "Concatenate just the output values of the tasks, separated by two newlines: {{ parent_outputs }}", + id="END", + **task_args, + ) + for artifact in image_artifacts: + task = PromptTask([prompt_text, artifact], **task_args) + task.add_child(end_task) + tasks.append(task) + workflow = Workflow(tasks=[*tasks, end_task]) + result = workflow.run() + else: + result = agent.run([prompt_text, *image_artifacts]) output = result.output_task.output.value else: output = "No image provided"