fix: retrieval stuck at generating response (#1988)

electron/main.ts

@@ -25,25 +25,11 @@ import { migrateExtensions } from './utils/migration'
 import { cleanUpAndQuit } from './utils/clean'
 import { setupExtensions } from './utils/extension'
 import { setupCore } from './utils/setup'
+import { setupReactDevTool } from './utils/dev'
 
 app
   .whenReady()
-  .then(async () => {
-    if (!app.isPackaged) {
-      // Which means you're running from source code
-      const { default: installExtension, REACT_DEVELOPER_TOOLS } = await import(
-        'electron-devtools-installer'
-      ) // Don't use import on top level, since the installer package is dev-only
-      try {
-        const name = installExtension(REACT_DEVELOPER_TOOLS)
-        console.log(`Added Extension: ${name}`)
-      } catch (err) {
-        console.log('An error occurred while installing devtools:')
-        console.error(err)
-        // Only log the error and don't throw it because it's not critical
-      }
-    }
-  })
+  .then(setupReactDevTool)
   .then(setupCore)
   .then(createUserSpace)
   .then(migrateExtensions)

electron/utils/dev.ts

@@ -0,0 +1,18 @@
+import { app } from 'electron'
+
+export const setupReactDevTool = async () => {
+  if (!app.isPackaged) {
+    // Which means you're running from source code
+    const { default: installExtension, REACT_DEVELOPER_TOOLS } = await import(
+      'electron-devtools-installer'
+    ) // Don't use import on top level, since the installer package is dev-only
+    try {
+      const name = installExtension(REACT_DEVELOPER_TOOLS)
+      console.log(`Added Extension: ${name}`)
+    } catch (err) {
+      console.log('An error occurred while installing devtools:')
+      console.error(err)
+      // Only log the error and don't throw it because it's not critical
+    }
+  }
+}

extensions/assistant-extension/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@janhq/assistant-extension",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "This extension enables assistants, including Jan, a default assistant that can call all downloaded models",
   "main": "dist/index.js",
   "node": "dist/node/index.js",

extensions/assistant-extension/src/index.ts

@@ -14,6 +14,7 @@ import {
 
 export default class JanAssistantExtension extends AssistantExtension {
   private static readonly _homeDir = "file://assistants";
+  private static readonly _threadDir = "file://threads";
 
   controller = new AbortController();
   isCancelled = false;
@@ -64,6 +65,8 @@ export default class JanAssistantExtension extends AssistantExtension {
     if (
       data.model?.engine !== InferenceEngine.tool_retrieval_enabled ||
       !data.messages ||
+      // TODO: Since the engine is defined, its unsafe to assume that assistant tools are defined
+      // That could lead to an issue where thread stuck at generating response
       !data.thread?.assistants[0]?.tools
     ) {
       return;
@@ -71,11 +74,12 @@ export default class JanAssistantExtension extends AssistantExtension {
 
     const latestMessage = data.messages[data.messages.length - 1];
 
-    // Ingest the document if needed
+    // 1. Ingest the document if needed
     if (
       latestMessage &&
       latestMessage.content &&
-      typeof latestMessage.content !== "string"
+      typeof latestMessage.content !== "string" &&
+      latestMessage.content.length > 1
     ) {
       const docFile = latestMessage.content[1]?.doc_url?.url;
       if (docFile) {
@@ -86,9 +90,29 @@ export default class JanAssistantExtension extends AssistantExtension {
           data.model?.proxyEngine
         );
       }
+    } else if (
+      // Check whether we need to ingest document or not
+      // Otherwise wrong context will be sent
+      !(await fs.existsSync(
+        await joinPath([
+          JanAssistantExtension._threadDir,
+          data.threadId,
+          "memory",
+        ])
+      ))
+    ) {
+      // No document ingested, reroute the result to inference engine
+      const output = {
+        ...data,
+        model: {
+          ...data.model,
+          engine: data.model.proxyEngine,
+        },
+      };
+      events.emit(MessageEvent.OnMessageSent, output);
+      return;
     }
-
-    // Load agent on thread changed
+    // 2. Load agent on thread changed
     if (instance.retrievalThreadId !== data.threadId) {
       await executeOnMain(NODE, "toolRetrievalLoadThreadMemory", data.threadId);
 
@@ -103,22 +127,22 @@ export default class JanAssistantExtension extends AssistantExtension {
       );
     }
 
+    // 3. Using the retrieval template with the result and query
     if (latestMessage.content) {
       const prompt =
         typeof latestMessage.content === "string"
           ? latestMessage.content
           : latestMessage.content[0].text;
       // Retrieve the result
-      console.debug("toolRetrievalQuery", latestMessage.content);
       const retrievalResult = await executeOnMain(
         NODE,
         "toolRetrievalQueryResult",
         prompt
       );
+      console.debug("toolRetrievalQueryResult", retrievalResult);
 
-      // Update the message content
-      // Using the retrieval template with the result and query
-      if (data.thread?.assistants[0].tools)
+      // Update message content
+      if (data.thread?.assistants[0]?.tools && retrievalResult)
         data.messages[data.messages.length - 1].content =
           data.thread.assistants[0].tools[0].settings?.retrieval_template
             ?.replace("{CONTEXT}", retrievalResult)
@@ -140,7 +164,7 @@ export default class JanAssistantExtension extends AssistantExtension {
       return message;
     });
 
-    // Reroute the result to inference engine
+    // 4. Reroute the result to inference engine
     const output = {
       ...data,
       model: {
@@ -248,12 +272,12 @@ export default class JanAssistantExtension extends AssistantExtension {
             chunk_size: 1024,
             chunk_overlap: 64,
             retrieval_template: `Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
-            ----------------
-            CONTEXT: {CONTEXT}
-            ----------------
-            QUESTION: {QUESTION}
-            ----------------
-            Helpful Answer:`,
+----------------
+CONTEXT: {CONTEXT}
+----------------
+QUESTION: {QUESTION}
+----------------
+Helpful Answer:`,
           },
         },
       ],

extensions/assistant-extension/src/node/index.ts

@@ -1,39 +1,39 @@
 import { getJanDataFolderPath, normalizeFilePath } from "@janhq/core/node";
-import { Retrieval } from "./tools/retrieval";
+import { retrieval } from "./tools/retrieval";
 import path from "path";
 
-const retrieval = new Retrieval();
-
-export async function toolRetrievalUpdateTextSplitter(
+export function toolRetrievalUpdateTextSplitter(
   chunkSize: number,
-  chunkOverlap: number,
+  chunkOverlap: number
 ) {
   retrieval.updateTextSplitter(chunkSize, chunkOverlap);
-  return Promise.resolve();
 }
 export async function toolRetrievalIngestNewDocument(
   file: string,
-  engine: string,
+  engine: string
 ) {
   const filePath = path.join(getJanDataFolderPath(), normalizeFilePath(file));
   const threadPath = path.dirname(filePath.replace("files", ""));
   retrieval.updateEmbeddingEngine(engine);
-  await retrieval.ingestAgentKnowledge(filePath, `${threadPath}/memory`);
-  return Promise.resolve();
+  return retrieval
+    .ingestAgentKnowledge(filePath, `${threadPath}/memory`)
+    .catch((err) => {
+      console.error(err);
+    });
 }
 
 export async function toolRetrievalLoadThreadMemory(threadId: string) {
-  try {
-    await retrieval.loadRetrievalAgent(
-      path.join(getJanDataFolderPath(), "threads", threadId, "memory"),
-    );
-    return Promise.resolve();
-  } catch (err) {
-    console.debug(err);
-  }
+  return retrieval
+    .loadRetrievalAgent(
+      path.join(getJanDataFolderPath(), "threads", threadId, "memory")
+    )
+    .catch((err) => {
+      console.error(err);
+    });
 }
 
 export async function toolRetrievalQueryResult(query: string) {
-  const res = await retrieval.generateResult(query);
-  return Promise.resolve(res);
+  return retrieval.generateResult(query).catch((err) => {
+    console.error(err);
+  });
 }

extensions/assistant-extension/src/node/tools/retrieval/index.ts

@@ -35,6 +35,7 @@ export class Retrieval {
     if (engine === "nitro") {
       this.embeddingModel = new OpenAIEmbeddings(
         { openAIApiKey: "nitro-embedding" },
+        // TODO: Raw settings
         { basePath: "http://127.0.0.1:3928/v1" },
       );
     } else {
@@ -75,3 +76,5 @@ export class Retrieval {
     return Promise.resolve(serializedDoc);
   };
 }
+
+export const retrieval = new Retrieval();