From d256f49d078348c2e9424ea375033c81f9943ee4 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 5 Oct 2023 10:58:28 +0700
Subject: [PATCH] chore: add stop model action

---
 electron/core/plugins/inference-plugin/index.js |  5 +++++
 web/app/_components/ModelRow/index.tsx          |  4 +++-
 web/app/_hooks/useGetModels.ts                  |  9 ---------
 web/app/_hooks/useStartStopModel.ts             | 10 ++++++++--
 web/shared/coreService.ts                       |  1 +
 5 files changed, 17 insertions(+), 12 deletions(-)
 delete mode 100644 web/app/_hooks/useGetModels.ts

diff --git a/electron/core/plugins/inference-plugin/index.js b/electron/core/plugins/inference-plugin/index.js
index 16ee2ca4f..80281737c 100644
--- a/electron/core/plugins/inference-plugin/index.js
+++ b/electron/core/plugins/inference-plugin/index.js
@@ -19,9 +19,14 @@ const dispose = async () =>
   });
 const inferenceUrl = () => "http://localhost:3928/llama/chat_completion";
 
+const stopModel = () => {
+  window.electronAPI.invokePluginFunc(MODULE_PATH, "killSubprocess");
+};
+
 // Register all the above functions and objects with the relevant extension points
 export function init({ register }) {
   register("initModel", "initModel", initModel);
   register("inferenceUrl", "inferenceUrl", inferenceUrl);
   register("dispose", "dispose", dispose);
+  register("stopModel", "stopModel", stopModel);
 }
diff --git a/web/app/_components/ModelRow/index.tsx b/web/app/_components/ModelRow/index.tsx
index 7c98b6edd..d1fa97d5a 100644
--- a/web/app/_components/ModelRow/index.tsx
+++ b/web/app/_components/ModelRow/index.tsx
@@ -14,7 +14,7 @@ type Props = {
 };
 
 const ModelRow: React.FC<Props> = ({ model }) => {
-  const { startModel } = useStartStopModel();
+  const { startModel, stopModel } = useStartStopModel();
   const activeModel = useAtomValue(currentProductAtom);
   const { deleteModel } = useDeleteModel();
 
@@ -31,6 +31,8 @@ const ModelRow: React.FC<Props> = ({ model }) => {
   const onModelActionClick = (action: ModelActionType) => {
     if (action === ModelActionType.Start) {
       startModel(model.id);
+    } else {
+      stopModel(model.id);
     }
   };
 
diff --git a/web/app/_hooks/useGetModels.ts b/web/app/_hooks/useGetModels.ts
deleted file mode 100644
index 9aa6b0181..000000000
--- a/web/app/_hooks/useGetModels.ts
+++ /dev/null
@@ -1,9 +0,0 @@
-import React, { useState } from "react";
-
-export default function useGetModels() {
-  const [models, setModels] = useState<any[]>()
-
-  return {
-    models
-  };
-}
diff --git a/web/app/_hooks/useStartStopModel.ts b/web/app/_hooks/useStartStopModel.ts
index 2201e8fce..b04166a1b 100644
--- a/web/app/_hooks/useStartStopModel.ts
+++ b/web/app/_hooks/useStartStopModel.ts
@@ -1,9 +1,12 @@
 import { executeSerial } from "@/_services/pluginService";
-import { DataService } from "../../shared/coreService";
+import { DataService, InfereceService } from "../../shared/coreService";
 import useInitModel from "./useInitModel";
+import { useSetAtom } from "jotai";
+import { currentProductAtom } from "@/_helpers/atoms/Model.atom";
 
 export default function useStartStopModel() {
   const { initModel } = useInitModel();
+  const setActiveModel = useSetAtom(currentProductAtom);
 
   const startModel = async (modelId: string) => {
     const model = await executeSerial(DataService.GET_MODEL_BY_ID, modelId);
@@ -14,7 +17,10 @@ export default function useStartStopModel() {
     }
   };
 
-  const stopModel = async (modelId: string) => {};
+  const stopModel = async (modelId: string) => {
+    await executeSerial(InfereceService.STOP_MODEL, modelId);
+    setActiveModel(undefined);
+  };
 
   return { startModel, stopModel };
 }
diff --git a/web/shared/coreService.ts b/web/shared/coreService.ts
index b0f5554a1..ac2f23192 100644
--- a/web/shared/coreService.ts
+++ b/web/shared/coreService.ts
@@ -30,6 +30,7 @@ export enum ModelService {
 export enum InfereceService {
   INFERENCE_URL = "inferenceUrl",
   INIT_MODEL = "initModel",
+  STOP_MODEL = "stopModel",
 }
 
 export enum ModelManagementService {