File size: 1,206 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

// create a cloudclient class that takes in an api key and an optional database
// this should wrap ChromaClient and specify the auth scheme correctly

import { ChromaClient } from "./ChromaClient";

interface CloudClientParams {
    apiKey?: string;
    database?: string;
    cloudHost?: string;
    cloudPort?: string;
}

class CloudClient extends ChromaClient{

    constructor({apiKey, database, cloudHost, cloudPort}: CloudClientParams)  {
        // If no API key is provided, try to load it from the environment variable
        if (!apiKey) {
            apiKey = process.env.CHROMA_API_KEY;
        }
        if (!apiKey) {
            throw new Error("No API key provided");
        }

        cloudHost = cloudHost || "https://api.trychroma.com";
        cloudPort = cloudPort || "8000";

        const path = `${cloudHost}:${cloudPort}`;

        const auth =  {
                provider: "token",
                credentials: apiKey,
                providerOptions: { headerType: "X_CHROMA_TOKEN" },
            }

        return new ChromaClient({
            path: path,
            auth: auth,
            database: database,
        })

        super()
    }
}

export { CloudClient };