diff --git a/skyvern/forge/sdk/artifact/storage/base.py b/skyvern/forge/sdk/artifact/storage/base.py index 2a2554fb..c7f3f491 100644 --- a/skyvern/forge/sdk/artifact/storage/base.py +++ b/skyvern/forge/sdk/artifact/storage/base.py @@ -94,7 +94,7 @@ class BaseStorage(ABC): pass @abstractmethod - async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None: + async def get_streaming_file(self, organization_id: str, file_name: str) -> bytes | None: pass @abstractmethod diff --git a/skyvern/forge/sdk/artifact/storage/local.py b/skyvern/forge/sdk/artifact/storage/local.py index 4069f262..ea9e7639 100644 --- a/skyvern/forge/sdk/artifact/storage/local.py +++ b/skyvern/forge/sdk/artifact/storage/local.py @@ -116,10 +116,10 @@ class LocalStorage(BaseStorage): async def save_streaming_file(self, organization_id: str, file_name: str) -> None: return - async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None: - file_path = Path(f"{get_skyvern_temp_dir()}/skyvern_screenshot.png") - if not use_default: - file_path = Path(f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}") + async def get_streaming_file(self, organization_id: str, file_name: str) -> bytes | None: + # make the directory if it doesn't exist + Path(f"{get_skyvern_temp_dir()}/{organization_id}").mkdir(parents=True, exist_ok=True) + file_path = Path(f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}") try: with open(file_path, "rb") as f: return f.read() diff --git a/skyvern/utils/files.py b/skyvern/utils/files.py new file mode 100644 index 00000000..ce614540 --- /dev/null +++ b/skyvern/utils/files.py @@ -0,0 +1,29 @@ +import json +import os + +import aiofiles + +from skyvern.forge.sdk.api.files import get_skyvern_temp_dir + + +def get_skyvern_state_file_path() -> str: + return f"{get_skyvern_temp_dir()}/current.json" + + +async def initialize_skyvern_state_file( + task_id: str | None = None, workflow_run_id: str | None = None, organization_id: str | None = None +) -> None: + # create the file if it doesn't exist + async with aiofiles.open(get_skyvern_state_file_path(), "w") as json_file: + await json_file.write( + json.dumps({"task_id": task_id, "workflow_run_id": workflow_run_id, "organization_id": organization_id}) + ) + + +def get_json_from_file(file_path: str) -> dict[str, str]: + # check if file exists + if not os.path.exists(file_path): + return {} + + with open(file_path, "r") as json_file: + return json.load(json_file)