add reset observer cruise functions (#1537)

This commit is contained in:
Shuchang Zheng
2025-01-11 19:49:51 -08:00
committed by GitHub
parent e54977ef29
commit 29120b4508

View File

@@ -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,