#![cfg(feature = "smart-summary")]
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)]
pub(crate) struct OllamaRequest {
model: String,
prompt: String,
system: String,
}
#[derive(Debug, serde::Deserialize)]
pub(crate) struct OllamaChunk {
response: String,
done: bool,
}
#[derive(Debug, serde::Deserialize)]
pub(crate) 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)]
pub(crate) 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 {
task: Option<relm4::JoinHandle<()>>,
waiting: bool,
http: soup::Session,
}
impl SmartSummaryButton {
async fn summarize(
sender: relm4::Sender<Result<String, Error>>,
http: soup::Session,
text: String,
) {
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.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 mut prompt_suffix = settings.get::<String>("smart-summary-prompt-suffix");
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::with_capacity(2048);
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}")]
#[allow(private_interfaces)]
Ollama(#[from] OllamaError),
#[error("i/o error: {0}")]
Io(#[from] std::io::Error)
}
#[derive(Debug)]
pub(crate) enum Input {
#[doc(hidden)] ButtonPressed,
Text(String),
Cancel,
}
#[derive(Debug)]
pub(crate) enum Output {
Start,
Chunk(String),
Done,
Error(Error)
}
#[relm4::component(pub(crate))]
impl Component for SmartSummaryButton {
type Input = Input;
type Output = Output;
type Init = soup::Session;
type CommandOutput = Result<String, Error>;
view! {
#[root]
#[name = "button"]
gtk::Button {
connect_clicked => Input::ButtonPressed,
#[watch]
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.task.is_some() || model.waiting {
gtk::Spinner { set_spinning: true }
} else {
gtk::Label { set_markup: "✨" }
}
}
}
fn init(
init: Self::Init,
root: Self::Root,
sender: ComponentSender<Self>
) -> ComponentParts<Self> {
let model = Self {
http: init,
..Self::default()
};
let widgets = view_output!();
ComponentParts { model, widgets }
}
fn update(
&mut self,
msg: Self::Input,
sender: ComponentSender<Self>,
_root: &Self::Root
) {
match msg {
Input::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.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);
log::debug!("XDG_DATA_DIRS={:?}", std::env::var("XDG_DATA_DIRS"));
let sender = sender.command_sender().clone();
relm4::spawn_local(Self::summarize(
sender, self.http.clone(), text
));
}
}
}
fn update_cmd(&mut self, msg: Self::CommandOutput, sender: ComponentSender<Self>, _root: &Self::Root) {
match msg {
Ok(chunk) if chunk.is_empty() => {
self.task = None;
self.waiting = false;
let _ = sender.output(Output::Done);
},
Err(err) => {
self.task = None;
self.waiting = false;
let _ = sender.output(Output::Error(err));
}
Ok(chunk) => {
let _ = sender.output(Output::Chunk(chunk));
},
}
}
}