feat: EF move to embeddings sub-package

README.md

@@ -73,72 +73,72 @@ package tech.amikos;
 
 import tech.amikos.chromadb.*;
 import tech.amikos.chromadb.Collection;
+import tech.amikos.chromadb.embeddings.DefaultEmbeddingFunction;
 
 import java.util.*;
 
 public class Main {
-    public static void main(String[] args) {
-        try {
-            Client client = new Client(System.getenv("CHROMA_URL"));
-            client.reset();
-            EmbeddingFunction ef = new DefaultEmbeddingFunction();
-            Collection collection = client.createCollection("test-collection", null, true, ef);
-            List<Map<String, String>> metadata = new ArrayList<>();
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "scientist");
-            }});
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "spy");
-            }});
-            collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
-            Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
-            System.out.println(qr);
-        } catch (Exception e) {
-            System.out.println(e);
-        }
+  public static void main(String[] args) {
+    try {
+      Client client = new Client(System.getenv("CHROMA_URL"));
+      client.reset();
+      EmbeddingFunction ef = new DefaultEmbeddingFunction();
+      Collection collection = client.createCollection("test-collection", null, true, ef);
+      List<Map<String, String>> metadata = new ArrayList<>();
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "scientist");
+      }});
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "spy");
+      }});
+      collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
+      Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
+      System.out.println(qr);
+    } catch (Exception e) {
+      System.out.println(e);
     }
+  }
 }
 ```
 
 ### Example OpenAI Embedding Function
 
-In this example we rely on `tech.amikos.chromadb.OpenAIEmbeddingFunction` to generate embeddings for our documents.
+In this example we rely on `tech.amikos.chromadb.embeddings.openai.OpenAIEmbeddingFunction` to generate embeddings for our documents.
 
 | **Important**: Ensure you have `OPENAI_API_KEY` environment variable set
 
 ```java
 package tech.amikos;
 
-import com.google.gson.internal.LinkedTreeMap;
 import tech.amikos.chromadb.Client;
 import tech.amikos.chromadb.Collection;
 import tech.amikos.chromadb.EmbeddingFunction;
-import tech.amikos.chromadb.OpenAIEmbeddingFunction;
+import tech.amikos.chromadb.embeddings.openai.OpenAIEmbeddingFunction;
 
 import java.util.*;
 
 public class Main {
-    public static void main(String[] args) {
-        try {
-            Client client = new Client(System.getenv("CHROMA_URL"));
-            String apiKey = System.getenv("OPENAI_API_KEY");
-            EmbeddingFunction ef = new OpenAIEmbeddingFunction(apiKey, "text-embedding-3-small");
-            Collection collection = client.createCollection("test-collection", null, true, ef);
-            List<Map<String, String>> metadata = new ArrayList<>();
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "scientist");
-            }});
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "spy");
-            }});
-            collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
-            Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
-            System.out.println(qr);
-        } catch (Exception e) {
-            e.printStackTrace();
-            System.out.println(e);
-        }
+  public static void main(String[] args) {
+    try {
+      Client client = new Client(System.getenv("CHROMA_URL"));
+      String apiKey = System.getenv("OPENAI_API_KEY");
+      EmbeddingFunction ef = new OpenAIEmbeddingFunction(apiKey, "text-embedding-3-small");
+      Collection collection = client.createCollection("test-collection", null, true, ef);
+      List<Map<String, String>> metadata = new ArrayList<>();
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "scientist");
+      }});
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "spy");
+      }});
+      collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
+      Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
+      System.out.println(qr);
+    } catch (Exception e) {
+      e.printStackTrace();
+      System.out.println(e);
     }
