Client.java

package tech.amikos.chromadb;

import com.google.gson.internal.LinkedTreeMap;
import tech.amikos.chromadb.handler.ApiClient;
import tech.amikos.chromadb.handler.ApiException;
import tech.amikos.chromadb.handler.DefaultApi;
import tech.amikos.chromadb.model.CreateCollection;

import java.math.BigDecimal;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * ChromaDB Client
 */
public class Client {
    final ApiClient apiClient = new ApiClient();
    private int timeout = 60;
    DefaultApi api;

    public Client(String basePath) {
        apiClient.setBasePath(basePath);
        api = new DefaultApi(apiClient);
        apiClient.setHttpClient(apiClient.getHttpClient().newBuilder()
                .readTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
                .writeTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
                .build());
        api.getApiClient().setUserAgent("Chroma-JavaClient/0.1.x");
    }

    /**
     * Set the timeout for the client
     * @param timeout timeout in seconds
     */
    public void setTimeout(int timeout) {
        this.timeout = timeout;
        apiClient.setHttpClient(apiClient.getHttpClient().newBuilder()
                .readTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
                .writeTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
                .build());
    }

    public Collection getCollection(String collectionName, EmbeddingFunction embeddingFunction) throws ApiException {
        return new Collection(api, collectionName, embeddingFunction).fetch();
    }

    public Map<String, BigDecimal> heartbeat() throws ApiException {
        return api.heartbeat();
    }

    public Collection createCollection(String collectionName, Map<String, String> metadata, Boolean createOrGet, EmbeddingFunction embeddingFunction) throws ApiException {
        CreateCollection req = new CreateCollection();
        req.setName(collectionName);
        Map<String, String> _metadata = metadata;
        if (metadata == null || metadata.isEmpty() || !metadata.containsKey("embedding_function")) {
            _metadata = new LinkedTreeMap<>();
            _metadata.put("embedding_function", embeddingFunction.getClass().getName());
        }
        req.setMetadata(_metadata);
        req.setGetOrCreate(createOrGet);
        LinkedTreeMap resp = (LinkedTreeMap) api.createCollection(req);
        return new Collection(api, (String) resp.get("name"), embeddingFunction).fetch();
    }

    public Collection deleteCollection(String collectionName) throws ApiException {
        Collection collection = Collection.getInstance(api, collectionName);
        api.deleteCollection(collectionName);
        return collection;
    }

    public Collection upsert(String collectionName, EmbeddingFunction ef) throws ApiException {
        Collection collection = getCollection(collectionName, ef);
//        collection.upsert();
        return collection;
    }

    public Boolean reset() throws ApiException {
        return api.reset();
    }

    public List<Collection> listCollections() throws ApiException {
        List<LinkedTreeMap> apiResponse = (List<LinkedTreeMap>) api.listCollections();
        return apiResponse.stream().map((LinkedTreeMap m) -> {
            try {
                return getCollection((String) m.get("name"), null);
            } catch (ApiException e) {
                e.printStackTrace(); //this is not great as we're swallowing the exception
            }
            return null;
        }).collect(Collectors.toList());
    }

    public String version() throws ApiException {
        return api.version();
    }
}