add workflow run block screenshot and observer thought screenshots (#1444)

skyvern/forge/sdk/artifact/manager.py

@@ -9,6 +9,7 @@ from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityT
 from skyvern.forge.sdk.db.id import generate_artifact_id
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
+from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
 
 LOG = structlog.get_logger(__name__)
 
@@ -151,6 +152,27 @@ class ArtifactManager:
             path=path,
         )
 
+    async def create_workflow_run_block_artifact(
+        self,
+        workflow_run_block: WorkflowRunBlock,
+        artifact_type: ArtifactType,
+        data: bytes | None = None,
+        path: str | None = None,
+    ) -> str:
+        artifact_id = generate_artifact_id()
+        uri = app.STORAGE.build_workflow_run_block_uri(artifact_id, workflow_run_block, artifact_type)
+        return await self._create_artifact(
+            aio_task_primary_key=workflow_run_block.workflow_run_block_id,
+            artifact_id=artifact_id,
+            artifact_type=artifact_type,
+            uri=uri,
+            workflow_run_block_id=workflow_run_block.workflow_run_block_id,
+            workflow_run_id=workflow_run_block.workflow_run_id,
+            organization_id=workflow_run_block.organization_id,
+            data=data,
+            path=path,
+        )
+
     async def create_llm_artifact(
         self,
         data: bytes,

skyvern/forge/sdk/artifact/storage/base.py

@@ -3,6 +3,7 @@ from abc import ABC, abstractmethod
 from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityType
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
+from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
 
 # TODO: This should be a part of the ArtifactType model
 FILE_EXTENTSION_MAP: dict[ArtifactType, str] = {
@@ -52,6 +53,12 @@ class BaseStorage(ABC):
     ) -> str:
         pass
 
+    @abstractmethod
+    def build_workflow_run_block_uri(
+        self, artifact_id: str, workflow_run_block: WorkflowRunBlock, artifact_type: ArtifactType
+    ) -> str:
+        pass
+
     @abstractmethod
     async def store_artifact(self, artifact: Artifact, data: bytes) -> None:
         pass

skyvern/forge/sdk/artifact/storage/local.py

@@ -12,6 +12,7 @@ from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityT
 from skyvern.forge.sdk.artifact.storage.base import FILE_EXTENTSION_MAP, BaseStorage
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
+from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
 
 LOG = structlog.get_logger()
 
@@ -40,6 +41,12 @@ class LocalStorage(BaseStorage):
         file_ext = FILE_EXTENTSION_MAP[artifact_type]
         return f"file://{self.artifact_path}/{settings.ENV}/observers/{observer_cruise.observer_cruise_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
 
+    def build_workflow_run_block_uri(
+        self, artifact_id: str, workflow_run_block: WorkflowRunBlock, artifact_type: ArtifactType
+    ) -> str:
+        file_ext = FILE_EXTENTSION_MAP[artifact_type]
+        return f"file://{self.artifact_path}/{settings.ENV}/workflow_runs/{workflow_run_block.workflow_run_id}/{workflow_run_block.workflow_run_block_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
+
     async def store_artifact(self, artifact: Artifact, data: bytes) -> None:
         file_path = None
         try:

skyvern/forge/sdk/artifact/storage/s3.py

@@ -16,6 +16,7 @@ from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityT
 from skyvern.forge.sdk.artifact.storage.base import FILE_EXTENTSION_MAP, BaseStorage
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
+from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
 
 
 class S3Storage(BaseStorage):
@@ -43,6 +44,12 @@ class S3Storage(BaseStorage):
         file_ext = FILE_EXTENTSION_MAP[artifact_type]
         return f"s3://{self.bucket}/{settings.ENV}/observers/{observer_cruise.observer_cruise_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
 
+    def build_workflow_run_block_uri(
+        self, artifact_id: str, workflow_run_block: WorkflowRunBlock, artifact_type: ArtifactType
+    ) -> str:
+        file_ext = FILE_EXTENTSION_MAP[artifact_type]
+        return f"s3://{self.bucket}/{settings.ENV}/workflow_runs/{workflow_run_block.workflow_run_id}/{workflow_run_block.workflow_run_block_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
+
     async def store_artifact(self, artifact: Artifact, data: bytes) -> None:
         await self.async_client.upload_file(artifact.uri, data)