From 479186bbef815330aa984466fc126ddc9668b2a3 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Wed, 15 May 2024 00:16:15 -0700
Subject: [PATCH] Quick fix for taskoutput change bug in forloopblock (#323)

---
 skyvern/forge/sdk/workflow/context_manager.py |  5 +++++
 skyvern/forge/sdk/workflow/models/block.py    | 15 +++++++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py
index 32eb498e..521f8acb 100644
--- a/skyvern/forge/sdk/workflow/context_manager.py
+++ b/skyvern/forge/sdk/workflow/context_manager.py
@@ -232,6 +232,11 @@ class WorkflowRunContext:
                 and isinstance(parameter.source, OutputParameter)
                 and parameter.source.key == output_parameter.key
             ):
+                value = (
+                    value["extracted_information"]
+                    if isinstance(value, dict) and "extracted_information" in value
+                    else value
+                )
                 if parameter.value:
                     LOG.warning(
                         f"Context parameter {parameter.key} already has a value, overwriting",
diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py
index 67a3c9cf..d42fc529 100644
--- a/skyvern/forge/sdk/workflow/models/block.py
+++ b/skyvern/forge/sdk/workflow/models/block.py
@@ -339,14 +339,25 @@ class ForLoopBlock(Block):
         return context_parameters
 
     def get_loop_over_parameter_values(self, workflow_run_context: WorkflowRunContext) -> list[Any]:
-        if isinstance(self.loop_over, WorkflowParameter) or isinstance(self.loop_over, OutputParameter):
+        if isinstance(self.loop_over, WorkflowParameter):
             parameter_value = workflow_run_context.get_value(self.loop_over.key)
+        elif isinstance(self.loop_over, OutputParameter):
+            # If the output parameter is for a TaskBlock, it will be a TaskOutput object. We need to extract the
+            # value from the TaskOutput object's extracted_information field.
+            output_parameter_value = workflow_run_context.get_value(self.loop_over.key)
+            if isinstance(output_parameter_value, dict) and "extracted_information" in output_parameter_value:
+                parameter_value = output_parameter_value["extracted_information"]
+            else:
+                parameter_value = output_parameter_value
         elif isinstance(self.loop_over, ContextParameter):
             parameter_value = self.loop_over.value
             if not parameter_value:
                 source_parameter_value = workflow_run_context.get_value(self.loop_over.source.key)
                 if isinstance(source_parameter_value, dict):
-                    parameter_value = source_parameter_value.get(self.loop_over.key)
+                    if "extracted_information" in source_parameter_value:
+                        parameter_value = source_parameter_value["extracted_information"].get(self.loop_over.key)
+                    else:
+                        parameter_value = source_parameter_value.get(self.loop_over.key)
                 else:
                     raise ValueError("ContextParameter source value should be a dict")
         else: