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

TaskRun pydantic model gets renamed to Run (#2042)

Browse Source
This commit is contained in:
Shuchang Zheng
2025-03-30 18:41:24 -07:00
committed by GitHub
parent 05e28931bc
commit 12ef2100b5
3 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -66,7 +66,7 @@ from skyvern.forge.sdk.schemas.credentials import Credential, CredentialType
from skyvern.forge.sdk.schemas.organization_bitwarden_collections import OrganizationBitwardenCollection from skyvern.forge.sdk.schemas.organization_bitwarden_collections import OrganizationBitwardenCollection
from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken
from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
from skyvern.forge.sdk.schemas.runs import TaskRun from skyvern.forge.sdk.schemas.runs import Run
from skyvern.forge.sdk.schemas.task_generations import TaskGeneration from skyvern.forge.sdk.schemas.task_generations import TaskGeneration
from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Status, Thought, ThoughtType from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Status, Thought, ThoughtType
from skyvern.forge.sdk.schemas.tasks import OrderBy, SortDirection, Task, TaskStatus from skyvern.forge.sdk.schemas.tasks import OrderBy, SortDirection, Task, TaskStatus
@@ -2789,7 +2789,7 @@ class AgentDB:
title: str | None = None, title: str | None = None,
url: str | None = None, url: str | None = None,
url_hash: str | None = None, url_hash: str | None = None,
) -> TaskRun: ) -> Run:
async with self.Session() as session: async with self.Session() as session:
task_run = TaskRunModel( task_run = TaskRunModel(
task_run_type=task_run_type, task_run_type=task_run_type,
@@ -2802,7 +2802,7 @@ class AgentDB:
session.add(task_run) session.add(task_run)
await session.commit() await session.commit()
await session.refresh(task_run) await session.refresh(task_run)
return TaskRun.model_validate(task_run) return Run.model_validate(task_run)
async def create_credential( async def create_credential(
self, self,
@@ -2916,7 +2916,7 @@ class AgentDB:
return OrganizationBitwardenCollection.model_validate(organization_bitwarden_collection) return OrganizationBitwardenCollection.model_validate(organization_bitwarden_collection)
return None return None
async def cache_task_run(self, run_id: str, organization_id: str | None = None) -> TaskRun: async def cache_task_run(self, run_id: str, organization_id: str | None = None) -> Run:
async with self.Session() as session: async with self.Session() as session:
task_run = ( task_run = (
await session.scalars( await session.scalars(
@@ -2927,12 +2927,12 @@ class AgentDB:
task_run.cached = True task_run.cached = True
await session.commit() await session.commit()
await session.refresh(task_run) await session.refresh(task_run)
return TaskRun.model_validate(task_run) return Run.model_validate(task_run)
raise NotFoundError(f"TaskRun {run_id} not found") raise NotFoundError(f"Run {run_id} not found")
async def get_cached_task_run( async def get_cached_task_run(
self, task_run_type: RunType, url_hash: str | None = None, organization_id: str | None = None self, task_run_type: RunType, url_hash: str | None = None, organization_id: str | None = None
) -> TaskRun | None: ) -> Run | None:
async with self.Session() as session: async with self.Session() as session:
query = select(TaskRunModel) query = select(TaskRunModel)
if task_run_type: if task_run_type:
@@ -2943,16 +2943,16 @@ class AgentDB:
query = query.filter_by(organization_id=organization_id) query = query.filter_by(organization_id=organization_id)
query = query.filter_by(cached=True).order_by(TaskRunModel.created_at.desc()) query = query.filter_by(cached=True).order_by(TaskRunModel.created_at.desc())
task_run = (await session.scalars(query)).first() task_run = (await session.scalars(query)).first()
return TaskRun.model_validate(task_run) if task_run else None return Run.model_validate(task_run) if task_run else None
async def get_task_run( async def get_task_run(
self, self,
run_id: str, run_id: str,
organization_id: str | None = None, organization_id: str | None = None,
) -> TaskRun | None: ) -> Run | None:
async with self.Session() as session: async with self.Session() as session:
query = select(TaskRunModel).filter_by(run_id=run_id) query = select(TaskRunModel).filter_by(run_id=run_id)
if organization_id: if organization_id:
query = query.filter_by(organization_id=organization_id) query = query.filter_by(organization_id=organization_id)
task_run = (await session.scalars(query)).first() task_run = (await session.scalars(query)).first()
return TaskRun.model_validate(task_run) if task_run else None return Run.model_validate(task_run) if task_run else None

2
skyvern/forge/sdk/schemas/runs.py
View File

@@ -5,7 +5,7 @@ from pydantic import BaseModel, ConfigDict
from skyvern.schemas.runs import RunType from skyvern.schemas.runs import RunType
class TaskRun(BaseModel): class Run(BaseModel):
model_config = ConfigDict(from_attributes=True) model_config = ConfigDict(from_attributes=True)
task_run_id: str task_run_id: str

4
skyvern/forge/sdk/services/task_run_service.py
View File

@@ -1,9 +1,9 @@
from skyvern.forge import app from skyvern.forge import app
from skyvern.forge.sdk.schemas.runs import TaskRun from skyvern.forge.sdk.schemas.runs import Run
from skyvern.schemas.runs import RunEngine, RunResponse, RunType from skyvern.schemas.runs import RunEngine, RunResponse, RunType
async def get_task_run(run_id: str, organization_id: str | None = None) -> TaskRun | None: async def get_task_run(run_id: str, organization_id: str | None = None) -> Run | None:
return await app.DATABASE.get_task_run(run_id, organization_id=organization_id) return await app.DATABASE.get_task_run(run_id, organization_id=organization_id)