using a workflow for multiple images

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"