mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2025-09-10 23:44:36 +00:00
Move the code over from private repository (#3)
This commit is contained in:
parent
32dd6d92a5
commit
9eddb3d812
93 changed files with 16798 additions and 0 deletions
220
skyvern/forge/sdk/db/utils.py
Normal file
220
skyvern/forge/sdk/db/utils.py
Normal file
|
@ -0,0 +1,220 @@
|
|||
import json
|
||||
import typing
|
||||
|
||||
import pydantic.json
|
||||
import structlog
|
||||
|
||||
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType
|
||||
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
|
||||
from skyvern.forge.sdk.db.models import (
|
||||
ArtifactModel,
|
||||
AWSSecretParameterModel,
|
||||
OrganizationAuthTokenModel,
|
||||
OrganizationModel,
|
||||
StepModel,
|
||||
TaskModel,
|
||||
WorkflowModel,
|
||||
WorkflowParameterModel,
|
||||
WorkflowRunModel,
|
||||
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.workflow import (
|
||||
Workflow,
|
||||
WorkflowDefinition,
|
||||
WorkflowRun,
|
||||
WorkflowRunParameter,
|
||||
WorkflowRunStatus,
|
||||
)
|
||||
|
||||
LOG = structlog.get_logger()
|
||||
|
||||
|
||||
@typing.no_type_check
|
||||
def _custom_json_serializer(*args, **kwargs) -> str:
|
||||
"""
|
||||
Encodes json in the same way that pydantic does.
|
||||
"""
|
||||
return json.dumps(*args, default=pydantic.json.pydantic_encoder, **kwargs)
|
||||
|
||||
|
||||
def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False) -> Task:
|
||||
if debug_enabled:
|
||||
LOG.debug("Converting TaskModel to Task", task_id=task_obj.task_id)
|
||||
task = Task(
|
||||
task_id=task_obj.task_id,
|
||||
status=TaskStatus(task_obj.status),
|
||||
created_at=task_obj.created_at,
|
||||
modified_at=task_obj.modified_at,
|
||||
url=task_obj.url,
|
||||
webhook_callback_url=task_obj.webhook_callback_url,
|
||||
navigation_goal=task_obj.navigation_goal,
|
||||
data_extraction_goal=task_obj.data_extraction_goal,
|
||||
navigation_payload=task_obj.navigation_payload,
|
||||
extracted_information=task_obj.extracted_information,
|
||||
failure_reason=task_obj.failure_reason,
|
||||
organization_id=task_obj.organization_id,
|
||||
proxy_location=ProxyLocation(task_obj.proxy_location) if task_obj.proxy_location else None,
|
||||
extracted_information_schema=task_obj.extracted_information_schema,
|
||||
workflow_run_id=task_obj.workflow_run_id,
|
||||
order=task_obj.order,
|
||||
retry=task_obj.retry,
|
||||
)
|
||||
return task
|
||||
|
||||
|
||||
def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step:
|
||||
if debug_enabled:
|
||||
LOG.debug("Converting StepModel to Step", step_id=step_model.step_id)
|
||||
return Step(
|
||||
task_id=step_model.task_id,
|
||||
step_id=step_model.step_id,
|
||||
created_at=step_model.created_at,
|
||||
modified_at=step_model.modified_at,
|
||||
status=StepStatus(step_model.status),
|
||||
output=step_model.output,
|
||||
order=step_model.order,
|
||||
is_last=step_model.is_last,
|
||||
retry_index=step_model.retry_index,
|
||||
organization_id=step_model.organization_id,
|
||||
input_token_count=step_model.input_token_count,
|
||||
output_token_count=step_model.output_token_count,
|
||||
step_cost=step_model.step_cost,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_organization(org_model: OrganizationModel) -> Organization:
|
||||
return Organization(
|
||||
organization_id=org_model.organization_id,
|
||||
organization_name=org_model.organization_name,
|
||||
webhook_callback_url=org_model.webhook_callback_url,
|
||||
max_steps_per_run=org_model.max_steps_per_run,
|
||||
created_at=org_model.created_at,
|
||||
modified_at=org_model.modified_at,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_organization_auth_token(org_auth_token: OrganizationAuthTokenModel) -> OrganizationAuthToken:
|
||||
return OrganizationAuthToken(
|
||||
id=org_auth_token.id,
|
||||
organization_id=org_auth_token.organization_id,
|
||||
token_type=OrganizationAuthTokenType(org_auth_token.token_type),
|
||||
token=org_auth_token.token,
|
||||
valid=org_auth_token.valid,
|
||||
created_at=org_auth_token.created_at,
|
||||
modified_at=org_auth_token.modified_at,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_artifact(artifact_model: ArtifactModel, debug_enabled: bool = False) -> Artifact:
|
||||
if debug_enabled:
|
||||
LOG.debug("Converting ArtifactModel to Artifact", artifact_id=artifact_model.artifact_id)
|
||||
|
||||
return Artifact(
|
||||
artifact_id=artifact_model.artifact_id,
|
||||
artifact_type=ArtifactType[artifact_model.artifact_type.upper()],
|
||||
uri=artifact_model.uri,
|
||||
task_id=artifact_model.task_id,
|
||||
step_id=artifact_model.step_id,
|
||||
created_at=artifact_model.created_at,
|
||||
modified_at=artifact_model.modified_at,
|
||||
organization_id=artifact_model.organization_id,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = False) -> Workflow:
|
||||
if debug_enabled:
|
||||
LOG.debug("Converting WorkflowModel to Workflow", workflow_id=workflow_model.workflow_id)
|
||||
|
||||
return Workflow(
|
||||
workflow_id=workflow_model.workflow_id,
|
||||
organization_id=workflow_model.organization_id,
|
||||
title=workflow_model.title,
|
||||
description=workflow_model.description,
|
||||
workflow_definition=WorkflowDefinition.model_validate(workflow_model.workflow_definition),
|
||||
created_at=workflow_model.created_at,
|
||||
modified_at=workflow_model.modified_at,
|
||||
deleted_at=workflow_model.deleted_at,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_workflow_run(workflow_run_model: WorkflowRunModel, debug_enabled: bool = False) -> WorkflowRun:
|
||||
if debug_enabled:
|
||||
LOG.debug("Converting WorkflowRunModel to WorkflowRun", workflow_run_id=workflow_run_model.workflow_run_id)
|
||||
|
||||
return WorkflowRun(
|
||||
workflow_run_id=workflow_run_model.workflow_run_id,
|
||||
workflow_id=workflow_run_model.workflow_id,
|
||||
status=WorkflowRunStatus[workflow_run_model.status],
|
||||
proxy_location=ProxyLocation(workflow_run_model.proxy_location) if workflow_run_model.proxy_location else None,
|
||||
webhook_callback_url=workflow_run_model.webhook_callback_url,
|
||||
created_at=workflow_run_model.created_at,
|
||||
modified_at=workflow_run_model.modified_at,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_workflow_parameter(
|
||||
workflow_parameter_model: WorkflowParameterModel, debug_enabled: bool = False
|
||||
) -> WorkflowParameter:
|
||||
if debug_enabled:
|
||||
LOG.debug(
|
||||
"Converting WorkflowParameterModel to WorkflowParameter",
|
||||
workflow_parameter_id=workflow_parameter_model.workflow_parameter_id,
|
||||
)
|
||||
|
||||
workflow_parameter_type = WorkflowParameterType[workflow_parameter_model.workflow_parameter_type.upper()]
|
||||
|
||||
return WorkflowParameter(
|
||||
workflow_parameter_id=workflow_parameter_model.workflow_parameter_id,
|
||||
workflow_parameter_type=workflow_parameter_type,
|
||||
workflow_id=workflow_parameter_model.workflow_id,
|
||||
default_value=workflow_parameter_type.convert_value(workflow_parameter_model.default_value),
|
||||
key=workflow_parameter_model.key,
|
||||
description=workflow_parameter_model.description,
|
||||
created_at=workflow_parameter_model.created_at,
|
||||
modified_at=workflow_parameter_model.modified_at,
|
||||
deleted_at=workflow_parameter_model.deleted_at,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_aws_secret_parameter(
|
||||
aws_secret_parameter_model: AWSSecretParameterModel, debug_enabled: bool = False
|
||||
) -> AWSSecretParameter:
|
||||
if debug_enabled:
|
||||
LOG.debug(
|
||||
"Converting AWSSecretParameterModel to AWSSecretParameter",
|
||||
aws_secret_parameter_id=aws_secret_parameter_model.id,
|
||||
)
|
||||
|
||||
return AWSSecretParameter(
|
||||
aws_secret_parameter_id=aws_secret_parameter_model.aws_secret_parameter_id,
|
||||
workflow_id=aws_secret_parameter_model.workflow_id,
|
||||
key=aws_secret_parameter_model.key,
|
||||
description=aws_secret_parameter_model.description,
|
||||
aws_key=aws_secret_parameter_model.aws_key,
|
||||
created_at=aws_secret_parameter_model.created_at,
|
||||
modified_at=aws_secret_parameter_model.modified_at,
|
||||
deleted_at=aws_secret_parameter_model.deleted_at,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_workflow_run_parameter(
|
||||
workflow_run_parameter_model: WorkflowRunParameterModel,
|
||||
workflow_parameter: WorkflowParameter,
|
||||
debug_enabled: bool = False,
|
||||
) -> WorkflowRunParameter:
|
||||
if debug_enabled:
|
||||
LOG.debug(
|
||||
"Converting WorkflowRunParameterModel to WorkflowRunParameter",
|
||||
workflow_run_id=workflow_run_parameter_model.workflow_run_id,
|
||||
workflow_parameter_id=workflow_run_parameter_model.workflow_parameter_id,
|
||||
)
|
||||
|
||||
return WorkflowRunParameter(
|
||||
workflow_run_id=workflow_run_parameter_model.workflow_run_id,
|
||||
workflow_parameter_id=workflow_run_parameter_model.workflow_parameter_id,
|
||||
value=workflow_parameter.workflow_parameter_type.convert_value(workflow_run_parameter_model.value),
|
||||
created_at=workflow_run_parameter_model.created_at,
|
||||
)
|
Loading…
Add table
Add a link
Reference in a new issue