Skip to main content

oxi_tui/content/
message.rs

1use std::time::Instant;
2
3pub type MessageId = u64;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub enum MessageRole {
7    User,
8    Assistant,
9    System,
10    Tool,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub enum ContentBlock {
15    Text(String),
16    Thinking(String),
17    ToolCall {
18        id: String,
19        name: String,
20        args: String,
21        status: ToolCallStatus,
22    },
23    ToolResult {
24        call_id: String,
25        content: String,
26        is_error: bool,
27    },
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
31pub enum ToolCallStatus {
32    Pending,
33    Running,
34    Completed,
35    Failed,
36}
37
38#[derive(Debug, Clone)]
39pub struct ChatMessage {
40    pub id: MessageId,
41    pub role: MessageRole,
42    pub blocks: Vec<ContentBlock>,
43    pub created_at: Instant,
44}
45
46impl ChatMessage {
47    #[must_use]
48    pub fn new(id: MessageId, role: MessageRole) -> Self {
49        Self {
50            id,
51            role,
52            blocks: Vec::new(),
53            created_at: Instant::now(),
54        }
55    }
56
57    #[must_use]
58    pub fn text_content(&self) -> Option<&str> {
59        self.blocks.iter().find_map(|block| match block {
60            ContentBlock::Text(text) => Some(text.as_str()),
61            _ => None,
62        })
63    }
64
65    pub fn append_text(&mut self, text: &str) {
66        if let Some(ContentBlock::Text(content)) = self.blocks.last_mut() {
67            content.push_str(text);
68        } else {
69            self.blocks.push(ContentBlock::Text(text.to_owned()));
70        }
71    }
72
73    #[must_use]
74    pub fn content_hash(&self) -> u64 {
75        use std::hash::{Hash, Hasher};
76
77        let mut hasher = std::collections::hash_map::DefaultHasher::new();
78        self.role.hash(&mut hasher);
79        self.blocks.hash(&mut hasher);
80
81        crate::widget::hash_combine(self.id, hasher.finish())
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::{ChatMessage, ContentBlock, MessageRole, ToolCallStatus};
88
89    #[test]
90    fn append_text_creates_a_text_block_when_empty() {
91        let mut message = ChatMessage::new(1, MessageRole::Assistant);
92        message.append_text("hello");
93        assert_eq!(message.blocks, [ContentBlock::Text("hello".into())]);
94        assert_eq!(message.text_content(), Some("hello"));
95    }
96
97    #[test]
98    fn append_text_extends_the_last_text_block() {
99        let mut message = ChatMessage::new(1, MessageRole::Assistant);
100        message.blocks.push(ContentBlock::Text("hel".into()));
101        message.append_text("lo");
102        assert_eq!(message.blocks, [ContentBlock::Text("hello".into())]);
103    }
104
105    #[test]
106    fn append_text_creates_a_block_after_non_text_content() {
107        let mut message = ChatMessage::new(1, MessageRole::Assistant);
108        message.blocks.push(ContentBlock::ToolCall {
109            id: "call-1".into(),
110            name: "search".into(),
111            args: "{}".into(),
112            status: ToolCallStatus::Completed,
113        });
114        message.append_text("done");
115        assert_eq!(message.blocks.len(), 2);
116        assert_eq!(message.blocks[1], ContentBlock::Text("done".into()));
117    }
118}