fix(ts-sdk): use textLemmatized for BM25 keyword search on Milvus, OpenSearch, and MongoDB (#6497)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

mem0-ts/src/oss/src/tests/milvus.test.ts

@@ -432,12 +432,18 @@ describe("Milvus vector store (TS OSS SDK)", () => {
       ["a"],
       [{ data: "hello world", text_lemmatized: "hello world lemma" }],
     );
+    await store.insert(
+      [[0.2, 0.3, 0.4]],
+      ["c"],
+      [{ data: "hello world", textLemmatized: "hello world camel" }],
+    );
     // Falls back to raw data when there is no lemmatized text.
     await store.insert([[0.4, 0.5, 0.6]], ["b"], [{ data: "just data" }]);
 
     const insertCalls = client.calls.filter((c) => c.method === "insert");
     expect(insertCalls[0].args.data[0].text).toBe("hello world lemma");
-    expect(insertCalls[1].args.data[0].text).toBe("just data");
+    expect(insertCalls[1].args.data[0].text).toBe("hello world camel");
+    expect(insertCalls[2].args.data[0].text).toBe("just data");
   });
 
   it("writes the BM25 text field on update for a BM25 collection", async () => {

mem0-ts/src/oss/src/tests/mongodb.test.ts

@@ -5,6 +5,7 @@ const mockFindOne = jest.fn();
 const mockUpdateOne = jest.fn();
 const mockListSearchIndexes = jest.fn();
 const mockCreateSearchIndex = jest.fn();
+const mockDropSearchIndex = jest.fn();
 const mockDrop = jest.fn();
 const mockToArray = jest.fn();
 const mockLimit = jest.fn().mockReturnThis();
@@ -26,6 +27,7 @@ const mockCollection = {
   updateOne: mockUpdateOne,
   listSearchIndexes: mockListSearchIndexes,
   createSearchIndex: mockCreateSearchIndex,
+  dropSearchIndex: mockDropSearchIndex,
   drop: mockDrop,
   find: mockFind,
   aggregate: mockAggregate,
@@ -74,6 +76,25 @@ describe("MongoDB Vector Store", () => {
     await store.close();
   });
 
+  const expectedTextSearchIndexDefinition = {
+    name: "test_col_text_search_index",
+    definition: {
+      mappings: {
+        dynamic: false,
+        fields: {
+          payload: {
+            type: "document",
+            fields: {
+              data: { type: "string" },
+              textLemmatized: { type: "string" },
+              text_lemmatized: { type: "string" },
+            },
+          },
+        },
+      },
+    },
+  };
+
   it("should initialize client and check/create collection and indexes", async () => {
     await store.initialize();
 
@@ -84,6 +105,83 @@ describe("MongoDB Vector Store", () => {
     });
     expect(mockCollection.deleteOne).toHaveBeenCalledWith({ _id: 0 });
     expect(mockCreateSearchIndex).toHaveBeenCalledTimes(2);
+    expect(mockCreateSearchIndex).toHaveBeenCalledWith(
+      expectedTextSearchIndexDefinition,
+    );
+  });
+
+  it("should drop and recreate a stale text search index on upgrade", async () => {
+    mockListCollections.mockReturnValue({
+      toArray: jest.fn().mockResolvedValue([{ name: "test_col" }]),
+    });
+    mockListSearchIndexes.mockReturnValue({
+      toArray: jest.fn().mockResolvedValue([
+        { name: "test_col_vector_index" },
+        {
+          name: "test_col_text_search_index",
+          definition: {
+            mappings: {
+              dynamic: false,
+              fields: {
+                payload: {
+                  type: "document",
+                  fields: {
+                    data: { type: "string" },
+                    text_lemmatized: { type: "string" },
+                  },
+                },
+              },
+            },
+          },
+        },
+      ]),
+    });
+
+    await store.initialize();
+
+    expect(mockDropSearchIndex).toHaveBeenCalledWith(
+      "test_col_text_search_index",
+    );
+    expect(mockCreateSearchIndex).toHaveBeenCalledTimes(1);
+    expect(mockCreateSearchIndex).toHaveBeenCalledWith(
+      expectedTextSearchIndexDefinition,
+    );
+  });
+
+  it("should not recreate a text search index that already has textLemmatized", async () => {
+    mockListCollections.mockReturnValue({
+      toArray: jest.fn().mockResolvedValue([{ name: "test_col" }]),
+    });
+    mockListSearchIndexes.mockReturnValue({
+      toArray: jest.fn().mockResolvedValue([
+        { name: "test_col_vector_index" },
+        {
+          name: "test_col_text_search_index",
+          latestDefinition: expectedTextSearchIndexDefinition.definition,
+        },
+      ]),
+    });
+
+    await store.initialize();
+
+    expect(mockDropSearchIndex).not.toHaveBeenCalled();
+    expect(mockCreateSearchIndex).not.toHaveBeenCalled();
+  });
+
+  it("should map payload.textLemmatized in the text search index", async () => {
+    await store.initialize();
+
+    const textIndexCall = mockCreateSearchIndex.mock.calls.find(
+      ([arg]: any[]) => arg.name === "test_col_text_search_index",
+    );
+    expect(textIndexCall).toBeDefined();
+    expect(textIndexCall![0].definition.mappings.fields.payload.fields).toEqual(
+      {
+        data: { type: "string" },
+        text_lemmatized: { type: "string" },
+        textLemmatized: { type: "string" },
+      },
+    );
   });
 
   it("should insert documents correctly", async () => {
@@ -197,7 +295,11 @@ describe("MongoDB Vector Store", () => {
           index: "test_col_text_search_index",
           text: {
             query: "test",
-            path: ["payload.data", "payload.text_lemmatized"],
+            path: [
+              "payload.data",
+              "payload.text_lemmatized",
+              "payload.textLemmatized",
+            ],
           },
         },
       },

mem0-ts/src/oss/src/vector_stores/milvus.ts

@@ -252,13 +252,13 @@ export class Milvus implements VectorStore {
   }
 
   /**
-   * Text fed to the BM25 sparse index for a payload. Prefers the lemmatized
+   * Text fed to the BM25 sparse index for a payload. Prefers `textLemmatized`,
-   * text, falls back to the raw memory `data`, and truncates to the VarChar
+   * then `text_lemmatized`, then raw `data`; truncates to the VarChar limit.
-   * limit (mirrors the Python provider).
    */
   private bm25Text(payload?: Record<string, any>): string {
     if (!payload) return "";
-    const raw = payload.text_lemmatized || payload.data || "";
+    const raw =
+      payload.textLemmatized || payload.text_lemmatized || payload.data || "";
     return String(raw).slice(0, 65535);
   }
 

mem0-ts/src/oss/src/vector_stores/mongodb.ts

@@ -62,6 +62,46 @@ export class MongoDB implements VectorStore {
     return this._initPromise;
   }
 
+  private textSearchIndexDefinition(textIndexName: string) {
+    return {
+      name: textIndexName,
+      definition: {
+        mappings: {
+          dynamic: false,
+          fields: {
+            payload: {
+              type: "document",
+              fields: {
+                data: { type: "string" },
+                textLemmatized: { type: "string" },
+                text_lemmatized: { type: "string" },
+              },
+            },
+          },
+        },
+      },
+    };
+  }
+
+  private textSearchIndexMappingIsCurrent(
+    index: Record<string, unknown>,
+  ): boolean {
+    const definition =
+      (index.latestDefinition as Record<string, unknown> | undefined) ??
+      (index.definition as Record<string, unknown> | undefined);
+    const payloadFields = (
+      (definition?.mappings as Record<string, unknown> | undefined)?.fields as
+        | Record<string, unknown>
+        | undefined
+    )?.payload as Record<string, unknown> | undefined;
+    const fields =
+      (payloadFields?.fields as Record<string, unknown> | undefined) ?? {};
+    const textLemmatized = fields.textLemmatized as
+      | { type?: string }
+      | undefined;
+    return textLemmatized?.type === "string";
+  }
+
   private async _doInitialize(): Promise<void> {
     await this.ensureClient();
     try {
@@ -111,32 +151,24 @@ export class MongoDB implements VectorStore {
       // Create Text Search Index for keywordSearch
       const textIndexName = `${this.collectionName}_text_search_index`;
       try {
-        let foundTextIndex = false;
+        let existingTextIndex: Record<string, unknown> | null = null;
         try {
           const indexes = await this.collection.listSearchIndexes().toArray();
-          foundTextIndex = indexes.some((idx) => idx.name === textIndexName);
+          existingTextIndex =
+            (indexes.find((idx) => idx.name === textIndexName) as
+              | Record<string, unknown>
+              | undefined) ?? null;
         } catch (e) {
           // ignore
         }
 
-        if (!foundTextIndex) {
+        const textSearchIndex = this.textSearchIndexDefinition(textIndexName);
-          await this.collection.createSearchIndex({
+
-            name: textIndexName,
+        if (!existingTextIndex) {
-            definition: {
+          await this.collection.createSearchIndex(textSearchIndex);
-              mappings: {
+        } else if (!this.textSearchIndexMappingIsCurrent(existingTextIndex)) {
-                dynamic: false,
+          await this.collection.dropSearchIndex(textIndexName);
-                fields: {
+          await this.collection.createSearchIndex(textSearchIndex);
-                  payload: {
-                    type: "document",
-                    fields: {
-                      data: { type: "string" },
-                      text_lemmatized: { type: "string" },
-                    },
-                  },
-                },
-              },
-            },
-          });
         }
       } catch (e: any) {
         console.warn(
@@ -286,7 +318,11 @@ export class MongoDB implements VectorStore {
             index: textIndexName,
             text: {
               query: query,
-              path: ["payload.data", "payload.text_lemmatized"],
+              path: [
+                "payload.data",
+                "payload.text_lemmatized",
+                "payload.textLemmatized",
+              ],
             },
           },
         },

mem0-ts/src/oss/src/vector_stores/opensearch.ts

@@ -276,6 +276,7 @@ export class OpenSearchDB implements VectorStore {
       should: [
         { match: { "payload.data": query } },
         { match: { "payload.text_lemmatized": query } },
+        { match: { "payload.textLemmatized": query } },
       ],
       minimum_should_match: 1,
     };

mem0-ts/src/oss/tests/opensearch.unit.test.ts

@@ -153,4 +153,26 @@ describe("OpenSearchDB", () => {
 
     await expect(store.get("missing")).resolves.toBeNull();
   });
+
+  it("keywordSearch queries lemmatized payload fields", async () => {
+    const client = createClient();
+    const store = await createStore(client);
+
+    await store.keywordSearch("stud french", 5, { user_id: "alice" });
+
+    const searchCall = client.search.mock.calls.find(
+      ([arg]: any[]) => arg.index === collectionName,
+    );
+    expect(searchCall).toBeDefined();
+    const should = searchCall![0].body.query.bool.should;
+    expect(should).toContainEqual({
+      match: { "payload.textLemmatized": "stud french" },
+    });
+    expect(should).toContainEqual({
+      match: { "payload.text_lemmatized": "stud french" },
+    });
+    expect(should).toContainEqual({
+      match: { "payload.data": "stud french" },
+    });
+  });
 });