From 3987fdc95b0e8c1cf56f0318b842f9ca466d52eb Mon Sep 17 00:00:00 2001
From: hiro <vuonghoainam.work@gmail.com>
Date: Tue, 5 Dec 2023 02:04:41 +0700
Subject: [PATCH] feat: Add nitro inference engine stop model handler

---
 .../inference-nitro-extension/src/index.ts    | 19 +++++++-----
 .../inference-nitro-extension/src/module.ts   | 30 +++++++++++--------
 2 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/extensions/inference-nitro-extension/src/index.ts b/extensions/inference-nitro-extension/src/index.ts
index 89d051f16..4d6d87dea 100644
--- a/extensions/inference-nitro-extension/src/index.ts
+++ b/extensions/inference-nitro-extension/src/index.ts
@@ -134,17 +134,19 @@ export default class JanInferenceNitroExtension implements InferenceExtension {
       model: model
     });
 
-    if (nitro_init_result.error) {
+    if (nitro_init_result.error === null) {
       events.emit(EventName.OnModelFail, model)
     }
     else{
+      JanInferenceNitroExtension._currentModel = model;
       events.emit(EventName.OnModelReady, model);
     }
   }
 
   private static async handleModelStop(model: Model) {
     if (model.engine !== 'nitro') { return }
-    else { 
+    else {
+      await executeOnMain(MODULE, "stopModel")
       events.emit(EventName.OnModelStopped, model)
     }
   }
@@ -168,9 +170,11 @@ export default class JanInferenceNitroExtension implements InferenceExtension {
     };
 
     return new Promise(async (resolve, reject) => {
-      requestInference(data.messages ?? [], 
+      requestInference(
+          data.messages ?? [], 
           JanInferenceNitroExtension._engineSettings, 
-          JanInferenceNitroExtension._currentModel)
+          JanInferenceNitroExtension._currentModel
+        )
         .subscribe({
         next: (_content) => {},
           complete: async () => {
@@ -212,11 +216,12 @@ export default class JanInferenceNitroExtension implements InferenceExtension {
     instance.isCancelled = false;
     instance.controller = new AbortController();
 
-    requestInference(data.messages ?? [], 
+    requestInference(
+        data.messages ?? [], 
         JanInferenceNitroExtension._engineSettings, 
         JanInferenceNitroExtension._currentModel, 
-        instance.controller)
-      .subscribe({
+        instance.controller
+      ).subscribe({
       next: (content) => {
         const messageContent: ThreadContent = {
           type: ContentType.Text,
diff --git a/extensions/inference-nitro-extension/src/module.ts b/extensions/inference-nitro-extension/src/module.ts
index 5b7a52c60..1323dca6d 100644
--- a/extensions/inference-nitro-extension/src/module.ts
+++ b/extensions/inference-nitro-extension/src/module.ts
@@ -20,22 +20,25 @@ let subprocess = null;
 let currentModelFile = null;
 
 /**
- * The response from the initModel function.
- * @property error - An error message if the model fails to load.
+ * Stops a Nitro subprocess.
+ * @param wrapper - The model wrapper.
+ * @returns A Promise that resolves when the subprocess is terminated successfully, or rejects with an error message if the subprocess fails to terminate.
  */
-interface InitModelResponse {
-  error?: any;
-  modelFile?: string;
+function stopModel(): Promise<ModelOperationResponse> {
+  return new Promise((resolve, reject) => {
+    checkAndUnloadNitro()
+    resolve({ error: undefined})
+  })
 }
 
 /**
  * Initializes a Nitro subprocess to load a machine learning model.
- * @param modelFile - The name of the machine learning model file.
+ * @param wrapper - The model wrapper.
  * @returns A Promise that resolves when the model is loaded successfully, or rejects with an error message if the model is not found or fails to load.
  * TODO: Should pass absolute of the model file instead of just the name - So we can modurize the module.ts to npm package
  * TODO: Should it be startModel instead?
  */
-function initModel(wrapper: any): Promise<InitModelResponse> {
+function initModel(wrapper: any): Promise<ModelOperationResponse> {
     currentModelFile = wrapper.modelFullPath;
     if (wrapper.model.engine !== "nitro") {
       return Promise.resolve({ error: "Not a nitro model" })
@@ -89,11 +92,11 @@ function loadLLMModel(settings): Promise<Response> {
 
 /**
  * Validates the status of a model.
- * @returns {Promise<InitModelResponse>} A promise that resolves to an object.
+ * @returns {Promise<ModelOperationResponse>} A promise that resolves to an object.
  * If the model is loaded successfully, the object is empty.
  * If the model is not loaded successfully, the object contains an error message.
  */
-async function validateModelStatus(): Promise<InitModelResponse> {
+async function validateModelStatus(): Promise<ModelOperationResponse> {
   // Send a GET request to the validation URL.
   // Retry the request up to 3 times if it fails, with a delay of 500 milliseconds between retries.
   return fetchRetry(NITRO_HTTP_VALIDATE_MODEL_URL, {
@@ -140,17 +143,18 @@ function killSubprocess(): Promise<void> {
  * Check port is used or not, if used, attempt to unload model
  * If unload failed, kill the port
  */
-function checkAndUnloadNitro() {
-  return tcpPortUsed.check(PORT, LOCAL_HOST).then((inUse) => {
+async function checkAndUnloadNitro() {
+  return tcpPortUsed.check(PORT, LOCAL_HOST).then(async (inUse) => {
     // If inUse - try unload or kill process, otherwise do nothing
     if (inUse) {
       // Attempt to unload model
-      return fetch(NITRO_HTTP_UNLOAD_MODEL_URL, {
+      return await fetch(NITRO_HTTP_UNLOAD_MODEL_URL, {
         method: "GET",
         headers: {
           "Content-Type": "application/json",
         },
-      }).catch((err) => {
+      })
+      .catch((err) => {
         console.error(err);
         // Fallback to kill the port
         return killSubprocess();