From 693796a68dae560c225836430cf43500d7b978a4 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Sun, 17 Nov 2024 14:36:52 +0700
Subject: [PATCH] fix: remove hacky retrieval settings and engine routing

---
 .../assistant-extension/src/node/engine.ts    | 38 -------------------
 .../assistant-extension/src/node/retrieval.ts | 25 ++++--------
 .../ThreadCenterPanel/ChatInput/index.tsx     | 35 ++++-------------
 3 files changed, 15 insertions(+), 83 deletions(-)
 delete mode 100644 extensions/assistant-extension/src/node/engine.ts

diff --git a/extensions/assistant-extension/src/node/engine.ts b/extensions/assistant-extension/src/node/engine.ts
deleted file mode 100644
index 05a380340..000000000
--- a/extensions/assistant-extension/src/node/engine.ts
+++ /dev/null
@@ -1,38 +0,0 @@
-import fs from 'fs'
-import path from 'path'
-import { SettingComponentProps, getJanDataFolderPath } from '@janhq/core/node'
-
-// Sec: Do not send engine settings over requests
-// Read it manually instead
-export const readEmbeddingEngine = (engineName: string) => {
-  if (engineName !== 'openai' && engineName !== 'groq') {
-    const engineSettings = fs.readFileSync(
-      path.join(getJanDataFolderPath(), 'engines', `${engineName}.json`),
-      'utf-8'
-    )
-    return JSON.parse(engineSettings)
-  } else {
-    const settingDirectoryPath = path.join(
-      getJanDataFolderPath(),
-      'settings',
-      '@janhq',
-      // TODO: James - To be removed
-      engineName === 'openai'
-        ? 'inference-openai-extension'
-        : 'inference-groq-extension',
-      'settings.json'
-    )
-
-    const content = fs.readFileSync(settingDirectoryPath, 'utf-8')
-    const settings: SettingComponentProps[] = JSON.parse(content)
-    const apiKeyId = engineName === 'openai' ? 'openai-api-key' : 'groq-api-key'
-    const keySetting = settings.find((setting) => setting.key === apiKeyId)
-
-    let apiKey = keySetting?.controllerProps.value
-    if (typeof apiKey !== 'string') apiKey = ''
-
-    return {
-      api_key: apiKey,
-    }
-  }
-}
diff --git a/extensions/assistant-extension/src/node/retrieval.ts b/extensions/assistant-extension/src/node/retrieval.ts
index 3386029fa..5804ff763 100644
--- a/extensions/assistant-extension/src/node/retrieval.ts
+++ b/extensions/assistant-extension/src/node/retrieval.ts
@@ -8,7 +8,6 @@ import { MemoryVectorStore } from 'langchain/vectorstores/memory'
 import { HNSWLib } from 'langchain/vectorstores/hnswlib'
 
 import { OpenAIEmbeddings } from 'langchain/embeddings/openai'
-import { readEmbeddingEngine } from './engine'
 
 export class Retrieval {
   public chunkSize: number = 100
@@ -28,8 +27,8 @@ export class Retrieval {
     // declare time-weighted retriever and storage
     this.timeWeightedVectorStore = new MemoryVectorStore(
       new OpenAIEmbeddings(
-        { openAIApiKey: 'nitro-embedding' },
-        { basePath: 'http://127.0.0.1:3928/v1' }
+        { openAIApiKey: 'cortex-embedding' },
+        { basePath: 'http://127.0.0.1:39291/v1' }
       )
     )
     this.timeWeightedretriever = new TimeWeightedVectorStoreRetriever({
@@ -49,21 +48,11 @@ export class Retrieval {
   }
 
   public updateEmbeddingEngine(model: string, engine: string): void {
-    // Engine settings are not compatible with the current embedding model params
-    // Switch case manually for now
-    if (engine === 'nitro') {
-      this.embeddingModel = new OpenAIEmbeddings(
-        { openAIApiKey: 'nitro-embedding', model },
-        // TODO: Raw settings
-        { basePath: 'http://127.0.0.1:3928/v1' },
-      )
-    } else {
-      // Fallback to OpenAI Settings
-      const settings = readEmbeddingEngine(engine)
-      this.embeddingModel = new OpenAIEmbeddings({
-        openAIApiKey: settings.api_key,
-      })
-    }
+    this.embeddingModel = new OpenAIEmbeddings(
+      { openAIApiKey: 'cortex-embedding', model },
+      // TODO: Raw settings
+      { basePath: 'http://127.0.0.1:39291/v1' }
+    )
 
     // update time-weighted embedding model
     this.timeWeightedVectorStore.embeddings = this.embeddingModel
diff --git a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx
index 066c93430..5662cd0c0 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx
@@ -249,19 +249,11 @@ const ChatInput = () => {
                   <li
                     className={twMerge(
                       'text-[hsla(var(--text-secondary)] hover:bg-secondary flex w-full cursor-pointer items-center space-x-2 px-4 py-2 hover:bg-[hsla(var(--dropdown-menu-hover-bg))]',
-                      activeThread?.assistants[0].model.settings?.text_model ===
-                        false
-                        ? 'cursor-not-allowed opacity-50'
-                        : 'cursor-pointer'
+                      'cursor-pointer'
                     )}
                     onClick={() => {
-                      if (
-                        activeThread?.assistants[0].model.settings
-                          ?.text_model !== false
-                      ) {
-                        fileInputRef.current?.click()
-                        setShowAttacmentMenus(false)
-                      }
+                      fileInputRef.current?.click()
+                      setShowAttacmentMenus(false)
                     }}
                   >
                     <FileTextIcon size={16} />
@@ -270,22 +262,11 @@ const ChatInput = () => {
                 }
                 content={
                   (!activeThread?.assistants[0].tools ||
-                    !activeThread?.assistants[0].tools[0]?.enabled ||
-                    activeThread?.assistants[0].model.settings?.text_model ===
-                      false) && (
-                    <>
-                      {activeThread?.assistants[0].model.settings
-                        ?.text_model === false ? (
-                        <span>
-                          This model does not support text-based retrieval.
-                        </span>
-                      ) : (
-                        <span>
-                          Turn on Retrieval in Assistant Settings to use this
-                          feature.
-                        </span>
-                      )}
-                    </>
+                    !activeThread?.assistants[0].tools[0]?.enabled) && (
+                    <span>
+                      Turn on Retrieval in Assistant Settings to use this
+                      feature.
+                    </span>
                   )
                 }
               />