New Features (#6)

* feat: Added support for reset

* fix: Improved delete Collection test case.

* feat: Cohere support

Refs: #1

* feat: Added integration tests in a separate workflow

- Release workflow now is only triggered upon new release and it bumps version to tag name

* fix: Added reset after each test to ensure test are not failing due to what previous tests have done.

* feat: List collections

- Improvements of how Collections are handled

Refs: #3

* feat: Full feature parity with the Python API (except for raw SQL execution)

Refs: #2, #3, #4

Author: tazarov

.github/workflows/integration-test.yml

@@ -0,0 +1,61 @@
+name: Integration test
+
+on:
+  push:
+    branches:
+      - develop
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  integration-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up JDK 8
+        uses: actions/setup-java@v3
+        with:
+          java-version: '8'
+          distribution: 'adopt'
+          cache: maven
+
+      - name: Install Helm
+        uses: azure/setup-helm@v1
+        with:
+          version: v3.4.0
+
+      - name: start minikube
+        id: minikube
+        uses: medyagh/setup-minikube@latest
+        with:
+          kubernetes-version: 1.27.3
+      - name: Add helm repo
+        run: |
+          set -e
+          helm repo add chromadb https://amikos-tech.github.io/chromadb-chart/
+          helm repo update
+      - name: Install chromadb
+        run: |
+          set -e
+          helm install chromadb chromadb/chromadb --set chromadb.allowReset=true
+      - name: Wait for deployment to be ready
+        id: wait-and-set
+        run: |
+          set -e
+          kubectl wait \
+          --for=condition=ready pod \
+          --selector=app.kubernetes.io/name=chromadb \
+          --timeout=120s
+          echo "chroma-url=$(minikube service chromadb --url)" >> $GITHUB_OUTPUT
+      - name: Hearthbeat
+        run: |
+          set -e
+          kubectl get svc -A
+          curl $(minikube service chromadb --url)/api/v1
+      - name: Test with Maven
+        run: mvn clean test
+        env:
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+          CHROMA_URL: ${{steps.wait-and-set.outputs.chroma-url}}

.github/workflows/release.yml

@@ -1,9 +1,8 @@
 name: release
 
 on:
-  push:
-    branches: [ "main" ]
-  pull_request:
+  release:
+    types: [created]
     branches: [ "main" ]
 
 jobs:
@@ -20,65 +19,37 @@ jobs:
           server-id: ossrh
           server-username: MAVEN_USERNAME
           server-password: MAVEN_PASSWORD
-
-      - name: Install Helm
-        uses: azure/setup-helm@v1
-        with:
-          version: v3.4.0
-
-      - name: start minikube
-        id: minikube
-        uses: medyagh/setup-minikube@latest
-        with:
-          kubernetes-version: 1.27.3
-      - name: Add helm repo
-        run: |
-          set -e
-          helm repo add chromadb https://amikos-tech.github.io/chromadb-chart/
-          helm repo update
-      - name: Install chromadb
-        run: |
-          set -e
-          helm install chromadb chromadb/chromadb
-      - name: Wait for deployment to be ready
-        id: wait-and-set
-        run: |
-          set -e
-          kubectl wait \
-          --for=condition=ready pod \
-          --selector=app.kubernetes.io/name=chromadb \
-          --timeout=120s
-          echo "chroma-url=$(minikube service chromadb --url)" >> $GITHUB_OUTPUT
-      - name: Hearthbeat
-        run: |
-          set -e
-          kubectl get svc -A
-          curl $(minikube service chromadb --url)/api/v1
-      - name: Build with Maven
-        run: mvn clean package
-        env:
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-          CHROMA_URL: ${{steps.wait-and-set.outputs.chroma-url}}
-#      - name: Publish package
-#        run: mvn --batch-mode deploy
-#        env:
-#          MAVEN_USERNAME: ${{ secrets.OSSRH_USERNAME }}
-#          MAVEN_PASSWORD: ${{ secrets.OSSRH_TOKEN }}
       - id: install-secret-key
         name: Install gpg secret key
         run: |
           # Install gpg secret key
-          cat <(echo -e "${{ secrets.AMIKOS_OSS_GPG_SECRET_KEY }}") | gpg --batch --import
+          cat <(echo -e "${GPG_KEY}") | gpg --batch --import
           # Verify gpg secret key
           gpg --list-secret-keys --keyid-format LONG
+        env:
+          GPG_KEY: ${{ secrets.AMIKOS_OSS_GPG_SECRET_KEY }}
+      - name: Version bump
+        run: |
+          mvn \
+                --no-transfer-progress \
+                --batch-mode \
+                -Dgpg.skip=true \
+                -DskipTests \
+                versions:set \
+                -DnewVersion=${{  github.ref_name }}
+          git config --local user.email "action@github.com"
+          git config --local user.name "GitHub Action"
+          git add pom.xml
+          git commit -m "Increment version to ${{  github.ref_name }}" -a
+          git push
       - name: Publish package
         run: |
           mvn \
                 --no-transfer-progress \
                 --batch-mode \
                 -Dgpg.passphrase=${GPG_PASSPHRASE} \
                 -DskipTests \
-                clean deploy
+                clean package deploy
         env:
 #          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           MAVEN_USERNAME: ${{ secrets.OSSRH_USERNAME }}

README.md

@@ -2,12 +2,43 @@
 
 This is a very basic/naive implementation in Java of the Chroma Vector Database API.
 
-| Important: For now it only supports OpenAI embeddings.
+This client works with Chroma Version `0.4.3`
+
+## Features
+
+### Embeddings Support
+
+- [x] OpenAI API
+- [x] Cohere API (including Multi-language support)
+- [ ] Sentence Transformers
+- [ ] PaLM API
+- [ ] Custom Embedding Function
+
+### Feature Parity with ChromaDB API
+
+- [x] Reset
+- [x] Heartbeat
+- [x] List Collections
+- [ ] Raw SQL
+- [x] Get Version
+- [x] Create Collection
+- [x] Delete Collection
+- [x] Collection Add
+- [x] Collection Get (partial without additional parameters)
+- [x] Collection Count
+- [x] Collection Query
+- [x] Collection Modify
+- [x] Collection Update
+- [x] Collection Upsert
+- [x] Collection Create Index
+- [x] Collection Delete - delete documents in collection
 
 ## TODO
 
 - [ ] Add support for other embedding functions
-- [ ] Push the package to Maven Central - https://docs.github.com/en/actions/publishing-packages/publishing-java-packages-with-maven
+- [ ] Push the package to Maven
+  Central - https://docs.github.com/en/actions/publishing-packages/publishing-java-packages-with-maven
+- [ ] Fluent API - make it easier for users to make use of the library
 
 ## Usage
 

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>io.github.amikos-tech</groupId>
     <artifactId>chromadb-java-client</artifactId>
-    <version>0.1.0-SNAPSHOT</version>
+    <version>0.1.1-SNAPSHOT</version>
     <description>
         Chroma Vector DB Java Client
     </description>
@@ -357,6 +357,11 @@
                     </execution>
                 </executions>
             </plugin>
+            <plugin>
+                <groupId>org.codehaus.mojo</groupId>
+                <artifactId>versions-maven-plugin</artifactId>
+                <version>2.8.1</version> <!-- Use the latest version available -->
+            </plugin>
         </plugins>
     </build>
 

src/main/java/tech/amikos/chromadb/Client.java

@@ -7,7 +7,9 @@
 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
@@ -22,8 +24,8 @@ public Client(String basePath) {
         api = new DefaultApi(apiClient);
     }
 
-    public Collection getCollection(String collectionName, EmbeddingFunction embeddingFunction) {
-        return new Collection(api, collectionName, embeddingFunction);
+    public Collection getCollection(String collectionName, EmbeddingFunction embeddingFunction) throws ApiException {
+        return new Collection(api, collectionName, embeddingFunction).fetch();
     }
 
     public Map<String, BigDecimal> heartbeat() throws ApiException {
@@ -33,21 +35,46 @@ public Map<String, BigDecimal> heartbeat() throws ApiException {
     public Collection createCollection(String collectionName, Map<String, String> metadata, Boolean createOrGet, EmbeddingFunction embeddingFunction) throws ApiException {
         CreateCollection req = new CreateCollection();
         req.setName(collectionName);
-        req.setMetadata(metadata);
+        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);
+        return new Collection(api, (String) resp.get("name"), embeddingFunction).fetch();
     }
 
     public Collection deleteCollection(String collectionName) throws ApiException {
-        Collection collection = Collection.getInstance(api,collectionName);
-        collection.delete();
+        Collection collection = Collection.getInstance(api, collectionName);
+        api.deleteCollection(collectionName);
         return collection;
     }
 
-    public Collection upsert(String collectionName, EmbeddingFunction ef) {
+    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();
+    }
 }

src/main/java/tech/amikos/chromadb/CohereEmbeddingFunction.java

@@ -0,0 +1,26 @@
+package tech.amikos.chromadb;
+
+import tech.amikos.cohere.CohereClient;
+import tech.amikos.cohere.CreateEmbeddingRequest;
+import tech.amikos.cohere.CreateEmbeddingResponse;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class CohereEmbeddingFunction implements EmbeddingFunction {
+
+    private final String cohereAPIKey;
+
+    public CohereEmbeddingFunction(String cohereAPIKey) {
+        this.cohereAPIKey = cohereAPIKey;
+
+    }
+
+    @Override
+    public List<List<Float>> createEmbedding(List<String> documents) {
+        CohereClient client = new CohereClient(this.cohereAPIKey);
+        CreateEmbeddingResponse response = client.createEmbedding(new CreateEmbeddingRequest().texts(documents.toArray(new String[0])));
+        return response.getEmbeddings();
+    }
+}