File size: 1,430 Bytes
2415446
a1bab2d
 
2415446
 
 
 
 
 
 
 
 
 
a1bab2d
2415446
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
"""Message tree node model."""
from __future__ import annotations


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()