Enforce navigation_goal in NavigationBlock and data_extraction_goal i…

…n ExtractionBlock (#1257)
Skyvern-AI · Nov 25, 2024 · f28a3bc · f28a3bc
1 parent d520254
commit f28a3bc
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py
@@ -1372,10 +1372,14 @@ async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult:
 class NavigationBlock(BaseTaskBlock):
     block_type: Literal[BlockType.NAVIGATION] = BlockType.NAVIGATION
 
+    navigation_goal: str
+
 
 class ExtractionBlock(BaseTaskBlock):
     block_type: Literal[BlockType.EXTRACTION] = BlockType.EXTRACTION
 
+    data_extraction_goal: str
+
 
 class LoginBlock(BaseTaskBlock):
     block_type: Literal[BlockType.LOGIN] = BlockType.LOGIN

diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py
@@ -253,9 +253,9 @@ class NavigationBlockYAML(BlockYAML):
 class ExtractionBlockYAML(BlockYAML):
     block_type: Literal[BlockType.EXTRACTION] = BlockType.EXTRACTION  # type: ignore
 
+    data_extraction_goal: str
     url: str | None = None
     title: str = ""
-    data_extraction_goal: str | None = None
     data_schema: dict[str, Any] | list | None = None
     max_retries: int = 0
     max_steps_per_run: int | None = None

diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
@@ -1207,7 +1207,7 @@ async def create_workflow_from_request(
             raise e
 
     @staticmethod
-    async def _create_output_parameter_for_block(workflow_id: str, block_yaml: BLOCK_YAML_TYPES) -> OutputParameter:
+    async def create_output_parameter_for_block(workflow_id: str, block_yaml: BLOCK_YAML_TYPES) -> OutputParameter:
         output_parameter_key = f"{block_yaml.label}_output"
         return await app.DATABASE.create_output_parameter(
             workflow_id=workflow_id,
@@ -1221,7 +1221,7 @@ async def _create_all_output_parameters_for_workflow(
     ) -> dict[str, OutputParameter]:
         output_parameters = {}
         for block_yaml in block_yamls:
-            output_parameter = await WorkflowService._create_output_parameter_for_block(
+            output_parameter = await WorkflowService.create_output_parameter_for_block(
                 workflow_id=workflow_id, block_yaml=block_yaml
             )
             output_parameters[block_yaml.label] = output_parameter