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.)

1 files changed, 171 insertions, 26 deletions

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<OllamaResult> for Result<OllamaChunk, OllamaError> {
+    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<relm4::JoinHandle<()>>,
+    waiting: bool,
+
+    http: soup::Session,
+}
+
+impl SmartSummaryButton {
+    async fn prompt_llm(
+        sender: relm4::Sender<Result<String, Error>>,
+        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<OllamaResult, serde_json::Error> = 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<String, Error>;
 
     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 `<b>` tags
             set_tooltip_markup: Some(gettext("<b>Smart Summary</b>\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<Self>
     ) -> ComponentParts<Self> {
-        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::<String>("llm-endpoint"),
+                    glib::UriFlags::NONE,
+                ).unwrap();
+                let model = settings.get::<String>("smart-summary-model");
+                let system_prompt = settings.get::<String>("smart-summary-system-prompt");
+                let prompt_prefix = settings.get::<String>("smart-summary-prompt-prefix");
+                let prompt_suffix = settings.get::<String>("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<Self>, _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) => {