oxi_tui/content/
chat_log.rs1use super::{ChatMessage, MessageId, MessageRole, StreamId, StreamingState};
14use crate::widget::{hash_combine, hash_str};
15
16#[derive(Debug, Clone)]
18pub struct ChatLog {
19 messages: Vec<ChatMessage>,
20 next_id: MessageId,
21 active_stream: StreamingState,
22 cached_hash: u64,
23}
24
25impl Default for ChatLog {
26 fn default() -> Self {
27 let mut log = Self {
28 messages: Vec::new(),
29 next_id: 0,
30 active_stream: StreamingState::default(),
31 cached_hash: 0,
32 };
33 log.refresh_hash();
34 log
35 }
36}
37
38impl ChatLog {
39 #[must_use]
41 pub fn new() -> Self {
42 Self::default()
43 }
44
45 #[must_use]
48 pub fn append_message(&mut self, role: MessageRole) -> MessageId {
49 let id = self.next_id;
50 self.next_id = self.next_id.wrapping_add(1);
51 self.messages.push(ChatMessage::new(id, role));
52
53 if role == MessageRole::Assistant {
54 self.active_stream.start(id);
55 }
56 self.refresh_hash();
57 id
58 }
59
60 pub fn append_token(&mut self, token: &str) {
63 let message_index = self.active_stream.active_stream.and_then(|stream_id| {
64 self.messages.iter().rposition(|message| {
65 message.id == stream_id && message.role == MessageRole::Assistant
66 })
67 });
68
69 let message_index = if let Some(index) = message_index {
70 index
71 } else {
72 let _ = self.append_message(MessageRole::Assistant);
73 self.messages.len().saturating_sub(1)
74 };
75
76 if let Some(message) = self.messages.get_mut(message_index) {
77 message.append_text(token);
78 }
79 self.active_stream.push_token(token);
80 self.refresh_hash();
81 }
82
83 pub fn finalize_stream(&mut self) {
85 self.active_stream.finalize();
86 self.refresh_hash();
87 }
88
89 #[must_use]
91 pub fn messages(&self) -> &[ChatMessage] {
92 &self.messages
93 }
94
95 #[must_use]
97 pub fn active_stream(&self) -> Option<StreamId> {
98 self.active_stream.active_stream
99 }
100
101 #[must_use]
103 pub const fn content_hash(&self) -> u64 {
104 self.cached_hash
105 }
106
107 fn refresh_hash(&mut self) {
108 let message_count = u64::try_from(self.messages.len()).unwrap_or(u64::MAX);
109 let last_message_hash = self.messages.last().map_or(0, ChatMessage::content_hash);
110 let stream_id_hash = match self.active_stream.active_stream {
111 Some(stream_id) => hash_combine(1, stream_id),
112 None => 0,
113 };
114 let partial_hash = hash_str(&self.active_stream.partial_text);
115 let messages_hash = hash_combine(message_count, last_message_hash);
116 let stream_hash = hash_combine(stream_id_hash, partial_hash);
117 self.cached_hash = hash_combine(messages_hash, stream_hash);
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use super::ChatLog;
124 use crate::content::{ContentBlock, MessageRole};
125
126 #[test]
127 fn append_message_assigns_incrementing_ids() {
128 let mut log = ChatLog::new();
129 let first = log.append_message(MessageRole::User);
130 let second = log.append_message(MessageRole::System);
131
132 assert_eq!(first, 0);
133 assert_eq!(second, 1);
134 assert_eq!(log.messages()[0].id, first);
135 assert_eq!(log.messages()[1].id, second);
136 }
137
138 #[test]
139 fn append_token_creates_assistant_if_none() {
140 let mut log = ChatLog::new();
141 log.append_token("hello");
142
143 assert_eq!(log.messages().len(), 1);
144 assert_eq!(log.messages()[0].role, MessageRole::Assistant);
145 assert_eq!(
146 log.messages()[0].blocks,
147 [ContentBlock::Text("hello".to_owned())]
148 );
149 }
150
151 #[test]
152 fn append_token_appends_to_last_text() {
153 let mut log = ChatLog::new();
154 let _ = log.append_message(MessageRole::Assistant);
155 log.append_token("hel");
156 log.append_token("lo");
157
158 assert_eq!(
159 log.messages()[0].blocks,
160 [ContentBlock::Text("hello".to_owned())]
161 );
162 }
163
164 #[test]
165 fn finalize_stream_clears() {
166 let mut log = ChatLog::new();
167 let id = log.append_message(MessageRole::Assistant);
168 assert_eq!(log.active_stream(), Some(id));
169
170 log.finalize_stream();
171
172 assert_eq!(log.active_stream(), None);
173 }
174
175 #[test]
176 fn content_hash_changes_on_append() {
177 let mut log = ChatLog::new();
178 let before = log.content_hash();
179
180 let _ = log.append_message(MessageRole::Assistant);
181 log.append_token("hello");
182
183 assert_ne!(before, log.content_hash());
184 }
185}