+  }
 }
 ```
 
@@ -174,7 +174,7 @@ curl http://localhost:11434/api/embeddings -d '{\n  "model": "llama2",\n  "promp
 
 ### Example Cohere Embedding Function
 
-In this example we rely on `tech.amikos.chromadb.CohereEmbeddingFunction` to generate embeddings for our documents.
+In this example we rely on `tech.amikos.chromadb.embeddings.cohere.CohereEmbeddingFunction` to generate embeddings for our documents.
 
 | **Important**: Ensure you have `COHERE_API_KEY` environment variable set
 
@@ -183,32 +183,33 @@ package tech.amikos;
 
 import tech.amikos.chromadb.*;
 import tech.amikos.chromadb.Collection;
+import tech.amikos.chromadb.embeddings.cohere.CohereEmbeddingFunction;
 
 import java.util.*;
 
 public class Main {
-    public static void main(String[] args) {
-        try {
-            Client client = new Client(System.getenv("CHROMA_URL"));
-            client.reset();
-            String apiKey = System.getenv("COHERE_API_KEY");
-            EmbeddingFunction ef = new CohereEmbeddingFunction(apiKey);
-            Collection collection = client.createCollection("test-collection", null, true, ef);
-            List<Map<String, String>> metadata = new ArrayList<>();
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "scientist");
-            }});
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "spy");
-            }});
-            collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
-            Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
-            System.out.println(qr);
-        } catch (Exception e) {
-            e.printStackTrace();
-            System.out.println(e);
-        }
+  public static void main(String[] args) {
+    try {
+      Client client = new Client(System.getenv("CHROMA_URL"));
+      client.reset();
+      String apiKey = System.getenv("COHERE_API_KEY");
+      EmbeddingFunction ef = new CohereEmbeddingFunction(apiKey);
+      Collection collection = client.createCollection("test-collection", null, true, ef);
+      List<Map<String, String>> metadata = new ArrayList<>();
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "scientist");
+      }});
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "spy");
+      }});
+      collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
+      Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
+      System.out.println(qr);
+    } catch (Exception e) {
+      e.printStackTrace();
+      System.out.println(e);
     }
+  }
 }
 ```
 
@@ -220,7 +221,7 @@ The above should output:
 
 ### Example Hugging Face Sentence Transformers Embedding Function
 
-In this example we rely on `tech.amikos.chromadb.HuggingFaceEmbeddingFunction` to generate embeddings for our documents.
+In this example we rely on `tech.amikos.chromadb.embeddings.hf.HuggingFaceEmbeddingFunction` to generate embeddings for our documents.
 
 | **Important**: Ensure you have `HF_API_KEY` environment variable set
 
@@ -229,31 +230,32 @@ package tech.amikos;
 
 import tech.amikos.chromadb.*;
 import tech.amikos.chromadb.Collection;
+import tech.amikos.chromadb.embeddings.hf.HuggingFaceEmbeddingFunction;
 
 import java.util.*;
 
 public class Main {
-    public static void main(String[] args) {
-        try {
-            Client client = new Client(System.getenv("CHROMA_URL"));
-            client.reset();
-            String apiKey = System.getenv("HF_API_KEY");
-            EmbeddingFunction ef = new HuggingFaceEmbeddingFunction(apiKey);
-            Collection collection = client.createCollection("test-collection", null, true, ef);
-            List<Map<String, String>> metadata = new ArrayList<>();
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "scientist");
-            }});
-            metadata.add(new HashMap<String, String>() {{
-                put("type", "spy");
-            }});
-            collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
-            Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
-            System.out.println(qr);
-        } catch (Exception e) {
-            System.out.println(e);
-        }
+  public static void main(String[] args) {
+    try {
+      Client client = new Client(System.getenv("CHROMA_URL"));
+      client.reset();
+      String apiKey = System.getenv("HF_API_KEY");
+      EmbeddingFunction ef = new HuggingFaceEmbeddingFunction(apiKey);
+      Collection collection = client.createCollection("test-collection", null, true, ef);
+      List<Map<String, String>> metadata = new ArrayList<>();
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "scientist");
+      }});
+      metadata.add(new HashMap<String, String>() {{
+        put("type", "spy");
+      }});
+      collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
+      Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
+      System.out.println(qr);
+    } catch (Exception e) {
+      System.out.println(e);
     }
