Spaces:
Running
Running
File size: 2,177 Bytes
287a0bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import { expect, test } from "@jest/globals";
import chroma from "./initClient";
import { IncludeEnum } from "../src/types";
import { IDS, DOCUMENTS, EMBEDDINGS, METADATAS } from "./data";
test("it should get embedding with matching documents", 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.get({
ids: ["test1"],
include: [
IncludeEnum.Embeddings,
IncludeEnum.Metadatas,
IncludeEnum.Documents,
]
});
expect(results).toBeDefined();
expect(results).toBeInstanceOf(Object);
expect(results.embeddings![0]).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
await collection.update({
ids: ["test1"],
embeddings: [[1, 2, 3, 4, 5, 6, 7, 8, 9, 11]],
metadatas: [{ test: "test1new" }],
documents: ["doc1new"]
});
const results2 = await collection.get({
ids: ["test1"],
include: [
IncludeEnum.Embeddings,
IncludeEnum.Metadatas,
IncludeEnum.Documents,
]
});
expect(results2).toBeDefined();
expect(results2).toBeInstanceOf(Object);
expect(results2.embeddings![0]).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 11]);
expect(results2.metadatas[0]).toEqual({ test: "test1new", float_value: -2 });
expect(results2.documents[0]).toEqual("doc1new");
});
// this currently fails
// test("it should update metadata or documents to array of Nones", async () => {
// await chroma.reset();
// const collection = await chroma.createCollection({ name: "test" });
// await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS });
// await collection.update({
// ids: ["test1"],
// metadatas: [undefined],
// });
// const results3 = await collection.get({
// ids: ["test1"],
// include: [
// IncludeEnum.Embeddings,
// IncludeEnum.Metadatas,
// IncludeEnum.Documents,
// ]
// });
// expect(results3).toBeDefined();
// expect(results3).toBeInstanceOf(Object);
// expect(results3.metadatas[0]).toEqual({});
// });
|