Merge pull request HamaWhiteGG#19 from HamaWhiteGG/dev

add embeddings for openai-client
tokuhirom · Jun 25, 2023 · 6daf274 · 6daf274
2 parents 8644700 + 30e833a
commit 6daf274
Show file tree

Hide file tree

Showing 6 changed files with 173 additions and 0 deletions.
diff --git a/openai-client/src/main/java/com/hw/openai/OpenAiClient.java b/openai-client/src/main/java/com/hw/openai/OpenAiClient.java
@@ -23,6 +23,8 @@
 import com.hw.openai.entity.chat.ChatCompletionResp;
 import com.hw.openai.entity.completions.Completion;
 import com.hw.openai.entity.completions.CompletionResp;
+import com.hw.openai.entity.embeddings.Embedding;
+import com.hw.openai.entity.embeddings.EmbeddingResp;
 import com.hw.openai.entity.models.Model;
 import com.hw.openai.entity.models.ModelResp;
 import com.hw.openai.service.OpenAiService;
@@ -217,4 +219,14 @@ public String chatCompletion(ChatCompletion chatCompletion) {
     public ChatCompletionResp create(ChatCompletion chatCompletion) {
         return service.chatCompletion(chatCompletion).blockingGet();
     }
+
+    /**
+     * Creates an embedding vector representing the input text.
+     *
+     * @param embedding The Embedding object containing the input text.
+     * @return The embedding vector response.
+     */
+    public EmbeddingResp embedding(Embedding embedding) {
+        return service.embedding(embedding).blockingGet();
+    }
 }
diff --git a/openai-client/src/main/java/com/hw/openai/entity/embeddings/Embedding.java b/openai-client/src/main/java/com/hw/openai/entity/embeddings/Embedding.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hw.openai.entity.embeddings;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+
+import jakarta.validation.constraints.NotBlank;
+import jakarta.validation.constraints.NotEmpty;
+import lombok.Builder;
+import lombok.Data;
+
+import java.io.Serializable;
+import java.util.List;
+
+/**
+ * @author HamaWhite
+ */
+@Data
+@Builder
+@JsonInclude(JsonInclude.Include.NON_NULL)
+public class Embedding implements Serializable {
+
+    /**
+     * ID of the model to use.
+     */
+    @NotBlank
+    private String model;
+
+    /**
+     * Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request,
+     * pass an array of strings or array of token arrays. Each input must not exceed the max input tokens for the
+     * model (8191 tokens for text-embedding-ada-002).
+     */
+    @NotEmpty
+    private List<String> input;
+
+    /**
+     * A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
+     */
+    private String user;
+}
diff --git a/openai-client/src/main/java/com/hw/openai/entity/embeddings/EmbeddingData.java b/openai-client/src/main/java/com/hw/openai/entity/embeddings/EmbeddingData.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hw.openai.entity.embeddings;
+
+import lombok.Data;
+
+import java.util.List;
+
+/**
+ * @author HamaWhite
+ */
+@Data
+public class EmbeddingData {
+
+    private String object;
+
+    private Integer index;
+
+    private List<Float> embedding;
+}
diff --git a/openai-client/src/main/java/com/hw/openai/entity/embeddings/EmbeddingResp.java b/openai-client/src/main/java/com/hw/openai/entity/embeddings/EmbeddingResp.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hw.openai.entity.embeddings;
+
+import com.hw.openai.entity.completions.Usage;
+
+import lombok.Data;
+
+import java.util.List;
+
+/**
+ * EmbeddingResp
+ * @author HamaWhite
+ */
+@Data
+public class EmbeddingResp {
+
+    private String object;
+
+    private List<EmbeddingData> data;
+
+    private String model;
+
+    private Usage usage;
+}
diff --git a/openai-client/src/main/java/com/hw/openai/service/OpenAiService.java b/openai-client/src/main/java/com/hw/openai/service/OpenAiService.java
@@ -22,6 +22,8 @@
 import com.hw.openai.entity.chat.ChatCompletionResp;
 import com.hw.openai.entity.completions.Completion;
 import com.hw.openai.entity.completions.CompletionResp;
+import com.hw.openai.entity.embeddings.Embedding;
+import com.hw.openai.entity.embeddings.EmbeddingResp;
 import com.hw.openai.entity.models.Model;
 import com.hw.openai.entity.models.ModelResp;
 
@@ -73,4 +75,14 @@ public interface OpenAiService {
      */
     @POST("chat/completions")
     Single<ChatCompletionResp> chatCompletion(@Body ChatCompletion chatCompletion);
+
+    /**
+     * Creates an embedding vector representing the input text.
+     *
+     * @param embedding The Embedding object containing the input text.
+     * @return A Single object that emits an EmbeddingResp, representing the response containing the embedding vector.
+     */
+    @POST("embeddings")
+    Single<EmbeddingResp> embedding(@Body Embedding embedding);
+
 }
diff --git a/openai-client/src/test/java/com/hw/openai/OpenAiClientTest.java b/openai-client/src/test/java/com/hw/openai/OpenAiClientTest.java
@@ -21,6 +21,7 @@
 import com.hw.openai.entity.chat.ChatCompletion;
 import com.hw.openai.entity.chat.Message;
 import com.hw.openai.entity.completions.Completion;
+import com.hw.openai.entity.embeddings.Embedding;
 import com.hw.openai.entity.models.Model;
 import com.hw.openai.entity.models.ModelResp;
 
@@ -103,4 +104,18 @@ void testChatCompletion() {
 
         assertThat(client.chatCompletion(chatCompletion)).isEqualTo("Hello there! How can I assist you today?");
     }
+
+    @Test
+    void testEmbeddings() {
+        var embedding = Embedding.builder()
+                .model("text-embedding-ada-002")
+                .input(List.of("The food was delicious and the waiter..."))
+                .build();
+
+        var response = client.embedding(embedding);
+
+        assertThat(response).as("Response should not be null").isNotNull();
+        assertThat(response.getData()).as("Data list should have size 1").hasSize(1);
+        assertThat(response.getData().get(0).getEmbedding()).as("Embedding should have size 1536").hasSize(1536);
+    }
 }