vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Workflow: Output Parameters & Code Blocks (#117)

Browse Source
This commit is contained in:
Kerem Yilmaz
2024-03-21 17:16:56 -07:00
committed by GitHub
parent d2ca6ca792
commit 066c2302b5
11 changed files with 556 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
skyvern/forge/sdk/db/client.py
View File

@@ -15,11 +15,13 @@ from skyvern.forge.sdk.db.models import (
AWSSecretParameterModel,
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
StepModel,
TaskModel,
WorkflowModel,
WorkflowParameterModel,
WorkflowRunModel,
WorkflowRunOutputParameterModel,
WorkflowRunParameterModel,
)
from skyvern.forge.sdk.db.utils import (
@@ -28,17 +30,30 @@ from skyvern.forge.sdk.db.utils import (
convert_to_aws_secret_parameter,
convert_to_organization,
convert_to_organization_auth_token,
convert_to_output_parameter,
convert_to_step,
convert_to_task,
convert_to_workflow,
convert_to_workflow_parameter,
convert_to_workflow_run,
convert_to_workflow_run_output_parameter,
convert_to_workflow_run_parameter,
)
from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step, StepStatus
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import AWSSecretParameter, WorkflowParameter, WorkflowParameterType
from skyvern.forge.sdk.workflow.models.workflow import Workflow, WorkflowRun, WorkflowRunParameter, WorkflowRunStatus
from skyvern.forge.sdk.workflow.models.parameter import (
AWSSecretParameter,
OutputParameter,
WorkflowParameter,
WorkflowParameterType,
)
from skyvern.forge.sdk.workflow.models.workflow import (
Workflow,
WorkflowRun,
WorkflowRunOutputParameter,
WorkflowRunParameter,
WorkflowRunStatus,
)
from skyvern.webeye.actions.models import AgentStepOutput
LOG = structlog.get_logger()
@@ -777,6 +792,68 @@ class AgentDB:
session.refresh(aws_secret_parameter)
return convert_to_aws_secret_parameter(aws_secret_parameter)
async def create_output_parameter(
self,
workflow_id: str,
key: str,
description: str | None = None,
) -> OutputParameter:
with self.Session() as session:
output_parameter = OutputParameterModel(
key=key,
description=description,
workflow_id=workflow_id,
)
session.add(output_parameter)
session.commit()
session.refresh(output_parameter)
return convert_to_output_parameter(output_parameter)
async def get_workflow_output_parameters(self, workflow_id: str) -> list[OutputParameter]:
try:
with self.Session() as session:
output_parameters = session.query(OutputParameterModel).filter_by(workflow_id=workflow_id).all()
return [convert_to_output_parameter(parameter) for parameter in output_parameters]
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def get_workflow_run_output_parameters(self, workflow_run_id: str) -> list[WorkflowRunOutputParameter]:
try:
with self.Session() as session:
workflow_run_output_parameters = (
session.query(WorkflowRunOutputParameterModel)
.filter_by(workflow_run_id=workflow_run_id)
.order_by(WorkflowRunOutputParameterModel.created_at)
.all()
)
return [
convert_to_workflow_run_output_parameter(parameter, self.debug_enabled)
for parameter in workflow_run_output_parameters
]
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def create_workflow_run_output_parameter(
self, workflow_run_id: str, output_parameter_id: str, value: dict[str, Any] | list | str | None
) -> WorkflowRunOutputParameter:
try:
with self.Session() as session:
workflow_run_output_parameter = WorkflowRunOutputParameterModel(
workflow_run_id=workflow_run_id,
output_parameter_id=output_parameter_id,
value=value,
)
session.add(workflow_run_output_parameter)
session.commit()
session.refresh(workflow_run_output_parameter)
return convert_to_workflow_run_output_parameter(workflow_run_output_parameter, self.debug_enabled)
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def get_workflow_parameters(self, workflow_id: str) -> list[WorkflowParameter]:
try:
with self.Session() as session:

6
skyvern/forge/sdk/db/id.py
View File

@@ -37,6 +37,7 @@ WORKFLOW_PREFIX = "w"
WORKFLOW_RUN_PREFIX = "wr"
WORKFLOW_PARAMETER_PREFIX = "wp"
AWS_SECRET_PARAMETER_PREFIX = "asp"
OUTPUT_PARAMETER_PREFIX = "op"
def generate_workflow_id() -> str:
@@ -59,6 +60,11 @@ def generate_workflow_parameter_id() -> str:
return f"{WORKFLOW_PARAMETER_PREFIX}_{int_id}"
def generate_output_parameter_id() -> str:
int_id = generate_id()
return f"{OUTPUT_PARAMETER_PREFIX}_{int_id}"
def generate_organization_auth_token_id() -> str:
int_id = generate_id()
return f"{ORGANIZATION_AUTH_TOKEN_PREFIX}_{int_id}"

24
skyvern/forge/sdk/db/models.py
View File

@@ -9,6 +9,7 @@ from skyvern.forge.sdk.db.id import (
generate_aws_secret_parameter_id,
generate_org_id,
generate_organization_auth_token_id,
generate_output_parameter_id,
generate_step_id,
generate_task_id,
generate_workflow_id,
@@ -150,6 +151,18 @@ class WorkflowParameterModel(Base):
deleted_at = Column(DateTime, nullable=True)
class OutputParameterModel(Base):
__tablename__ = "output_parameters"
output_parameter_id = Column(String, primary_key=True, index=True, default=generate_output_parameter_id)
key = Column(String, nullable=False)
description = Column(String, nullable=True)
workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False)
created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
deleted_at = Column(DateTime, nullable=True)
class AWSSecretParameterModel(Base):
__tablename__ = "aws_secret_parameters"
@@ -173,3 +186,14 @@ class WorkflowRunParameterModel(Base):
# Can be bool | int | float | str | dict | list depending on the workflow parameter type
value = Column(String, nullable=False)
created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
class WorkflowRunOutputParameterModel(Base):
__tablename__ = "workflow_run_output_parameters"
workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), primary_key=True, index=True)
output_parameter_id = Column(
String, ForeignKey("output_parameters.output_parameter_id"), primary_key=True, index=True
)
value = Column(JSON, nullable=False)
created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)

