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:
@@ -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 () => {
|
||||
|
||||
@@ -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",
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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" },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user