vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Implement get_latest_screenshots, add action_screenshots to TaskResponse (#148)

Browse Source
This commit is contained in:
Kerem Yilmaz
2024-04-02 14:43:29 -07:00
committed by GitHub
parent 5c0f482b05
commit 69c458bd7c
4 changed files with 65 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
skyvern/forge/sdk/db/client.py
View File

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