org id not None (#2696)

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,
+        thought_types: list[ThoughtType],
-        organization_id: str | None = None,
+        organization_id: str,
     ) -> list[Thought]:
         async with self.Session() as session:
             query = (

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(organization_id=current_org.organization_id, **params)  # type: ignore
-    artifacts = await app.DATABASE.get_artifacts_by_entity_id(**params)  # type: ignore
 
     if settings.ENV != "local" or settings.GENERATE_PRESIGNED_URLS:
         signed_urls = await app.ARTIFACT_MANAGER.get_share_links(artifacts)

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(
+async def get_thought_timelines(*, task_v2_id: str, organization_id: str) -> list[WorkflowRunTimeline]:
-    task_v2_id: str,
-    organization_id: str | None = None,
-) -> 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,