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

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Clement Antony
2026-07-23 00:39:18 +05:30
committed by GitHub
parent a58e0586ad
commit 9caffeaa7b
6 changed files with 194 additions and 27 deletions
+7 -1
View File
@@ -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 () => {
+103 -1
View File
@@ -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",
],
},
},
},
+4 -4
View File
@@ -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, 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);
}
+57 -21
View File
@@ -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) {
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",
],
},
},
},
@@ -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,
};
@@ -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" },
});
});
});