node.py

"""Message tree node model."""

from dataclasses import dataclass, field
from enum import Enum

from ..models import MessageScope


class MessageState(Enum):
    """State of a message node in the tree."""

    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    ERROR = "error"


class MessageReferenceKind(Enum):
    """Kind of platform message represented by a tree reference."""

    PROMPT = "prompt"
    STATUS = "status"


@dataclass
class MessageNode:
    """A single user prompt/status node in a messaging conversation tree."""

    node_id: str
    scope: MessageScope
    prompt: str
    status_message_id: str | None
    state: MessageState = MessageState.PENDING
    parent_id: str | None = None
    parent_reference_id: str | None = None
    session_id: str | None = None
    children_ids: list[str] = field(default_factory=list)

    def update_state(
        self,
        state: MessageState,
        *,
        session_id: str | None = None,
    ) -> None:
        self.state = state
        if session_id:
            self.session_id = session_id

    def mark_error(self) -> None:
        self.update_state(MessageState.ERROR)

    def clear_status(self) -> None:
        """Invalidate the response and resume point while retaining its prompt."""
        self.status_message_id = None
        self.session_id = None
        self.mark_error()