From 19c7c56af739596ed21c513b40be2d9e88822462 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Mon, 24 Mar 2025 15:15:21 -0700 Subject: [PATCH] migrate ProxyLocation to skyvern.schemas.runs (#2011) --- evaluation/core/__init__.py | 3 +- .../script/create_webvoyager_workflow.py | 2 +- skyvern/agent/client.py | 2 +- skyvern/forge/sdk/db/client.py | 3 +- skyvern/forge/sdk/db/utils.py | 3 +- skyvern/forge/sdk/schemas/task_runs.py | 2 +- skyvern/forge/sdk/schemas/task_v2.py | 2 +- skyvern/forge/sdk/schemas/tasks.py | 81 +------------------ skyvern/forge/sdk/services/task_v2_service.py | 2 +- skyvern/forge/sdk/workflow/models/workflow.py | 2 +- skyvern/forge/sdk/workflow/models/yaml.py | 2 +- skyvern/forge/sdk/workflow/service.py | 3 +- skyvern/schemas/runs.py | 81 +++++++++++++++++++ skyvern/webeye/browser_factory.py | 2 +- skyvern/webeye/browser_manager.py | 3 +- 15 files changed, 100 insertions(+), 93 deletions(-) create mode 100644 skyvern/schemas/runs.py diff --git a/evaluation/core/__init__.py b/evaluation/core/__init__.py index 90d498ae..ac72c71d 100644 --- a/evaluation/core/__init__.py +++ b/evaluation/core/__init__.py @@ -13,8 +13,9 @@ from skyvern.forge import app from skyvern.forge.prompts import prompt_engine from skyvern.forge.sdk.api.files import create_folder_if_not_exist from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Request -from skyvern.forge.sdk.schemas.tasks import ProxyLocation, TaskRequest, TaskResponse, TaskStatus +from skyvern.forge.sdk.schemas.tasks import TaskRequest, TaskResponse, TaskStatus from skyvern.forge.sdk.workflow.models.workflow import WorkflowRequestBody, WorkflowRunResponse, WorkflowRunStatus +from skyvern.schemas.runs import ProxyLocation class TaskOutput(BaseModel): diff --git a/evaluation/script/create_webvoyager_workflow.py b/evaluation/script/create_webvoyager_workflow.py index d25130f2..577c0f78 100644 --- a/evaluation/script/create_webvoyager_workflow.py +++ b/evaluation/script/create_webvoyager_workflow.py @@ -11,8 +11,8 @@ from evaluation.core import Evaluator, SkyvernClient from evaluation.core.utils import load_webvoyager_case_from_json from skyvern.forge import app from skyvern.forge.prompts import prompt_engine -from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.workflow.models.workflow import WorkflowRequestBody +from skyvern.schemas.runs import ProxyLocation load_dotenv() diff --git a/skyvern/agent/client.py b/skyvern/agent/client.py index f6f2c795..844597eb 100644 --- a/skyvern/agent/client.py +++ b/skyvern/agent/client.py @@ -6,8 +6,8 @@ import httpx from skyvern.config import settings from skyvern.exceptions import SkyvernClientException from skyvern.forge.sdk.schemas.task_runs import TaskRunResponse -from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.workflow.models.workflow import RunWorkflowResponse, WorkflowRunResponse +from skyvern.schemas.runs import ProxyLocation class RunEngine(StrEnum): diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 1fda6faa..5e397752 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -69,7 +69,7 @@ from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrow from skyvern.forge.sdk.schemas.task_generations import TaskGeneration from skyvern.forge.sdk.schemas.task_runs import TaskRun, TaskRunType from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Status, Thought, ThoughtType -from skyvern.forge.sdk.schemas.tasks import OrderBy, ProxyLocation, SortDirection, Task, TaskStatus +from skyvern.forge.sdk.schemas.tasks import OrderBy, SortDirection, Task, TaskStatus from skyvern.forge.sdk.schemas.totp_codes import TOTPCode from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock from skyvern.forge.sdk.workflow.models.block import BlockStatus, BlockType @@ -91,6 +91,7 @@ from skyvern.forge.sdk.workflow.models.workflow import ( WorkflowRunStatus, WorkflowStatus, ) +from skyvern.schemas.runs import ProxyLocation from skyvern.webeye.actions.actions import Action from skyvern.webeye.actions.models import AgentStepOutput diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index 7596c7ae..7c331644 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -25,7 +25,7 @@ from skyvern.forge.sdk.db.models import ( ) from skyvern.forge.sdk.models import Step, StepStatus from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken -from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus +from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock from skyvern.forge.sdk.workflow.models.block import BlockStatus, BlockType from skyvern.forge.sdk.workflow.models.parameter import ( @@ -45,6 +45,7 @@ from skyvern.forge.sdk.workflow.models.workflow import ( WorkflowRunStatus, WorkflowStatus, ) +from skyvern.schemas.runs import ProxyLocation LOG = structlog.get_logger() diff --git a/skyvern/forge/sdk/schemas/task_runs.py b/skyvern/forge/sdk/schemas/task_runs.py index 216046c1..8273f87e 100644 --- a/skyvern/forge/sdk/schemas/task_runs.py +++ b/skyvern/forge/sdk/schemas/task_runs.py @@ -3,7 +3,7 @@ from enum import StrEnum from pydantic import BaseModel, ConfigDict -from skyvern.forge.sdk.schemas.tasks import ProxyLocation +from skyvern.schemas.runs import ProxyLocation class TaskRunStatus(StrEnum): diff --git a/skyvern/forge/sdk/schemas/task_v2.py b/skyvern/forge/sdk/schemas/task_v2.py index f1162e7f..c2f4d24f 100644 --- a/skyvern/forge/sdk/schemas/task_v2.py +++ b/skyvern/forge/sdk/schemas/task_v2.py @@ -5,7 +5,7 @@ from typing import Any from pydantic import BaseModel, ConfigDict, Field, field_validator from skyvern.forge.sdk.core.validators import validate_url -from skyvern.forge.sdk.schemas.tasks import ProxyLocation +from skyvern.schemas.runs import ProxyLocation DEFAULT_WORKFLOW_TITLE = "New Workflow" diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py index aa1e4d13..ad5324a7 100644 --- a/skyvern/forge/sdk/schemas/tasks.py +++ b/skyvern/forge/sdk/schemas/tasks.py @@ -3,7 +3,6 @@ from __future__ import annotations from datetime import datetime from enum import StrEnum from typing import Any -from zoneinfo import ZoneInfo from pydantic import BaseModel, Field, field_validator @@ -11,85 +10,7 @@ from skyvern.exceptions import InvalidTaskStatusTransition, TaskAlreadyCanceled, from skyvern.forge.sdk.core.validators import validate_url from skyvern.forge.sdk.db.enums import TaskType from skyvern.forge.sdk.schemas.files import FileInfo - - -class ProxyLocation(StrEnum): - US_CA = "US-CA" - US_NY = "US-NY" - US_TX = "US-TX" - US_FL = "US-FL" - US_WA = "US-WA" - RESIDENTIAL = "RESIDENTIAL" - RESIDENTIAL_ES = "RESIDENTIAL_ES" - RESIDENTIAL_IE = "RESIDENTIAL_IE" - RESIDENTIAL_GB = "RESIDENTIAL_GB" - RESIDENTIAL_IN = "RESIDENTIAL_IN" - RESIDENTIAL_JP = "RESIDENTIAL_JP" - RESIDENTIAL_FR = "RESIDENTIAL_FR" - RESIDENTIAL_DE = "RESIDENTIAL_DE" - RESIDENTIAL_NZ = "RESIDENTIAL_NZ" - RESIDENTIAL_ZA = "RESIDENTIAL_ZA" - RESIDENTIAL_AR = "RESIDENTIAL_AR" - RESIDENTIAL_ISP = "RESIDENTIAL_ISP" - NONE = "NONE" - - -def get_tzinfo_from_proxy(proxy_location: ProxyLocation) -> ZoneInfo | None: - if proxy_location == ProxyLocation.NONE: - return None - - if proxy_location == ProxyLocation.US_CA: - return ZoneInfo("America/Los_Angeles") - - if proxy_location == ProxyLocation.US_NY: - return ZoneInfo("America/New_York") - - if proxy_location == ProxyLocation.US_TX: - return ZoneInfo("America/Chicago") - - if proxy_location == ProxyLocation.US_FL: - return ZoneInfo("America/New_York") - - if proxy_location == ProxyLocation.US_WA: - return ZoneInfo("America/New_York") - - if proxy_location == ProxyLocation.RESIDENTIAL: - return ZoneInfo("America/New_York") - - if proxy_location == ProxyLocation.RESIDENTIAL_ES: - return ZoneInfo("Europe/Madrid") - - if proxy_location == ProxyLocation.RESIDENTIAL_IE: - return ZoneInfo("Europe/Dublin") - - if proxy_location == ProxyLocation.RESIDENTIAL_GB: - return ZoneInfo("Europe/London") - - if proxy_location == ProxyLocation.RESIDENTIAL_IN: - return ZoneInfo("Asia/Kolkata") - - if proxy_location == ProxyLocation.RESIDENTIAL_JP: - return ZoneInfo("Asia/Tokyo") - - if proxy_location == ProxyLocation.RESIDENTIAL_FR: - return ZoneInfo("Europe/Paris") - - if proxy_location == ProxyLocation.RESIDENTIAL_DE: - return ZoneInfo("Europe/Berlin") - - if proxy_location == ProxyLocation.RESIDENTIAL_NZ: - return ZoneInfo("Pacific/Auckland") - - if proxy_location == ProxyLocation.RESIDENTIAL_ZA: - return ZoneInfo("Africa/Johannesburg") - - if proxy_location == ProxyLocation.RESIDENTIAL_AR: - return ZoneInfo("America/Argentina/Buenos_Aires") - - if proxy_location == ProxyLocation.RESIDENTIAL_ISP: - return ZoneInfo("America/New_York") - - return None +from skyvern.schemas.runs import ProxyLocation class TaskBase(BaseModel): diff --git a/skyvern/forge/sdk/services/task_v2_service.py b/skyvern/forge/sdk/services/task_v2_service.py index 91709696..eedef7b5 100644 --- a/skyvern/forge/sdk/services/task_v2_service.py +++ b/skyvern/forge/sdk/services/task_v2_service.py @@ -22,7 +22,6 @@ from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.task_runs import TaskRunType from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Metadata, TaskV2Status, ThoughtScenario, ThoughtType -from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunTimeline, WorkflowRunTimelineType from skyvern.forge.sdk.workflow.models.block import ( BlockResult, @@ -54,6 +53,7 @@ from skyvern.forge.sdk.workflow.models.yaml import ( WorkflowCreateYAMLRequest, WorkflowDefinitionYAML, ) +from skyvern.schemas.runs import ProxyLocation from skyvern.webeye.browser_factory import BrowserState from skyvern.webeye.scraper.scraper import ElementTreeFormat, ScrapedPage, scrape_website from skyvern.webeye.utils.page import SkyvernFrame diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py index 4efeb48b..80afc2f2 100644 --- a/skyvern/forge/sdk/workflow/models/workflow.py +++ b/skyvern/forge/sdk/workflow/models/workflow.py @@ -7,10 +7,10 @@ from pydantic import BaseModel, field_validator from skyvern.forge.sdk.core.validators import validate_url from skyvern.forge.sdk.schemas.files import FileInfo from skyvern.forge.sdk.schemas.task_v2 import TaskV2 -from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.workflow.exceptions import WorkflowDefinitionHasDuplicateBlockLabels from skyvern.forge.sdk.workflow.models.block import BlockTypeVar from skyvern.forge.sdk.workflow.models.parameter import PARAMETER_TYPE +from skyvern.schemas.runs import ProxyLocation class WorkflowRequestBody(BaseModel): diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py index 2dd42b8f..7212d74c 100644 --- a/skyvern/forge/sdk/workflow/models/yaml.py +++ b/skyvern/forge/sdk/workflow/models/yaml.py @@ -4,11 +4,11 @@ from typing import Annotated, Any, Literal from pydantic import BaseModel, Field from skyvern.config import settings -from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.workflow.models.block import BlockType, FileType from skyvern.forge.sdk.workflow.models.constants import FileStorageType from skyvern.forge.sdk.workflow.models.parameter import ParameterType, WorkflowParameterType from skyvern.forge.sdk.workflow.models.workflow import WorkflowStatus +from skyvern.schemas.runs import ProxyLocation class ParameterYAML(BaseModel, abc.ABC): diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 5cd3fd56..34fc2fd2 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -25,7 +25,7 @@ from skyvern.forge.sdk.db.enums import TaskType from skyvern.forge.sdk.models import Step, StepStatus from skyvern.forge.sdk.schemas.files import FileInfo from skyvern.forge.sdk.schemas.organizations import Organization -from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task +from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock, WorkflowRunTimeline, WorkflowRunTimelineType from skyvern.forge.sdk.workflow.exceptions import ( ContextParameterSourceNotDefined, @@ -90,6 +90,7 @@ from skyvern.forge.sdk.workflow.models.yaml import ( WorkflowCreateYAMLRequest, WorkflowDefinitionYAML, ) +from skyvern.schemas.runs import ProxyLocation from skyvern.webeye.browser_factory import BrowserState LOG = structlog.get_logger() diff --git a/skyvern/schemas/runs.py b/skyvern/schemas/runs.py new file mode 100644 index 00000000..3bb718fb --- /dev/null +++ b/skyvern/schemas/runs.py @@ -0,0 +1,81 @@ +from enum import StrEnum +from zoneinfo import ZoneInfo + + +class ProxyLocation(StrEnum): + US_CA = "US-CA" + US_NY = "US-NY" + US_TX = "US-TX" + US_FL = "US-FL" + US_WA = "US-WA" + RESIDENTIAL = "RESIDENTIAL" + RESIDENTIAL_ES = "RESIDENTIAL_ES" + RESIDENTIAL_IE = "RESIDENTIAL_IE" + RESIDENTIAL_GB = "RESIDENTIAL_GB" + RESIDENTIAL_IN = "RESIDENTIAL_IN" + RESIDENTIAL_JP = "RESIDENTIAL_JP" + RESIDENTIAL_FR = "RESIDENTIAL_FR" + RESIDENTIAL_DE = "RESIDENTIAL_DE" + RESIDENTIAL_NZ = "RESIDENTIAL_NZ" + RESIDENTIAL_ZA = "RESIDENTIAL_ZA" + RESIDENTIAL_AR = "RESIDENTIAL_AR" + RESIDENTIAL_ISP = "RESIDENTIAL_ISP" + NONE = "NONE" + + +def get_tzinfo_from_proxy(proxy_location: ProxyLocation) -> ZoneInfo | None: + if proxy_location == ProxyLocation.NONE: + return None + + if proxy_location == ProxyLocation.US_CA: + return ZoneInfo("America/Los_Angeles") + + if proxy_location == ProxyLocation.US_NY: + return ZoneInfo("America/New_York") + + if proxy_location == ProxyLocation.US_TX: + return ZoneInfo("America/Chicago") + + if proxy_location == ProxyLocation.US_FL: + return ZoneInfo("America/New_York") + + if proxy_location == ProxyLocation.US_WA: + return ZoneInfo("America/New_York") + + if proxy_location == ProxyLocation.RESIDENTIAL: + return ZoneInfo("America/New_York") + + if proxy_location == ProxyLocation.RESIDENTIAL_ES: + return ZoneInfo("Europe/Madrid") + + if proxy_location == ProxyLocation.RESIDENTIAL_IE: + return ZoneInfo("Europe/Dublin") + + if proxy_location == ProxyLocation.RESIDENTIAL_GB: + return ZoneInfo("Europe/London") + + if proxy_location == ProxyLocation.RESIDENTIAL_IN: + return ZoneInfo("Asia/Kolkata") + + if proxy_location == ProxyLocation.RESIDENTIAL_JP: + return ZoneInfo("Asia/Tokyo") + + if proxy_location == ProxyLocation.RESIDENTIAL_FR: + return ZoneInfo("Europe/Paris") + + if proxy_location == ProxyLocation.RESIDENTIAL_DE: + return ZoneInfo("Europe/Berlin") + + if proxy_location == ProxyLocation.RESIDENTIAL_NZ: + return ZoneInfo("Pacific/Auckland") + + if proxy_location == ProxyLocation.RESIDENTIAL_ZA: + return ZoneInfo("Africa/Johannesburg") + + if proxy_location == ProxyLocation.RESIDENTIAL_AR: + return ZoneInfo("America/Argentina/Buenos_Aires") + + if proxy_location == ProxyLocation.RESIDENTIAL_ISP: + return ZoneInfo("America/New_York") + + return None diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py index 4945bdbb..66db458d 100644 --- a/skyvern/webeye/browser_factory.py +++ b/skyvern/webeye/browser_factory.py @@ -26,7 +26,7 @@ from skyvern.exceptions import ( ) from skyvern.forge.sdk.api.files import get_download_dir, make_temp_directory from skyvern.forge.sdk.core.skyvern_context import current, ensure_context -from skyvern.forge.sdk.schemas.tasks import ProxyLocation, get_tzinfo_from_proxy +from skyvern.schemas.runs import ProxyLocation, get_tzinfo_from_proxy from skyvern.webeye.utils.page import SkyvernFrame LOG = structlog.get_logger() diff --git a/skyvern/webeye/browser_manager.py b/skyvern/webeye/browser_manager.py index 0a733010..6e2e02ed 100644 --- a/skyvern/webeye/browser_manager.py +++ b/skyvern/webeye/browser_manager.py @@ -7,8 +7,9 @@ from playwright.async_api import async_playwright from skyvern.exceptions import MissingBrowserState from skyvern.forge import app -from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task +from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.workflow.models.workflow import WorkflowRun +from skyvern.schemas.runs import ProxyLocation from skyvern.webeye.browser_factory import BrowserContextFactory, BrowserState, VideoArtifact LOG = structlog.get_logger()