fix(pgvector, server): add rich filter operators and fix /search 502 (#5263)
This commit is contained in:
Kartik
@@ -28,6 +28,119 @@ function escapeFilterKey(key: string): string {
|
||||
return key;
|
||||
}
|
||||
|
||||
interface FilterResult {
|
||||
conditions: string[];
|
||||
values: any[];
|
||||
paramIndex: number;
|
||||
}
|
||||
|
||||
const OPERATOR_SQL_MAP: Record<string, { template: string; numeric: boolean }> =
|
||||
{
|
||||
eq: { template: "payload->>'%KEY%' = $%IDX%", numeric: false },
|
||||
ne: { template: "payload->>'%KEY%' != $%IDX%", numeric: false },
|
||||
gt: { template: "(payload->>'%KEY%')::numeric > $%IDX%", numeric: true },
|
||||
gte: { template: "(payload->>'%KEY%')::numeric >= $%IDX%", numeric: true },
|
||||
lt: { template: "(payload->>'%KEY%')::numeric < $%IDX%", numeric: true },
|
||||
lte: { template: "(payload->>'%KEY%')::numeric <= $%IDX%", numeric: true },
|
||||
in: { template: "payload->>'%KEY%' = ANY($%IDX%::text[])", numeric: false },
|
||||
nin: {
|
||||
template: "NOT (payload->>'%KEY%' = ANY($%IDX%::text[]))",
|
||||
numeric: false,
|
||||
},
|
||||
contains: { template: "payload->>'%KEY%' LIKE $%IDX%", numeric: false },
|
||||
icontains: { template: "payload->>'%KEY%' ILIKE $%IDX%", numeric: false },
|
||||
};
|
||||
|
||||
export function buildFilterConditions(
|
||||
filters: Record<string, any> | undefined,
|
||||
startIndex: number,
|
||||
): FilterResult {
|
||||
const conditions: string[] = [];
|
||||
const values: any[] = [];
|
||||
let paramIndex = startIndex;
|
||||
|
||||
if (!filters) {
|
||||
return { conditions, values, paramIndex };
|
||||
}
|
||||
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
if (key === "$or") {
|
||||
const orGroups: string[] = [];
|
||||
for (const orFilter of value as Record<string, any>[]) {
|
||||
const sub = buildFilterConditions(orFilter, paramIndex);
|
||||
if (sub.conditions.length > 0) {
|
||||
orGroups.push("(" + sub.conditions.join(" AND ") + ")");
|
||||
values.push(...sub.values);
|
||||
paramIndex = sub.paramIndex;
|
||||
}
|
||||
}
|
||||
if (orGroups.length > 0) {
|
||||
conditions.push("(" + orGroups.join(" OR ") + ")");
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (key === "$not") {
|
||||
const notGroups: string[] = [];
|
||||
for (const notFilter of value as Record<string, any>[]) {
|
||||
const sub = buildFilterConditions(notFilter, paramIndex);
|
||||
if (sub.conditions.length > 0) {
|
||||
notGroups.push("(" + sub.conditions.join(" AND ") + ")");
|
||||
values.push(...sub.values);
|
||||
paramIndex = sub.paramIndex;
|
||||
}
|
||||
}
|
||||
if (notGroups.length > 0) {
|
||||
conditions.push("NOT (" + notGroups.join(" OR ") + ")");
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const safeKey = escapeFilterKey(key);
|
||||
|
||||
if (value === "*") {
|
||||
conditions.push(`payload ? $${paramIndex}`);
|
||||
values.push(key);
|
||||
paramIndex++;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (typeof value === "object" && value !== null && !Array.isArray(value)) {
|
||||
for (const [op, opValue] of Object.entries(value)) {
|
||||
const mapping = OPERATOR_SQL_MAP[op];
|
||||
if (!mapping) {
|
||||
throw new Error(`Unsupported filter operator: ${op}`);
|
||||
}
|
||||
const clause = mapping.template
|
||||
.replace("%KEY%", safeKey)
|
||||
.replace("%IDX%", String(paramIndex));
|
||||
conditions.push(clause);
|
||||
|
||||
if (op === "in" || op === "nin") {
|
||||
values.push((opValue as any[]).map(String));
|
||||
} else if (op === "contains" || op === "icontains") {
|
||||
values.push(`%${opValue}%`);
|
||||
} else if (mapping.numeric) {
|
||||
values.push(Number(opValue));
|
||||
} else {
|
||||
values.push(String(opValue));
|
||||
}
|
||||
paramIndex++;
|
||||
}
|
||||
} else if (Array.isArray(value)) {
|
||||
conditions.push(`payload->>'${safeKey}' = ANY($${paramIndex}::text[])`);
|
||||
values.push(value.map(String));
|
||||
paramIndex++;
|
||||
} else {
|
||||
conditions.push(`payload->>'${safeKey}' = $${paramIndex}`);
|
||||
values.push(value);
|
||||
paramIndex++;
|
||||
}
|
||||
}
|
||||
|
||||
return { conditions, values, paramIndex };
|
||||
}
|
||||
|
||||
interface PGVectorConfig extends VectorStoreConfig {
|
||||
dbname?: string;
|
||||
user: string;
|
||||
@@ -203,23 +316,15 @@ export class PGVector implements VectorStore {
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[] | null> {
|
||||
try {
|
||||
const filterConditions: string[] = [];
|
||||
const filterValues: any[] = [query, topK];
|
||||
let filterIndex = 3;
|
||||
|
||||
if (filters) {
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
const safeKey = escapeFilterKey(key);
|
||||
filterConditions.push(`payload->>'${safeKey}' = $${filterIndex}`);
|
||||
filterValues.push(value);
|
||||
filterIndex++;
|
||||
}
|
||||
}
|
||||
const {
|
||||
conditions,
|
||||
values,
|
||||
paramIndex: _,
|
||||
} = buildFilterConditions(filters, 3);
|
||||
const filterValues: any[] = [query, topK, ...values];
|
||||
|
||||
const filterClause =
|
||||
filterConditions.length > 0
|
||||
? "AND " + filterConditions.join(" AND ")
|
||||
: "";
|
||||
conditions.length > 0 ? "AND " + conditions.join(" AND ") : "";
|
||||
|
||||
const searchQuery = `
|
||||
SELECT id, ts_rank_cd(to_tsvector('simple', payload->>'textLemmatized'), plainto_tsquery('simple', $1)) AS score, payload
|
||||
@@ -248,24 +353,16 @@ export class PGVector implements VectorStore {
|
||||
topK: number = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[]> {
|
||||
const filterConditions: string[] = [];
|
||||
const queryVector = `[${query.join(",")}]`;
|
||||
const filterValues: any[] = [queryVector, topK];
|
||||
let filterIndex = 3;
|
||||
|
||||
if (filters) {
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
const safeKey = escapeFilterKey(key);
|
||||
filterConditions.push(`payload->>'${safeKey}' = $${filterIndex}`);
|
||||
filterValues.push(value);
|
||||
filterIndex++;
|
||||
}
|
||||
}
|
||||
const {
|
||||
conditions,
|
||||
values,
|
||||
paramIndex: _,
|
||||
} = buildFilterConditions(filters, 3);
|
||||
const filterValues: any[] = [queryVector, topK, ...values];
|
||||
|
||||
const filterClause =
|
||||
filterConditions.length > 0
|
||||
? "WHERE " + filterConditions.join(" AND ")
|
||||
: "";
|
||||
conditions.length > 0 ? "WHERE " + conditions.join(" AND ") : "";
|
||||
|
||||
const searchQuery = `
|
||||
SELECT id, vector <=> $1::vector AS distance, payload
|
||||
@@ -337,23 +434,14 @@ export class PGVector implements VectorStore {
|
||||
filters?: SearchFilters,
|
||||
topK: number = 100,
|
||||
): Promise<[VectorStoreResult[], number]> {
|
||||
const filterConditions: string[] = [];
|
||||
const filterValues: any[] = [];
|
||||
let paramIndex = 1;
|
||||
|
||||
if (filters) {
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
const safeKey = escapeFilterKey(key);
|
||||
filterConditions.push(`payload->>'${safeKey}' = $${paramIndex}`);
|
||||
filterValues.push(value);
|
||||
paramIndex++;
|
||||
}
|
||||
}
|
||||
const {
|
||||
conditions,
|
||||
values: filterValues,
|
||||
paramIndex,
|
||||
} = buildFilterConditions(filters, 1);
|
||||
|
||||
const filterClause =
|
||||
filterConditions.length > 0
|
||||
? "WHERE " + filterConditions.join(" AND ")
|
||||
: "";
|
||||
conditions.length > 0 ? "WHERE " + conditions.join(" AND ") : "";
|
||||
|
||||
const listQuery = `
|
||||
SELECT id, payload
|
||||
@@ -368,11 +456,11 @@ export class PGVector implements VectorStore {
|
||||
${filterClause}
|
||||
`;
|
||||
|
||||
filterValues.push(topK); // Add limit as the last parameter
|
||||
const listValues = [...filterValues, topK];
|
||||
|
||||
const [listResult, countResult] = await Promise.all([
|
||||
this.client.query(listQuery, filterValues),
|
||||
this.client.query(countQuery, filterValues.slice(0, -1)), // Remove limit parameter for count query
|
||||
this.client.query(listQuery, listValues),
|
||||
this.client.query(countQuery, filterValues),
|
||||
]);
|
||||
|
||||
const results = listResult.rows.map((row) => ({
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
/// <reference types="jest" />
|
||||
|
||||
jest.mock("pg", () => {
|
||||
const Client = jest.fn().mockImplementation(() => ({
|
||||
connect: jest.fn().mockResolvedValue(undefined),
|
||||
end: jest.fn().mockResolvedValue(undefined),
|
||||
query: jest.fn().mockResolvedValue({ rows: [] }),
|
||||
}));
|
||||
const escapeIdentifier = (str: string) => `"${str.replace(/"/g, '""')}"`;
|
||||
return {
|
||||
__esModule: true,
|
||||
default: { Client, escapeIdentifier },
|
||||
Client,
|
||||
escapeIdentifier,
|
||||
};
|
||||
});
|
||||
|
||||
import { buildFilterConditions } from "../src/vector_stores/pgvector";
|
||||
|
||||
describe("buildFilterConditions", () => {
|
||||
test("returns empty for undefined filters", () => {
|
||||
const result = buildFilterConditions(undefined, 1);
|
||||
expect(result.conditions).toEqual([]);
|
||||
expect(result.values).toEqual([]);
|
||||
expect(result.paramIndex).toBe(1);
|
||||
});
|
||||
|
||||
test("returns empty for empty filters", () => {
|
||||
const result = buildFilterConditions({}, 1);
|
||||
expect(result.conditions).toEqual([]);
|
||||
expect(result.values).toEqual([]);
|
||||
});
|
||||
|
||||
test("simple equality", () => {
|
||||
const result = buildFilterConditions({ user_id: "alice" }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("payload->>'user_id' = $1");
|
||||
expect(result.values).toEqual(["alice"]);
|
||||
expect(result.paramIndex).toBe(2);
|
||||
});
|
||||
|
||||
test("multiple equalities", () => {
|
||||
const result = buildFilterConditions(
|
||||
{ user_id: "alice", agent_id: "bot1" },
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(2);
|
||||
expect(result.values).toEqual(["alice", "bot1"]);
|
||||
expect(result.paramIndex).toBe(3);
|
||||
});
|
||||
|
||||
test("eq operator", () => {
|
||||
const result = buildFilterConditions({ status: { eq: "active" } }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("= $1");
|
||||
expect(result.values).toEqual(["active"]);
|
||||
});
|
||||
|
||||
test("ne operator", () => {
|
||||
const result = buildFilterConditions({ status: { ne: "deleted" } }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("!= $1");
|
||||
expect(result.values).toEqual(["deleted"]);
|
||||
});
|
||||
|
||||
test("gt operator", () => {
|
||||
const result = buildFilterConditions({ price: { gt: 100 } }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("::numeric > $1");
|
||||
expect(result.values).toEqual([100]);
|
||||
});
|
||||
|
||||
test("gte operator", () => {
|
||||
const result = buildFilterConditions({ price: { gte: 100 } }, 1);
|
||||
expect(result.conditions[0]).toContain("::numeric >= $1");
|
||||
expect(result.values).toEqual([100]);
|
||||
});
|
||||
|
||||
test("lt operator", () => {
|
||||
const result = buildFilterConditions({ price: { lt: 50 } }, 1);
|
||||
expect(result.conditions[0]).toContain("::numeric < $1");
|
||||
expect(result.values).toEqual([50]);
|
||||
});
|
||||
|
||||
test("lte operator", () => {
|
||||
const result = buildFilterConditions({ price: { lte: 50 } }, 1);
|
||||
expect(result.conditions[0]).toContain("::numeric <= $1");
|
||||
expect(result.values).toEqual([50]);
|
||||
});
|
||||
|
||||
test("range combination (gte + lte)", () => {
|
||||
const result = buildFilterConditions({ score: { gte: 1, lte: 10 } }, 1);
|
||||
expect(result.conditions).toHaveLength(2);
|
||||
expect(result.conditions[0]).toContain("::numeric >= $1");
|
||||
expect(result.conditions[1]).toContain("::numeric <= $2");
|
||||
expect(result.values).toEqual([1, 10]);
|
||||
});
|
||||
|
||||
test("in operator", () => {
|
||||
const result = buildFilterConditions(
|
||||
{ status: { in: ["active", "pending"] } },
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("= ANY($1::text[])");
|
||||
expect(result.values).toEqual([["active", "pending"]]);
|
||||
});
|
||||
|
||||
test("nin operator", () => {
|
||||
const result = buildFilterConditions(
|
||||
{ status: { nin: ["deleted", "archived"] } },
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("NOT");
|
||||
expect(result.conditions[0]).toContain("= ANY($1::text[])");
|
||||
expect(result.values).toEqual([["deleted", "archived"]]);
|
||||
});
|
||||
|
||||
test("contains operator", () => {
|
||||
const result = buildFilterConditions({ name: { contains: "alice" } }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("LIKE $1");
|
||||
expect(result.values).toEqual(["%alice%"]);
|
||||
});
|
||||
|
||||
test("icontains operator", () => {
|
||||
const result = buildFilterConditions({ name: { icontains: "Alice" } }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("ILIKE $1");
|
||||
expect(result.values).toEqual(["%Alice%"]);
|
||||
});
|
||||
|
||||
test("wildcard value", () => {
|
||||
const result = buildFilterConditions({ metadata_key: "*" }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("payload ? $1");
|
||||
expect(result.values).toEqual(["metadata_key"]);
|
||||
});
|
||||
|
||||
test("list shorthand (array value)", () => {
|
||||
const result = buildFilterConditions({ tags: ["a", "b", "c"] }, 1);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain("= ANY($1::text[])");
|
||||
expect(result.values).toEqual([["a", "b", "c"]]);
|
||||
});
|
||||
|
||||
test("$or operator", () => {
|
||||
const result = buildFilterConditions(
|
||||
{
|
||||
$or: [{ user_id: "alice" }, { user_id: "bob" }],
|
||||
},
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain(" OR ");
|
||||
expect(result.conditions[0]).toMatch(/^\(/);
|
||||
expect(result.values).toEqual(["alice", "bob"]);
|
||||
});
|
||||
|
||||
test("$not operator", () => {
|
||||
const result = buildFilterConditions(
|
||||
{
|
||||
$not: [{ status: "deleted" }],
|
||||
},
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toMatch(/^NOT/);
|
||||
expect(result.values).toEqual(["deleted"]);
|
||||
});
|
||||
|
||||
test("$or with operators", () => {
|
||||
const result = buildFilterConditions(
|
||||
{
|
||||
$or: [{ price: { gt: 100 } }, { price: { lt: 10 } }],
|
||||
},
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(1);
|
||||
expect(result.conditions[0]).toContain(" OR ");
|
||||
expect(result.values).toEqual([100, 10]);
|
||||
});
|
||||
|
||||
test("mixed simple and operator filters", () => {
|
||||
const result = buildFilterConditions(
|
||||
{
|
||||
user_id: "alice",
|
||||
score: { gte: 5 },
|
||||
},
|
||||
1,
|
||||
);
|
||||
expect(result.conditions).toHaveLength(2);
|
||||
expect(result.values[0]).toBe("alice");
|
||||
expect(result.values[1]).toBe(5);
|
||||
});
|
||||
|
||||
test("unsupported operator throws", () => {
|
||||
expect(() => buildFilterConditions({ x: { badop: 1 } }, 1)).toThrow(
|
||||
"Unsupported filter operator",
|
||||
);
|
||||
});
|
||||
|
||||
test("in with numeric values converts to strings", () => {
|
||||
const result = buildFilterConditions({ priority: { in: [1, 2, 3] } }, 1);
|
||||
expect(result.values).toEqual([["1", "2", "3"]]);
|
||||
});
|
||||
|
||||
test("paramIndex increments correctly across multiple fields", () => {
|
||||
const result = buildFilterConditions(
|
||||
{ a: "x", b: { gt: 5 }, c: { in: [1, 2] } },
|
||||
3,
|
||||
);
|
||||
expect(result.conditions[0]).toContain("$3");
|
||||
expect(result.conditions[1]).toContain("$4");
|
||||
expect(result.conditions[2]).toContain("$5");
|
||||
expect(result.paramIndex).toBe(6);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user