From 848f9c5eed790ad8b56df1c2c8c226edef027690 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Sat, 9 Nov 2024 12:54:06 +0700
Subject: [PATCH] fix: migrate model settings and parameters but still allow
 model.yml editing

---
 extensions/model-extension/src/cortex.ts   |  4 +--
 extensions/model-extension/src/index.ts    |  8 +++++-
 web/containers/Providers/EventListener.tsx | 29 +++++++++++++++++++---
 web/hooks/useModels.ts                     |  4 ---
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/extensions/model-extension/src/cortex.ts b/extensions/model-extension/src/cortex.ts
index 5b7d1e36b..b7111c859 100644
--- a/extensions/model-extension/src/cortex.ts
+++ b/extensions/model-extension/src/cortex.ts
@@ -123,10 +123,10 @@ export class CortexAPI implements ICortexAPI {
    * @param model
    * @returns
    */
-  updateModel(model: object): Promise<void> {
+  updateModel(model: Partial<Model>): Promise<void> {
     return this.queue.add(() =>
       ky
-        .patch(`${API_URL}/v1/models/${model}`, { json: { model } })
+        .patch(`${API_URL}/v1/models/${model.id}`, { json: { ...model } })
         .json()
         .then()
     )
diff --git a/extensions/model-extension/src/index.ts b/extensions/model-extension/src/index.ts
index 8f50bd5d0..e5fc7ccf6 100644
--- a/extensions/model-extension/src/index.ts
+++ b/extensions/model-extension/src/index.ts
@@ -193,7 +193,13 @@ export default class JanModelExtension extends ModelExtension {
                 ]) // Copied models
               : model.sources[0].url, // Symlink models,
             model.name
-          )
+          ).then((e) => {
+            this.updateModel({
+              id: model.id,
+              ...model.settings,
+              ...model.parameters,
+            } as Partial<Model>)
+          })
         )
       )
 
diff --git a/web/containers/Providers/EventListener.tsx b/web/containers/Providers/EventListener.tsx
index 9535bbfa6..5cb0debab 100644
--- a/web/containers/Providers/EventListener.tsx
+++ b/web/containers/Providers/EventListener.tsx
@@ -2,7 +2,17 @@ import { PropsWithChildren, useCallback, useEffect } from 'react'
 
 import React from 'react'
 
-import { DownloadEvent, events, DownloadState, ModelEvent } from '@janhq/core'
+import {
+  DownloadEvent,
+  events,
+  DownloadState,
+  ModelEvent,
+  ExtensionTypeEnum,
+  ModelExtension,
+  ModelManager,
+  Model,
+} from '@janhq/core'
+
 import { useSetAtom } from 'jotai'
 
 import { setDownloadStateAtom } from '@/hooks/useDownloadState'
@@ -18,6 +28,7 @@ import EventHandler from './EventHandler'
 import ModelImportListener from './ModelImportListener'
 import QuickAskListener from './QuickAskListener'
 
+import { extensionManager } from '@/extension'
 import {
   InstallingExtensionState,
   removeInstallingExtensionAtom,
@@ -83,12 +94,24 @@ const EventListenerWrapper = ({ children }: PropsWithChildren) => {
   )
 
   const onFileDownloadSuccess = useCallback(
-    (state: DownloadState) => {
+    async (state: DownloadState) => {
       console.debug('onFileDownloadSuccess', state)
       if (state.downloadType !== 'extension') {
+        // Update model metadata accordingly
+        const model = ModelManager.instance().models.get(state.modelId)
+        if (model) {
+          await extensionManager
+            .get<ModelExtension>(ExtensionTypeEnum.Model)
+            ?.updateModel({
+              id: model.id,
+              ...model.settings,
+              ...model.parameters,
+            } as Partial<Model>)
+            .catch((e) => console.debug(e))
+        }
         state.downloadState = 'end'
         setDownloadState(state)
-        if (state.percent !== 0) removeDownloadingModel(state.modelId)
+        removeDownloadingModel(state.modelId)
       }
       events.emit(ModelEvent.OnModelsUpdate, {})
     },
diff --git a/web/hooks/useModels.ts b/web/hooks/useModels.ts
index b8b680715..400e02793 100644
--- a/web/hooks/useModels.ts
+++ b/web/hooks/useModels.ts
@@ -35,10 +35,6 @@ const useModels = () => {
       const localModels = (await getModels()).map((e) => ({
         ...e,
         name: ModelManager.instance().models.get(e.id)?.name ?? e.id,
-        settings:
-          ModelManager.instance().models.get(e.id)?.settings ?? e.settings,
-        parameters:
-          ModelManager.instance().models.get(e.id)?.parameters ?? e.parameters,
         metadata:
           ModelManager.instance().models.get(e.id)?.metadata ?? e.metadata,
       }))