+  }
 }
 ```
 

pom.xml

@@ -228,7 +228,7 @@
                             <!--                                            <additionalProperty>ignoreImportMappings=false</additionalProperty>-->
                             <!--                                        </additionalProperties>-->
                             <importMappings>
-                                <importMapping>CreateEmbeddingRequest=tech.amikos.openai.CreateEmbeddingRequest
+                                <importMapping>CreateEmbeddingRequest=tech.amikos.chromadb.embeddings.openai.CreateEmbeddingRequest
                                 </importMapping>
                             </importMappings>
                             <generateApiTests>false</generateApiTests>

src/main/java/tech/amikos/chromadb/DefaultEmbeddingFunction.java → src/main/java/tech/amikos/chromadb/embeddings/DefaultEmbeddingFunction.java

@@ -1,4 +1,4 @@
-package tech.amikos.chromadb;
+package tech.amikos.chromadb.embeddings;
 
 import ai.djl.huggingface.tokenizers.Encoding;
 import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
@@ -11,6 +11,8 @@
 import org.nd4j.linalg.api.ops.impl.transforms.clip.ClipByValue;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.shade.guava.primitives.Floats;
+import tech.amikos.chromadb.EFException;
+import tech.amikos.chromadb.EmbeddingFunction;
 
 import java.io.*;
 import java.net.URL;

src/main/java/tech/amikos/cohere/CohereClient.java → src/main/java/tech/amikos/chromadb/embeddings/cohere/CohereClient.java

@@ -1,4 +1,4 @@
-package tech.amikos.cohere;
+package tech.amikos.chromadb.embeddings.cohere;
 
 import com.google.gson.Gson;
 import okhttp3.*;

src/main/java/tech/amikos/chromadb/CohereEmbeddingFunction.java → src/main/java/tech/amikos/chromadb/embeddings/cohere/CohereEmbeddingFunction.java

@@ -1,12 +1,8 @@
-package tech.amikos.chromadb;
+package tech.amikos.chromadb.embeddings.cohere;
 
-import tech.amikos.cohere.CohereClient;
-import tech.amikos.cohere.CreateEmbeddingRequest;
-import tech.amikos.cohere.CreateEmbeddingResponse;
+import tech.amikos.chromadb.EmbeddingFunction;
 
-import java.util.Arrays;
 import java.util.List;
-import java.util.stream.Collectors;
 
 public class CohereEmbeddingFunction implements EmbeddingFunction {
 

src/main/java/tech/amikos/cohere/CreateEmbeddingRequest.java → src/main/java/tech/amikos/chromadb/embeddings/cohere/CreateEmbeddingRequest.java

@@ -1,4 +1,4 @@
-package tech.amikos.cohere;
+package tech.amikos.chromadb.embeddings.cohere;
 
 import com.google.gson.*;
 import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/cohere/CreateEmbeddingResponse.java → src/main/java/tech/amikos/chromadb/embeddings/cohere/CreateEmbeddingResponse.java

@@ -1,4 +1,4 @@
-package tech.amikos.cohere;
+package tech.amikos.chromadb.embeddings.cohere;
 
 import com.google.gson.Gson;
 import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/hf/CreateEmbeddingRequest.java → src/main/java/tech/amikos/chromadb/embeddings/hf/CreateEmbeddingRequest.java

@@ -1,4 +1,4 @@
-package tech.amikos.hf;
+package tech.amikos.chromadb.embeddings.hf;
 
 import com.google.gson.Gson;
 import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/hf/CreateEmbeddingResponse.java → src/main/java/tech/amikos/chromadb/embeddings/hf/CreateEmbeddingResponse.java

@@ -1,4 +1,4 @@
-package tech.amikos.hf;
+package tech.amikos.chromadb.embeddings.hf;
 
 import com.google.gson.Gson;
 

src/main/java/tech/amikos/hf/HuggingFaceClient.java → src/main/java/tech/amikos/chromadb/embeddings/hf/HuggingFaceClient.java

@@ -1,4 +1,4 @@
-package tech.amikos.hf;
+package tech.amikos.chromadb.embeddings.hf;
 
 import com.google.gson.Gson;
 import okhttp3.*;

src/main/java/tech/amikos/chromadb/HuggingFaceEmbeddingFunction.java → src/main/java/tech/amikos/chromadb/embeddings/hf/HuggingFaceEmbeddingFunction.java

@@ -1,31 +1,27 @@
-package tech.amikos.chromadb;
+package tech.amikos.chromadb.embeddings.hf;
 
 
-import tech.amikos.hf.CreateEmbeddingRequest;
-import tech.amikos.hf.CreateEmbeddingResponse;
-import tech.amikos.hf.HuggingFaceClient;
+import tech.amikos.chromadb.EmbeddingFunction;
 
 import java.util.List;
 
 public class HuggingFaceEmbeddingFunction implements EmbeddingFunction {
-
     private final String hfAPIKey;
-
+    private final HuggingFaceClient client;
     public HuggingFaceEmbeddingFunction(String hfAPIKey) {
         this.hfAPIKey = hfAPIKey;
-
+         this.client = new HuggingFaceClient(this.hfAPIKey);
     }
 
     @Override
     public List<List<Float>> createEmbedding(List<String> documents) {
-        HuggingFaceClient client = new HuggingFaceClient(this.hfAPIKey);
-        CreateEmbeddingResponse response = client.createEmbedding(new CreateEmbeddingRequest().inputs(documents.toArray(new String[0])));
+
+        CreateEmbeddingResponse response = this.client.createEmbedding(new CreateEmbeddingRequest().inputs(documents.toArray(new String[0])));
         return response.getEmbeddings();
     }
 
     @Override
     public List<List<Float>> createEmbedding(List<String> documents, String model) {
-        HuggingFaceClient client = new HuggingFaceClient(this.hfAPIKey);
         client.modelId(model);
         CreateEmbeddingResponse response = client.createEmbedding(new CreateEmbeddingRequest().inputs(documents.toArray(new String[0])));
         return response.getEmbeddings();

src/main/java/tech/amikos/openai/CreateEmbeddingRequest.java → src/main/java/tech/amikos/chromadb/embeddings/openai/CreateEmbeddingRequest.java

@@ -1,4 +1,4 @@
-package tech.amikos.openai;
+package tech.amikos.chromadb.embeddings.openai;
 
 import com.google.gson.*;
 import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/openai/CreateEmbeddingResponse.java → src/main/java/tech/amikos/chromadb/embeddings/openai/CreateEmbeddingResponse.java

@@ -1,4 +1,4 @@
-package tech.amikos.openai;
+package tech.amikos.chromadb.embeddings.openai;
 
 import com.google.gson.Gson;
 import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/openai/OpenAIClient.java → src/main/java/tech/amikos/chromadb/embeddings/openai/OpenAIClient.java

@@ -1,4 +1,4 @@
-package tech.amikos.openai;
+package tech.amikos.chromadb.embeddings.openai;
 
 import com.google.gson.Gson;
 import okhttp3.*;

src/main/java/tech/amikos/chromadb/OpenAIEmbeddingFunction.java → src/main/java/tech/amikos/chromadb/embeddings/openai/OpenAIEmbeddingFunction.java

@@ -1,11 +1,8 @@
-package tech.amikos.chromadb;
+package tech.amikos.chromadb.embeddings.openai;
 
-import tech.amikos.openai.CreateEmbeddingRequest;
-import tech.amikos.openai.CreateEmbeddingResponse;
-import tech.amikos.openai.OpenAIClient;
+import tech.amikos.chromadb.EmbeddingFunction;
 
 import java.util.List;
-import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 public class OpenAIEmbeddingFunction implements EmbeddingFunction {