From cb6d9edb2bc40d20314b7428e734d97faa242047 Mon Sep 17 00:00:00 2001
From: Itai Smith <itaismith@gmail.com>
Date: Wed, 20 Nov 2024 15:01:06 -0800
Subject: [PATCH] Export Collection type directly

---
 clients/js/src/ChromaClient.ts            |  9 +++--
 clients/js/src/index.ts                   |  2 +-
 clients/js/src/types.ts                   |  7 ----
 clients/js/test/collection.client.test.ts | 46 ++---------------------
 4 files changed, 10 insertions(+), 54 deletions(-)

diff --git a/clients/js/src/ChromaClient.ts b/clients/js/src/ChromaClient.ts
index 415d72193a1..ac6c9bc170d 100644
--- a/clients/js/src/ChromaClient.ts
+++ b/clients/js/src/ChromaClient.ts
@@ -290,7 +290,7 @@ export class ChromaClient {
   /**
    * Lists all collections.
    *
-   * @returns {Promise<CollectionType[]>} A promise that resolves to a list of collection names.
+   * @returns {Promise<string[]>} A promise that resolves to a list of collection names.
    * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get.
    * @param {PositiveInteger} [params.offset] - Optional offset on the items to get.
    * @throws {Error} If there is an issue listing the collections.
@@ -304,16 +304,17 @@ export class ChromaClient {
    * ```
    */
   async listCollections({ limit, offset }: ListCollectionsParams = {}): Promise<
-    CollectionParams[]
+    string[]
   > {
     await this.init();
-    return (await this.api.listCollections(
+    const collections = (await this.api.listCollections(
       this.tenant,
       this.database,
       limit,
       offset,
       this.api.options,
-    )) as CollectionParams[];
+    )) as Collection[];
+    return collections.map((collection: Collection) => collection.name);
   }
 
   /**
diff --git a/clients/js/src/index.ts b/clients/js/src/index.ts
index 36c439ead7b..fe9e64e2f18 100644
--- a/clients/js/src/index.ts
+++ b/clients/js/src/index.ts
@@ -1,6 +1,7 @@
 export { ChromaClient } from "./ChromaClient";
 export { AdminClient } from "./AdminClient";
 export { CloudClient } from "./CloudClient";
+export { Collection } from "./Collection";
 export type { IEmbeddingFunction } from "./embeddings/IEmbeddingFunction";
 export { OpenAIEmbeddingFunction } from "./embeddings/OpenAIEmbeddingFunction";
 export { CohereEmbeddingFunction } from "./embeddings/CohereEmbeddingFunction";
@@ -14,7 +15,6 @@ export { OllamaEmbeddingFunction } from "./embeddings/OllamaEmbeddingFunction";
 export type {
   IncludeEnum,
   GetParams,
-  CollectionType,
   CollectionMetadata,
   Embedding,
   Embeddings,
diff --git a/clients/js/src/types.ts b/clients/js/src/types.ts
index 88ffa29da92..67487b73049 100644
--- a/clients/js/src/types.ts
+++ b/clients/js/src/types.ts
@@ -54,13 +54,6 @@ export type WhereDocument = {
     | WhereDocument[];
 };
 
-export type CollectionType = {
-  name: string;
-  id: string;
-  metadata: Metadata | null;
-  configuration_json: any;
-};
-
 export type MultiGetResponse = {
   ids: IDs;
   embeddings: Embeddings | null;
diff --git a/clients/js/test/collection.client.test.ts b/clients/js/test/collection.client.test.ts
index 22b3a3cbb89..22cb63edcec 100644
--- a/clients/js/test/collection.client.test.ts
+++ b/clients/js/test/collection.client.test.ts
@@ -3,13 +3,9 @@ import {
   test,
   beforeEach,
   describe,
-  afterAll,
-  beforeAll,
 } from "@jest/globals";
-import { DefaultEmbeddingFunction } from "../src/embeddings/DefaultEmbeddingFunction";
-import { StartedTestContainer } from "testcontainers";
-import { ChromaClient } from "../src/ChromaClient";
-import { startChromaContainer } from "./startChromaContainer";
+import { DefaultEmbeddingFunction } from "../src";
+import { ChromaClient } from "../src";
 
 describe("collection operations", () => {
   // connects to the unauthenticated chroma instance started in
@@ -42,23 +38,7 @@ describe("collection operations", () => {
 
     const [returnedCollection] = collections;
 
-    expect({
-      ...returnedCollection,
-      configuration_json: undefined,
-      id: undefined,
-    }).toMatchInlineSnapshot(`
-    {
-      "configuration_json": undefined,
-      "database": "default_database",
-      "dimension": null,
-      "id": undefined,
-      "log_position": 0,
-      "metadata": null,
-      "name": "test",
-      "tenant": "default_tenant",
-      "version": 0,
-    }
-  `);
+    expect(returnedCollection).toEqual("test")
 
     expect([{ name: "test2", metadata: null }]).not.toEqual(
       expect.arrayContaining(collections),
@@ -79,25 +59,7 @@ describe("collection operations", () => {
     const collections2 = await client.listCollections();
     expect(collections2).toHaveLength(1);
     const [returnedCollection2] = collections2;
-    expect({
-      ...returnedCollection2,
-      configuration_json: undefined,
-      id: undefined,
-    }).toMatchInlineSnapshot(`
-    {
-      "configuration_json": undefined,
-      "database": "default_database",
-      "dimension": null,
-      "id": undefined,
-      "log_position": 0,
-      "metadata": {
-        "test": "test",
-      },
-      "name": "test2",
-      "tenant": "default_tenant",
-      "version": 0,
-    }
-  `);
+    expect(returnedCollection2).toEqual("test2");
   });
 
   test("it should get a collection", async () => {