From 9daee14167dc0d3c2c5ecb186fd7d522744a2539 Mon Sep 17 00:00:00 2001
From: hiro <vuonghoainam.work@gmail.com>
Date: Tue, 5 Dec 2023 22:48:04 +0700
Subject: [PATCH] fix: Add hack waiting for model loading

---
 .../inference-nitro-extension/src/index.ts    |  1 -
 web/hooks/useSendChatMessage.ts               | 28 +++++++++++++++----
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/extensions/inference-nitro-extension/src/index.ts b/extensions/inference-nitro-extension/src/index.ts
index 4d6d87dea..f8bbd2a09 100644
--- a/extensions/inference-nitro-extension/src/index.ts
+++ b/extensions/inference-nitro-extension/src/index.ts
@@ -198,7 +198,6 @@ export default class JanInferenceNitroExtension implements InferenceExtension {
     instance: JanInferenceNitroExtension
   ) {
     if (data.model.engine !== 'nitro') { return }
-
     const timestamp = Date.now();
     const message: ThreadMessage = {
       id: ulid(),
diff --git a/web/hooks/useSendChatMessage.ts b/web/hooks/useSendChatMessage.ts
index 4ed080438..2a9e7a102 100644
--- a/web/hooks/useSendChatMessage.ts
+++ b/web/hooks/useSendChatMessage.ts
@@ -50,7 +50,6 @@ export default function useSendChatMessage() {
   const [queuedMessage, setQueuedMessage] = useState(false)
 
   const modelRef = useRef<Model | undefined>()
-
   useEffect(() => {
     modelRef.current = activeModel
   }, [activeModel])
@@ -91,19 +90,35 @@ export default function useSendChatMessage() {
       id: ulid(),
       messages: messages,
       threadId: activeThread.id,
-      model: activeThread.assistants[0].model??selectedModel,
+      model: activeThread.assistants[0].model ?? selectedModel,
     }
 
     const modelId = selectedModel?.id ?? activeThread.assistants[0].model.id
 
     if (activeModel?.id !== modelId) {
       setQueuedMessage(true)
-      await startModel(modelId)
+      startModel(modelId)
+      await WaitForModelStarting(modelId)
       setQueuedMessage(false)
     }
     events.emit(EventName.OnMessageSent, messageRequest)
   }
 
+  // TODO: Refactor @louis
+  const WaitForModelStarting = async (modelId: string) => {
+    return new Promise<void>((resolve) => {
+      setTimeout(async () => {
+        if (modelRef.current?.id !== modelId) {
+          console.log('waiting for model to start')
+          await WaitForModelStarting(modelId)
+          resolve()
+        } else {
+          resolve()
+        }
+      }, 200)
+    })
+  }
+
   const sendChatMessage = async () => {
     if (!currentPrompt || currentPrompt.trim().length === 0) {
       return
@@ -180,8 +195,7 @@ export default function useSendChatMessage() {
       id: msgId,
       threadId: activeThread.id,
       messages,
-      parameters: activeThread.assistants[0].model.parameters,
-      model: selectedModel??activeThread.assistants[0].model,
+      model: selectedModel ?? activeThread.assistants[0].model,
     }
     const timestamp = Date.now()
     const threadMessage: ThreadMessage = {
@@ -213,9 +227,11 @@ export default function useSendChatMessage() {
 
     if (activeModel?.id !== modelId) {
       setQueuedMessage(true)
-      await startModel(modelId)
+      startModel(modelId)
+      await WaitForModelStarting(modelId)
       setQueuedMessage(false)
     }
+    console.log('messageRequest', messageRequest)
     events.emit(EventName.OnMessageSent, messageRequest)
   }