Move the code over from private repository (#3)

skyvern/forge/sdk/workflow/context_manager.py

@@ -0,0 +1,79 @@
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from skyvern.forge.sdk.api.aws import AsyncAWSClient
+from skyvern.forge.sdk.workflow.models.parameter import PARAMETER_TYPE, Parameter, ParameterType, WorkflowParameter
+
+if TYPE_CHECKING:
+    from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunParameter
+
+
+LOG = structlog.get_logger()
+
+
+class ContextManager:
+    aws_client: AsyncAWSClient
+    parameters: dict[str, PARAMETER_TYPE]
+    values: dict[str, Any]
+
+    def __init__(self, workflow_parameter_tuples: list[tuple[WorkflowParameter, "WorkflowRunParameter"]]) -> None:
+        self.aws_client = AsyncAWSClient()
+        self.parameters = {}
+        self.values = {}
+        for parameter, run_parameter in workflow_parameter_tuples:
+            if parameter.key in self.parameters:
+                prev_value = self.parameters[parameter.key]
+                new_value = run_parameter.value
+                LOG.error(
+                    f"Duplicate parameter key {parameter.key} found while initializing context manager, previous value: {prev_value}, new value: {new_value}. Using new value."
+                )
+
+            self.parameters[parameter.key] = parameter
+            self.values[parameter.key] = run_parameter.value
+
+    async def register_parameter_value(
+        self,
+        parameter: PARAMETER_TYPE,
+    ) -> None:
+        if parameter.parameter_type == ParameterType.WORKFLOW:
+            LOG.error(f"Workflow parameters are set while initializing context manager. Parameter key: {parameter.key}")
+            raise ValueError(
+                f"Workflow parameters are set while initializing context manager. Parameter key: {parameter.key}"
+            )
+        elif parameter.parameter_type == ParameterType.AWS_SECRET:
+            secret_value = await self.aws_client.get_secret(parameter.aws_key)
+            if secret_value is not None:
+                self.values[parameter.key] = secret_value
+        else:
+            # ContextParameter values will be set within the blocks
+            return None
+
+    async def register_block_parameters(
+        self,
+        parameters: list[PARAMETER_TYPE],
+    ) -> None:
+        for parameter in parameters:
+            if parameter.key in self.parameters:
+                LOG.debug(f"Parameter {parameter.key} already registered, skipping")
+                continue
+
+            if parameter.parameter_type == ParameterType.WORKFLOW:
+                LOG.error(
+                    f"Workflow parameter {parameter.key} should have already been set through workflow run parameters"
+                )
+                raise ValueError(
+                    f"Workflow parameter {parameter.key} should have already been set through workflow run parameters"
+                )
+
+            self.parameters[parameter.key] = parameter
+            await self.register_parameter_value(parameter)
+
+    def get_parameter(self, key: str) -> Parameter:
+        return self.parameters[key]
+
+    def get_value(self, key: str) -> Any:
+        return self.values[key]
+
+    def set_value(self, key: str, value: Any) -> None:
+        self.values[key] = value