Stage 3

Cargo.toml

@@ -35,10 +35,13 @@ async-openai       = { version = "0.23" }
 color-eyre         = { version = "0.6" }
 eframe             = { version = "0.27", features = ["persistence"] }
 egui_commonmark    = { version = "0.16" }
+egui_extras        = { version = "0.27", features = ["svg"] }
 enigo              = { version = "0.2" }
 futures            = { version = "0.3" }
 get-selected-text  = { version = "0.1" }
 global-hotkey      = { version = "0.5" }
+# TODO?: remove it since it requires libxml2.
+llm_utils          = { version = "0.0.6" }
 reqwew             = { version = "0.2" }
 serde              = { version = "1.0", features = ["derive"] }
 thiserror          = { version = "1.0" }

src/air.rs

@@ -1,3 +1,5 @@
+// TODO: use unwrap here and return result in other places.
+
 // std
 use std::{
 	sync::{atomic::Ordering, Once},
@@ -12,6 +14,7 @@ use crate::{
 	os::{AppKit, Os},
 	prelude::Result,
 	service::Services,
+	state::State,
 	ui::Uis,
 };
 
@@ -20,20 +23,22 @@ struct AiR {
 	once: Once,
 	runtime: Runtime,
 	components: Components,
+	state: State,
 	services: Services,
 	uis: Uis,
 }
 impl AiR {
-	fn init(ctx: &Context) -> Self {
-		Self::set_fonts(ctx);
+	fn init(ctx: Context) -> Self {
+		Self::set_fonts(&ctx);
 
 		let once = Once::new();
-		let runtime = Runtime::new().expect("runtime must be created");
-		let components = Components::init();
-		let services = Services::init(ctx);
+		let runtime = Runtime::new().unwrap();
+		let mut components = Components::init().unwrap();
+		let state = Default::default();
+		let services = Services::init(&ctx, &runtime, &mut components, &state).unwrap();
 		let uis = Uis::init();
 
-		Self { once, runtime, components, services, uis }
+		Self { once, runtime, components, state, services, uis }
 	}
 
 	fn set_fonts(ctx: &Context) {
@@ -62,36 +67,37 @@ impl AiR {
 	}
 
 	fn try_unhide(&mut self, ctx: &Context) {
-		let to_hidden = self.services.to_hidden.load(Ordering::SeqCst);
+		let to_hide = self.state.to_hide.load(Ordering::SeqCst);
 		let focused = ctx.input(|i| i.focused);
 
-		if to_hidden && !focused {
+		if to_hide && !focused {
 			self.components.active_timer.refresh();
-			self.services.to_hidden.store(true, Ordering::SeqCst);
+			self.state.to_hide.store(true, Ordering::SeqCst);
 
 			// TODO: https://github.com/emilk/egui/discussions/4635.
 			// ctx.send_viewport_cmd(ViewportCommand::Minimized(true));
 			Os::hide();
-		} else if !to_hidden && focused {
+		} else if !to_hide && focused {
 			// TODO: find a better place to initialize this.
 			self.once.call_once(Os::set_move_to_active_space);
-			self.services.to_hidden.store(true, Ordering::SeqCst);
+			self.state.to_hide.store(true, Ordering::SeqCst);
 		}
 	}
 }
 impl App for AiR {
 	fn update(&mut self, ctx: &Context, _: &mut Frame) {
 		let air_ctx = AiRContext {
 			egui_ctx: ctx,
-			runtime: &self.runtime,
 			components: &mut self.components,
-			services: &mut self.services,
+			state: &self.state,
+			services: &self.services,
 		};
 
 		self.uis.draw(air_ctx);
-		// TODO?: these will be called multiple times, move to focus service.
+		// TODO: these will be called multiple times, move to focus service.
 		self.try_unhide(ctx);
 
+		// TODO: move to timer service.
 		if self.components.active_timer.duration() > Duration::from_secs(15) {
 			self.components.active_timer.refresh();
 
@@ -102,16 +108,16 @@ impl App for AiR {
 	}
 
 	fn save(&mut self, _: &mut dyn Storage) {
-		self.components.setting.save().expect("setting must be saved");
+		self.components.setting.save().unwrap();
 	}
 }
 
 #[derive(Debug)]
 pub struct AiRContext<'a> {
 	pub egui_ctx: &'a Context,
-	pub runtime: &'a Runtime,
 	pub components: &'a mut Components,
-	pub services: &'a mut Services,
+	pub state: &'a State,
+	pub services: &'a Services,
 }
 
 pub fn launch() -> Result<()> {
@@ -121,15 +127,14 @@ pub fn launch() -> Result<()> {
 			viewport: ViewportBuilder::default()
 				.with_icon(
 					icon_data::from_png_bytes(include_bytes!("../asset/icon.png").as_slice())
-						.expect("icon must be valid"),
+						.unwrap(),
 				)
 				.with_inner_size((720., 360.))
 				.with_min_inner_size((720., 360.))
 				.with_transparent(true),
-			follow_system_theme: true,
 			..Default::default()
 		},
-		Box::new(|c| Box::new(AiR::init(&c.egui_ctx))),
+		Box::new(|c| Box::new(AiR::init(c.egui_ctx.clone()))),
 	)?;
 
 	Ok(())

src/component.rs

@@ -1,5 +1,8 @@
 // TODO?: refresh trait.
 
+pub mod tokenizer;
+use tokenizer::Tokenizer;
+
 pub mod function;
 
 pub mod keyboard;
@@ -20,20 +23,28 @@ use timer::Timer;
 
 pub mod util;
 
+// std
+use std::sync::Arc;
+// self
+use crate::prelude::*;
+
 #[derive(Debug)]
 pub struct Components {
 	pub active_timer: Timer,
 	pub setting: Setting,
 	pub quote: Quoter,
-	pub openai: OpenAi,
+	pub tokenizer: Tokenizer,
+	pub openai: Arc<OpenAi>,
 }
 impl Components {
-	pub fn init() -> Self {
+	pub fn init() -> Result<Self> {
 		let active_timer = Timer::default();
-		let setting = Setting::load().expect("setting must be loaded");
+		let setting = Setting::load()?;
 		let quote = Quoter::default();
-		let openai = OpenAi::new(setting.ai.clone());
+		let tokenizer = Tokenizer::new(setting.ai.model.as_str());
+		// TODO: no clone.
+		let openai = Arc::new(OpenAi::new(setting.ai.clone()));
 
-		Self { active_timer, setting, quote, openai }
+		Ok(Self { active_timer, setting, quote, tokenizer, openai })
 	}
 }

src/component/function.rs

@@ -1,20 +1,28 @@
-// TODO:detect input language.
+// self
+use super::setting::Translation;
 
 #[derive(Debug)]
 pub enum Function {
-	Polish,
-	// Translate,
+	Rewrite,
+	RewriteDirectly,
+	Translate,
+	TranslateDirectly,
 }
 impl Function {
-	pub fn prompt(&self) -> &'static str {
+	pub fn prompt(&self, setting: &Translation) -> String {
 		match self {
-			Self::Polish =>
+			Self::Rewrite | Self::RewriteDirectly =>
 				"As an English professor, assist me in refining this text. \
 				Amend any grammatical errors and enhance the language to sound more like a native speaker.\
-				Provide the refined text only, without any other things.",
-			// Self::Translate => "\
-			// 	Translate the following text into English.\
-			// 	",
+				Provide the refined text only, without any other things."
+					.into(),
+			Self::Translate | Self::TranslateDirectly => format!(
+				"As a language professor, assist me in translate this text from {} to {}. \
+				Amend any grammatical errors and enhance the language to sound more like a native speaker.\
+				Provide the translated text only, without any other things.",
+				setting.source.as_str(),
+				setting.target.as_str()
+			),
 		}
 	}
 }

src/component/keyboard.rs

@@ -1,15 +1,17 @@
 // crates.io
-use enigo::{Enigo, Settings};
+use enigo::{Enigo, Keyboard as _, Settings};
 // self
 use crate::prelude::*;
 
 #[derive(Debug)]
 pub struct Keyboard {
-	enigo: Enigo,
+	pub enigo: Enigo,
 }
 impl Keyboard {
 	pub fn init() -> Result<Self> {
-		Ok(Self { enigo: Enigo::new(&Settings::default()).map_err(EnigoError::NewCon)? })
+		let enigo = Enigo::new(&Settings::default()).map_err(EnigoError::NewCon)?;
+
+		Ok(Self { enigo })
 	}
 
 	// pub fn copy(&mut self) -> Result<()> {
@@ -19,4 +21,8 @@ impl Keyboard {
 	//
 	// 	Ok(())
 	// }
+
+	pub fn text(&mut self, text: &str) -> Result<()> {
+		Ok(self.enigo.text(text).map_err(EnigoError::Input)?)
+	}
 }

src/component/openai.rs

@@ -1,44 +1,43 @@
 // std
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 // crates.io
 use async_openai::{
 	config::OpenAIConfig,
 	types::{
-		// ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestMessage,
-		ChatCompletionRequestSystemMessageArgs,
-		ChatCompletionRequestUserMessageArgs,
-		CreateChatCompletionRequestArgs,
+		ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestUserMessageArgs,
+		ChatCompletionResponseStream, CreateChatCompletionRequestArgs,
 	},
 	Client,
 };
 use eframe::egui::WidgetText;
-use futures::StreamExt;
 use serde::{Deserialize, Serialize};
-use tokio::runtime::Runtime;
+use tokio::sync::RwLock as TokioRwLock;
 // self
-use crate::{component::setting::Ai, prelude::*};
+use super::setting::Ai;
+use crate::prelude::*;
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct OpenAi {
-	pub client: Arc<Client<OpenAIConfig>>,
-	// TODO: use Mutex.
+	pub client: Arc<TokioRwLock<Client<OpenAIConfig>>>,
 	pub setting: Ai,
-	pub output: Arc<Mutex<String>>,
 }
 impl OpenAi {
 	pub fn new(setting: Ai) -> Self {
-		Self {
-			client: Arc::new(Client::with_config(
-				OpenAIConfig::new()
-					.with_api_base("https://aihubmix.com/v1")
-					.with_api_key(&setting.api_key),
-			)),
-			setting,
-			output: Arc::new(Mutex::new(Default::default())),
-		}
+		let client = Arc::new(TokioRwLock::new(Client::with_config(
+			OpenAIConfig::new().with_api_base(&setting.api_base).with_api_key(&setting.api_key),
+		)));
+
+		Self { client, setting }
+	}
+
+	pub fn reload(&mut self, setting: Ai) {
+		*self.client.blocking_write() = Client::with_config(
+			OpenAIConfig::new().with_api_base(&setting.api_base).with_api_key(&setting.api_key),
+		);
+		self.setting = setting;
 	}
 
-	pub fn chat(&self, runtime: &Runtime, prompt: &str, content: &str) -> Result<()> {
+	pub async fn chat(&self, prompt: &str, content: &str) -> Result<ChatCompletionResponseStream> {
 		let msg = [
 			ChatCompletionRequestSystemMessageArgs::default().content(prompt).build()?.into(),
 			ChatCompletionRequestUserMessageArgs::default().content(content).build()?.into(),
@@ -49,34 +48,9 @@ impl OpenAi {
 			.max_tokens(4_096_u16)
 			.messages(&msg)
 			.build()?;
-		let c = self.client.clone();
-		let o = self.output.clone();
-
-		runtime.spawn(async move {
-			match c.chat().create_stream(req).await {
-				Ok(mut s) => {
-					o.lock().expect("output must be available").clear();
-
-					while let Some(r) = s.next().await {
-						match r {
-							Ok(resp) => {
-								resp.choices.iter().for_each(|c| {
-									if let Some(c) = &c.delta.content {
-										o.lock().expect("output must be available").push_str(c);
-									}
-								});
-							},
-							Err(e) => println!("failed to receive chat response: {e}"),
-						}
-					}
-				},
-				Err(e) => {
-					tracing::error!("failed to create chat stream: {e}");
-				},
-			}
-		});
+		let stream = self.client.read().await.chat().create_stream(req).await?;
 
-		Ok(())
+		Ok(stream)
 	}
 }