refactor loopblock value (#1379)

skyvern/forge/sdk/workflow/exceptions.py

@@ -110,5 +110,17 @@ class WorkflowParameterMissingRequiredValue(BaseWorkflowHTTPException):
 
 
 class InvalidWaitBlockTime(SkyvernException):
-    def __init__(self, max_sec: int):
+    def __init__(self, max_sec: int) -> None:
         super().__init__(f"Invalid wait time for wait block, it should be a number between 0 and {max_sec}.")
+
+
+class FailedToFormatJinjaStyleParameter(SkyvernException):
+    def __init__(self, template: str, msg: str) -> None:
+        super().__init__(
+            f"Failed to format Jinja style parameter {template}. Please make sure the variable reference is correct. reason: {msg}"
+        )
+
+
+class NoIterableValueFound(SkyvernException):
+    def __init__(self) -> None:
+        super().__init__("No iterable value found for the loop block")

skyvern/forge/sdk/workflow/models/block.py

@@ -47,8 +47,10 @@ from skyvern.forge.sdk.db.enums import TaskType
 from skyvern.forge.sdk.schemas.tasks import Task, TaskOutput, TaskStatus
 from skyvern.forge.sdk.workflow.context_manager import BlockMetadata, WorkflowRunContext
 from skyvern.forge.sdk.workflow.exceptions import (
+    FailedToFormatJinjaStyleParameter,
     InvalidEmailClientConfiguration,
     InvalidFileType,
+    NoIterableValueFound,
     NoValidEmailRecipient,
 )
 from skyvern.forge.sdk.workflow.models.parameter import (
@@ -576,14 +578,17 @@ class LoopBlockExecutedResult(BaseModel):
 class ForLoopBlock(Block):
     block_type: Literal[BlockType.FOR_LOOP] = BlockType.FOR_LOOP
 
-    loop_over: PARAMETER_TYPE
+    loop_over: PARAMETER_TYPE | None
+    loop_variable_reference: str | None
     loop_blocks: list[BlockTypeVar]
 
     def get_all_parameters(
         self,
         workflow_run_id: str,
     ) -> list[PARAMETER_TYPE]:
-        parameters = {self.loop_over}
+        parameters = set()
+        if self.loop_over is not None:
+            parameters.add(self.loop_over)
 
         for loop_block in self.loop_blocks:
             for parameter in loop_block.get_all_parameters(workflow_run_id):
@@ -600,6 +605,9 @@ class ForLoopBlock(Block):
                 if isinstance(parameter, ContextParameter):
                     context_parameters.append(parameter)
 
+        if self.loop_over is None:
+            return context_parameters
+
         for context_parameter in context_parameters:
             if context_parameter.source.key != self.loop_over.key:
                 continue
@@ -620,6 +628,18 @@ class ForLoopBlock(Block):
         return context_parameters
 
     def get_loop_over_parameter_values(self, workflow_run_context: WorkflowRunContext) -> list[Any]:
+        # parse the value from self.loop_variable_reference and then from self.loop_over
+        if self.loop_variable_reference:
+            value_template = f'{{{{ {self.loop_variable_reference.strip(" {}")} | tojson }}}}'
+            try:
+                value_json = self.format_block_parameter_template_from_workflow_run_context(
+                    value_template, workflow_run_context
+                )
+            except Exception as e:
+                raise FailedToFormatJinjaStyleParameter(value_template, str(e))
+            parameter_value = json.loads(value_json)
+
+        elif self.loop_over is not None:
             if isinstance(self.loop_over, WorkflowParameter):
                 parameter_value = workflow_run_context.get_value(self.loop_over.key)
             elif isinstance(self.loop_over, OutputParameter):
@@ -642,7 +662,10 @@ class ForLoopBlock(Block):
                     else:
                         raise ValueError("ContextParameter source value should be a dict")
             else:
-            raise NotImplementedError
+                raise NotImplementedError()
+
+        else:
+            raise NoIterableValueFound()
 
         if isinstance(parameter_value, list):
             return parameter_value
@@ -725,7 +748,15 @@ class ForLoopBlock(Block):
 
     async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult:
         workflow_run_context = self.get_workflow_run_context(workflow_run_id)
+        try:
             loop_over_values = self.get_loop_over_parameter_values(workflow_run_context)
+        except Exception as e:
+            return self.build_block_result(
+                success=False,
+                failure_reason=f"failed to get loop values: {str(e)}",
+                status=BlockStatus.failed,
+            )
+
         LOG.info(
             f"Number of loop_over values: {len(loop_over_values)}",
             block_type=self.block_type,

skyvern/forge/sdk/workflow/models/yaml.py

@@ -142,6 +142,7 @@ class ForLoopBlockYAML(BlockYAML):
 
     loop_over_parameter_key: str
     loop_blocks: list["BLOCK_YAML_SUBCLASSES"]
+    loop_variable_reference: str | None = None
 
 
 class CodeBlockYAML(BlockYAML):

skyvern/forge/sdk/workflow/service.py

@@ -1319,10 +1319,27 @@ class WorkflowService:
                 await WorkflowService.block_yaml_to_block(workflow, loop_block, parameters)
                 for loop_block in block_yaml.loop_blocks
             ]
+
+            loop_over_parameter: Parameter | None = None
+            if block_yaml.loop_over_parameter_key:
                 loop_over_parameter = parameters[block_yaml.loop_over_parameter_key]
+
+            if block_yaml.loop_variable_reference:
+                # it's backaward compatible with jinja style parameter and context paramter
+                # we trim the format like {{ loop_key }} into loop_key to initialize the context parater,
+                # otherwise it might break the context parameter initialization chain, blow up the worklofw parameters
+                # TODO: consider remove this if we totally give up context parameter
+                trimmed_key = block_yaml.loop_variable_reference.strip(" {}")
+                if trimmed_key in parameters:
+                    loop_over_parameter = parameters[trimmed_key]
+
+            if loop_over_parameter is None and not block_yaml.loop_variable_reference:
+                raise Exception("empty loop value parameter")
+
             return ForLoopBlock(
                 label=block_yaml.label,
                 loop_over=loop_over_parameter,
+                loop_variable_reference=block_yaml.loop_variable_reference,
                 loop_blocks=loop_blocks,
                 output_parameter=output_parameter,
                 continue_on_failure=block_yaml.continue_on_failure,