chain-service/chain_service/database/models/progress_chain.py
2024-02-28 16:33:51 +07:00

53 lines
1.6 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=TypeAdapter(List[Action]).dump_python(chain.actions),
)