add stt model getter and setter

Author: yujonglee

Cargo.lock

@@ -15,11 +15,13 @@ export default function LocalAI() {
   const sttRunning = useQuery({
     queryKey: ["local-stt", "running"],
     queryFn: async () => localSttCommands.isServerRunning(),
+    refetchInterval: 3000,
   });
 
   const llmRunning = useQuery({
     queryKey: ["local-llm", "running"],
     queryFn: async () => localLlmCommands.isServerRunning(),
+    refetchInterval: 3000,
   });
 
   return (
@@ -44,6 +46,20 @@ function SpeechToTextDetails(
     },
   });
 
+  // const currentModel = useQuery({
+  //   queryKey: ["local-stt", "current-model"],
+  //   queryFn: () => localSttCommands.getCurrentModel(),
+  // });
+
+  // const setCurrentModel = useMutation({
+  //   mutationFn: async (model: SupportedModel) => {
+  //     await localSttCommands.setCurrentModel(model);
+  //   },
+  //   onSuccess: () => {
+  //     queryClient.invalidateQueries({ queryKey: ["local-stt", "current-model"] });
+  //   },
+  // });
+
   return (
     <div className="space-y-4">
       <div className="flex items-center justify-between rounded-lg border p-4">

apps/desktop/src/components/settings/views/ai.tsx

@@ -29,11 +29,13 @@ hypr-file = { workspace = true }
 hypr-listener-interface = { workspace = true }
 
 tauri = { workspace = true, features = ["test"] }
+tauri-plugin-store2 = { workspace = true }
 tauri-specta = { workspace = true, features = ["derive", "typescript"] }
 
 serde = { workspace = true }
 serde_json = { workspace = true }
 specta = { workspace = true }
+strum = { workspace = true, features = ["derive"] }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 

plugins/local-stt/Cargo.toml

@@ -4,6 +4,8 @@ const COMMANDS: &[&str] = &[
     "download_model",
     "start_server",
     "stop_server",
+    "get_current_model",
+    "set_current_model",
 ];
 
 fn main() {

plugins/local-stt/build.rs

@@ -21,6 +21,12 @@ async startServer() : Promise<null> {
 },
 async stopServer() : Promise<null> {
     return await TAURI_INVOKE("plugin:local-stt|stop_server");
+},
+async getCurrentModel() : Promise<SupportedModel> {
+    return await TAURI_INVOKE("plugin:local-stt|get_current_model");
+},
+async setCurrentModel(model: SupportedModel) : Promise<null> {
+    return await TAURI_INVOKE("plugin:local-stt|set_current_model", { model });
 }
 }
 
@@ -34,6 +40,7 @@ async stopServer() : Promise<null> {
 
 /** user-defined types **/
 
+export type SupportedModel = "QuantizedTiny" | "QuantizedTinyEn" | "QuantizedLargeV3Turbo"
 export type TAURI_CHANNEL<TSend> = null
 
 /** tauri-specta globals **/

plugins/local-stt/js/bindings.gen.ts

@@ -0,0 +1,13 @@
+# Automatically generated - DO NOT EDIT!
+
+"$schema" = "../../schemas/schema.json"
+
+[[permission]]
+identifier = "allow-get-current-model"
+description = "Enables the get_current_model command without any pre-configured scope."
+commands.allow = ["get_current_model"]
+
+[[permission]]
+identifier = "deny-get-current-model"
+description = "Denies the get_current_model command without any pre-configured scope."
+commands.deny = ["get_current_model"]

plugins/local-stt/permissions/autogenerated/commands/get_current_model.toml

@@ -0,0 +1,13 @@
+# Automatically generated - DO NOT EDIT!
+
+"$schema" = "../../schemas/schema.json"
+
+[[permission]]
+identifier = "allow-set-current-model"
+description = "Enables the set_current_model command without any pre-configured scope."
+commands.allow = ["set_current_model"]
+
+[[permission]]
+identifier = "deny-set-current-model"
+description = "Denies the set_current_model command without any pre-configured scope."
+commands.deny = ["set_current_model"]

plugins/local-stt/permissions/autogenerated/commands/set_current_model.toml

@@ -7,6 +7,8 @@ Default permissions for the plugin
 - `allow-download-model`
 - `allow-start-server`
 - `allow-stop-server`
+- `allow-get-current-model`
+- `allow-set-current-model`
 
 ## Permission Table
 
@@ -46,6 +48,32 @@ Denies the download_model command without any pre-configured scope.
 <tr>
 <td>
 
+`local-stt:allow-get-current-model`
+
+</td>
+<td>
+
+Enables the get_current_model command without any pre-configured scope.
+
+</td>
+</tr>
+
+<tr>
+<td>
+
+`local-stt:deny-get-current-model`
+
+</td>
+<td>
+
+Denies the get_current_model command without any pre-configured scope.
+
+</td>
+</tr>
+
+<tr>
+<td>
+
 `local-stt:allow-get-status`
 
 </td>
@@ -124,6 +152,32 @@ Denies the is_server_running command without any pre-configured scope.
 <tr>
 <td>
 
+`local-stt:allow-set-current-model`
+
+</td>
+<td>
+
+Enables the set_current_model command without any pre-configured scope.
+
+</td>
+</tr>
+
+<tr>
+<td>
+
+`local-stt:deny-set-current-model`
+
+</td>
+<td>
+
+Denies the set_current_model command without any pre-configured scope.
+
+</td>
+</tr>
+
+<tr>
+<td>
+
 `local-stt:allow-start-server`
 
 </td>

plugins/local-stt/permissions/autogenerated/reference.md

@@ -6,4 +6,6 @@ permissions = [
     "allow-download-model",
     "allow-start-server",
     "allow-stop-server",
+    "allow-get-current-model",
+    "allow-set-current-model",
 ]

plugins/local-stt/permissions/default.toml

@@ -304,6 +304,16 @@
           "type": "string",
           "const": "deny-download-model"
         },
+        {
+          "description": "Enables the get_current_model command without any pre-configured scope.",
+          "type": "string",
+          "const": "allow-get-current-model"
+        },
+        {
+          "description": "Denies the get_current_model command without any pre-configured scope.",
+          "type": "string",
+          "const": "deny-get-current-model"
+        },
         {
           "description": "Enables the get_status command without any pre-configured scope.",
           "type": "string",
@@ -334,6 +344,16 @@
           "type": "string",
           "const": "deny-is-server-running"
         },
+        {
+          "description": "Enables the set_current_model command without any pre-configured scope.",
+          "type": "string",
+          "const": "allow-set-current-model"
+        },
+        {
+          "description": "Denies the set_current_model command without any pre-configured scope.",
+          "type": "string",
+          "const": "deny-set-current-model"
+        },
         {
           "description": "Enables the start_server command without any pre-configured scope.",
           "type": "string",

plugins/local-stt/permissions/schemas/schema.json

@@ -38,3 +38,20 @@ pub async fn start_server<R: tauri::Runtime>(app: tauri::AppHandle<R>) -> Result
 pub async fn stop_server<R: tauri::Runtime>(app: tauri::AppHandle<R>) -> Result<(), String> {
     app.stop_server().await.map_err(|e| e.to_string())
 }
+
+#[tauri::command]
+#[specta::specta]
+pub fn get_current_model<R: tauri::Runtime>(
+    app: tauri::AppHandle<R>,
+) -> Result<crate::SupportedModel, String> {
+    app.get_current_model().map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+#[specta::specta]
+pub fn set_current_model<R: tauri::Runtime>(
+    app: tauri::AppHandle<R>,
+    model: crate::SupportedModel,
+) -> Result<(), String> {
+    app.set_current_model(model).map_err(|e| e.to_string())
+}

plugins/local-stt/src/commands.rs

@@ -10,6 +10,8 @@ pub enum Error {
     TauriError(#[from] tauri::Error),
     #[error(transparent)]
     IoError(#[from] std::io::Error),
+    #[error(transparent)]
+    StoreError(#[from] tauri_plugin_store2::Error),
 }
 
 impl Serialize for Error {

plugins/local-stt/src/error.rs

@@ -2,6 +2,7 @@ use std::future::Future;
 use std::path::PathBuf;
 
 use tauri::{ipc::Channel, Manager, Runtime};
+use tauri_plugin_store2::StorePluginExt;
 
 #[derive(Debug, Clone)]
 pub struct ModelConfig {
@@ -10,6 +11,7 @@ pub struct ModelConfig {
 }
 
 pub trait LocalSttPluginExt<R: Runtime> {
+    fn local_stt_store(&self) -> tauri_plugin_store2::ScopedStore<R, crate::StoreKey>;
     fn api_base(&self) -> impl Future<Output = Option<String>>;
     fn is_model_downloaded(&self) -> impl Future<Output = Result<bool, crate::Error>>;
     fn is_server_running(&self) -> impl Future<Output = bool>;
@@ -18,9 +20,15 @@ pub trait LocalSttPluginExt<R: Runtime> {
     fn download_config(&self) -> impl Future<Output = Result<(), crate::Error>>;
     fn download_tokenizer(&self) -> impl Future<Output = Result<(), crate::Error>>;
     fn download_model(&self, c: Channel<u8>) -> impl Future<Output = Result<(), crate::Error>>;
+    fn get_current_model(&self) -> Result<crate::SupportedModel, crate::Error>;
+    fn set_current_model(&self, model: crate::SupportedModel) -> Result<(), crate::Error>;
 }
 
 impl<R: Runtime, T: Manager<R>> LocalSttPluginExt<R> for T {
+    fn local_stt_store(&self) -> tauri_plugin_store2::ScopedStore<R, crate::StoreKey> {
+        self.scoped_store(crate::PLUGIN_NAME).unwrap()
+    }
+
     #[tracing::instrument(skip_all)]
     async fn api_base(&self) -> Option<String> {
         let state = self.state::<crate::SharedState>();
@@ -186,4 +194,18 @@ impl<R: Runtime, T: Manager<R>> LocalSttPluginExt<R> for T {
 
         Ok(())
     }
+
+    #[tracing::instrument(skip_all)]
+    fn get_current_model(&self) -> Result<crate::SupportedModel, crate::Error> {
+        let store = self.local_stt_store();
+        let model = store.get(crate::StoreKey::DefaultModel)?;
+        Ok(model.unwrap_or(crate::SupportedModel::QuantizedLargeV3Turbo))
+    }
+
+    #[tracing::instrument(skip_all)]
+    fn set_current_model(&self, model: crate::SupportedModel) -> Result<(), crate::Error> {
+        let store = self.local_stt_store();
+        store.set(crate::StoreKey::DefaultModel, model)?;
+        Ok(())
+    }
 }

plugins/local-stt/src/ext.rs

@@ -6,11 +6,13 @@ mod error;
 mod ext;
 mod model;
 mod server;
+mod store;
 
 pub use error::*;
 pub use ext::*;
 use model::*;
 use server::*;
+use store::*;
 
 pub type SharedState = std::sync::Arc<tokio::sync::Mutex<State>>;
 
@@ -41,6 +43,8 @@ fn make_specta_builder<R: tauri::Runtime>() -> tauri_specta::Builder<R> {
             commands::download_model::<Wry>,
             commands::start_server::<Wry>,
             commands::stop_server::<Wry>,
+            commands::get_current_model::<Wry>,
+            commands::set_current_model::<Wry>,
         ])
         .error_handling(tauri_specta::ErrorHandlingMode::Throw)
 }
@@ -51,6 +55,11 @@ pub fn init<R: tauri::Runtime>() -> tauri::plugin::TauriPlugin<R> {
     tauri::plugin::Builder::new(PLUGIN_NAME)
         .invoke_handler(specta_builder.invoke_handler())
         .setup(|app, _api| {
+            let mut state = State::default();
+            state.model = app
+                .get_current_model()
+                .unwrap_or(SupportedModel::QuantizedLargeV3Turbo);
+
             app.manage(SharedState::default());
             Ok(())
         })

plugins/local-stt/src/lib.rs

@@ -1,4 +1,4 @@
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, specta::Type)]
 pub enum SupportedModel {
     QuantizedTiny,
     QuantizedTinyEn,

plugins/local-stt/src/model.rs

@@ -0,0 +1,8 @@
+use tauri_plugin_store2::ScopedStoreKey;
+
+#[derive(serde::Deserialize, specta::Type, PartialEq, Eq, Hash, strum::Display)]
+pub enum StoreKey {
+    DefaultModel,
+}
+
+impl ScopedStoreKey for StoreKey {}