From 9caffeaa7bd4dcbccc8b8b54dc1484e3b256dd5b Mon Sep 17 00:00:00 2001 From: Clement Antony <167181878+clementantonyk@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:39:18 +0530 Subject: [PATCH] fix(ts-sdk): use textLemmatized for BM25 keyword search on Milvus, OpenSearch, and MongoDB (#6497) Co-authored-by: kartik-mem0 --- mem0-ts/src/oss/src/tests/milvus.test.ts | 8 +- mem0-ts/src/oss/src/tests/mongodb.test.ts | 104 +++++++++++++++++- mem0-ts/src/oss/src/vector_stores/milvus.ts | 8 +- mem0-ts/src/oss/src/vector_stores/mongodb.ts | 78 +++++++++---- .../src/oss/src/vector_stores/opensearch.ts | 1 + mem0-ts/src/oss/tests/opensearch.unit.test.ts | 22 ++++ 6 files changed, 194 insertions(+), 27 deletions(-) diff --git a/mem0-ts/src/oss/src/tests/milvus.test.ts b/mem0-ts/src/oss/src/tests/milvus.test.ts index 9f1a92625..7643b1055 100644 --- a/mem0-ts/src/oss/src/tests/milvus.test.ts +++ b/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 () => { diff --git a/mem0-ts/src/oss/src/tests/mongodb.test.ts b/mem0-ts/src/oss/src/tests/mongodb.test.ts index d339de582..fe847b56c 100644 --- a/mem0-ts/src/oss/src/tests/mongodb.test.ts +++ b/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", + ], }, }, }, diff --git a/mem0-ts/src/oss/src/vector_stores/milvus.ts b/mem0-ts/src/oss/src/vector_stores/milvus.ts index 9d63b0023..284e9e08a 100644 --- a/mem0-ts/src/oss/src/vector_stores/milvus.ts +++ b/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, falls back to the raw memory `data`, and truncates to the VarChar - * limit (mirrors the Python provider). + * Text fed to the BM25 sparse index for a payload. Prefers `textLemmatized`, + * then `text_lemmatized`, then raw `data`; truncates to the VarChar limit. */ private bm25Text(payload?: Record): 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); } diff --git a/mem0-ts/src/oss/src/vector_stores/mongodb.ts b/mem0-ts/src/oss/src/vector_stores/mongodb.ts index f2cdfd309..a3bcecf21 100644 --- a/mem0-ts/src/oss/src/vector_stores/mongodb.ts +++ b/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, + ): boolean { + const definition = + (index.latestDefinition as Record | undefined) ?? + (index.definition as Record | undefined); + const payloadFields = ( + (definition?.mappings as Record | undefined)?.fields as + | Record + | undefined + )?.payload as Record | undefined; + const fields = + (payloadFields?.fields as Record | undefined) ?? {}; + const textLemmatized = fields.textLemmatized as + | { type?: string } + | undefined; + return textLemmatized?.type === "string"; + } + private async _doInitialize(): Promise { 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 | 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 + | undefined) ?? null; } catch (e) { // ignore } - if (!foundTextIndex) { - await this.collection.createSearchIndex({ - name: textIndexName, - definition: { - mappings: { - dynamic: false, - fields: { - payload: { - type: "document", - fields: { - data: { type: "string" }, - text_lemmatized: { type: "string" }, - }, - }, - }, - }, - }, - }); + const textSearchIndex = this.textSearchIndexDefinition(textIndexName); + + if (!existingTextIndex) { + await this.collection.createSearchIndex(textSearchIndex); + } else if (!this.textSearchIndexMappingIsCurrent(existingTextIndex)) { + await this.collection.dropSearchIndex(textIndexName); + await this.collection.createSearchIndex(textSearchIndex); } } 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", + ], }, }, }, diff --git a/mem0-ts/src/oss/src/vector_stores/opensearch.ts b/mem0-ts/src/oss/src/vector_stores/opensearch.ts index 12964d9f8..10f42e46a 100644 --- a/mem0-ts/src/oss/src/vector_stores/opensearch.ts +++ b/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, }; diff --git a/mem0-ts/src/oss/tests/opensearch.unit.test.ts b/mem0-ts/src/oss/tests/opensearch.unit.test.ts index 11682c860..e53a1dbc0 100644 --- a/mem0-ts/src/oss/tests/opensearch.unit.test.ts +++ b/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" }, + }); + }); });