ravenrs_ai/
assistant.rs

1//! A model in use: answers away from the UI thread, and lets go of its memory when
2//! nobody has asked anything for a while.
3
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::mpsc::{self, RecvTimeoutError, Sender};
6use std::sync::{Arc, Mutex};
7use std::time::Duration;
8
9use ravenrs_widgets::UiWaker;
10
11use crate::tools::json_closed;
12use crate::{LocalEngine, LocalModel, ModelFolder};
13
14/// How long a model stays in memory after answering.
15const STAYS_LOADED: Duration = Duration::from_secs(120);
16const MAX_TOKENS: usize = 160;
17/// How much a choice can take: a tool and its arguments are short.
18const CHOICE_TOKENS: usize = 80;
19/// How long someone waits, counting the loading: for a whole choice, or for the first
20/// word of an answer. A slow machine should not leave them looking at it forever.
21const WAIT_LIMIT: Duration = Duration::from_secs(60);
22
23struct Job {
24    system: String,
25    user: String,
26    text: Arc<Mutex<String>>,
27    waker: UiWaker,
28    if_it_fails: String,
29    done: Option<Arc<AtomicBool>>,
30    /// Set when the wait ended first: nobody wants this answer any more.
31    given_up: Arc<AtomicBool>,
32    /// Letting go of it ends the wait.
33    waiting: Option<Sender<()>>,
34}
35
36impl Job {
37    /// Takes in a piece of the answer, and says whether to go on.
38    fn take(&mut self, piece: &str) -> bool {
39        let Ok(mut text) = self.text.lock() else { return false };
40        if self.given_up.load(Ordering::Acquire) {
41            return false;
42        }
43        text.push_str(piece);
44        let go_on = match self.done {
45            // A choice ends with its JSON, whatever a model goes on to say.
46            Some(_) => !json_closed(&text),
47            None => {
48                self.waiting = None;
49                true
50            }
51        };
52        drop(text);
53        self.waker.wake();
54        go_on
55    }
56
57    /// Says it is over, with `if_it_fails` when nothing came.
58    fn finish(&self, answered: bool) {
59        if !answered && let Ok(mut text) = self.text.lock() {
60            text.clone_from(&self.if_it_fails);
61        }
62        if let Some(done) = &self.done {
63            done.store(true, Ordering::Release);
64        }
65        self.waker.wake();
66    }
67}
68
69/// A model that answers on its own thread. Dropping it ends the thread and frees
70/// the memory.
71#[derive(Clone, Debug)]
72pub struct LocalAssistant {
73    jobs: Sender<Job>,
74    too_slow: Arc<AtomicBool>,
75    wait_limit: Duration,
76}
77
78impl Job {
79    /// What every job starts with.
80    fn blank() -> Self {
81        Self {
82            system: String::new(),
83            user: String::new(),
84            text: Arc::default(),
85            waker: UiWaker::default(),
86            if_it_fails: String::new(),
87            done: None,
88            given_up: Arc::default(),
89            waiting: None,
90        }
91    }
92}
93
94impl std::fmt::Debug for Job {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("Job").finish_non_exhaustive()
97    }
98}
99
100impl LocalAssistant {
101    /// An assistant over `model`. Nothing is loaded until the first question.
102    pub fn new(folder: ModelFolder, model: LocalModel) -> Self {
103        let (jobs, inbox) = mpsc::channel::<Job>();
104        std::thread::spawn(move || {
105            let mut engine: Option<LocalEngine> = None;
106            loop {
107                let mut job = if engine.is_some() {
108                    match inbox.recv_timeout(STAYS_LOADED) {
109                        Ok(job) => job,
110                        Err(RecvTimeoutError::Timeout) => {
111                            engine = None;
112                            continue;
113                        }
114                        Err(RecvTimeoutError::Disconnected) => return,
115                    }
116                } else {
117                    match inbox.recv() {
118                        Ok(job) => job,
119                        Err(_) => return,
120                    }
121                };
122                if job.given_up.load(Ordering::Acquire) {
123                    continue;
124                }
125                if engine.is_none() {
126                    engine = LocalEngine::load(&folder, &model).ok();
127                }
128                let limit = if job.done.is_some() { CHOICE_TOKENS } else { MAX_TOKENS };
129                let (system, user) = (job.system.clone(), job.user.clone());
130                let answered = engine.as_mut().is_some_and(|engine| {
131                    let wrote = engine.generate(&system, &user, limit, &mut |piece| job.take(piece));
132                    wrote.is_ok_and(|generated| generated.tokens > 0)
133                });
134                if !job.given_up.load(Ordering::Acquire) {
135                    job.finish(answered);
136                }
137            }
138        });
139        Self { jobs, too_slow: Arc::new(AtomicBool::new(false)), wait_limit: WAIT_LIMIT }
140    }
141
142    /// Asks it to answer `user`, following `system`, writing to `text` as it goes.
143    /// `if_it_fails` is what `text` becomes if the model cannot answer.
144    pub fn write(
145        &self,
146        system: String,
147        user: String,
148        text: Arc<Mutex<String>>,
149        waker: UiWaker,
150        if_it_fails: String,
151    ) {
152        self.send(Job { system, user, text, waker, if_it_fails, done: None, ..Job::blank() });
153    }
154
155    /// Asks it to answer `user`, following `system`, and says when it is over: `done`
156    /// is set after `text` has all of it, or stays empty if the model could not answer.
157    pub fn answer(
158        &self,
159        system: String,
160        user: String,
161        text: Arc<Mutex<String>>,
162        done: Arc<AtomicBool>,
163        waker: UiWaker,
164    ) {
165        let if_it_fails = String::new();
166        self.send(Job { system, user, text, waker, if_it_fails, done: Some(done), ..Job::blank() });
167    }
168
169    /// Whether a question waited too long on this machine. From then on it answers
170    /// nothing, at once, so nobody waits for it again until another assistant is made.
171    pub fn too_slow(&self) -> bool {
172        self.too_slow.load(Ordering::Acquire)
173    }
174
175    fn send(&self, mut job: Job) {
176        if self.too_slow() {
177            job.finish(false);
178            return;
179        }
180        let (waiting, ended) = mpsc::channel();
181        job.waiting = Some(waiting);
182        let watched = Job {
183            system: String::new(),
184            user: String::new(),
185            text: job.text.clone(),
186            waker: job.waker.clone(),
187            if_it_fails: job.if_it_fails.clone(),
188            done: job.done.clone(),
189            given_up: job.given_up.clone(),
190            waiting: None,
191        };
192        let (too_slow, limit) = (self.too_slow.clone(), self.wait_limit);
193        std::thread::spawn(move || {
194            if matches!(ended.recv_timeout(limit), Err(RecvTimeoutError::Timeout)) {
195                too_slow.store(true, Ordering::Release);
196                watched.given_up.store(true, Ordering::Release);
197                watched.finish(false);
198            }
199        });
200        // The thread ends only with this handle, so this does not fail.
201        let _ = self.jobs.send(job);
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    /// An assistant whose model never gets to answer.
210    fn stuck() -> (LocalAssistant, mpsc::Receiver<Job>) {
211        let (jobs, inbox) = mpsc::channel();
212        let wait_limit = Duration::from_millis(20);
213        (LocalAssistant { jobs, too_slow: Arc::default(), wait_limit }, inbox)
214    }
215
216    fn wait_for(done: &AtomicBool) {
217        let started = std::time::Instant::now();
218        while !done.load(Ordering::Acquire) {
219            assert!(started.elapsed() < Duration::from_secs(5), "it waited forever");
220            std::thread::sleep(Duration::from_millis(5));
221        }
222    }
223
224    #[test]
225    fn a_choice_that_takes_too_long_ends_empty() {
226        let (assistant, _inbox) = stuck();
227        let (text, done) = (Arc::new(Mutex::new("half".to_owned())), Arc::new(AtomicBool::new(false)));
228        assistant.answer("rules".into(), "question".into(), text.clone(), done.clone(), UiWaker::default());
229        wait_for(&done);
230        assert!(assistant.too_slow());
231        assert_eq!(*text.lock().expect("text"), "");
232    }
233
234    #[test]
235    fn after_that_nobody_waits_again() {
236        let (assistant, inbox) = stuck();
237        let (text, done) = (Arc::new(Mutex::new(String::new())), Arc::new(AtomicBool::new(false)));
238        assistant.answer("rules".into(), "question".into(), text.clone(), done.clone(), UiWaker::default());
239        wait_for(&done);
240        drop(inbox.try_recv());
241        let words = Arc::new(Mutex::new(String::new()));
242        assistant.write("rules".into(), "question".into(), words.clone(), UiWaker::default(), "found".into());
243        assert_eq!(*words.lock().expect("text"), "found");
244        assert!(inbox.try_recv().is_err(), "it did not reach the model");
245    }
246
247    #[test]
248    fn nothing_arrives_after_the_wait_is_over() {
249        let given_up = Arc::new(AtomicBool::new(true));
250        let text = Arc::new(Mutex::new(String::new()));
251        let mut job = Job { text: text.clone(), given_up, ..Job::blank() };
252        assert!(!job.take("late"));
253        assert_eq!(*text.lock().expect("text"), "");
254    }
255
256    #[test]
257    fn a_choice_stops_when_its_json_closes() {
258        let done = Some(Arc::new(AtomicBool::new(false)));
259        let mut job = Job { done, ..Job::blank() };
260        assert!(job.take(r#"{"tool":"a","args":{"#));
261        assert!(!job.take(r#"}}"#));
262    }
263
264    #[test]
265    fn the_first_word_ends_the_wait_for_an_answer() {
266        let (waiting, ended) = mpsc::channel::<()>();
267        let mut job = Job { waiting: Some(waiting), ..Job::blank() };
268        assert!(job.take("A"));
269        assert!(matches!(ended.try_recv(), Err(mpsc::TryRecvError::Disconnected)));
270    }
271}
ravenrs_ai/assistant.rs

ravenrs_ai/
assistant.rs
