shu/test discriminated run response schema (#2046)

2025-03-30 21:14:52 -07:00
parent 8253738c7b
commit ea67502161
6 changed files with 158 additions and 112 deletions
--- a/skyvern/services/run_service.py
+++ b/skyvern/services/run_service.py
@@ -0,0 +1,71 @@
+from skyvern.forge import app
+from skyvern.schemas.runs import RunEngine, RunResponse, RunType, TaskRunRequest, TaskRunResponse
+
+
+async def get_run_response(run_id: str, organization_id: str | None = None) -> RunResponse | None:
+    run = await app.DATABASE.get_run(run_id, organization_id=organization_id)
+    if not run:
+        return None
+
+    if run.task_run_type == RunType.task_v1:
+        # fetch task v1 from db and transform to task run response
+        task_v1 = await app.DATABASE.get_task(run.task_v1_id, organization_id=organization_id)
+        if not task_v1:
+            return None
+        return TaskRunResponse(
+            run_id=run.run_id,
+            run_type=run.task_run_type,
+            status=str(task_v1.status),
+            output=task_v1.extracted_information,
+            failure_reason=task_v1.failure_reason,
+            created_at=task_v1.created_at,
+            modified_at=task_v1.modified_at,
+            run_request=TaskRunRequest(
+                engine=RunEngine.skyvern_v1,
+                prompt=task_v1.navigation_goal,
+                url=task_v1.url,
+                webhook_url=task_v1.webhook_callback_url,
+                totp_identifier=task_v1.totp_identifier,
+                totp_url=task_v1.totp_verification_url,
+                proxy_location=task_v1.proxy_location,
+                max_steps=task_v1.max_steps_per_run,
+                data_extraction_schema=task_v1.extracted_information_schema,
+                error_code_mapping=task_v1.error_code_mapping,
+            ),
+        )
+    elif run.task_run_type == RunType.task_v2:
+        task_v2 = await app.DATABASE.get_task_v2(run.task_v2_id, organization_id=organization_id)
+        if not task_v2:
+            return None
+        return TaskRunResponse(
+            run_id=run.run_id,
+            run_type=run.task_run_type,
+            status=task_v2.status,
+            output=task_v2.output,
+            failure_reason=task_v2.failure_reason,
+            created_at=task_v2.created_at,
+            modified_at=task_v2.modified_at,
+            run_request=TaskRunRequest(
+                engine=RunEngine.skyvern_v2,
+                prompt=task_v2.prompt,
+                url=task_v2.url,
+                webhook_url=task_v2.webhook_callback_url,
+                totp_identifier=task_v2.totp_identifier,
+                totp_url=task_v2.totp_verification_url,
+                proxy_location=task_v2.proxy_location,
+                data_extraction_schema=task_v2.data_extraction_schema,
+                error_code_mapping=task_v2.error_code_mapping,
+            ),
+        )
+    elif run.task_run_type == RunType.workflow_run:
+        raise NotImplementedError("Workflow run response not implemented")
+        # return WorkflowRunResponse(
+        #     run_id=run.run_id,
+        #     run_type=run.task_run_type,
+        #     status=run.status,
+        #     output=run.output,
+        #     parameters=None,
+        #     created_at=run.created_at,
+        #     modified_at=run.modified_at,
+        # )
+    raise ValueError(f"Invalid task run type: {run.task_run_type}")