mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2025-09-02 18:50:24 +00:00
79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
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
|