diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 04e08e28..0f035b8b 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -359,8 +359,9 @@ class AgentDB: async def get_total_unique_step_order_count_by_task_ids( self, + *, task_ids: list[str], - organization_id: str | None = None, + organization_id: str, ) -> int: """ Get the total count of unique (step.task_id, step.order) pairs of StepModel for the given task ids @@ -1009,6 +1010,8 @@ class AgentDB: async def get_artifacts_by_entity_id( self, + *, + organization_id: str, artifact_type: ArtifactType | None = None, task_id: str | None = None, step_id: str | None = None, @@ -1016,7 +1019,6 @@ class AgentDB: workflow_run_block_id: str | None = None, thought_id: str | None = None, task_v2_id: str | None = None, - organization_id: str | None = None, ) -> list[Artifact]: try: async with self.Session() as session: @@ -1053,16 +1055,18 @@ class AgentDB: async def get_artifact_by_entity_id( self, + *, artifact_type: ArtifactType, + organization_id: str, task_id: str | None = None, step_id: str | None = None, workflow_run_id: str | None = None, workflow_run_block_id: str | None = None, thought_id: str | None = None, task_v2_id: str | None = None, - organization_id: str | None = None, ) -> Artifact | None: artifacts = await self.get_artifacts_by_entity_id( + organization_id=organization_id, artifact_type=artifact_type, task_id=task_id, step_id=step_id, @@ -1070,7 +1074,6 @@ class AgentDB: workflow_run_block_id=workflow_run_block_id, thought_id=thought_id, task_v2_id=task_v2_id, - organization_id=organization_id, ) return artifacts[0] if artifacts else None @@ -2401,9 +2404,10 @@ class AgentDB: async def get_thoughts( self, + *, task_v2_id: str, - thought_types: list[ThoughtType] | None = None, - organization_id: str | None = None, + thought_types: list[ThoughtType], + organization_id: str, ) -> list[Thought]: async with self.Session() as session: query = ( diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index 9512574e..e1441796 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -1135,13 +1135,10 @@ async def get_artifacts( ) analytics.capture("skyvern-oss-agent-entity-artifacts-get") - params = { - "organization_id": current_org.organization_id, entity_type_to_param[entity_type]: entity_id, } - - artifacts = await app.DATABASE.get_artifacts_by_entity_id(**params) # type: ignore + artifacts = await app.DATABASE.get_artifacts_by_entity_id(organization_id=current_org.organization_id, **params) # type: ignore if settings.ENV != "local" or settings.GENERATE_PRESIGNED_URLS: signed_urls = await app.ARTIFACT_MANAGER.get_share_links(artifacts) diff --git a/skyvern/services/task_v2_service.py b/skyvern/services/task_v2_service.py index 8c4a72e0..1b6bc705 100644 --- a/skyvern/services/task_v2_service.py +++ b/skyvern/services/task_v2_service.py @@ -1351,12 +1351,9 @@ def _generate_random_string(length: int = 5) -> str: return "".join(random.choices(RANDOM_STRING_POOL, k=length)) -async def get_thought_timelines( - task_v2_id: str, - organization_id: str | None = None, -) -> list[WorkflowRunTimeline]: +async def get_thought_timelines(*, task_v2_id: str, organization_id: str) -> list[WorkflowRunTimeline]: thoughts = await app.DATABASE.get_thoughts( - task_v2_id, + task_v2_id=task_v2_id, organization_id=organization_id, thought_types=[ ThoughtType.plan,