added generation api (#1)

Author: dzlab

src/main/java/com/github/llmjava/cohere4j/AuthorizationInterceptor.java

@@ -0,0 +1,28 @@
+package com.github.llmjava.cohere4j;
+
+
+import okhttp3.Interceptor;
+import okhttp3.Request;
+import okhttp3.Response;
+
+import java.io.IOException;
+
+public class AuthorizationInterceptor implements Interceptor {
+
+    private final String apiKey;
+
+    AuthorizationInterceptor(String apiKey) {
+        this.apiKey = apiKey;
+    }
+
+    @Override
+    public Response intercept(Chain chain) throws IOException {
+
+        Request request = chain.request()
+                .newBuilder()
+                .addHeader("Authorization", "Bearer " + apiKey)
+                .build();
+
+        return chain.proceed(request);
+    }
+}

src/main/java/com/github/llmjava/cohere4j/CohereApi.java

@@ -0,0 +1,25 @@
+package com.github.llmjava.cohere4j;
+
+import com.github.llmjava.cohere4j.request.CompletionRequest;
+import com.github.llmjava.cohere4j.response.CompletionResponse;
+import com.github.llmjava.cohere4j.response.streaming.StreamingCompletionResponse;
+import retrofit2.Call;
+import retrofit2.http.Body;
+import retrofit2.http.Headers;
+import retrofit2.http.POST;
+import retrofit2.http.Streaming;
+
+public interface CohereApi {
+
+    @POST("/v1/generate")
+    @Headers({"accept: application/json", "content-type: application/json"})
+    Call<CompletionResponse>
+    generate(@Body CompletionRequest request);
+
+    @Streaming
+    @POST("/v1/generate")
+    @Headers({"accept: application/stream+json", "content-type: application/json"})
+    Call<CompletionResponse>
+    generateStream(@Body CompletionRequest request);
+
+}

src/main/java/com/github/llmjava/cohere4j/CohereApiFactory.java

@@ -0,0 +1,43 @@
+package com.github.llmjava.cohere4j;
+
+import com.google.gson.FieldNamingPolicy;
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import okhttp3.OkHttpClient;
+import retrofit2.Retrofit;
+import retrofit2.converter.gson.GsonConverterFactory;
+
+import java.time.Duration;
+
+public class CohereApiFactory {
+
+    public CohereApi build(CohereConfig config) {
+        String apiKey = config.getApiKey();
+        Duration timeout = config.getTimeout();
+        CohereApi api = buildApi(apiKey, timeout);
+        return api;
+    }
+
+    CohereApi buildApi(String apiKey, Duration timeout) {
+        OkHttpClient okHttpClient = new OkHttpClient.Builder()
+                .addInterceptor(new AuthorizationInterceptor(apiKey))
+                .callTimeout(timeout)
+                .connectTimeout(timeout)
+                .readTimeout(timeout)
+                .writeTimeout(timeout)
+                .build();
+
+        Gson gson = new GsonBuilder()
+                .setFieldNamingPolicy(FieldNamingPolicy.LOWER_CASE_WITH_UNDERSCORES)
+                .create();
+
+        Retrofit retrofit = new Retrofit.Builder()
+                .baseUrl(CohereConfig.BASE_URL)
+                .client(okHttpClient)
+                .addConverterFactory(GsonConverterFactory.create(gson))
+                .build();
+
+        return retrofit.create(CohereApi.class);
+    }
+
+}

src/main/java/com/github/llmjava/cohere4j/CohereClient.java

@@ -0,0 +1,104 @@
+package com.github.llmjava.cohere4j;
+
+import com.github.llmjava.cohere4j.callback.AsyncCallback;
+import com.github.llmjava.cohere4j.callback.StreamingCallback;
+import com.github.llmjava.cohere4j.request.CompletionRequest;
+import com.github.llmjava.cohere4j.response.CompletionResponse;
+import com.github.llmjava.cohere4j.response.streaming.StreamingCompletionResponse;
+import retrofit2.Call;
+import retrofit2.Response;
+
+import java.io.IOException;
+
+public class CohereClient {
+    private final CohereApi api;
+    private final CohereConfig config;
+
+    CohereClient(Builder builder) {
+        this.api = builder.api;
+        this.config = builder.config;
+    }
+
+    public String generate(CompletionRequest request) {
+        try {
+            Response<CompletionResponse> response = api.generate(request).execute();
+            if (response.isSuccessful()) {
+                return response.body().getTexts().get(0);
+            } else  {
+                throw newException(response);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public void generateAsync(CompletionRequest request, AsyncCallback<String> callback) {
+        api.generate(request).enqueue(new retrofit2.Callback<CompletionResponse>() {
+            @Override
+            public void onResponse(Call<CompletionResponse> call, Response<CompletionResponse> response) {
+                if (response.isSuccessful()) {
+                    callback.onSuccess(response.body().getTexts().get(0));
+                } else  {
+                    callback.onFailure(newException(response));
+                }
+            }
+
+            @Override
+            public void onFailure(Call<CompletionResponse> call, Throwable throwable) {
+                callback.onFailure(throwable);
+            }
+        });
+    }
+
+    public void generateStream(CompletionRequest request, StreamingCallback<String> callback) {
+        api.generateStream(request).enqueue(new retrofit2.Callback<CompletionResponse>() {
+            @Override
+            public void onResponse(Call<CompletionResponse> call, Response<CompletionResponse> response) {
+                if (response.isSuccessful()) {
+                    CompletionResponse resp = response.body();
+                    callback.onPart(resp.getTexts().get(0));
+                    callback.onComplete();
+
+                } else  {
+                    callback.onFailure(newException(response));
+                }
+            }
+
+            @Override
+            public void onFailure(Call<CompletionResponse> call, Throwable throwable) {
+                callback.onFailure(throwable);
+            }
+        });
+    }
+
+    /**
+     * Parse exceptions:
+     * status code: 429; body: {"message":"You are using a Trial key, which is limited to 5 API calls / minute. You can continue to use the Trial key for free or upgrade to a Production key with higher rate limits at 'https://dashboard.cohere.ai/api-keys'. Contact us on 'https://discord.gg/XW44jPfYJu' or email us at support@cohere.com with any questions"}
+     */
+    private static RuntimeException newException(retrofit2.Response<?> response) {
+        try {
+            int code = response.code();
+            String body = response.errorBody().string();
+            String errorMessage = String.format("status code: %s; body: %s", code, body);
+            return new RuntimeException(errorMessage);
+        } catch (IOException e) {
+            return new RuntimeException(e);
+        }
+    }
+
+
+    public static class Builder {
+        private CohereApi api;
+        private CohereConfig config;
+
+        public Builder withConfig(CohereConfig config) {
+            this.config = config;
+            this.api = new CohereApiFactory().build(config);
+            return this;
+        }
+
+        public CohereClient build() {
+            return new CohereClient(this);
+        }
+    }
+}

src/main/java/com/github/llmjava/cohere4j/CohereConfig.java

@@ -0,0 +1,38 @@
+package com.github.llmjava.cohere4j;
+
+import org.apache.commons.configuration2.Configuration;
+import org.apache.commons.configuration2.builder.fluent.Configurations;
+import org.apache.commons.configuration2.ex.ConfigurationException;
+
+import java.time.Duration;
+
+public class CohereConfig {
+    private final Configuration config;
+
+    public CohereConfig(Configuration config) {
+        this.config = config;
+    }
+
+    public static final String BASE_URL = "https://api.cohere.ai/";
+
+    public static final String API_KEY = "cohere.apiKey";
+
+    public static final String TIMEOUT = "timeout";
+
+    public static final Long DEFAULT_TIMEOUT_MILLIS = 10*1000l;
+
+    public String getApiKey() {
+        return config.getString(API_KEY);
+    }
+
+    public Duration getTimeout() {
+        Long timeout = config.getLong(CohereConfig.TIMEOUT, CohereConfig.DEFAULT_TIMEOUT_MILLIS);
+        return Duration.ofMillis(timeout);
+    }
+
+    public static CohereConfig fromProperties(String path) throws ConfigurationException {
+        Configuration baseConfig = new Configurations().properties(path);
+        return new CohereConfig(baseConfig);
+    }
+
+}

src/main/java/com/github/llmjava/cohere4j/callback/AsyncCallback.java

@@ -0,0 +1,7 @@
+package com.github.llmjava.cohere4j.callback;
+
+public interface AsyncCallback<T> {
+
+    void onSuccess(T response);
+    void onFailure(Throwable throwable);
+}

src/main/java/com/github/llmjava/cohere4j/callback/StreamingCallback.java

@@ -0,0 +1,8 @@
+package com.github.llmjava.cohere4j.callback;
+
+public interface StreamingCallback<S> {
+
+    void onPart(S response);
+    void onComplete();
+    void onFailure(Throwable throwable);
+}

src/main/java/com/github/llmjava/cohere4j/request/CompletionRequest.java

@@ -0,0 +1,80 @@
+package com.github.llmjava.cohere4j.request;
+
+public class CompletionRequest {
+
+    /**
+     * max_tokens
+     * integer
+     * The maximum number of tokens the model will generate as part of the response. Note: Setting a low value may result in incomplete generations.
+     * Defaults to 20. See BPE Tokens for more details.
+     *
+     * Can only be set to 0 if return_likelihoods is set to ALL to get the likelihood of the prompt.
+     */
+    private Integer max_tokens;
+
+    /**
+     * One of NONE|START|END to specify how the API will handle inputs longer than the maximum token length.
+     *
+     * Passing START will discard the start of the input. END will discard the end of the input. In both cases, input is discarded until the remaining input is exactly the maximum input token length for the model.
+     *
+     * If NONE is selected, when the input exceeds the maximum input token length an error will be returned.
+     *
+     * Default: END
+     */
+    private String truncate;
+
+    /**
+     * One of GENERATION|ALL|NONE to specify how and if the token likelihoods are returned with the response. Defaults to NONE.
+     *
+     * If GENERATION is selected, the token likelihoods will only be provided for generated text.
+     *
+     * If ALL is selected, the token likelihoods will be provided both for the prompt and the generated text.
+     *
+     * Default: NONE
+     */
+    private String return_likelihoods;
+
+    /**
+     * The input text that serves as the starting point for generating the response.
+     * Note: The prompt will be pre-processed and modified before reaching the model.
+     */
+    private String prompt;
+
+    CompletionRequest(Builder builder) {
+        this.max_tokens = builder.max_tokens;
+        this.truncate = builder.truncate;
+        this.return_likelihoods = builder.return_likelihoods;
+        this.prompt = builder.prompt;
+    }
+
+    public static class Builder {
+        private Integer max_tokens;
+        private String truncate;
+        private String return_likelihoods;
+        private String prompt;
+
+        public Builder withMaxTokens(Integer maxTokens) {
+            this.max_tokens = maxTokens;
+            return this;
+        }
+
+        public Builder withTruncate(String truncate) {
+            this.truncate = truncate;
+            return this;
+        }
+
+        public Builder withLikelihoods(String likelihoods) {
+            this.return_likelihoods = likelihoods;
+            return this;
+        }
+
+        public Builder withPrompt(String prompt) {
+            this.prompt = prompt;
+            return this;
+        }
+
+        public CompletionRequest build() {
+            return new CompletionRequest(this);
+        }
+    }
+}

src/main/java/com/github/llmjava/cohere4j/response/CompletionResponse.java

@@ -0,0 +1,22 @@
+package com.github.llmjava.cohere4j.response;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class CompletionResponse {
+    private String id;
+    private String prompt;
+    private List<Generation> generations;
+
+    public String getPrompt() {
+        return prompt;
+    }
+
+    public List<String> getTexts() {
+        List<String> texts = new ArrayList<>();
+        for(Generation generation: generations) {
+            texts.add(generation.getText());
+        }
+        return texts;
+    }
+}

src/main/java/com/github/llmjava/cohere4j/response/Generation.java

@@ -0,0 +1,15 @@
+package com.github.llmjava.cohere4j.response;
+
+import java.util.List;
+
+public class Generation {
+    private String id;
+    private String text;
+    private Integer index;
+    private Double likelihood;
+    private List<Likelihood> token_likelihoods;
+
+    public String getText() {
+        return text;
+    }
+}