From ff642f44ee1ff7b5379d72108ff0c383f0bc54f7 Mon Sep 17 00:00:00 2001 From: lawyzheng Date: Mon, 16 Dec 2024 11:19:50 +0800 Subject: [PATCH] add timezone support --- skyvern/forge/agent.py | 3 +- .../forge/prompts/skyvern/custom-select.j2 | 4 +- .../skyvern/decisive-criterion-validate.j2 | 4 +- .../forge/prompts/skyvern/extract-action.j2 | 4 +- .../prompts/skyvern/extract-information.j2 | 4 +- .../prompts/skyvern/single-click-action.j2 | 4 +- .../prompts/skyvern/single-input-action.j2 | 4 +- .../prompts/skyvern/single-select-action.j2 | 4 +- .../prompts/skyvern/single-upload-action.j2 | 4 +- skyvern/forge/sdk/core/skyvern_context.py | 2 + skyvern/forge/sdk/schemas/tasks.py | 41 +++++++++++++++++++ skyvern/webeye/actions/handler.py | 8 +++- skyvern/webeye/browser_factory.py | 28 +++++++++---- 13 files changed, 86 insertions(+), 28 deletions(-) diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index 5c9eb7f68..75586cfb2 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -1310,6 +1310,7 @@ async def _build_extract_action_prompt( if not template: raise UnsupportedTaskType(task_type=task_type) + context = skyvern_context.ensure_context() return prompt_engine.load_prompt( template=template, navigation_goal=navigation_goal, @@ -1320,7 +1321,7 @@ async def _build_extract_action_prompt( data_extraction_goal=task.data_extraction_goal, action_history=actions_and_results_str, error_code_mapping_str=(json.dumps(task.error_code_mapping) if task.error_code_mapping else None), - utc_datetime=datetime.utcnow().strftime("%Y-%m-%d %H:%M"), + local_datetime=datetime.now(context.tz_info).isoformat(), verification_code_check=verification_code_check, complete_criterion=task.complete_criterion, terminate_criterion=task.terminate_criterion, diff --git a/skyvern/forge/prompts/skyvern/custom-select.j2 b/skyvern/forge/prompts/skyvern/custom-select.j2 index 56e63071c..98cadabb6 100644 --- a/skyvern/forge/prompts/skyvern/custom-select.j2 +++ b/skyvern/forge/prompts/skyvern/custom-select.j2 @@ -51,7 +51,7 @@ Select History: {{ select_history }} ``` {% endif %} -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` \ No newline at end of file diff --git a/skyvern/forge/prompts/skyvern/decisive-criterion-validate.j2 b/skyvern/forge/prompts/skyvern/decisive-criterion-validate.j2 index 68978ad0c..ab939e189 100644 --- a/skyvern/forge/prompts/skyvern/decisive-criterion-validate.j2 +++ b/skyvern/forge/prompts/skyvern/decisive-criterion-validate.j2 @@ -49,7 +49,7 @@ User details: {{ navigation_payload_str }} ``` -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` diff --git a/skyvern/forge/prompts/skyvern/extract-action.j2 b/skyvern/forge/prompts/skyvern/extract-action.j2 index 011be699a..cb5809b1e 100644 --- a/skyvern/forge/prompts/skyvern/extract-action.j2 +++ b/skyvern/forge/prompts/skyvern/extract-action.j2 @@ -73,7 +73,7 @@ User details: Action results from previous steps: (note: even if the action history suggests goal is achieved, check the screenshot and the DOM elements to make sure the goal is achieved) {{ action_history }} {% endif %} -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` diff --git a/skyvern/forge/prompts/skyvern/extract-information.j2 b/skyvern/forge/prompts/skyvern/extract-information.j2 index 7eb64b1c4..b1d9a8ac6 100644 --- a/skyvern/forge/prompts/skyvern/extract-information.j2 +++ b/skyvern/forge/prompts/skyvern/extract-information.j2 @@ -26,7 +26,7 @@ Text extracted from the webpage: {{ extracted_text }} User Navigation Payload: {{ navigation_payload }} -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` \ No newline at end of file diff --git a/skyvern/forge/prompts/skyvern/single-click-action.j2 b/skyvern/forge/prompts/skyvern/single-click-action.j2 index d53d5ebfd..f9c0e18d7 100644 --- a/skyvern/forge/prompts/skyvern/single-click-action.j2 +++ b/skyvern/forge/prompts/skyvern/single-click-action.j2 @@ -35,7 +35,7 @@ User details: {{ navigation_payload_str }} ``` -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` diff --git a/skyvern/forge/prompts/skyvern/single-input-action.j2 b/skyvern/forge/prompts/skyvern/single-input-action.j2 index 7bff5867a..bcf6cec7f 100644 --- a/skyvern/forge/prompts/skyvern/single-input-action.j2 +++ b/skyvern/forge/prompts/skyvern/single-input-action.j2 @@ -37,7 +37,7 @@ User details: {{ navigation_payload_str }} ``` -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` diff --git a/skyvern/forge/prompts/skyvern/single-select-action.j2 b/skyvern/forge/prompts/skyvern/single-select-action.j2 index 34de0ff70..89ae123b0 100644 --- a/skyvern/forge/prompts/skyvern/single-select-action.j2 +++ b/skyvern/forge/prompts/skyvern/single-select-action.j2 @@ -39,7 +39,7 @@ User details: {{ navigation_payload_str }} ``` -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` diff --git a/skyvern/forge/prompts/skyvern/single-upload-action.j2 b/skyvern/forge/prompts/skyvern/single-upload-action.j2 index 9bbcf0344..af6c4cb5e 100644 --- a/skyvern/forge/prompts/skyvern/single-upload-action.j2 +++ b/skyvern/forge/prompts/skyvern/single-upload-action.j2 @@ -35,7 +35,7 @@ User details: {{ navigation_payload_str }} ``` -Current datetime in UTC, YYYY-MM-DD HH:MM format: +Current datetime, ISO format: ``` -{{ utc_datetime }} +{{ local_datetime }} ``` diff --git a/skyvern/forge/sdk/core/skyvern_context.py b/skyvern/forge/sdk/core/skyvern_context.py index 11a615fc8..2a6a29be2 100644 --- a/skyvern/forge/sdk/core/skyvern_context.py +++ b/skyvern/forge/sdk/core/skyvern_context.py @@ -1,5 +1,6 @@ from contextvars import ContextVar from dataclasses import dataclass, field +from zoneinfo import ZoneInfo @dataclass @@ -10,6 +11,7 @@ class SkyvernContext: workflow_id: str | None = None workflow_run_id: str | None = None max_steps_override: int | None = None + tz_info: ZoneInfo | None = None totp_codes: dict[str, str | None] = field(default_factory=dict) def __repr__(self) -> str: diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py index abe90b137..41fdece34 100644 --- a/skyvern/forge/sdk/schemas/tasks.py +++ b/skyvern/forge/sdk/schemas/tasks.py @@ -3,6 +3,7 @@ from datetime import datetime from enum import StrEnum from typing import Any +from zoneinfo import ZoneInfo from pydantic import BaseModel, Field, field_validator @@ -26,6 +27,46 @@ class ProxyLocation(StrEnum): 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/Kolkata") + + return None + + class TaskBase(BaseModel): title: str | None = Field( default=None, diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py index 70226da39..07f3ceb56 100644 --- a/skyvern/webeye/actions/handler.py +++ b/skyvern/webeye/actions/handler.py @@ -52,6 +52,7 @@ from skyvern.forge.sdk.api.files import download_file, get_download_dir, list_files_in_directory from skyvern.forge.sdk.core.aiohttp_helper import aiohttp_post from skyvern.forge.sdk.core.security import generate_skyvern_signature +from skyvern.forge.sdk.core.skyvern_context import ensure_context from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType from skyvern.forge.sdk.models import Step from skyvern.forge.sdk.schemas.tasks import Task @@ -1877,6 +1878,7 @@ async def select_from_dropdown( html = incremental_scraped.build_html_tree(element_tree=trimmed_element_tree) + skyvern_context = ensure_context() prompt = prompt_engine.load_prompt( "custom-select", field_information=context.field, @@ -1886,7 +1888,7 @@ async def select_from_dropdown( navigation_payload_str=json.dumps(task.navigation_payload), elements=html, select_history=json.dumps(build_sequential_select_history(select_history)) if select_history else "", - utc_datetime=datetime.utcnow().strftime("%Y-%m-%d %H:%M"), + local_datetime=datetime.now(skyvern_context.tz_info).isoformat(), ) LOG.info( @@ -2449,6 +2451,8 @@ async def extract_information_for_navigation_goal( element_tree_in_prompt: str = scraped_page.build_element_tree(element_tree_format) scraped_page_refreshed = await scraped_page.refresh() + + context = ensure_context() extract_information_prompt = prompt_engine.load_prompt( prompt_template, navigation_goal=task.navigation_goal, @@ -2459,7 +2463,7 @@ async def extract_information_for_navigation_goal( current_url=scraped_page_refreshed.url, extracted_text=scraped_page_refreshed.extracted_text, error_code_mapping_str=(json.dumps(task.error_code_mapping) if task.error_code_mapping else None), - utc_datetime=datetime.utcnow().strftime("%Y-%m-%d %H:%M"), + local_datetime=datetime.now(context.tz_info).isoformat(), ) json_response = await app.LLM_API_HANDLER( diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py index 9fed7b079..3de70e362 100644 --- a/skyvern/webeye/browser_factory.py +++ b/skyvern/webeye/browser_factory.py @@ -25,7 +25,7 @@ ) 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 +from skyvern.forge.sdk.schemas.tasks import ProxyLocation, get_tzinfo_from_proxy from skyvern.webeye.utils.page import SkyvernFrame LOG = structlog.get_logger() @@ -128,7 +128,7 @@ def initialize_download_dir() -> str: class BrowserContextCreator(Protocol): def __call__( - self, playwright: Playwright, **kwargs: dict[str, Any] + self, playwright: Playwright, proxy_location: ProxyLocation | None = None, **kwargs: dict[str, Any] ) -> Awaitable[tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]]: ... @@ -162,14 +162,13 @@ def update_chromium_browser_preferences(user_data_dir: str, download_dir: str) - f.write(preference_file_content) @staticmethod - def build_browser_args() -> dict[str, Any]: + def build_browser_args(proxy_location: ProxyLocation | None = None) -> dict[str, Any]: video_dir = f"{settings.VIDEO_PATH}/{datetime.utcnow().strftime('%Y-%m-%d')}" har_dir = ( f"{settings.HAR_PATH}/{datetime.utcnow().strftime('%Y-%m-%d')}/{BrowserContextFactory.get_subdir()}.har" ) - return { + args = { "locale": settings.BROWSER_LOCALE, - "timezone_id": settings.BROWSER_TIMEZONE, "color_scheme": "no-preference", "args": [ "--disable-blink-features=AutomationControlled", @@ -188,6 +187,11 @@ def build_browser_args() -> dict[str, Any]: }, } + if proxy_location: + if tz_info := get_tzinfo_from_proxy(proxy_location=proxy_location): + args["timezone_id"] = tz_info.key + return args + @staticmethod def build_browser_artifacts( video_artifacts: list[VideoArtifact] | None = None, @@ -221,6 +225,12 @@ async def create_browser_context( browser_context, browser_artifacts, cleanup_func = await creator(playwright, **kwargs) set_browser_console_log(browser_context=browser_context, browser_artifacts=browser_artifacts) set_download_file_listener(browser_context=browser_context, **kwargs) + + proxy_location: ProxyLocation | None = kwargs.get("proxy_location") + if proxy_location is not None: + context = ensure_context() + context.tz_info = get_tzinfo_from_proxy(proxy_location) + return browser_context, browser_artifacts, cleanup_func except Exception as e: if browser_context is not None: @@ -279,7 +289,7 @@ async def read_browser_console_log(self) -> bytes: async def _create_headless_chromium( - playwright: Playwright, **kwargs: dict + playwright: Playwright, proxy_location: ProxyLocation | None = None, **kwargs: dict ) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]: user_data_dir = make_temp_directory(prefix="skyvern_browser_") download_dir = initialize_download_dir() @@ -287,7 +297,7 @@ async def _create_headless_chromium( user_data_dir=user_data_dir, download_dir=download_dir, ) - browser_args = BrowserContextFactory.build_browser_args() + browser_args = BrowserContextFactory.build_browser_args(proxy_location=proxy_location) browser_args.update( { "user_data_dir": user_data_dir, @@ -301,7 +311,7 @@ async def _create_headless_chromium( async def _create_headful_chromium( - playwright: Playwright, **kwargs: dict + playwright: Playwright, proxy_location: ProxyLocation | None = None, **kwargs: dict ) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]: user_data_dir = make_temp_directory(prefix="skyvern_browser_") download_dir = initialize_download_dir() @@ -309,7 +319,7 @@ async def _create_headful_chromium( user_data_dir=user_data_dir, download_dir=download_dir, ) - browser_args = BrowserContextFactory.build_browser_args() + browser_args = BrowserContextFactory.build_browser_args(proxy_location=proxy_location) browser_args.update( { "user_data_dir": user_data_dir,