refactored

chain_service/database/models/progress_chain.py

@@ -1,35 +1,49 @@
 from .chain import Chain
 from .base import BaseConfig, BaseMongoModel
 
+from enum import Enum
 from datetime import datetime
 from pydantic import BaseModel, Field
 from typing import Literal, Union, Annotated, Optional, List
 
 
-class WaitAction(BaseModel):
-    action_type: Literal["wait"]
-    wait_for: int
-    done: Annotated[bool, Field(default=False)]
+class ProgressActionStatusEnum(str, Enum):
+    PENDING = "pending"
+    DONE = "done"
+    FAILED = "failed"
+
+
+class BaseProgressAction(BaseModel):
+    action_type: Literal["wait", "comment"]
+
+    status: Annotated[
+        ProgressActionStatusEnum, Field(default=ProgressActionStatusEnum.PENDING)
+    ]
+
     started_at: Annotated[Optional[datetime], Field(default=None)]
     finished_at: Annotated[Optional[datetime], Field(default=None)]
+    error_text: Annotated[Optional[str], Field(default=None)]
+
+
+class WaitProgressAction(BaseProgressAction):
+    action_type: Literal["wait"]
+    wait_for: int
 
     class Config(BaseConfig):
         pass
 
 
-class TextCommentAction(BaseModel):
-    action_type: Literal["text_comment"]
+class CommentProgressAction(BaseProgressAction):
+    action_type: Literal["comment"]
     text: str
-    done: Annotated[bool, Field(default=False)]
-    started_at: Annotated[Optional[datetime], Field(default=None)]
-    finished_at: Annotated[Optional[datetime], Field(default=None)]
 
     class Config(BaseConfig):
         pass
 
 
 Action = Annotated[
-    Union[WaitAction, TextCommentAction], Field(description="action_type")
+    Union[WaitProgressAction, CommentProgressAction],
+    Field(description="action_type"),
 ]
 
 
@@ -50,3 +64,15 @@ class ProgressChain(BaseMongoModel):
             name=chain.name,
             actions=map(Chain.model_dump, chain.actions),
         )
+
+    @property
+    def is_finished(self):
+        return all(
+            action.status is ProgressActionStatusEnum.DONE for action in self.actions
+        )
+
+    @property
+    def has_failed(self):
+        return any(
+            action.status is ProgressActionStatusEnum.FAILED for action in self.actions
+        )