From cec9b7e69127ff4314be11edbcde5a016c38bd72 Mon Sep 17 00:00:00 2001
From: Logan Yang <logancyang@gmail.com>
Date: Tue, 7 Jan 2025 14:13:50 -0800
Subject: [PATCH] Use reranking on NaN chunks (#1018)

---
 src/search/findRelevantNotes.ts |  2 +-
 src/search/hybridRetriever.ts   | 34 +++++++++++++++++++++++----------
 src/utils.ts                    |  5 +++++
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/search/findRelevantNotes.ts b/src/search/findRelevantNotes.ts
index b00986bf..38d66b2d 100644
--- a/src/search/findRelevantNotes.ts
+++ b/src/search/findRelevantNotes.ts
@@ -27,7 +27,7 @@ async function getNoteEmbeddings(notePath: string, db: Orama<any>): Promise<numb
 
   const embeddings: number[][] = [];
   for (const hit of hits) {
-    if (!hit.document.embedding) {
+    if (!hit?.document?.embedding) {
       if (debug) {
         console.log("No embedding found for note:", notePath);
       }
diff --git a/src/search/hybridRetriever.ts b/src/search/hybridRetriever.ts
index eb96f9ec..9992470a 100644
--- a/src/search/hybridRetriever.ts
+++ b/src/search/hybridRetriever.ts
@@ -24,7 +24,7 @@ export class HybridRetriever extends BaseRetriever {
       timeRange?: { startTime: number; endTime: number };
       textWeight?: number;
       returnAll?: boolean;
-      useRerankerThreshold?: number;
+      useRerankerThreshold?: number; // reranking API is only called with this set
     }
   ) {
     super();
@@ -57,16 +57,30 @@ export class HybridRetriever extends BaseRetriever {
     const combinedChunks = this.filterAndFormatChunks(oramaChunks, explicitChunks);
 
     let finalChunks = combinedChunks;
-    const maxOramaScore = combinedChunks.reduce(
-      (max, chunk) => Math.max(max, chunk.metadata.score ?? 0),
-      0
+
+    // Add check for empty array
+    if (combinedChunks.length === 0) {
+      if (getSettings().debug) {
+        console.log("No chunks found for query:", query);
+      }
+      return finalChunks;
+    }
+
+    const maxOramaScore = combinedChunks.reduce((max, chunk) => {
+      const score = chunk.metadata.score;
+      const isValidScore = typeof score === "number" && !isNaN(score);
+      return isValidScore ? Math.max(max, score) : max;
+    }, 0);
+
+    const allScoresAreNaN = combinedChunks.every(
+      (chunk) => typeof chunk.metadata.score !== "number" || isNaN(chunk.metadata.score)
     );
-    // Apply reranking if max score is below the threshold
-    if (
+
+    const shouldRerank =
       this.options.useRerankerThreshold &&
-      maxOramaScore < this.options.useRerankerThreshold &&
-      maxOramaScore > 0
-    ) {
+      (maxOramaScore < this.options.useRerankerThreshold || allScoresAreNaN);
+    // Apply reranking if max score is below the threshold or all scores are NaN
+    if (shouldRerank) {
       const rerankResponse = await BrevilabsClient.getInstance().rerank(
         query,
         // Limit the context length to 3000 characters to avoid overflowing the reranker
@@ -95,7 +109,7 @@ export class HybridRetriever extends BaseRetriever {
       console.log("Orama Chunks: ", oramaChunks);
       console.log("Combined Chunks: ", combinedChunks);
       console.log("Max Orama Score: ", maxOramaScore);
-      if (this.options.useRerankerThreshold && maxOramaScore < this.options.useRerankerThreshold) {
+      if (shouldRerank) {
         console.log("Reranked Chunks: ", finalChunks);
       } else {
         console.log("No reranking applied.");
diff --git a/src/utils.ts b/src/utils.ts
index f94b1923..b8bc3e89 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -452,6 +452,11 @@ export function extractChatHistory(memoryVariables: MemoryVariables): [string, s
   return chatHistory;
 }
 
+// TODO: Deprecate this. Note mentions should be an object with title and path (optional).
+// When user input `[[` the popup should show title and path for selection.
+// The selected item has path to avoid duplicate titles. If user manually types
+// the full title, path can still be missing. In that case title is used to retrieve
+// the note.
 export function extractNoteTitles(query: string): string[] {
   // Use a regular expression to extract note titles wrapped in [[]]
   const regex = /\[\[(.*?)\]\]/g;