Skip to content

Commit

Permalink
fix: do not return false immediately with negative operators
Browse files Browse the repository at this point in the history
  • Loading branch information
@thucpn
thucpn committed Jul 19, 2024
1 parent c69e740 commit c1e82bd
Show file tree
Hide file tree
Showing 2 changed files with 197 additions and 19 deletions.
28 changes: 22 additions & 6 deletions packages/llamaindex/src/storage/vectorStore/SimpleVectorStore.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,23 +62,32 @@ const OPERATOR_TO_FILTER: {
return parseArrayValue(value).every((v) => metadata[key].includes(v));
},
[FilterOperator.TEXT_MATCH]: ({ key, value }, metadata) => {
return metadata[key].includes(parsePrimitiveValue(value));
if (typeof metadata[key] !== "string") return false;
return metadata[key].includes(parsePrimitiveValue(value).toString());
},
[FilterOperator.CONTAINS]: ({ key, value }, metadata) => {
if (!Array.isArray(metadata[key])) return false;
return !!parseArrayValue(metadata[key]).find((v) => v === value);
},
[FilterOperator.GT]: ({ key, value }, metadata) => {
return metadata[key] > parsePrimitiveValue(value);
const val = metadata[key];
if (typeof val !== "string" && typeof val !== "number") return false;
return val > parsePrimitiveValue(value);
},
[FilterOperator.LT]: ({ key, value }, metadata) => {
return metadata[key] < parsePrimitiveValue(value);
const val = metadata[key];
if (typeof val !== "string" && typeof val !== "number") return false;
return val < parsePrimitiveValue(value);
},
[FilterOperator.GTE]: ({ key, value }, metadata) => {
return metadata[key] >= parsePrimitiveValue(value);
const val = metadata[key];
if (typeof val !== "string" && typeof val !== "number") return false;
return val >= parsePrimitiveValue(value);
},
[FilterOperator.LTE]: ({ key, value }, metadata) => {
return metadata[key] <= parsePrimitiveValue(value);
const val = metadata[key];
if (typeof val !== "string" && typeof val !== "number") return false;
return val <= parsePrimitiveValue(value);
},
};

Expand All @@ -94,7 +103,14 @@ const buildFilterFn = (
const queryCondition = condition || "and"; // default to and

const itemFilterFn = (filter: MetadataFilter): boolean => {
if (metadata[filter.key] === undefined) return false; // always return false if the metadata key is not present
// for all operators except != and nin, if the metadata key is not present, return false
if (
metadata[filter.key] === undefined &&
filter.operator !== FilterOperator.NE &&
filter.operator !== FilterOperator.NIN
) {
return false;
}
const metadataLookupFn = OPERATOR_TO_FILTER[filter.operator];
if (!metadataLookupFn)
throw new Error(`Unsupported operator: ${filter.operator}`);
Expand Down
188 changes: 175 additions & 13 deletions packages/llamaindex/tests/vectorStores/SimpleVectorStore.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ describe("SimpleVectorStore", () => {
private: "true",
weight: 1.2,
type: ["husky", "puppy"],
height: 50,
},
}),
new TextNode({
Expand Down Expand Up @@ -87,19 +88,6 @@ describe("SimpleVectorStore", () => {
title: "No filter",
expected: 3,
},
{
title: "Filter with non-exist key",
filters: {
filters: [
{
key: "non-exist-key",
value: "cat",
operator: "==",
},
],
},
expected: 0,
},
{
title: "Filter EQ",
filters: {
Expand Down Expand Up @@ -309,4 +297,178 @@ describe("SimpleVectorStore", () => {
});
});
});

describe("[SimpleVectorStore] query nodes with optional key in metadata", () => {
const testcases: FilterTestCase[] = [
{
title: "Filter EQ with an optional key",
filters: {
filters: [
{
key: "height",
value: 50,
operator: "==",
},
],
},
expected: 1,
},
{
title: "Filter NE with an optional key",
filters: {
filters: [
{
key: "height",
value: 50,
operator: "!=",
},
],
},
expected: 2,
},
{
title: "Filter GT with an optional key",
filters: {
filters: [
{
key: "height",
value: 48,
operator: ">",
},
],
},
expected: 1,
},
{
title: "Filter GTE with an optional key",
filters: {
filters: [
{
key: "height",
value: 50,
operator: ">=",
},
],
},
expected: 1,
},
{
title: "Filter LT with an optional key",
filters: {
filters: [
{
key: "height",
value: 50,
operator: "<",
},
],
},
expected: 0,
},
{
title: "Filter LTE with an optional key",
filters: {
filters: [
{
key: "height",
value: 50,
operator: "<=",
},
],
},
expected: 1,
},
{
title: "Filter IN with an optional key",
filters: {
filters: [
{
key: "non-existing-key",
value: ["a", "b"],
operator: "in",
},
],
},
expected: 0,
},
{
title: "Filter NIN with an optional key",
filters: {
filters: [
{
key: "non-existing-key",
value: ["a", "b"],
operator: "nin",
},
],
},
expected: 3,
},
{
title: "Filter ANY with an optional key",
filters: {
filters: [
{
key: "non-existing-key",
value: ["a", "b"],
operator: "any",
},
],
},
expected: 0,
},
{
title: "Filter ALL with an optional key",
filters: {
filters: [
{
key: "non-existing-key",
value: ["a", "b"],
operator: "all",
},
],
},
expected: 0,
},
{
title: "Filter CONTAINS with an optional key",
filters: {
filters: [
{
key: "non-existing-key",
value: "a",
operator: "contains",
},
],
},
expected: 0,
},
{
title: "Filter TEXT_MATCH with an optional key",
filters: {
filters: [
{
key: "non-existing-key",
value: "a",
operator: "text_match",
},
],
},
expected: 0,
},
];

testcases.forEach((tc) => {
it(`[${tc.title}] should return ${tc.expected} nodes`, async () => {
await store.add(nodes);
const result = await store.query({
queryEmbedding: [0.1, 0.2],
similarityTopK: 3,
mode: VectorStoreQueryMode.DEFAULT,
filters: tc.filters,
});
expect(result.ids).length(tc.expected);
});
});
});
});

0 comments on commit c1e82bd

Please sign in to comment.