Forráskód Böngészése

fix: Replace generic exceptions with specific error classes in task p… (#12036)

Signed-off-by: -LAN- <laipz8200@outlook.com>
-LAN- 4 hónapja
szülő
commit
82134a1d50

+ 1 - 1
api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py

@@ -177,7 +177,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline, MessageCycleMan
             else:
                 continue
 
-        raise Exception("Queue listening stopped unexpectedly.")
+        raise RuntimeError("queue listening stopped unexpectedly.")
 
     def _to_stream_response(
         self, generator: Generator[StreamResponse, None, None]

+ 17 - 0
api/core/app/task_pipeline/exc.py

@@ -0,0 +1,17 @@
+class TaskPipilineError(ValueError):
+    pass
+
+
+class RecordNotFoundError(TaskPipilineError):
+    def __init__(self, record_name: str, record_id: str):
+        super().__init__(f"{record_name} with id {record_id} not found")
+
+
+class WorkflowRunNotFoundError(RecordNotFoundError):
+    def __init__(self, workflow_run_id: str):
+        super().__init__("WorkflowRun", workflow_run_id)
+
+
+class WorkflowNodeExecutionNotFoundError(RecordNotFoundError):
+    def __init__(self, workflow_node_execution_id: str):
+        super().__init__("WorkflowNodeExecution", workflow_node_execution_id)

+ 4 - 2
api/core/app/task_pipeline/workflow_cycle_manage.py

@@ -58,6 +58,8 @@ from models.workflow import (
     WorkflowRunStatus,
 )
 
+from .exc import WorkflowNodeExecutionNotFoundError, WorkflowRunNotFoundError
+
 
 class WorkflowCycleManage:
     _application_generate_entity: Union[AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity]
@@ -898,7 +900,7 @@ class WorkflowCycleManage:
         workflow_run = db.session.query(WorkflowRun).filter(WorkflowRun.id == workflow_run_id).first()
 
         if not workflow_run:
-            raise Exception(f"Workflow run not found: {workflow_run_id}")
+            raise WorkflowRunNotFoundError(workflow_run_id)
 
         return workflow_run
 
@@ -911,6 +913,6 @@ class WorkflowCycleManage:
         workflow_node_execution = self._wip_workflow_node_executions.get(node_execution_id)
 
         if not workflow_node_execution:
-            raise Exception(f"Workflow node execution not found: {node_execution_id}")
+            raise WorkflowNodeExecutionNotFoundError(node_execution_id)
 
         return workflow_node_execution