51
skyvern/forge/sdk/db/utils.py
View File

@@ -11,20 +11,28 @@ from skyvern.forge.sdk.db.models import (
AWSSecretParameterModel,
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
StepModel,
TaskModel,
WorkflowModel,
WorkflowParameterModel,
WorkflowRunModel,
WorkflowRunOutputParameterModel,
WorkflowRunParameterModel,
)
from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step, StepStatus
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import AWSSecretParameter, WorkflowParameter, WorkflowParameterType
from skyvern.forge.sdk.workflow.models.parameter import (
AWSSecretParameter,
OutputParameter,
WorkflowParameter,
WorkflowParameterType,
)
from skyvern.forge.sdk.workflow.models.workflow import (
Workflow,
WorkflowDefinition,
WorkflowRun,
WorkflowRunOutputParameter,
WorkflowRunParameter,
WorkflowRunStatus,
)
@@ -188,7 +196,7 @@ def convert_to_aws_secret_parameter(
if debug_enabled:
LOG.debug(
"Converting AWSSecretParameterModel to AWSSecretParameter",
aws_secret_parameter_id=aws_secret_parameter_model.id,
aws_secret_parameter_id=aws_secret_parameter_model.aws_secret_parameter_id,
)
return AWSSecretParameter(
@@ -203,6 +211,45 @@ def convert_to_aws_secret_parameter(
)
def convert_to_output_parameter(
output_parameter_model: OutputParameterModel, debug_enabled: bool = False
) -> OutputParameter:
if debug_enabled:
LOG.debug(
"Converting OutputParameterModel to OutputParameter",
output_parameter_id=output_parameter_model.output_parameter_id,
)
return OutputParameter(
output_parameter_id=output_parameter_model.output_parameter_id,
key=output_parameter_model.key,
description=output_parameter_model.description,
workflow_id=output_parameter_model.workflow_id,
created_at=output_parameter_model.created_at,
modified_at=output_parameter_model.modified_at,
deleted_at=output_parameter_model.deleted_at,
)
def convert_to_workflow_run_output_parameter(
workflow_run_output_parameter_model: WorkflowRunOutputParameterModel,
debug_enabled: bool = False,
) -> WorkflowRunOutputParameter:
if debug_enabled:
LOG.debug(
"Converting WorkflowRunOutputParameterModel to WorkflowRunOutputParameter",
workflow_run_id=workflow_run_output_parameter_model.workflow_run_id,
output_parameter_id=workflow_run_output_parameter_model.output_parameter_id,
)
return WorkflowRunOutputParameter(
workflow_run_id=workflow_run_output_parameter_model.workflow_run_id,
output_parameter_id=workflow_run_output_parameter_model.output_parameter_id,
value=workflow_run_output_parameter_model.value,
created_at=workflow_run_output_parameter_model.created_at,
)
def convert_to_workflow_run_parameter(
workflow_run_parameter_model: WorkflowRunParameterModel,
workflow_parameter: WorkflowParameter,