Skip to content

Commit

Permalink
Merge pull request #17 from upstash/dx-730-vector-sdk-metadata-filter…
Browse files Browse the repository at this point in the history
…ing-support

Add: Metadata Filtering Support
  • Loading branch information
ogzhanolguncu authored Mar 14, 2024
2 parents 181f8f0 + 1c97f95 commit 1cbbb6e
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 1 deletion.
83 changes: 83 additions & 0 deletions src/commands/client/query/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,4 +28,87 @@ describe("QUERY", () => {
},
]);
});

test("should query records filtered with metadata filter", async () => {
const initialVector = [6.6, 7.7];
const initialData = {
id: 34,
vector: initialVector,
metadata: {
city: "Istanbul",
population: 1546000,
geography: {
continent: "Asia",
},
},
};

await new UpsertCommand(initialData).exec(client);
//This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
await sleep(2000);
const res = await new QueryCommand<{
city: string;
population: number;
geography: { continent: string };
}>({
vector: initialVector,
topK: 1,
filter: "population > 1000000 AND geography.continent = 'Asia'",
includeVectors: true,
includeMetadata: true,
}).exec(client);

expect(res).toEqual([
{
id: "34",
score: 1,
vector: [6.6, 7.7],
metadata: { city: "Istanbul", population: 1546000, geography: { continent: "Asia" } },
},
]);
});

test("should narrow down the query results with filter", async () => {
const exampleVector = [6.6, 7.7];
const initialData = [
{
id: 1,
vector: exampleVector,
metadata: {
animal: "elephant",
tags: ["mammal"],
diet: "herbivore",
},
},
{
id: 2,
vector: exampleVector,
metadata: {
animal: "tiger",
tags: ["mammal"],
diet: "carnivore",
},
},
];

await new UpsertCommand(initialData).exec(client);
//This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
await sleep(2000);
const res = await new QueryCommand<{ animal: string; tags: string[]; diet: string }>({
vector: exampleVector,
topK: 1,
filter: "tags[0] = 'mammal' AND diet = 'carnivore'",
includeVectors: true,
includeMetadata: true,
}).exec(client);

expect(res).toEqual([
{
id: "2",
score: 1,
vector: [6.6, 7.7],
metadata: { animal: "tiger", tags: ["mammal"], diet: "carnivore" },
},
]);
});
});
1 change: 1 addition & 0 deletions src/commands/client/query/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { Command } from "@commands/command";
type QueryCommandPayload = {
vector: number[];
topK: number;
filter?: string;
includeVectors?: boolean;
includeMetadata?: boolean;
};
Expand Down
7 changes: 6 additions & 1 deletion src/vector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,13 +52,18 @@ export class Index<TIndexMetadata extends Record<string, unknown> = Record<strin
*
* @example
* ```js
* await index.query({ topK: 3, vector: [ 0.22, 0.66 ]})
* await index.query({
* topK: 3,
* vector: [ 0.22, 0.66 ],
* filter: "age >= 23 and (type = \'turtle\' OR type = \'cat\')"
* });
* ```
*
* @param {Object} args - The arguments for the query command.
* @param {number[]} args.vector - An array of numbers representing the feature vector for the query.
* This vector is utilized to find the most relevant items in the index.
* @param {number} args.topK - The desired number of top results to be returned, based on relevance or similarity to the query vector.
* @param {string} [args.filter] - An optional filter string to be used in the query. The filter string is used to narrow down the query results.
* @param {boolean} [args.includeVectors=false] - When set to true, includes the feature vectors of the returned items in the response.
* @param {boolean} [args.includeMetadata=false] - When set to true, includes additional metadata of the returned items in the response.
*
Expand Down

0 comments on commit 1cbbb6e

Please sign in to comment.