Refactor chat

README.md

@@ -16,8 +16,7 @@
 Built upon [egui](https://github.com/emilk/egui), a fast and cross-platform GUI toolkit written in pure Rust.
 
 ### Components
-These items either have their own `refresh` logic or do not require frequent refreshing.
-They are not time-sensitive, and their `refresh` method will be called at specific intervals (e.g., every 15 seconds).
+These items are static and they used to be called by other stuffs.
 
 ### OS
 Provides wrapped APIs to interact with the operating system.
@@ -26,5 +25,8 @@ Provides wrapped APIs to interact with the operating system.
 These items are time-sensitive and require frequent checking or updating.
 They will be spawned as separate threads and run in the background.
 
+### State
+Mutable version of the components. Usually, they are `Arc<Mutex<Components>>` in order to sync the state between service and UI.
+
 ### UI
 The user interface components.

src/air.rs

@@ -16,17 +16,17 @@ struct AiR {
 	uis: Uis,
 }
 impl AiR {
-	fn init(ctx: &Context) -> Result<Self> {
+	fn new(ctx: &Context) -> Result<Self> {
 		Self::set_fonts(ctx);
 
 		// To enable SVG.
 		egui_extras::install_image_loaders(ctx);
 
 		let once = Once::new();
-		let components = Components::init()?;
+		let components = Components::new()?;
 		let state = Default::default();
-		let services = Services::init(ctx, &components, &state)?;
-		let uis = Uis::init();
+		let services = Services::new(ctx, &components, &state)?;
+		let uis = Uis::new();
 
 		Ok(Self { once, components, state, services, uis })
 	}
@@ -62,7 +62,7 @@ impl App for AiR {
 			egui_ctx: ctx,
 			components: &mut self.components,
 			state: &self.state,
-			services: &self.services,
+			services: &mut self.services,
 		};
 
 		self.uis.draw(air_ctx);
@@ -108,7 +108,7 @@ pub struct AiRContext<'a> {
 	pub egui_ctx: &'a Context,
 	pub components: &'a mut Components,
 	pub state: &'a State,
-	pub services: &'a Services,
+	pub services: &'a mut Services,
 }
 
 pub fn launch() -> Result<()> {
@@ -125,7 +125,7 @@ pub fn launch() -> Result<()> {
 				.with_transparent(true),
 			..Default::default()
 		},
-		Box::new(|c| Ok(Box::new(AiR::init(&c.egui_ctx).unwrap()))),
+		Box::new(|c| Ok(Box::new(AiR::new(&c.egui_ctx).unwrap()))),
 	)?;
 
 	Ok(())

src/component.rs

@@ -8,7 +8,6 @@ pub mod keyboard;
 pub mod net;
 
 pub mod openai;
-use openai::OpenAi;
 
 pub mod quote;
 
@@ -17,46 +16,29 @@ use setting::Setting;
 
 pub mod util;
 
-// std
-use std::sync::Arc;
-// crates.io
-use tokio::sync::Mutex;
 // self
 use crate::prelude::*;
 
 #[derive(Debug)]
 pub struct Components {
 	pub setting: Setting,
-	// Keyboard didn't implement `Send`, can't use it between threads.
-	// pub keyboard: Arc<Mutex<Keyboard>>,
-	// TODO?: move the lock to somewhere else.
-	pub openai: Arc<Mutex<OpenAi>>,
 	#[cfg(feature = "tokenizer")]
 	pub tokenizer: Tokenizer,
 }
 impl Components {
-	pub fn init() -> Result<Self> {
+	pub fn new() -> Result<Self> {
 		let setting = Setting::load()?;
 
 		// TODO: https://github.com/emilk/egui/discussions/4670.
 		debug_assert_eq!(setting.ai.temperature, setting.ai.temperature * 10. / 10.);
 
-		let openai = Arc::new(Mutex::new(OpenAi::new(setting.ai.clone())));
 		#[cfg(feature = "tokenizer")]
 		let tokenizer = Tokenizer::new(setting.ai.model.as_str());
 
 		Ok(Self {
 			setting,
-			openai,
 			#[cfg(feature = "tokenizer")]
 			tokenizer,
 		})
 	}
-
-	// TODO?: move to somewhere else.
-	pub fn reload_openai(&self) {
-		tracing::info!("reloading openai component");
-
-		self.openai.blocking_lock().reload(self.setting.ai.clone());
-	}
 }

src/component/keyboard.rs

@@ -7,7 +7,7 @@ use crate::prelude::*;
 #[derive(Debug)]
 pub struct Keyboard(pub Enigo);
 impl Keyboard {
-	pub fn init() -> Result<Self> {
+	pub fn new() -> Result<Self> {
 		Ok(Self(Enigo::new(&Settings::default()).map_err(EnigoError::NewCon)?))
 	}
 

src/component/openai.rs

@@ -16,22 +16,16 @@ use crate::prelude::*;
 #[derive(Debug)]
 pub struct OpenAi {
 	pub client: Client<OpenAIConfig>,
-	pub setting: Ai,
+	pub model: Model,
+	pub temperature: f32,
 }
 impl OpenAi {
 	pub fn new(setting: Ai) -> Self {
-		let client = Client::with_config(
-			OpenAIConfig::new().with_api_base(&setting.api_base).with_api_key(&setting.api_key),
-		);
+		let Ai { api_base, api_key, model, temperature } = setting;
+		let client =
+			Client::with_config(OpenAIConfig::new().with_api_base(api_base).with_api_key(api_key));
 
-		Self { client, setting }
-	}
-
-	pub fn reload(&mut self, setting: Ai) {
-		self.client = Client::with_config(
-			OpenAIConfig::new().with_api_base(&setting.api_base).with_api_key(&setting.api_key),
-		);
-		self.setting = setting;
+		Self { client, model, temperature }
 	}
 
 	pub async fn chat(&self, prompt: &str, content: &str) -> Result<ChatCompletionResponseStream> {
@@ -40,8 +34,8 @@ impl OpenAi {
 			ChatCompletionRequestUserMessageArgs::default().content(content).build()?.into(),
 		];
 		let req = CreateChatCompletionRequestArgs::default()
-			.model(self.setting.model.as_str())
-			.temperature(self.setting.temperature)
+			.model(self.model.as_str())
+			.temperature(self.temperature)
 			.max_tokens(4_096_u16)
 			.messages(&msg)
 			.build()?;

src/component/setting.rs

@@ -102,7 +102,7 @@ impl Default for Rewrite {
 		Self {
 			prompt: "As language professor, assist me in refining this text. \
 				Amend any grammatical errors and enhance the language to sound more like a native speaker.\
-				Just provide the refined text only, without any other things."
+				Just provide the refined text only, without any other things:"
 				.into(),
 		}
 	}
@@ -128,7 +128,7 @@ impl Default for Translation {
 	fn default() -> Self {
 		Self {
 			prompt: "As a language professor, amend any grammatical errors and enhance the language to sound more like a native speaker. \
-				Provide the translated text only, without any other things.".into(),
+				Provide the translated text only, without any other things:".into(),
 			a: Language::ZhCn,
 			b: Language::EnGb,
 		}
@@ -144,7 +144,7 @@ pub enum Language {
 	EnGb,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(rename_all = "kebab-case")]
 pub struct Hotkeys {
 	pub rewrite: HotKey,

src/service.rs

@@ -1,3 +1,6 @@
+mod chat;
+use chat::Chat;
+
 mod hotkey;
 use hotkey::Hotkey;
 
@@ -7,6 +10,11 @@ use keyboard::Keyboard;
 mod quoter;
 use quoter::Quoter;
 
+// std
+use std::sync::{
+	atomic::{AtomicBool, Ordering},
+	Arc,
+};
 // crates.io
 use eframe::egui::Context;
 use tokio::runtime::Runtime;
@@ -18,21 +26,39 @@ pub struct Services {
 	pub keyboard: Keyboard,
 	pub rt: Option<Runtime>,
 	pub quoter: Quoter,
+	pub is_chatting: Arc<AtomicBool>,
+	pub chat: Chat,
 	pub hotkey: Hotkey,
 }
 impl Services {
-	pub fn init(ctx: &Context, components: &Components, state: &State) -> Result<Self> {
-		let keyboard = Keyboard::init();
+	pub fn new(ctx: &Context, components: &Components, state: &State) -> Result<Self> {
+		let keyboard = Keyboard::new();
 		let rt = Runtime::new()?;
-		let quoter = Quoter::init(&rt, state.chat.quote.clone());
-		let hotkey = Hotkey::init(ctx, keyboard.clone(), &rt, components, state)?;
+		let quoter = Quoter::new(&rt, state.chat.quote.clone());
+		let is_chatting = Arc::new(AtomicBool::new(false));
+		let chat = Chat::new(
+			keyboard.clone(),
+			&rt,
+			is_chatting.clone(),
+			components.setting.ai.clone(),
+			components.setting.chat.clone(),
+			state.chat.input.clone(),
+			state.chat.output.clone(),
+		);
+		let hotkey =
+			Hotkey::new(ctx, keyboard.clone(), &rt, &components.setting.hotkeys, chat.tx.clone())?;
+
+		Ok(Self { keyboard, rt: Some(rt), quoter, is_chatting, chat, hotkey })
+	}
 
-		Ok(Self { keyboard, rt: Some(rt), quoter, hotkey })
+	pub fn is_chatting(&self) -> bool {
+		self.is_chatting.load(Ordering::SeqCst)
 	}
 
 	pub fn abort(&mut self) {
 		self.keyboard.abort();
 		self.quoter.abort();
+		self.chat.abort();
 		self.hotkey.abort();
 
 		if let Some(rt) = self.rt.take() {

src/service/chat.rs

@@ -0,0 +1,100 @@
+// std
+use std::{
+	sync::{
+		atomic::{AtomicBool, Ordering},
+		mpsc::{self, Sender},
+		Arc,
+	},
+	time::Duration,
+};
+// crates.io
+use futures::StreamExt;
+use parking_lot::RwLock;
+use tokio::{runtime::Runtime, task::AbortHandle, time};
+// self
+use super::keyboard::Keyboard;
+use crate::component::{
+	function::Function,
+	openai::OpenAi,
+	setting::{Ai, Chat as ChatSetting},
+};
+
+pub type ChatArgs = (Function, String, bool);
+
+#[derive(Debug)]
+pub struct Chat {
+	pub tx: Sender<ChatArgs>,
+	abort_handle: AbortHandle,
+}
+impl Chat {
+	pub fn new(
+		keyboard: Keyboard,
+		rt: &Runtime,
+		is_chatting: Arc<AtomicBool>,
+		ai_setting: Ai,
+		chat_setting: ChatSetting,
+		input: Arc<RwLock<String>>,
+		output: Arc<RwLock<String>>,
+	) -> Self {
+		let openai = OpenAi::new(ai_setting);
+		let (tx, rx) = mpsc::channel();
+		// TODO: handle the error.
+		let abort_handle = rt
+			.spawn(async move {
+				loop {
+					let (func, content, type_in): ChatArgs = rx.recv().unwrap();
+
+					is_chatting.store(true, Ordering::SeqCst);
+
+					tracing::info!("func: {func:?}");
+					tracing::debug!("content: {content}");
+
+					input.write().clone_from(&content);
+					output.write().clear();
+
+					let mut stream =
+						openai.chat(&func.prompt(&chat_setting), &content).await.unwrap();
+
+					while let Some(r) = stream.next().await {
+						for s in r.unwrap().choices.into_iter().filter_map(|c| c.delta.content) {
+							output.write().push_str(&s);
+
+							// TODO?: move to outside of the loop.
+							if type_in {
+								keyboard.text(s);
+							}
+						}
+					}
+
+					// Allow the UI a moment to refresh the content.
+					time::sleep(Duration::from_millis(50)).await;
+
+					is_chatting.store(false, Ordering::SeqCst);
+				}
+			})
+			.abort_handle();
+
+		Self { abort_handle, tx }
+	}
+
+	pub fn abort(&self) {
+		self.abort_handle.abort();
+	}
+
+	// TODO: fix clippy.
+	#[allow(clippy::too_many_arguments)]
+	pub fn renew(
+		&mut self,
+		keyboard: Keyboard,
+		rt: &Runtime,
+		is_chatting: Arc<AtomicBool>,
+		ai_setting: Ai,
+		chat_setting: ChatSetting,
+		input: Arc<RwLock<String>>,
+		output: Arc<RwLock<String>>,
+	) {
+		self.abort();
+
+		*self = Self::new(keyboard, rt, is_chatting, ai_setting, chat_setting, input, output);
+	}
+}