From 29120b45084e7cc2af232685c1cb40205c8435fc Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Sat, 11 Jan 2025 19:49:51 -0800 Subject: [PATCH] add reset observer cruise functions (#1537) --- skyvern/forge/sdk/db/client.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 2bc9ed25..33d0b953 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -1744,6 +1744,19 @@ class AgentDB: await session.execute(stmt) await session.commit() + async def delete_observer_cruise_artifacts( + self, observer_cruise_id: str, organization_id: str | None = None + ) -> None: + async with self.Session() as session: + stmt = delete(ArtifactModel).where( + and_( + ArtifactModel.observer_cruise_id == observer_cruise_id, + ArtifactModel.organization_id == organization_id, + ) + ) + await session.execute(stmt) + await session.commit() + async def delete_task_steps(self, organization_id: str, task_id: str) -> None: async with self.Session() as session: # delete artifacts by filtering organization_id and task_id @@ -1933,6 +1946,19 @@ class AgentDB: return ObserverCruise.model_validate(observer_cruise) return None + async def delete_observer_thoughts_for_cruise( + self, observer_cruise_id: str, organization_id: str | None = None + ) -> None: + async with self.Session() as session: + stmt = delete(ObserverThoughtModel).where( + and_( + ObserverThoughtModel.observer_cruise_id == observer_cruise_id, + ObserverThoughtModel.organization_id == organization_id, + ) + ) + await session.execute(stmt) + await session.commit() + async def get_observer_cruise_by_workflow_run_id( self, workflow_run_id: str,