618 lines
19 KiB
TypeScript
618 lines
19 KiB
TypeScript
import {
|
|
AutoBuildPolicyType,
|
|
FieldType,
|
|
IndexType,
|
|
InvertedIndexFieldAttribute,
|
|
MetricType,
|
|
PartitionType,
|
|
ServerErrCode,
|
|
TableState,
|
|
} from "@mochow/mochow-sdk-node";
|
|
import { BaiduDB } from "../src/vector_stores/baidu";
|
|
|
|
// No jest.mock() here on purpose: the real SDK supplies the enums and the search request
|
|
// classes (which carry an internal `set` map the client reads, so they cannot be hand-rolled
|
|
// as plain literals). Only the network-facing MochowClient is faked, via the `client` config.
|
|
|
|
const OK = { code: 0, msg: "" };
|
|
const DIMS = 1536;
|
|
|
|
const normalTable = (schema: unknown = { fields: [], indexes: [] }) => ({
|
|
...OK,
|
|
table: { state: TableState.Normal, schema },
|
|
});
|
|
|
|
const CORE_FIELDS = [
|
|
{ fieldName: "id", fieldType: FieldType.String },
|
|
{ fieldName: "data", fieldType: FieldType.Text },
|
|
{ fieldName: "vector", fieldType: FieldType.FloatVector, dimension: DIMS },
|
|
{ fieldName: "metadata", fieldType: "JSON" },
|
|
];
|
|
|
|
const BM25_FIELDS = [
|
|
...CORE_FIELDS,
|
|
{ fieldName: "textLemmatized", fieldType: FieldType.Text },
|
|
];
|
|
|
|
/** Records call order, so ordering regressions (deleteCol vs. in-flight init) are visible. */
|
|
function fakeClient(overrides: Record<string, (...args: any[]) => any> = {}) {
|
|
const calls: string[] = [];
|
|
const track =
|
|
(name: string, impl: (...args: any[]) => any) =>
|
|
(...args: any[]) => {
|
|
calls.push(name);
|
|
return impl(...args);
|
|
};
|
|
|
|
const client: any = {
|
|
calls,
|
|
createDatabase: jest.fn(track("createDatabase", async () => OK)),
|
|
createTable: jest.fn(track("createTable", async () => OK)),
|
|
dropTable: jest.fn(track("dropTable", async () => OK)),
|
|
descTable: jest.fn(track("descTable", async () => normalTable())),
|
|
upsert: jest.fn(async () => OK),
|
|
delete: jest.fn(async () => OK),
|
|
query: jest.fn(),
|
|
select: jest.fn(),
|
|
vectorSearch: jest.fn(),
|
|
bm25Search: jest.fn(),
|
|
};
|
|
|
|
for (const [name, impl] of Object.entries(overrides)) {
|
|
client[name] = jest.fn(track(name, impl));
|
|
}
|
|
return client;
|
|
}
|
|
|
|
const makeStore = (client: any, extra: Record<string, unknown> = {}) =>
|
|
new BaiduDB({
|
|
endpoint: "http://127.0.0.1:5287",
|
|
account: "root",
|
|
apiKey: "test-key",
|
|
databaseName: "mem0_db",
|
|
tableName: "mem0",
|
|
embeddingModelDims: DIMS,
|
|
client,
|
|
...extra,
|
|
} as any);
|
|
|
|
/** Run the poll loop's setTimeout inline so tests never wait the real 2s interval. */
|
|
const runTimersInline = () =>
|
|
jest.spyOn(global, "setTimeout").mockImplementation(((fn: () => void) => {
|
|
fn();
|
|
return 0;
|
|
}) as any);
|
|
|
|
beforeEach(() => {
|
|
jest.spyOn(console, "warn").mockImplementation(() => {});
|
|
jest.spyOn(console, "error").mockImplementation(() => {});
|
|
});
|
|
|
|
afterEach(() => jest.restoreAllMocks());
|
|
|
|
describe("BaiduDB config", () => {
|
|
it("rejects a missing required field", () => {
|
|
expect(() => makeStore(fakeClient(), { tableName: "" })).toThrow(
|
|
/non-empty 'tableName'/,
|
|
);
|
|
});
|
|
|
|
it("does not require endpoint credentials when a client is injected", () => {
|
|
expect(() =>
|
|
makeStore(fakeClient(), { endpoint: "", account: "", apiKey: "" }),
|
|
).not.toThrow();
|
|
});
|
|
});
|
|
|
|
describe("BaiduDB table provisioning", () => {
|
|
it("creates the table with the schema mem0 needs", async () => {
|
|
const client = fakeClient();
|
|
await makeStore(client).initialize();
|
|
|
|
const spec = client.createTable.mock.calls[0][0];
|
|
expect(client.createDatabase).toHaveBeenCalledWith("mem0_db");
|
|
expect(spec.database).toBe("mem0_db");
|
|
expect(spec.table).toBe("mem0");
|
|
expect(spec.enableDynamicField).toBe(false);
|
|
// Mochow rejects a partition without partitionType.
|
|
expect(spec.partition).toEqual({
|
|
partitionType: PartitionType.HASH,
|
|
partitionNum: 1,
|
|
});
|
|
|
|
const fields = spec.schema.fields.map((f: any) => [
|
|
f.fieldName,
|
|
f.fieldType,
|
|
]);
|
|
expect(fields).toEqual([
|
|
["id", FieldType.String],
|
|
["data", FieldType.Text],
|
|
["vector", FieldType.FloatVector],
|
|
["textLemmatized", FieldType.Text],
|
|
["metadata", "JSON"],
|
|
]);
|
|
expect(spec.schema.fields[0]).toMatchObject({
|
|
primaryKey: true,
|
|
partitionKey: true,
|
|
notNull: true,
|
|
});
|
|
expect(spec.schema.fields[2].dimension).toBe(DIMS);
|
|
});
|
|
|
|
it("builds a vector index with a genuine row-count-increment auto-build policy", async () => {
|
|
const client = fakeClient();
|
|
await makeStore(client).initialize();
|
|
|
|
const [vectorIndex] = client.createTable.mock.calls[0][0].schema.indexes;
|
|
expect(vectorIndex).toMatchObject({
|
|
indexName: "vector_idx",
|
|
indexType: IndexType.HNSW,
|
|
field: "vector",
|
|
metricType: MetricType.L2,
|
|
params: { M: 16, efConstruction: 200 },
|
|
autoBuild: true,
|
|
});
|
|
// Regression guard: sdk.AutoBuildIncrement() stamps policyType "TIMING" in 2.1.5.
|
|
expect(vectorIndex.autoBuildPolicy).toEqual({
|
|
policyType: AutoBuildPolicyType.Increment,
|
|
rowCountIncrement: 10000,
|
|
});
|
|
expect(AutoBuildPolicyType.Increment).not.toBe(AutoBuildPolicyType.Timing);
|
|
});
|
|
|
|
it("declares the filtering and BM25 indexes with an indexType", async () => {
|
|
const client = fakeClient();
|
|
await makeStore(client).initialize();
|
|
|
|
const [, filtering, bm25] =
|
|
client.createTable.mock.calls[0][0].schema.indexes;
|
|
expect(filtering).toEqual({
|
|
indexName: "metadata_filtering_idx",
|
|
indexType: IndexType.FilteringIndex,
|
|
fields: ["metadata"],
|
|
});
|
|
// Memory.search() hands keywordSearch() an already-lemmatized query, so raw `data` is
|
|
// not worth indexing — only the lemmatized column is. The index is therefore named for
|
|
// that column and must never be called "data_bm25_idx": that is the name Python's
|
|
// keyword_search() queries with a raw, unstemmed query, and it must keep missing (and so
|
|
// falling back to vector search) rather than half-matching this stemmed index.
|
|
expect(bm25).toMatchObject({
|
|
indexName: "text_lemmatized_bm25_idx",
|
|
indexType: IndexType.InvertedIndex,
|
|
fields: ["textLemmatized"],
|
|
fieldAttributes: [InvertedIndexFieldAttribute.Analyzed],
|
|
});
|
|
});
|
|
|
|
it("honours a configured metric type", async () => {
|
|
const client = fakeClient();
|
|
await makeStore(client, { metricType: "COSINE" }).initialize();
|
|
const [vectorIndex] = client.createTable.mock.calls[0][0].schema.indexes;
|
|
expect(vectorIndex.metricType).toBe(MetricType.COSINE);
|
|
});
|
|
|
|
it("tolerates an existing database and table", async () => {
|
|
const client = fakeClient({
|
|
createDatabase: async () => ({
|
|
code: ServerErrCode.DBAlreadyExist,
|
|
msg: "db exists",
|
|
}),
|
|
createTable: async () => ({
|
|
code: ServerErrCode.TableAlreadyExist,
|
|
msg: "table exists",
|
|
}),
|
|
descTable: async () => normalTable({ fields: BM25_FIELDS, indexes: [] }),
|
|
});
|
|
await expect(makeStore(client).initialize()).resolves.toBeUndefined();
|
|
});
|
|
|
|
it("waits for a CREATING table to become NORMAL", async () => {
|
|
runTimersInline();
|
|
const states = [
|
|
TableState.Creating,
|
|
TableState.Creating,
|
|
TableState.Normal,
|
|
];
|
|
const client = fakeClient({
|
|
descTable: async () => ({
|
|
...OK,
|
|
table: { state: states.shift(), schema: { fields: [], indexes: [] } },
|
|
}),
|
|
});
|
|
|
|
await makeStore(client).initialize();
|
|
expect(client.descTable).toHaveBeenCalledTimes(3);
|
|
});
|
|
|
|
it("surfaces a non-zero envelope as an error rather than succeeding", async () => {
|
|
const client = fakeClient({
|
|
createTable: async () => ({
|
|
code: ServerErrCode.InvalidTableSchema,
|
|
msg: "bad schema",
|
|
}),
|
|
});
|
|
await expect(makeStore(client).initialize()).rejects.toThrow(
|
|
/createTable 'mem0' failed \(code 60\): bad schema/,
|
|
);
|
|
});
|
|
|
|
it("rejects an existing table whose vector dimension disagrees", async () => {
|
|
const client = fakeClient({
|
|
createTable: async () => ({
|
|
code: ServerErrCode.TableAlreadyExist,
|
|
msg: "",
|
|
}),
|
|
descTable: async () =>
|
|
normalTable({
|
|
fields: [
|
|
CORE_FIELDS[0],
|
|
CORE_FIELDS[1],
|
|
{
|
|
fieldName: "vector",
|
|
fieldType: FieldType.FloatVector,
|
|
dimension: 768,
|
|
},
|
|
CORE_FIELDS[3],
|
|
],
|
|
indexes: [],
|
|
}),
|
|
});
|
|
await expect(makeStore(client).initialize()).rejects.toThrow(
|
|
/stores 768-dimensional vectors, but 'embeddingModelDims' is 1536/,
|
|
);
|
|
});
|
|
|
|
it("rejects an existing table missing the core schema", async () => {
|
|
const client = fakeClient({
|
|
createTable: async () => ({
|
|
code: ServerErrCode.TableAlreadyExist,
|
|
msg: "",
|
|
}),
|
|
descTable: async () =>
|
|
normalTable({ fields: [CORE_FIELDS[0]], indexes: [] }),
|
|
});
|
|
await expect(makeStore(client).initialize()).rejects.toThrow(
|
|
/missing the id\/data\/vector\/metadata schema/,
|
|
);
|
|
});
|
|
});
|
|
|
|
describe("BaiduDB keyword search support detection", () => {
|
|
it("fails closed when an existing table has no inverted index", async () => {
|
|
const client = fakeClient({
|
|
createTable: async () => ({
|
|
code: ServerErrCode.TableAlreadyExist,
|
|
msg: "",
|
|
}),
|
|
descTable: async () => normalTable({ fields: CORE_FIELDS, indexes: [] }),
|
|
});
|
|
const store = makeStore(client);
|
|
|
|
await expect(store.keywordSearch("hello")).resolves.toBeNull();
|
|
expect(client.bm25Search).not.toHaveBeenCalled();
|
|
expect(console.warn).toHaveBeenCalledWith(
|
|
expect.stringContaining("text_lemmatized_bm25_idx"),
|
|
);
|
|
});
|
|
|
|
it("enables keyword search when the existing table carries the BM25 index", async () => {
|
|
const client = fakeClient({
|
|
createTable: async () => ({
|
|
code: ServerErrCode.TableAlreadyExist,
|
|
msg: "",
|
|
}),
|
|
descTable: async () =>
|
|
normalTable({
|
|
fields: BM25_FIELDS,
|
|
indexes: [{ indexName: "text_lemmatized_bm25_idx" }],
|
|
}),
|
|
});
|
|
client.bm25Search.mockResolvedValue({ ...OK, rows: [] });
|
|
|
|
await expect(makeStore(client).keywordSearch("hello")).resolves.toEqual([]);
|
|
expect(client.bm25Search).toHaveBeenCalled();
|
|
});
|
|
|
|
it("queries the inverted index with the caller's already-lemmatized text", async () => {
|
|
const client = fakeClient();
|
|
client.bm25Search.mockResolvedValue({
|
|
...OK,
|
|
rows: [
|
|
{ row: { id: "m1", data: "loves pizza", metadata: {} }, score: 3.5 },
|
|
],
|
|
});
|
|
|
|
const results = await makeStore(client).keywordSearch("love pizza", 7, {
|
|
userId: "alice",
|
|
});
|
|
expect(results).toEqual([
|
|
{ id: "m1", payload: { data: "loves pizza" }, score: 3.5 },
|
|
]);
|
|
|
|
const { request, ...ns } = client.bm25Search.mock.calls[0][0];
|
|
expect(ns).toEqual({ database: "mem0_db", table: "mem0" });
|
|
expect(request.indexName).toBe("text_lemmatized_bm25_idx");
|
|
expect(request.searchText).toBe("love pizza");
|
|
expect(request.limit).toBe(7);
|
|
expect(request.filter).toBe('metadata["userId"] = "alice"');
|
|
});
|
|
});
|
|
|
|
describe("BaiduDB writes", () => {
|
|
it("upserts the whole batch in one call and mirrors textLemmatized out of the payload", async () => {
|
|
const client = fakeClient();
|
|
|
|
await makeStore(client).insert(
|
|
[
|
|
[1, 2],
|
|
[3, 4],
|
|
],
|
|
["a", "b"],
|
|
[
|
|
{ data: "loves pizza", textLemmatized: "love pizza" },
|
|
{ data: "runs daily" },
|
|
],
|
|
);
|
|
|
|
expect(client.upsert).toHaveBeenCalledTimes(1);
|
|
expect(client.upsert.mock.calls[0][0]).toEqual({
|
|
database: "mem0_db",
|
|
table: "mem0",
|
|
rows: [
|
|
{
|
|
id: "a",
|
|
data: "loves pizza",
|
|
vector: [1, 2],
|
|
textLemmatized: "love pizza",
|
|
metadata: {},
|
|
},
|
|
// Falls back to `data` when the caller did not lemmatize.
|
|
{
|
|
id: "b",
|
|
data: "runs daily",
|
|
vector: [3, 4],
|
|
textLemmatized: "runs daily",
|
|
metadata: {},
|
|
},
|
|
],
|
|
});
|
|
});
|
|
|
|
it("refuses a ragged batch instead of silently truncating it", async () => {
|
|
await expect(
|
|
makeStore(fakeClient()).insert([[1]], ["a", "b"], [{}]),
|
|
).rejects.toThrow(/equal length \(got 1\/2\/1\)/);
|
|
});
|
|
|
|
it("updates and deletes by primary key", async () => {
|
|
const client = fakeClient();
|
|
const store = makeStore(client);
|
|
|
|
await store.update("m1", [9], { data: "new" });
|
|
expect(client.upsert.mock.calls[0][0].rows).toEqual([
|
|
{
|
|
id: "m1",
|
|
data: "new",
|
|
vector: [9],
|
|
textLemmatized: "new",
|
|
metadata: {},
|
|
},
|
|
]);
|
|
|
|
await store.delete("m1");
|
|
expect(client.delete).toHaveBeenCalledWith({
|
|
database: "mem0_db",
|
|
table: "mem0",
|
|
primaryKey: { id: "m1" },
|
|
});
|
|
});
|
|
|
|
it("throws when the server rejects an upsert", async () => {
|
|
const client = fakeClient();
|
|
client.upsert.mockResolvedValue({ code: 100, msg: "duplicate key" });
|
|
|
|
await expect(makeStore(client).insert([[1]], ["a"], [{}])).rejects.toThrow(
|
|
/upsert failed \(code 100\): duplicate key/,
|
|
);
|
|
});
|
|
});
|
|
|
|
describe("BaiduDB reads", () => {
|
|
it("maps vector search hits out of the nested row envelope", async () => {
|
|
const client = fakeClient();
|
|
client.vectorSearch.mockResolvedValue({
|
|
...OK,
|
|
rows: [
|
|
{
|
|
row: { id: "m1", data: "x", metadata: {} },
|
|
distance: 0.2,
|
|
score: 0.8,
|
|
},
|
|
],
|
|
});
|
|
|
|
const results = await makeStore(client).search([1, 2, 3], 5, {
|
|
userId: "alice",
|
|
});
|
|
expect(results).toEqual([{ id: "m1", payload: { data: "x" }, score: 0.8 }]);
|
|
|
|
const { request } = client.vectorSearch.mock.calls[0][0];
|
|
expect(request.vectorField).toBe("vector");
|
|
expect(request.vector).toEqual({ vector: [1, 2, 3] });
|
|
expect(request.limit).toBe(5);
|
|
expect(request.filter).toBe('metadata["userId"] = "alice"');
|
|
expect(request.projections).toEqual(["id", "data", "metadata"]);
|
|
expect(request.config.params).toEqual({ ef: 200 });
|
|
});
|
|
|
|
it("omits the filter when no filters are supplied", async () => {
|
|
const client = fakeClient();
|
|
client.vectorSearch.mockResolvedValue({ ...OK, rows: [] });
|
|
await makeStore(client).search([1], 5);
|
|
expect(client.vectorSearch.mock.calls[0][0].request.filter).toBeUndefined();
|
|
});
|
|
|
|
it("escapes quotes and rejects unsafe filter keys and values", async () => {
|
|
const client = fakeClient();
|
|
client.vectorSearch.mockResolvedValue({ ...OK, rows: [] });
|
|
const store = makeStore(client);
|
|
|
|
await store.search([1], 5, { userId: 'a"b', runId: 3, agentId: true });
|
|
expect(client.vectorSearch.mock.calls[0][0].request.filter).toBe(
|
|
'metadata["userId"] = "a\\"b" AND metadata["runId"] = 3 AND metadata["agentId"] = true',
|
|
);
|
|
|
|
await expect(store.search([1], 5, { "bad key": "x" })).rejects.toThrow(
|
|
/Invalid filter key/,
|
|
);
|
|
await expect(
|
|
store.search([1], 5, { userId: ["a"] as any }),
|
|
).rejects.toThrow(/must be str, int, float, or bool, got array/);
|
|
});
|
|
|
|
it("returns null for a missing id and throws on a real query failure", async () => {
|
|
const client = fakeClient();
|
|
const store = makeStore(client);
|
|
|
|
client.query.mockResolvedValue({ ...OK, row: {} });
|
|
await expect(store.get("nope")).resolves.toBeNull();
|
|
|
|
client.query.mockResolvedValue({
|
|
...OK,
|
|
row: { id: "m1", data: "stored text", metadata: { a: 1 } },
|
|
});
|
|
await expect(store.get("m1")).resolves.toEqual({
|
|
id: "m1",
|
|
payload: { a: 1, data: "stored text" },
|
|
});
|
|
|
|
client.query.mockResolvedValue({ code: 2, msg: "invalid parameter" });
|
|
await expect(store.get("m1")).rejects.toThrow(
|
|
/query 'm1' failed \(code 2\): invalid parameter/,
|
|
);
|
|
});
|
|
|
|
// The server signals a missing primary key with code 101; pymochow and mochow-sdk-go both
|
|
// name it (ROW_KEY_NOT_FOUND / RowKeyNotFound). The Node SDK's ServerErrCode stops at 100,
|
|
// so it has to be spelled out. Memory.get()/update()/delete() all branch on a null here.
|
|
it("returns null when the server reports the row key is missing", async () => {
|
|
const client = fakeClient();
|
|
const store = makeStore(client);
|
|
|
|
client.query.mockResolvedValue({ code: 101, msg: "row key not found" });
|
|
await expect(store.get("nope")).resolves.toBeNull();
|
|
});
|
|
|
|
it("lists flat select rows and reports how many came back", async () => {
|
|
const client = fakeClient();
|
|
client.select.mockResolvedValue({
|
|
...OK,
|
|
isTruncated: false,
|
|
nextMarker: "",
|
|
rows: [{ id: "m1", data: "x", metadata: {} }, { id: "m2" }],
|
|
});
|
|
|
|
await expect(
|
|
makeStore(client).list({ userId: "alice" }, 50),
|
|
).resolves.toEqual([
|
|
[
|
|
{ id: "m1", payload: { data: "x" } },
|
|
{ id: "m2", payload: {} },
|
|
],
|
|
2,
|
|
]);
|
|
expect(client.select).toHaveBeenCalledWith({
|
|
database: "mem0_db",
|
|
table: "mem0",
|
|
filter: 'metadata["userId"] = "alice"',
|
|
projections: ["id", "data", "metadata"],
|
|
limit: 50,
|
|
});
|
|
});
|
|
});
|
|
|
|
describe("BaiduDB deleteCol", () => {
|
|
it("waits for the drop to land before returning", async () => {
|
|
runTimersInline();
|
|
const client = fakeClient();
|
|
const store = makeStore(client);
|
|
await store.initialize();
|
|
|
|
client.descTable
|
|
.mockResolvedValueOnce({ ...OK, table: { state: TableState.Deleting } })
|
|
.mockResolvedValueOnce({
|
|
code: ServerErrCode.TableNotExist,
|
|
msg: "gone",
|
|
});
|
|
|
|
await store.deleteCol();
|
|
expect(client.dropTable).toHaveBeenCalledWith("mem0_db", "mem0");
|
|
expect(client.descTable).toHaveBeenCalledTimes(3); // 1 from initialize + 2 polls
|
|
});
|
|
|
|
it("is a no-op when the table is already gone", async () => {
|
|
const client = fakeClient();
|
|
const store = makeStore(client);
|
|
await store.initialize();
|
|
|
|
client.dropTable.mockResolvedValue({
|
|
code: ServerErrCode.TableNotExist,
|
|
msg: "gone",
|
|
});
|
|
await expect(store.deleteCol()).resolves.toBeUndefined();
|
|
});
|
|
|
|
// Regression: deleteCol() used to run alongside the fire-and-forget initialize() the
|
|
// constructor starts, so the in-flight createTable landed *after* dropTable and the table
|
|
// survived reset().
|
|
it("does not race the initialize() the constructor kicks off", async () => {
|
|
runTimersInline();
|
|
let release: () => void = () => {};
|
|
const gate = new Promise<void>((resolve) => {
|
|
release = resolve;
|
|
});
|
|
|
|
// The table exists after the first init, is gone once dropTable lands, then exists again.
|
|
const descQueue: unknown[] = [
|
|
normalTable(),
|
|
{ code: ServerErrCode.TableNotExist, msg: "gone" },
|
|
];
|
|
const client = fakeClient({
|
|
createDatabase: async () => {
|
|
await gate;
|
|
return OK;
|
|
},
|
|
descTable: async () => descQueue.shift() ?? normalTable(),
|
|
});
|
|
|
|
const store = makeStore(client); // initialize() is now in flight, parked on `gate`
|
|
const resetting = store.reset();
|
|
release();
|
|
await resetting;
|
|
|
|
expect(client.calls).toEqual([
|
|
"createDatabase",
|
|
"createTable",
|
|
"descTable",
|
|
"dropTable",
|
|
"descTable",
|
|
"createDatabase",
|
|
"createTable",
|
|
"descTable",
|
|
]);
|
|
expect(client.calls.indexOf("dropTable")).toBeGreaterThan(
|
|
client.calls.indexOf("createTable"),
|
|
);
|
|
expect(client.createTable).toHaveBeenCalledTimes(2);
|
|
});
|
|
});
|
|
|
|
describe("BaiduDB user id", () => {
|
|
it("round-trips the store user id", async () => {
|
|
const store = makeStore(fakeClient());
|
|
await expect(store.getUserId()).resolves.toBe("anonymous-baidu-user");
|
|
await store.setUserId("alice");
|
|
await expect(store.getUserId()).resolves.toBe("alice");
|
|
});
|
|
});
|