From 57267ab8118262191bbcbf6199d3b9d567cb76c3 Mon Sep 17 00:00:00 2001
From: jacoblee93 <jacoblee93@gmail.com>
Date: Thu, 21 Nov 2024 21:13:47 -0800
Subject: [PATCH] Make score inclusion toggleable

---
 .../src/retrievers/bm25.ts                    | 24 +++++++++++++------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/libs/langchain-community/src/retrievers/bm25.ts b/libs/langchain-community/src/retrievers/bm25.ts
index d8024babe495..716027cf684b 100644
--- a/libs/langchain-community/src/retrievers/bm25.ts
+++ b/libs/langchain-community/src/retrievers/bm25.ts
@@ -6,6 +6,7 @@ import { BM25 } from "../utils/@furkantoprak/bm25/BM25.js";
 export type BM25RetrieverOptions = {
   docs: Document[];
   k: number;
+  includeScore?: boolean;
 } & BaseRetrieverInput;
 
 /**
@@ -14,6 +15,8 @@ export type BM25RetrieverOptions = {
  * The k parameter determines the number of documents to return for each query.
  */
 export class BM25Retriever extends BaseRetriever {
+  includeScore = false;
+
   static lc_name() {
     return "BM25Retriever";
   }
@@ -35,6 +38,7 @@ export class BM25Retriever extends BaseRetriever {
     super(options);
     this.docs = options.docs;
     this.k = options.k;
+    this.includeScore = options.includeScore ?? this.includeScore;
   }
 
   private preprocessFunc(text: string): string[] {
@@ -53,13 +57,19 @@ export class BM25Retriever extends BaseRetriever {
 
     scoredDocs.sort((a, b) => b.score - a.score);
 
-    return scoredDocs.slice(0, this.k).map((item) => new Document({
-      ...item.document.id && { id: item.document.id },
-      pageContent: item.document.pageContent,
-      metadata: {
-        bm25Score: item.score,
-        ...item.document.metadata,
+    return scoredDocs.slice(0, this.k).map((item) => {
+      if (this.includeScore) {
+        return new Document({
+          ...(item.document.id && { id: item.document.id }),
+          pageContent: item.document.pageContent,
+          metadata: {
+            bm25Score: item.score,
+            ...item.document.metadata,
+          },
+        });
+      } else {
+        return item.document;
       }
-    }));
+    });
   }
 }