@ondeinference / cli
Code Commits Issues Pulls Cloud Agent CI/CD
development
development main
rs 343 lines 12.8 KB
Raw
1 //! Background GGUF chat worker for the onde-cli TUI.
2 //!
3 //! Loads a local GGUF file from an absolute path on disk and handles
4 //! multi-turn chat via tokio mpsc channels. No HuggingFace Hub access is
5 //! required — the parent directory of the file is used as the model root,
6 //! and `TokenSource::None` bypasses all HF authentication.
7
8 use std::path::PathBuf;
9 use std::time::Instant;
10
11 use onde::mistralrs::{GgufModelBuilder, RequestBuilder, Response, TextMessageRole, TokenSource};
12 use tokio::sync::mpsc;
13
14 // ── Public types ──────────────────────────────────────────────────────────────
15
16 /// Progress events emitted by the background chat worker.
17 #[derive(Debug, Clone)]
18 pub enum ChatProgress {
19 /// The worker has started loading the model from disk.
20 LoadingModel,
21 /// The model is loaded and ready to accept messages.
22 Ready { model_name: String },
23 /// Inference is running; waiting for the reply.
24 Thinking,
25 /// A complete assistant reply has been produced.
26 Reply {
27 _text: String,
28 duration_display: String,
29 },
30 /// A partial token arrived from the streaming response.
31 StreamDelta(String),
32 /// An unrecoverable error occurred. The worker exits after sending this.
33 Error(String),
34 }
35
36 /// A single turn in the conversation history.
37 #[derive(Debug, Clone)]
38 pub struct ChatMessage {
39 pub role: ChatRole,
40 pub content: String,
41 }
42
43 impl ChatMessage {
44 pub fn user(content: impl Into<String>) -> Self {
45 Self {
46 role: ChatRole::User,
47 content: content.into(),
48 }
49 }
50
51 pub fn assistant(content: impl Into<String>) -> Self {
52 Self {
53 role: ChatRole::Assistant,
54 content: content.into(),
55 }
56 }
57 }
58
59 /// The role of a participant in the conversation.
60 #[derive(Debug, Clone, PartialEq, Eq)]
61 pub enum ChatRole {
62 User,
63 Assistant,
64 }
65
66 /// Commands sent to the background chat worker.
67 #[derive(Debug)]
68 pub enum ChatCommand {
69 /// Send a user message and await a reply.
70 SendMessage(String),
71 /// Shut down the worker cleanly.
72 Quit,
73 }
74
75 // ── Duration formatting ───────────────────────────────────────────────────────
76
77 /// Format an elapsed [`std::time::Duration`] as a compact human-readable string.
78 ///
79 /// - Under 60 s → `"4.2s"`
80 /// - 60 s or more → `"1m 4.2s"`
81 fn format_duration(elapsed: std::time::Duration) -> String {
82 let total_secs = elapsed.as_secs_f64();
83 let mins = (total_secs / 60.0).floor() as u64;
84 let secs = total_secs - (mins as f64 * 60.0);
85 if mins > 0 {
86 format!("{}m {:.1}s", mins, secs)
87 } else {
88 format!("{:.1}s", total_secs)
89 }
90 }
91
92 // ── History helpers ───────────────────────────────────────────────────────────
93
94 /// Map a [`ChatRole`] to the corresponding [`TextMessageRole`] expected by
95 /// `mistralrs`.
96 #[cfg(test)]
97 fn to_mistral_role(role: &ChatRole) -> TextMessageRole {
98 match role {
99 ChatRole::User => TextMessageRole::User,
100 ChatRole::Assistant => TextMessageRole::Assistant,
101 }
102 }
103
104 /// Build a [`RequestBuilder`] from the accumulated history plus the new user
105 /// message. A fixed sampling preset (temperature 0.7, max 512 tokens) is
106 /// applied so callers do not need to thread sampling config through the
107 /// command channel.
108 fn build_request(history: &[(TextMessageRole, String)], user_message: &str) -> RequestBuilder {
109 let mut req = RequestBuilder::new()
110 .set_sampler_temperature(0.7)
111 .set_sampler_max_len(512);
112
113 for (role, content) in history {
114 req = req.add_message(role.clone(), content);
115 }
116
117 req = req.add_message(TextMessageRole::User, user_message);
118 req
119 }
120
121 // ── Worker entry point ────────────────────────────────────────────────────────
122
123 /// Start the background GGUF chat worker.
124 ///
125 /// The worker:
126 /// 1. Sends [`ChatProgress::LoadingModel`].
127 /// 2. Loads the GGUF file at `gguf_path` using `mistralrs::GgufModelBuilder`.
128 /// The parent directory is used as the model root; no HF token is required.
129 /// 3. Sends [`ChatProgress::Ready`] on success, or [`ChatProgress::Error`] and
130 /// returns on failure.
131 /// 4. Loops on `command_rx`:
132 /// - [`ChatCommand::SendMessage`] → runs inference and sends
133 /// [`ChatProgress::Thinking`] then [`ChatProgress::Reply`] (or
134 /// [`ChatProgress::Error`] on failure).
135 /// - [`ChatCommand::Quit`] → breaks the loop and returns.
136 ///
137 /// Conversation history is kept in-memory as a `Vec<(TextMessageRole, String)>`
138 /// so every request includes the full multi-turn context.
139 pub async fn start_chat(
140 gguf_path: PathBuf,
141 progress_tx: mpsc::UnboundedSender<ChatProgress>,
142 mut command_rx: mpsc::UnboundedReceiver<ChatCommand>,
143 ) {
144 // ── 1. Announce that loading has started ─────────────────────────────────
145 let _ = progress_tx.send(ChatProgress::LoadingModel);
146
147 // ── 2. Derive model_id (parent dir) and filename from the absolute path ──
148 let model_dir = match gguf_path.parent() {
149 Some(p) => p.to_string_lossy().to_string(),
150 None => {
151 let _ = progress_tx.send(ChatProgress::Error(format!(
152 "Cannot determine parent directory of: {}",
153 gguf_path.display()
154 )));
155 return;
156 }
157 };
158
159 let file_name = match gguf_path.file_name() {
160 Some(n) => n.to_string_lossy().to_string(),
161 None => {
162 let _ = progress_tx.send(ChatProgress::Error(format!(
163 "Cannot determine file name of: {}",
164 gguf_path.display()
165 )));
166 return;
167 }
168 };
169
170 // ── 3. Build and load the model ──────────────────────────────────────────
171 log::info!("[chat] loading GGUF from dir={model_dir:?} file={file_name:?}");
172 let model = match GgufModelBuilder::new(&model_dir, vec![file_name.clone()])
173 .with_token_source(TokenSource::None)
174 .with_logging()
175 .build()
176 .await
177 {
178 Ok(m) => {
179 log::info!("[chat] model loaded successfully");
180 m
181 }
182 Err(e) => {
183 log::error!("[chat] model load failed: {e}");
184 let _ = progress_tx.send(ChatProgress::Error(format!(
185 "Failed to load model \"{file_name}\": {e}"
186 )));
187 return;
188 }
189 };
190
191 let _ = progress_tx.send(ChatProgress::Ready {
192 model_name: file_name,
193 });
194
195 // ── 4. Conversation history (mistralrs roles + content) ──────────────────
196 //
197 // We store the history as raw (TextMessageRole, String) pairs rather than
198 // as `ChatMessage` values so we can hand slices directly to `build_request`
199 // without any extra mapping at inference time.
200 let mut history: Vec<(TextMessageRole, String)> = Vec::new();
201
202 // ── 5. Command loop ───────────────────────────────────────────────────────
203 loop {
204 let command = match command_rx.recv().await {
205 Some(cmd) => cmd,
206 // Sender dropped — treat as a quit signal.
207 None => break,
208 };
209
210 match command {
211 ChatCommand::Quit => break,
212
213 ChatCommand::SendMessage(user_text) => {
214 let _ = progress_tx.send(ChatProgress::Thinking);
215
216 let request = build_request(&history, &user_text);
217
218 let start = Instant::now();
219 let mut stream = match model.stream_chat_request(request).await {
220 Ok(s) => s,
221 Err(e) => {
222 let _ = progress_tx.send(ChatProgress::Error(format!(
223 "Failed to start streaming: {e}"
224 )));
225 continue;
226 }
227 };
228
229 let mut reply_buf = String::new();
230
231 while let Some(response) = stream.next().await {
232 match response {
233 Response::Chunk(chunk) => {
234 if let Some(choice) = chunk.choices.first()
235 && let Some(ref delta) = choice.delta.content
236 {
237 reply_buf.push_str(delta);
238 let _ = progress_tx.send(ChatProgress::StreamDelta(delta.clone()));
239 }
240 }
241 Response::Done(_) => break,
242 Response::ModelError(msg, _) => {
243 let _ = progress_tx
244 .send(ChatProgress::Error(format!("Model error: {msg}")));
245 break;
246 }
247 Response::InternalError(e) => {
248 let _ = progress_tx
249 .send(ChatProgress::Error(format!("Internal error: {e}")));
250 break;
251 }
252 Response::ValidationError(e) => {
253 let _ = progress_tx
254 .send(ChatProgress::Error(format!("Validation error: {e}")));
255 break;
256 }
257 Response::CompletionModelError(msg, _) => {
258 let _ = progress_tx
259 .send(ChatProgress::Error(format!("Completion error: {msg}")));
260 break;
261 }
262 Response::CompletionDone(_) => break,
263 Response::CompletionChunk(_) => {}
264 Response::ImageGeneration(_) => {}
265 Response::Speech { .. } => {}
266 Response::Raw { .. } => {}
267 Response::Embeddings { .. } => {}
268 // Agentic, diffusion and file events. This is a plain
269 // chat request with no tools configured, so none of
270 // them can arrive here.
271 Response::AgenticToolCallProgress { .. } => {}
272 Response::AgenticToolApprovalRequired { .. } => {}
273 Response::BlockDenoisingProgress(_) => {}
274 Response::File(_) => {}
275 }
276 }
277
278 let elapsed = start.elapsed();
279
280 if !reply_buf.is_empty() {
281 history.push((TextMessageRole::User, user_text));
282 history.push((TextMessageRole::Assistant, reply_buf.trim().to_string()));
283
284 let _ = progress_tx.send(ChatProgress::Reply {
285 _text: reply_buf.trim().to_string(),
286 duration_display: format_duration(elapsed),
287 });
288 }
289 }
290 }
291 }
292 }
293
294 // ── Tests ─────────────────────────────────────────────────────────────────────
295
296 #[cfg(test)]
297 mod tests {
298 use super::*;
299
300 #[test]
301 fn format_duration_under_one_minute() {
302 let d = std::time::Duration::from_secs_f64(4.567);
303 assert_eq!(format_duration(d), "4.6s");
304 }
305
306 #[test]
307 fn format_duration_exactly_one_minute() {
308 let d = std::time::Duration::from_secs(60);
309 assert_eq!(format_duration(d), "1m 0.0s");
310 }
311
312 #[test]
313 fn format_duration_over_one_minute() {
314 let d = std::time::Duration::from_secs_f64(125.3);
315 assert_eq!(format_duration(d), "2m 5.3s");
316 }
317
318 #[test]
319 fn chat_message_user_role() {
320 let msg = ChatMessage::user("hello");
321 assert_eq!(msg.role, ChatRole::User);
322 assert_eq!(msg.content, "hello");
323 }
324
325 #[test]
326 fn chat_message_assistant_role() {
327 let msg = ChatMessage::assistant("hi there");
328 assert_eq!(msg.role, ChatRole::Assistant);
329 assert_eq!(msg.content, "hi there");
330 }
331
332 #[test]
333 fn to_mistral_role_maps_correctly() {
334 assert!(matches!(
335 to_mistral_role(&ChatRole::User),
336 TextMessageRole::User
337 ));
338 assert!(matches!(
339 to_mistral_role(&ChatRole::Assistant),
340 TextMessageRole::Assistant
341 ));
342 }
343 }