53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
from .chain import Chain
|
|
from .base import BaseConfig, BaseMongoModel
|
|
|
|
from datetime import datetime
|
|
from pydantic import BaseModel, Field, TypeAdapter
|
|
from typing import Literal, Union, Annotated, Optional, List
|
|
|
|
|
|
class WaitAction(BaseModel):
|
|
action_type: Literal["wait"]
|
|
wait_for: int
|
|
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
|
|
|
|
|
|
class TextCommentAction(BaseModel):
|
|
action_type: Literal["text_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")
|
|
]
|
|
|
|
|
|
class ProgressChain(BaseMongoModel):
|
|
task_id: int
|
|
recipients: Annotated[Optional[List[int]], Field(default=[])]
|
|
name: Annotated[Optional[str], Field(default=None)]
|
|
actions: Annotated[Optional[List[Action]], Field(default=[])]
|
|
created_at: Annotated[datetime, Field(default_factory=datetime.utcnow)]
|
|
|
|
@classmethod
|
|
def create_from_chain(
|
|
cls, chain: Chain, task_id: int, recipients: Optional[List[int]] = []
|
|
):
|
|
return ProgressChain(
|
|
task_id=task_id,
|
|
recipients=recipients,
|
|
name=chain.name,
|
|
actions=map(Chain.model_dump, chain.actions),
|
|
)
|