Return action screenshots from get task endpoint as well (#149)

Skyvern-AI · Apr 2, 2024 · 999eda9 · 999eda9
1 parent 69c458b
commit 999eda9
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
@@ -787,7 +787,7 @@ async def execute_task_webhook(self, task: Task, last_step: Step, api_key: str |
         if recording_artifact:
             recording_url = await app.ARTIFACT_MANAGER.get_share_link(recording_artifact)
 
-        # get the artifact of the last  screenshot and get the screenshot_url
+        # get the artifact of the last TASK_RESPONSE_ACTION_SCREENSHOT_COUNT screenshots and get the screenshot_url
         latest_action_screenshot_artifacts = await app.DATABASE.get_latest_n_artifacts(
             task_id=task.task_id,
             organization_id=task.organization_id,

diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -218,6 +218,27 @@ async def get_task(
     if recording_artifact:
         recording_url = await app.ARTIFACT_MANAGER.get_share_link(recording_artifact)
 
+    # get the artifact of the last  screenshot and get the screenshot_url
+    latest_action_screenshot_artifacts = await app.DATABASE.get_latest_n_artifacts(
+        task_id=task_obj.task_id,
+        organization_id=task_obj.organization_id,
+        artifact_types=[ArtifactType.SCREENSHOT_ACTION],
+        n=SettingsManager.get_settings().TASK_RESPONSE_ACTION_SCREENSHOT_COUNT,
+    )
+    latest_action_screenshot_urls = []
+    if latest_action_screenshot_artifacts:
+        for artifact in latest_action_screenshot_artifacts:
+            screenshot_url = await app.ARTIFACT_MANAGER.get_share_link(artifact)
+            if screenshot_url:
+                latest_action_screenshot_urls.append(screenshot_url)
+            else:
+                LOG.error(
+                    "Failed to get share link for action screenshot",
+                    artifact_id=artifact.artifact_id,
+                )
+    else:
+        LOG.error("Failed to get latest action screenshots")
+
     failure_reason = None
     if task_obj.status == TaskStatus.failed and (latest_step.output or task_obj.failure_reason):
         failure_reason = ""
@@ -231,6 +252,7 @@ async def get_task(
             )
 
     return task_obj.to_task_response(
+        action_screenshot_urls=latest_action_screenshot_urls,
         screenshot_url=screenshot_url,
         recording_url=recording_url,
         failure_reason=failure_reason,