Skip to content

Commit

Permalink
add timezone support
Browse files Browse the repository at this point in the history
  • Loading branch information
LawyZheng committed Dec 16, 2024
1 parent bc57dc1 commit ff642f4
Show file tree
Hide file tree
Showing 13 changed files with 86 additions and 28 deletions.
3 changes: 2 additions & 1 deletion skyvern/forge/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/custom-select.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/decisive-criterion-validate.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/extract-action.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/extract-information.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/single-click-action.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/single-input-action.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/single-select-action.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
4 changes: 2 additions & 2 deletions skyvern/forge/prompts/skyvern/single-upload-action.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
```
2 changes: 2 additions & 0 deletions skyvern/forge/sdk/core/skyvern_context.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from contextvars import ContextVar
from dataclasses import dataclass, field
from zoneinfo import ZoneInfo


@dataclass
Expand All @@ -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:
Expand Down
41 changes: 41 additions & 0 deletions skyvern/forge/sdk/schemas/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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,
Expand Down
8 changes: 6 additions & 2 deletions skyvern/webeye/actions/handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down
28 changes: 19 additions & 9 deletions skyvern/webeye/browser_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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]]: ...


Expand Down Expand Up @@ -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",
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -279,15 +289,15 @@ 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()
BrowserContextFactory.update_chromium_browser_preferences(
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,
Expand All @@ -301,15 +311,15 @@ 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()
BrowserContextFactory.update_chromium_browser_preferences(
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,
Expand Down

0 comments on commit ff642f4

Please sign in to comment.