implement delete model function

Signed-off-by: cbh778899 <[email protected]>
SkywardAI · Oct 22, 2024 · af19865 · af19865
1 parent f659d17
commit af19865
Showing 1 changed file with 77 additions and 16 deletions.
diff --git a/src/components/settings/LlamaSettings.jsx b/src/components/settings/LlamaSettings.jsx
@@ -6,12 +6,16 @@ import useIDB from "../../utils/idb";
 import { getPlatformSettings, updatePlatformSettings } from "../../utils/general_settings";
 import { DEFAULT_LLAMA_CPP_MODEL_URL } from "../../utils/types";
 import DropdownComponent from "./components/DropdownComponent";
+import ButtonComponent from "./components/ButtonComponent";
+import ConfirmationDialog from "../ConfirmationDialog";
 
 export default function LlamaSettings({ trigger, enabled, updateEnabled, openDownloadProtector, updateState }) {
 
     const [model_download_link, setModelDownloadLink] = useState('');
+    const [selected_model, setSelectedModel] = useState({});
     const [reset_everytime, setResetEveryTime] = useState(false);
     const [downloaded_models, setDownloadedModels] = useState([])
+    const [delete_confirm_opened, setDeleteConfirmOpenStatus] = useState(false);
     const idb = useIDB();
 
     async function saveSettings() {
@@ -41,25 +45,58 @@ export default function LlamaSettings({ trigger, enabled, updateEnabled, openDow
                                 platform: 'Llama'
                             }
                             await idb.insert("downloaded-models", stored_model)
-                            setDownloadedModels([...downloaded_models, { title: stored_model['model-name'], value: stored_model.url }])
+                            setDownloadedModels([...downloaded_models, { title: stored_model['model-name'], value: JSON.stringify(stored_model) }])
+                            setSelectedModel(stored_model)
                         }
                     }
                 )
             }
-            await openDownloadProtector(
-                'Loading model...',
-                `Loading model ${stored_model['model-name']}`,
-                async callback => {
-                    callback(100, false);
-                    // load model using the model name retrieved
-                    await window['node-llama-cpp'].loadModel(stored_model['model-name'])
-                    updateState();
-                    callback(100, true)
-                }
-            )
+            await loadModel(stored_model);
         }
     }
 
+    async function loadModel(model) {
+        await openDownloadProtector(
+            'Loading model...',
+            `Loading model ${model['model-name']}`,
+            async callback => {
+                callback(100, false);
+                // load model using the model name retrieved
+                await window['node-llama-cpp'].loadModel(model['model-name'])
+                updateState();
+                callback(100, true)
+            }
+        )
+    }
+
+    async function deleteModel() {
+        let load_after_delete = {};
+        await openDownloadProtector(
+            'Please wait while we deleteing the model...',
+            `Loading model ${selected_model['model-name']}`,
+            async callback => {
+                window['node-llama-cpp'].deleteModel(selected_model['model-name']);
+                callback(80, false);
+                await idb.deleteOne("downloaded-models", [selected_model]);
+                callback(90, false);
+
+                const models = await idb.getAll("downloaded-models", {
+                    where: [{'platform': 'Llama'}]
+                });
+                setDownloadedModels(models.map(e=>{return { title: e['model-name'], value: JSON.stringify(e) }}))
+                load_after_delete = models.pop() || {}
+                const url = load_after_delete.url || '';
+                setSelectedModel(load_after_delete);
+                setModelDownloadLink(url);
+                updatePlatformSettings({
+                    llama_model_url: url
+                })
+                callback(100, true);
+            }
+        )
+        await loadModel(load_after_delete);
+    }
+
     useEffect(()=>{
         trigger && saveSettings();
     // eslint-disable-next-line
@@ -72,10 +109,10 @@ export default function LlamaSettings({ trigger, enabled, updateEnabled, openDow
             setResetEveryTime(llama_reset_everytime);
 
             const models = await idb.getAll("downloaded-models", {
-                where: [{'platform': 'Llama'}],
-                select: ["model-name", "url"]
+                where: [{'platform': 'Llama'}]
             });
-            setDownloadedModels(models.map(e=>{return { title: e['model-name'], value: e.url }}))
+            setDownloadedModels(models.map(e=>{return { title: e['model-name'], value: JSON.stringify(e) }}))
+            setSelectedModel(models.filter(e=>e.url === llama_model_url).pop())
         })()
     // eslint-disable-next-line
     }, [])
@@ -90,7 +127,12 @@ export default function LlamaSettings({ trigger, enabled, updateEnabled, openDow
             <DropdownComponent 
                 title={"Select a downloaded model to use"}
                 description={"If there's no models, you might want to download a new one by enter url below."}
-                value={downloaded_models} cb={setModelDownloadLink}
+                value={downloaded_models} cb={model=>{
+                    model = JSON.parse(model)
+                    setSelectedModel(model)
+                    setModelDownloadLink(model.url)
+                }}
+                selected={JSON.stringify(selected_model)}
             />
             <TextComponent 
                 title={"Or enter the link of model you want to use"}
@@ -103,6 +145,25 @@ export default function LlamaSettings({ trigger, enabled, updateEnabled, openDow
                 description={"Reset the conversation, only keeps the latest message and the system instruction, this is useful for many one-shot operations."}
                 value={reset_everytime} cb={setResetEveryTime}
             />
+            <ButtonComponent 
+                title={"Delete selected model"}
+                description={"This will delete the downloaded model, make sure before you choose to delete!"}
+                className="dangerous" disabled={!model_download_link}
+                value={"Delete Model"} cb={()=>{setDeleteConfirmOpenStatus(true)}}
+            />
+            <ConfirmationDialog
+                open_status={delete_confirm_opened}
+                setOpenStatus={setDeleteConfirmOpenStatus}
+                callback={cb=>{
+                    cb && deleteModel();
+                    setDeleteConfirmOpenStatus(false);
+                }}
+            >
+                <div>
+                    Are you sure you want to delete model<br/>
+                    <strong>{selected_model['model-name']}</strong>?
+                </div>
+            </ConfirmationDialog>
         </SettingSection>
     )
 }