From a26610044483ddd323479d748af401382b8df210 Mon Sep 17 00:00:00 2001 From: Vika Date: Wed, 4 Sep 2024 19:49:26 +0300 Subject: Smart Summary is now working! There's no preferences dialog, so you can't really adjust the prompt or the model it uses. The default settings work well for me. You may want to tweak them depending on your model preferences and compute budget. (Not many can afford to run Llama3-8B at high quantization. Conversely, you might have a better GPU than me and wish to run a 27B model or bigger.) --- src/components/post_editor.rs | 10 +- src/components/smart_summary.rs | 197 ++++++++++++++++++++++++++++++++++------ src/lib.rs | 2 +- 3 files changed, 179 insertions(+), 30 deletions(-) (limited to 'src') diff --git a/src/components/post_editor.rs b/src/components/post_editor.rs index 8c81b6b..f2268ad 100644 --- a/src/components/post_editor.rs +++ b/src/components/post_editor.rs @@ -115,7 +115,7 @@ pub enum Input { #[relm4::component(pub)] impl Component for PostEditor { - type Init = Option; + type Init = (::Init, Option); type Output = Option; type Input = Input; type CommandOutput = (); @@ -330,7 +330,11 @@ impl Component for Post } } - fn init(init: Self::Init, root: Self::Root, sender: ComponentSender) -> ComponentParts { + fn init( + (http, init): Self::Init, + root: Self::Root, + sender: ComponentSender + ) -> ComponentParts { let mut model = Self { smart_summary_busy_guard: None, sending: false, @@ -356,7 +360,7 @@ impl Component for Post wide_layout: gtk::GridLayout::new(), smart_summary: components::SmartSummaryButton::builder() - .launch(()) + .launch(http) .forward(sender.input_sender(), Input::SmartSummary), tracker: Default::default(), diff --git a/src/components/smart_summary.rs b/src/components/smart_summary.rs index 9da67af..050a52c 100644 --- a/src/components/smart_summary.rs +++ b/src/components/smart_summary.rs @@ -1,15 +1,144 @@ +use futures::AsyncBufReadExt; +use gio::prelude::SettingsExtManual; +use soup::prelude::*; use adw::prelude::*; use gettextrs::*; use relm4::{gtk, prelude::{Component, ComponentParts}, ComponentSender}; +// All of this is incredibly minimalist. +// This should be expanded later. +#[derive(Debug, serde::Serialize)] +struct OllamaRequest { + model: String, + prompt: String, + system: String, +} + +#[derive(Debug, serde::Deserialize)] +struct OllamaChunk { + response: String, + done: bool, +} + +#[derive(Debug, serde::Deserialize)] +struct OllamaError { + error: String +} +impl std::error::Error for OllamaError {} +impl std::fmt::Display for OllamaError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.error.fmt(f) + } +} + +#[derive(serde::Deserialize)] +#[serde(untagged)] +enum OllamaResult { + Ok(OllamaChunk), + Err(OllamaError), +} + +impl From for Result { + fn from(val: OllamaResult) -> Self { + match val { + OllamaResult::Ok(chunk) => Ok(chunk), + OllamaResult::Err(err) => Err(err) + } + } +} + + #[derive(Debug, Default)] pub(crate) struct SmartSummaryButton { - busy: bool, + task: Option>, + waiting: bool, + + http: soup::Session, +} + +impl SmartSummaryButton { + async fn prompt_llm( + sender: relm4::Sender>, + http: soup::Session, + endpoint: glib::Uri, + model: String, + system_prompt: String, + prompt_prefix: String, + mut prompt_suffix: String, + text: String, + ) { + let endpoint = endpoint.parse_relative("./api/generate", glib::UriFlags::NONE).unwrap(); + log::debug!("endpoint: {}, model: {}", endpoint, model); + log::debug!("system prompt: {}", system_prompt); + + let msg = soup::Message::from_uri( + "POST", + &endpoint + ); + + if !prompt_suffix.is_empty() { + prompt_suffix = String::from("\n\n") + &prompt_suffix; + } + msg.set_request_body_from_bytes(Some("application/json"), + Some(&glib::Bytes::from_owned(serde_json::to_vec(&OllamaRequest { + model, system: system_prompt, prompt: format!("{}\n\n{}{}", prompt_prefix, text, prompt_suffix), + }).unwrap())) + ); + + let mut stream = match http.send_future(&msg, glib::Priority::DEFAULT).await { + Ok(stream) => stream.into_async_buf_read(128), + Err(err) => { + let _ = sender.send(Err(err.into())); + return + } + }; + log::debug!("response: {:?} ({})", msg.status(), msg.reason_phrase().unwrap_or_default()); + let mut buffer = Vec::new(); + const DELIM: u8 = b'\n'; + loop { + let len = match stream.read_until(DELIM, &mut buffer).await { + Ok(len) => len, + Err(err) => { + let _ = sender.send(Err(err.into())); + return + } + }; + log::debug!("Got chunk ({} bytes): {}", len, String::from_utf8_lossy(&buffer)); + let response: Result = serde_json::from_slice(&buffer[..len]); + match response.map(Result::from) { + Ok(Ok(OllamaChunk { response: chunk, done })) => { + if !chunk.is_empty() { + sender.emit(Ok(chunk)); + } + if done { + sender.emit(Ok(String::new())); + return + } + }, + Ok(Err(err)) => { + sender.emit(Err(err.into())); + return + } + Err(err) => { + sender.emit(Err(err.into())); + return + } + } + buffer.truncate(0); + } + } } #[derive(Debug, thiserror::Error)] pub(crate) enum Error { - + #[error("glib error: {0}")] + Glib(#[from] glib::Error), + #[error("json error: {0}")] + Json(#[from] serde_json::Error), + #[error("ollama error: {0}")] + Ollama(#[from] OllamaError), + #[error("i/o error: {0}")] + Io(#[from] std::io::Error) } #[derive(Debug)] @@ -33,7 +162,7 @@ impl Component for SmartSummaryButton { type Input = Input; type Output = Output; - type Init = (); + type Init = soup::Session; type CommandOutput = Result; view! { @@ -42,11 +171,11 @@ impl Component for SmartSummaryButton { gtk::Button { connect_clicked => Input::ButtonPressed, #[watch] - set_sensitive: !model.busy, + set_sensitive: !(model.task.is_some() || model.waiting), // TRANSLATORS: please keep the newline and `` tags set_tooltip_markup: Some(gettext("Smart Summary\nAsk a language model for a single-sentence summary.")).as_deref(), - if model.busy { + if model.task.is_some() || model.waiting { gtk::Spinner { set_spinning: true } } else { gtk::Label { set_markup: "✨" } @@ -56,11 +185,14 @@ impl Component for SmartSummaryButton { } fn init( - _init: Self::Init, + init: Self::Init, root: Self::Root, sender: ComponentSender ) -> ComponentParts { - let model = Self::default(); + let model = Self { + http: init, + ..Self::default() + }; let widgets = view_output!(); ComponentParts { model, widgets } @@ -74,30 +206,41 @@ impl Component for SmartSummaryButton { ) { match msg { Input::Cancel => { - self.busy = false; - log::debug!("Parent component asked us to cancel."); + self.waiting = false; + if let Some(task) = self.task.take() { + log::debug!("Parent component asked us to cancel."); + task.abort(); + } else { + log::warn!("Parent component asked us to cancel, but we're not running a task."); + } }, Input::ButtonPressed => if let Ok(()) = sender.output(Output::Start) { - self.busy = true; + self.waiting = true; log::debug!("Requesting text to summarize from parent component..."); + // TODO: set timeout in case parent component never replies + // This shouldn't happen, but I feel like we should handle this case. }, Input::Text(text) => { log::debug!("Would generate summary for the following text:\n{}", text); - sender.command(|sender, shutdown| shutdown.register(async move { - tokio::time::sleep(std::time::Duration::from_millis(450)).await; - - for i in ["I'", "m ", "sorry,", " I", " am ", "unable", " to", " ", "write", " you ", "a summary.", " I", " am", " not ", "really ", "an ", "LLM."] { - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - - if sender.send(Ok(i.to_string())).is_err() { - return - }; - } - - log::debug!("Done with the summary."); - let _ = sender.send(Ok(String::new())); - }).drop_on_shutdown()); + log::debug!("XDG_DATA_DIRS={:?}", std::env::var("XDG_DATA_DIRS")); + let settings = gio::Settings::new(crate::APPLICATION_ID); + // We shouldn't let the user record a bad setting anyway. + let endpoint = glib::Uri::parse( + &settings.get::("llm-endpoint"), + glib::UriFlags::NONE, + ).unwrap(); + let model = settings.get::("smart-summary-model"); + let system_prompt = settings.get::("smart-summary-system-prompt"); + let prompt_prefix = settings.get::("smart-summary-prompt-prefix"); + let prompt_suffix = settings.get::("smart-summary-prompt-suffix"); + let sender = sender.command_sender().clone(); + relm4::spawn_local(Self::prompt_llm( + sender, self.http.clone(), + endpoint, model, system_prompt, + prompt_prefix, prompt_suffix, + text + )); } } } @@ -105,11 +248,13 @@ impl Component for SmartSummaryButton { fn update_cmd(&mut self, msg: Self::CommandOutput, sender: ComponentSender, _root: &Self::Root) { match msg { Ok(chunk) if chunk.is_empty() => { - self.busy = false; + self.task = None; + self.waiting = false; let _ = sender.output(Output::Done); }, Err(err) => { - self.busy = false; + self.task = None; + self.waiting = false; let _ = sender.output(Output::Error(err)); } Ok(chunk) => { diff --git a/src/lib.rs b/src/lib.rs index 7f2a78d..30cba6f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -255,7 +255,7 @@ impl AsyncComponent for App { secret_schema, post_editor: components::PostEditor::builder() - .launch(None) + .launch((http.clone(), None)) .forward(sender.input_sender(), Self::Input::PostEditor), signin: components::SignIn::builder() .launch((glib::Uri::parse( -- cgit 1.4.1