1
//! The agent handle.
2
3
+use std::future::Future;
4
use std::time::Duration;
5
6
use log::{error, info};
63
&self.engine
64
}
65
65
- /// Load a GGUF model, reporting each transition.
66
+ /// Run any load, reporting the transitions around it.
67
///
67
- /// Emits [`EngineStatus::Loading`] before starting and then either
68
+ /// Emits [`EngineStatus::Loading`] before `load` starts and then either
69
/// [`Ready`](EngineStatus::Ready) or [`Error`](EngineStatus::Error). That
69
- /// sequence is the whole reason this method exists: calling
70
- /// [`ChatEngine::load_gguf_model`] directly leaves the UI with no way to
71
- /// know a multi-minute download has begun.
70
+ /// sequence is the whole reason this method exists: calling a
71
+ /// [`ChatEngine`] load path directly leaves the UI with no way to know a
72
+ /// multi-minute download has begun.
73
///
73
- /// Downloads the weights first if they aren't in the local cache, so the
74
- /// `Loading` state can last a long time on a cold start.
75
- pub async fn load(
76
- &self,
77
- config: GgufModelConfig,
78
- system_prompt: Option<String>,
79
- sampling: Option<SamplingConfig>,
80
- ) -> Result<Duration, InferenceError> {
81
- let model_name = config.display_name.clone();
74
+ /// Ed takes a closure rather than wrapping each load path by name because
75
+ /// `onde`'s load APIs take types Ed has no business in its public
76
+ /// signature — `ChatEngine::load_assigned_model` takes an `Environment`
77
+ /// from the smbCloud SDK, and *which* crate that type comes from changed
78
+ /// between `onde` releases. Naming it here would drag that dependency into
79
+ /// every host and pin them all to Ed's version of it. Passing a closure
80
+ /// keeps Ed's dependency surface to `onde` alone and works for load paths
81
+ /// that don't exist yet.
82
+ ///
83
+ /// ```no_run
84
+ /// # use ed_agent::{Ed, GgufModelConfig};
85
+ /// # async fn example(ed: Ed) {
86
+ /// ed.load_with("Qwen 2.5", |engine| {
87
+ /// engine.load_gguf_model(GgufModelConfig::qwen25_1_5b(), None, None)
88
+ /// })
89
+ /// .await
90
+ /// .expect("load");
91
+ /// # }
92
+ /// ```
93
+ ///
94
+ /// `model_name` is what the host wants the UI to show while loading; it is
95
+ /// reported with each transition and needn't match what the engine ends up
96
+ /// calling the model.
97
+ pub async fn load_with<'a, F, Fut>(
98
+ &'a self,
99
+ model_name: &str,
100
+ load: F,
101
+ ) -> Result<Duration, InferenceError>
102
+ where
103
+ F: FnOnce(&'a ChatEngine) -> Fut,
104
+ Fut: Future<Output = Result<Duration, InferenceError>>,
105
+ {
106
info!("ed: loading model {model_name}");
107
self.sink
84
- .status_changed(EngineStatus::Loading, Some(&model_name), None);
108
+ .status_changed(EngineStatus::Loading, Some(model_name), None);
109
86
- match self
87
- .engine
88
- .load_gguf_model(config, system_prompt, sampling)
89
- .await
90
- {
110
+ match load(&self.engine).await {
111
Ok(elapsed) => {
112
info!("ed: loaded {model_name} in {}", format_duration(elapsed));
113
self.sink
94
- .status_changed(EngineStatus::Ready, Some(&model_name), None);
114
+ .status_changed(EngineStatus::Ready, Some(model_name), None);
115
Ok(elapsed)
116
}
117
Err(err) => {
118
error!("ed: failed to load {model_name}: {err}");
119
let message = err.to_string();
120
self.sink
101
- .status_changed(EngineStatus::Error, Some(&model_name), Some(&message));
121
+ .status_changed(EngineStatus::Error, Some(model_name), Some(&message));
122
Err(err)
123
}
124
}
125
}
126
127
+ /// Load a GGUF model, reporting each transition.
128
+ ///
129
+ /// A convenience over [`load_with`](Ed::load_with) for the most common
130
+ /// path. Downloads the weights first if they aren't in the local cache, so
131
+ /// the `Loading` state can last a long time on a cold start.
132
+ pub async fn load(
133
+ &self,
134
+ config: GgufModelConfig,
135
+ system_prompt: Option<String>,
136
+ sampling: Option<SamplingConfig>,
137
+ ) -> Result<Duration, InferenceError> {
138
+ let model_name = config.display_name.clone();
139
+ self.load_with(&model_name, |engine| {
140
+ engine.load_gguf_model(config, system_prompt, sampling)
141
+ })
142
+ .await
143
+ }
144
+
145
/// Unload the current model, returning its display name if one was loaded.
146
pub async fn unload(&self) -> Option<String> {
147
let unloaded = self.engine.unload_model().await;
208
#[derive(Default)]
209
struct RecordingSink {
210
statuses: Mutex<Vec<(EngineStatus, Option<String>)>>,
211
+ errors: Mutex<Vec<Option<String>>>,
212
replies: Mutex<Vec<ChatReply>>,
213
}
214
216
fn statuses(&self) -> Vec<(EngineStatus, Option<String>)> {
217
self.statuses.lock().unwrap().clone()
218
}
219
+
220
+ fn errors(&self) -> Vec<Option<String>> {
221
+ self.errors.lock().unwrap().clone()
222
+ }
223
}
224
225
impl StatusSink for &'static RecordingSink {
183
- fn status_changed(&self, status: EngineStatus, model_name: Option<&str>, _: Option<&str>) {
226
+ fn status_changed(
227
+ &self,
228
+ status: EngineStatus,
229
+ model_name: Option<&str>,
230
+ error: Option<&str>,
231
+ ) {
232
self.statuses
233
.lock()
234
.unwrap()
235
.push((status, model_name.map(str::to_owned)));
236
+ self.errors.lock().unwrap().push(error.map(str::to_owned));
237
}
238
239
fn replied(&self, reply: &ChatReply) {
283
ed.unload().await;
284
assert_eq!(ed.info().await.status, EngineStatus::Unloaded);
285
}
286
+
287
+ #[tokio::test]
288
+ async fn a_successful_load_reports_loading_then_ready() {
289
+ // The ordering is the feature: a UI that only ever sees the final
290
+ // state has no way to put up a spinner for a load that can take
291
+ // minutes.
292
+ let sink = sink();
293
+ let ed = Ed::with_sink(sink);
294
+
295
+ let elapsed = ed
296
+ .load_with("Test model", |_| async { Ok(Duration::from_secs(2)) })
297
+ .await
298
+ .expect("the closure succeeded");
299
+
300
+ assert_eq!(elapsed, Duration::from_secs(2));
301
+ assert_eq!(
302
+ sink.statuses(),
303
+ vec![
304
+ (EngineStatus::Loading, Some("Test model".to_owned())),
305
+ (EngineStatus::Ready, Some("Test model".to_owned())),
306
+ ]
307
+ );
308
+ }
309
+
310
+ #[tokio::test]
311
+ async fn a_failed_load_reports_loading_then_error_with_the_reason() {
312
+ let sink = sink();
313
+ let ed = Ed::with_sink(sink);
314
+
315
+ let err = ed
316
+ .load_with("Test model", |_| async {
317
+ Err(InferenceError::Other {
318
+ reason: "no weights".to_owned(),
319
+ })
320
+ })
321
+ .await
322
+ .expect_err("the closure failed");
323
+
324
+ assert!(err.to_string().contains("no weights"));
325
+ assert_eq!(
326
+ sink.statuses(),
327
+ vec![
328
+ (EngineStatus::Loading, Some("Test model".to_owned())),
329
+ (EngineStatus::Error, Some("Test model".to_owned())),
330
+ ]
331
+ );
332
+ assert_eq!(sink.errors(), vec![None, Some("no weights".to_owned())]);
333
+ }
334
}