Skip to content

Commit

Permalink
[ENH]: Support for $in and $nin metadata filters (chroma-core#1151)
Browse files Browse the repository at this point in the history
Refs: chroma-core#1105

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
	 - JS Client support for $in and $nin

## Test plan
*How are these changes tested?*

- [x] Tests pass locally `yarn test` for js

## Documentation Changes
TBD
  • Loading branch information
tazarov authored Sep 20, 2023
1 parent 8968222 commit 5436bd5
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 3 deletions.
8 changes: 5 additions & 3 deletions clients/js/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,15 @@ export type IDs = ID[];

export type PositiveInteger = number;

type LiteralValue = string | number;
type LiteralValue = string | number | boolean;
type ListLiteralValue = LiteralValue[];
type LiteralNumber = number;
type LogicalOperator = "$and" | "$or";
type InclusionOperator = "$in" | "$nin";
type WhereOperator = "$gt" | "$gte" | "$lt" | "$lte" | "$ne" | "$eq";

type OperatorExpression = {
[key in WhereOperator | LogicalOperator]?: LiteralValue | LiteralNumber;
[key in WhereOperator | InclusionOperator | LogicalOperator ]?: LiteralValue | ListLiteralValue;
};

type BaseWhere = {
Expand Down Expand Up @@ -77,4 +79,4 @@ export type CollectionMetadata = Record<string, unknown>;
// see all options here: https://www.jsdocs.io/package/@types/node-fetch#RequestInit
export type ConfigOptions = {
options?: RequestInit;
};
};
68 changes: 68 additions & 0 deletions clients/js/test/query.collection.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,3 +86,71 @@ test("it should query a collection with text", async () => {
expect.arrayContaining(results.documents[0])
);
})


test("it should query a collection with text and where", async () => {
await chroma.reset();
let embeddingFunction = new TestEmbeddingFunction();
const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction });
await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS });

const results = await collection.query({
queryTexts: ["test"],
nResults: 3,
where: { "float_value" : 2 }
});

expect(results).toBeDefined();
expect(results).toBeInstanceOf(Object);
expect(results.ids.length).toBe(1);
expect(["test3"]).toEqual(expect.arrayContaining(results.ids[0]));
expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0]));
expect(["This is a third test"]).toEqual(
expect.arrayContaining(results.documents[0])
);
})


test("it should query a collection with text and where in", async () => {
await chroma.reset();
let embeddingFunction = new TestEmbeddingFunction();
const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction });
await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS });

const results = await collection.query({
queryTexts: ["test"],
nResults: 3,
where: { "float_value" : { '$in': [2,5,10] }}
});

expect(results).toBeDefined();
expect(results).toBeInstanceOf(Object);
expect(results.ids.length).toBe(1);
expect(["test3"]).toEqual(expect.arrayContaining(results.ids[0]));
expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0]));
expect(["This is a third test"]).toEqual(
expect.arrayContaining(results.documents[0])
);
})

test("it should query a collection with text and where nin", async () => {
await chroma.reset();
let embeddingFunction = new TestEmbeddingFunction();
const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction });
await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS });

const results = await collection.query({
queryTexts: ["test"],
nResults: 3,
where: { "float_value" : { '$nin': [-2,0] }}
});

expect(results).toBeDefined();
expect(results).toBeInstanceOf(Object);
expect(results.ids.length).toBe(1);
expect(["test3"]).toEqual(expect.arrayContaining(results.ids[0]));
expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0]));
expect(["This is a third test"]).toEqual(
expect.arrayContaining(results.documents[0])
);
})

0 comments on commit 5436bd5

Please sign in to comment.