link actions to their screenshots - backend (#4404)

skyvern/forge/sdk/db/agent_db.py

@@ -3871,6 +3871,7 @@ class AgentDB(BaseAlchemyDB):
                 element_id=action.element_id,
                 skyvern_element_hash=action.skyvern_element_hash,
                 skyvern_element_data=action.skyvern_element_data,
+                screenshot_artifact_id=action.screenshot_artifact_id,
                 action_json=action.model_dump(),
                 confidence_float=action.confidence_float,
                 created_by=action.created_by,
@@ -3880,6 +3881,17 @@ class AgentDB(BaseAlchemyDB):
             await session.refresh(new_action)
             return Action.model_validate(new_action)
 
+    async def update_action_screenshot_artifact_id(
+        self, *, organization_id: str, action_id: str, screenshot_artifact_id: str
+    ) -> None:
+        async with self.Session() as session:
+            await session.execute(
+                update(ActionModel)
+                .where(ActionModel.action_id == action_id, ActionModel.organization_id == organization_id)
+                .values(screenshot_artifact_id=screenshot_artifact_id)
+            )
+            await session.commit()
+
     async def update_action_reasoning(
         self,
         organization_id: str,

skyvern/forge/sdk/db/models.py

@@ -680,6 +680,7 @@ class ActionModel(Base):
     action_json = Column(JSON, nullable=True)
     input_or_select_context = Column(JSON, nullable=True)
     confidence_float = Column(Numeric, nullable=True)
+    screenshot_artifact_id = Column(String, nullable=True)
 
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
     modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)

skyvern/forge/sdk/db/utils.py

@@ -709,6 +709,7 @@ def hydrate_action(action_model: ActionModel, empty_element_id: bool = False) ->
         "element_id": element_id,
         "skyvern_element_hash": action_model.skyvern_element_hash,
         "skyvern_element_data": action_model.skyvern_element_data,
+        "screenshot_artifact_id": action_model.screenshot_artifact_id,
         "created_at": action_model.created_at,
         "modified_at": action_model.modified_at,
     }