Skip to content

Commit

Permalink
[BUG]: Added missing $not_contains operator for WhereDocument (#1781)
Browse files Browse the repository at this point in the history
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
	 - Added missing `$not_contains` WhereDocument op

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

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

## Documentation Changes
N/A
  • Loading branch information
tazarov authored Mar 6, 2024
1 parent c51b230 commit 276226b
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 1 deletion.
2 changes: 1 addition & 1 deletion clients/js/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ type LogicalWhere = {

export type Where = BaseWhere | LogicalWhere;

type WhereDocumentOperator = "$contains" | LogicalOperator;
type WhereDocumentOperator = "$contains" | "$not_contains" | LogicalOperator;

export type WhereDocument = {
[key in WhereDocumentOperator]?: LiteralValue | LiteralNumber | WhereDocument[];
Expand Down
11 changes: 11 additions & 0 deletions clients/js/test/get.collection.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,17 @@ test("it should get embedding with matching documents", async () => {
expect(["test1"]).toEqual(expect.arrayContaining(results2.ids));
});

test("it should get records not matching", async () => {
await chroma.reset();
const collection = await chroma.createCollection({ name: "test" });
await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS });
const results2 = await collection.get({ whereDocument: { $not_contains: "This is another" } });
expect(results2).toBeDefined();
expect(results2).toBeInstanceOf(Object);
expect(results2.ids.length).toBe(2);
expect(["test1","test3"]).toEqual(expect.arrayContaining(results2.ids));
});

test("test gt, lt, in a simple small way", async () => {
await chroma.reset();
const collection = await chroma.createCollection({ name: "test" });
Expand Down
19 changes: 19 additions & 0 deletions clients/js/test/query.collection.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,25 @@ test("it should get embedding with matching documents", async () => {
});


test("it should exclude documents matching - not_contains", async () => {
await chroma.reset();
const collection = await chroma.createCollection({ name: "test" });
await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS });

const results = await collection.query({
queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
nResults: 3,
whereDocument: { $not_contains: "This is a test" }
});

// it should only return doc1
expect(results).toBeDefined();
expect(results).toBeInstanceOf(Object);
expect(results.ids.length).toBe(1);
expect(["test2","test3"]).toEqual(expect.arrayContaining(results.ids[0]));
});


// test queryTexts
test("it should query a collection with text", async () => {
await chroma.reset();
Expand Down

0 comments on commit 276226b

Please sign in to comment.