fix(ts-oss/turbopuffer): apply all filter operators, not just gte/lte (#6578)
This commit is contained in:
@@ -247,23 +247,54 @@ export class TurbopufferDB implements VectorStore {
|
||||
}
|
||||
}
|
||||
|
||||
// Maps mem0's universal filter operators to Turbopuffer's filter tokens.
|
||||
private static readonly OPERATOR_MAP: Record<string, string> = {
|
||||
eq: "Eq",
|
||||
ne: "NotEq",
|
||||
gt: "Gt",
|
||||
gte: "Gte",
|
||||
lt: "Lt",
|
||||
lte: "Lte",
|
||||
in: "In",
|
||||
nin: "NotIn",
|
||||
};
|
||||
|
||||
private convertFilters(filters?: SearchFilters): any {
|
||||
if (!filters || Object.keys(filters).length === 0) return null;
|
||||
|
||||
const conditions: any[] = [];
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
if (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
!Array.isArray(value)
|
||||
) {
|
||||
if ("gte" in value) conditions.push([key, "Gte", value.gte]);
|
||||
if ("lte" in value) conditions.push([key, "Lte", value.lte]);
|
||||
if ("gt" in value) conditions.push([key, "Gt", value.gt]);
|
||||
if ("lt" in value) conditions.push([key, "Lt", value.lt]);
|
||||
} else {
|
||||
conditions.push([key, "Eq", value]);
|
||||
// "*" is a match-any wildcard: it must not constrain the query. The old
|
||||
// code turned it into `[key, "Eq", "*"]`, matching nothing.
|
||||
if (value === "*") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Array shorthand: { key: [a, b] } means "in".
|
||||
if (Array.isArray(value)) {
|
||||
conditions.push([key, "In", value]);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (typeof value === "object" && value !== null) {
|
||||
// Operator dict: every operator present must hold. Previously only
|
||||
// `gte` and `lte` were read, so `gt`/`lt`/`ne`/`eq`/`in`/`nin` were
|
||||
// silently dropped and the filter returned unfiltered results.
|
||||
for (const [op, operand] of Object.entries(value)) {
|
||||
const token = TurbopufferDB.OPERATOR_MAP[op];
|
||||
if (!token) {
|
||||
throw new Error(
|
||||
`Unsupported Turbopuffer filter operator '${op}' for field '${key}'. ` +
|
||||
`Supported operators: ${Object.keys(TurbopufferDB.OPERATOR_MAP).join(", ")}.`,
|
||||
);
|
||||
}
|
||||
conditions.push([key, token, operand]);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Scalar shorthand: equality.
|
||||
conditions.push([key, "Eq", value]);
|
||||
}
|
||||
|
||||
if (conditions.length === 0) return null;
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
/// <reference types="jest" />
|
||||
/**
|
||||
* Turbopuffer vector store — filter translation unit tests.
|
||||
*
|
||||
* Drives the private convertFilters() through the public search() API with a
|
||||
* virtually-mocked @turbopuffer/turbopuffer peer, and asserts the filter tuple
|
||||
* handed to ns.query().
|
||||
*/
|
||||
|
||||
const mockQuery = jest.fn().mockResolvedValue({ rows: [] });
|
||||
|
||||
// The peer is an optional dependency and may not be installed; mock it
|
||||
// virtually. createClient() does `new sdk.default({...})`, whose namespace()
|
||||
// returns the object search() calls query() on.
|
||||
jest.mock(
|
||||
"@turbopuffer/turbopuffer",
|
||||
() => ({
|
||||
__esModule: true,
|
||||
default: class {
|
||||
namespace() {
|
||||
return { query: mockQuery };
|
||||
}
|
||||
},
|
||||
}),
|
||||
{ virtual: true },
|
||||
);
|
||||
|
||||
import { TurbopufferDB } from "../src/vector_stores/turbopuffer";
|
||||
|
||||
function makeStore() {
|
||||
return new TurbopufferDB({
|
||||
apiKey: "test-key",
|
||||
collectionName: "mem0",
|
||||
} as any);
|
||||
}
|
||||
|
||||
async function filterFor(filters: any): Promise<any> {
|
||||
mockQuery.mockClear();
|
||||
await makeStore().search([0.1, 0.2, 0.3], 5, filters);
|
||||
return mockQuery.mock.calls[0][0].filters;
|
||||
}
|
||||
|
||||
describe("TurbopufferDB convertFilters", () => {
|
||||
it("maps every operator, not just gte/lte", async () => {
|
||||
expect(await filterFor({ age: { gt: 18 } })).toEqual(["age", "Gt", 18]);
|
||||
expect(await filterFor({ age: { lt: 65 } })).toEqual(["age", "Lt", 65]);
|
||||
expect(await filterFor({ age: { ne: 40 } })).toEqual(["age", "NotEq", 40]);
|
||||
expect(await filterFor({ role: { eq: "admin" } })).toEqual([
|
||||
"role",
|
||||
"Eq",
|
||||
"admin",
|
||||
]);
|
||||
expect(await filterFor({ tier: { in: ["a", "b"] } })).toEqual([
|
||||
"tier",
|
||||
"In",
|
||||
["a", "b"],
|
||||
]);
|
||||
expect(await filterFor({ tier: { nin: ["x"] } })).toEqual([
|
||||
"tier",
|
||||
"NotIn",
|
||||
["x"],
|
||||
]);
|
||||
});
|
||||
|
||||
it("applies every operator in a compound range (AND), not just the first", async () => {
|
||||
const filter = await filterFor({ age: { gt: 18, lt: 65 } });
|
||||
expect(filter[0]).toBe("And");
|
||||
expect(filter[1]).toEqual(
|
||||
expect.arrayContaining([
|
||||
["age", "Gt", 18],
|
||||
["age", "Lt", 65],
|
||||
]),
|
||||
);
|
||||
});
|
||||
|
||||
it("treats a bare array value as an 'in' filter", async () => {
|
||||
expect(await filterFor({ tier: ["gold", "silver"] })).toEqual([
|
||||
"tier",
|
||||
"In",
|
||||
["gold", "silver"],
|
||||
]);
|
||||
});
|
||||
|
||||
it("keeps scalar equality working", async () => {
|
||||
expect(await filterFor({ user_id: "u1" })).toEqual(["user_id", "Eq", "u1"]);
|
||||
});
|
||||
|
||||
it("skips a '*' wildcard value instead of matching it literally", async () => {
|
||||
// Only the real agent_id clause survives; user_id: "*" contributes nothing.
|
||||
expect(await filterFor({ user_id: "*", agent_id: "a1" })).toEqual([
|
||||
"agent_id",
|
||||
"Eq",
|
||||
"a1",
|
||||
]);
|
||||
});
|
||||
|
||||
it("throws on an unsupported operator rather than silently dropping it", async () => {
|
||||
await expect(
|
||||
makeStore().search([0.1, 0.2, 0.3], 5, { name: { startsWith: "a" } }),
|
||||
).rejects.toThrow(/Unsupported Turbopuffer filter operator 'startsWith'/);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user