Implement get_latest_screenshots, add action_screenshots to TaskResponse (#148)

skyvern/forge/sdk/db/client.py

@@ -612,6 +612,32 @@ class AgentDB:
         artifact_types: list[ArtifactType] | None = None,
         organization_id: str | None = None,
     ) -> Artifact | None:
+        try:
+            artifacts = await self.get_latest_n_artifacts(
+                task_id=task_id,
+                step_id=step_id,
+                artifact_types=artifact_types,
+                organization_id=organization_id,
+                n=1,
+            )
+            if artifacts:
+                return artifacts[0]
+            return None
+        except SQLAlchemyError:
+            LOG.exception("SQLAlchemyError", exc_info=True)
+            raise
+        except Exception:
+            LOG.exception("UnexpectedError", exc_info=True)
+            raise
+
+    async def get_latest_n_artifacts(
+        self,
+        task_id: str,
+        step_id: str | None = None,
+        artifact_types: list[ArtifactType] | None = None,
+        organization_id: str | None = None,
+        n: int = 1,
+    ) -> list[Artifact] | None:
         try:
             async with self.Session() as session:
                 artifact_query = select(ArtifactModel).filter_by(task_id=task_id)
@@ -622,9 +648,11 @@ class AgentDB:
                 if artifact_types:
                     artifact_query = artifact_query.filter(ArtifactModel.artifact_type.in_(artifact_types))
 
-                artifact = (await session.scalars(artifact_query.order_by(ArtifactModel.created_at.desc()))).first()
-                if artifact:
-                    return convert_to_artifact(artifact, self.debug_enabled)
+                artifacts = (await session.scalars(artifact_query.order_by(ArtifactModel.created_at.desc()))).fetchmany(
+                    n
+                )
+                if artifacts:
+                    return [convert_to_artifact(artifact, self.debug_enabled) for artifact in artifacts]
                 return None
         except SQLAlchemyError:
             LOG.exception("SQLAlchemyError", exc_info=True)

skyvern/forge/sdk/schemas/tasks.py

@@ -166,7 +166,11 @@ class Task(TaskRequest):
             raise ValueError(f"cant_override_failure_reason({self.task_id})")
 
     def to_task_response(
-        self, screenshot_url: str | None = None, recording_url: str | None = None, failure_reason: str | None = None
+        self,
+        action_screenshot_urls: list[str] | None = None,
+        screenshot_url: str | None = None,
+        recording_url: str | None = None,
+        failure_reason: str | None = None,
     ) -> TaskResponse:
         return TaskResponse(
             request=self,
@@ -176,6 +180,7 @@ class Task(TaskRequest):
             modified_at=self.modified_at,
             extracted_information=self.extracted_information,
             failure_reason=failure_reason or self.failure_reason,
+            action_screenshot_urls=action_screenshot_urls,
             screenshot_url=screenshot_url,
             recording_url=recording_url,
             errors=self.errors,
@@ -189,6 +194,7 @@ class TaskResponse(BaseModel):
     created_at: datetime
     modified_at: datetime
     extracted_information: list | dict[str, Any] | str | None = None
+    action_screenshot_urls: list[str] | None = None
     screenshot_url: str | None = None
     recording_url: str | None = None
     failure_reason: str | None = None