add support for chat

dzlab · dzlab · commit ebfa2b130b33 · 2023-08-13T16:43:53.000+01:00
diff --git a/llm4j-api/src/main/java/org/llm4j/api/ChatHistory.java b/llm4j-api/src/main/java/org/llm4j/api/ChatHistory.java
@@ -0,0 +1,82 @@
+package org.llm4j.api;
+
+import org.apache.commons.lang3.tuple.ImmutablePair;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class ChatHistory {
+
+    String context;
+    List<Map.Entry<Message, Message>> exampleList = new ArrayList<>();
+    List<Message> messageList = new ArrayList<>();
+
+    public ChatHistory setContext(String context) {
+        this.context = context;
+        return this;
+    }
+
+    public ChatHistory addExample(String text1, String text2) {
+        Map.Entry<Message, Message> pair = new ImmutablePair<>(new Message(text1), new Message(text1));
+        exampleList.add(pair);
+        return this;
+    }
+
+    public ChatHistory addMessage(String content) {
+        messageList.add(new Message(content));
+        return this;
+    }
+
+    public ChatHistory addMessage(String author, String content) {
+        messageList.add(new Message(author, content));
+        return this;
+    }
+
+    public String getContext() {
+        return context;
+    }
+
+    public List<Map.Entry<Message, Message>> getExampleList() {
+        return exampleList;
+    }
+
+    public List<Message> getMessageList() {
+        return messageList;
+    }
+
+    public static class Message {
+
+        private final String author;
+        private final String content;
+
+        public Message(String author, String content) {
+            this.author = author;
+            this.content = content;
+        }
+
+        public Message(String content) {
+            this(null, content);
+        }
+
+        public String getAuthor() {
+            return author;
+        }
+
+        public String getContent() {
+            return content;
+        }
+
+        @Override
+        public String toString() {
+            String text = "";
+            if(author!=null && !author.isEmpty()) {
+                text += author;
+            }
+            if(content!=null && !content.isEmpty()) {
+                text += text.isEmpty() ? content: ":" + content;
+            }
+            return text;
+        }
+    }
+}
diff --git a/llm4j-api/src/main/java/org/llm4j/api/LanguageModel.java b/llm4j-api/src/main/java/org/llm4j/api/LanguageModel.java
@@ -3,5 +3,7 @@
 public interface LanguageModel {
 
     String process(String text);
+
+    String process(ChatHistory history);
 }
 
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/HFApi.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/HFApi.java
@@ -1,5 +1,8 @@
 package org.llm4j.huggingface;
 
+import org.llm4j.huggingface.request.TextEmbeddingRequest;
+import org.llm4j.huggingface.request.TextGenerationRequest;
+import org.llm4j.huggingface.request.TextGenerationResponse;
 import retrofit2.Call;
 import retrofit2.http.Body;
 import retrofit2.http.Headers;
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/HFApiClient.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/HFApiClient.java
@@ -4,7 +4,9 @@
 import com.google.gson.GsonBuilder;
 import okhttp3.OkHttpClient;
 import org.apache.commons.configuration2.Configuration;
-import retrofit2.Response;
+import org.llm4j.huggingface.request.TextEmbeddingRequest;
+import org.llm4j.huggingface.request.TextGenerationRequest;
+import org.llm4j.huggingface.request.TextGenerationResponse;
 import retrofit2.Retrofit;
 import retrofit2.converter.gson.GsonConverterFactory;
 
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/HFLanguageModel.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/HFLanguageModel.java
@@ -1,8 +1,11 @@
 package org.llm4j.huggingface;
 
 import org.apache.commons.configuration2.Configuration;
+import org.llm4j.api.ChatHistory;
 import org.llm4j.api.LanguageModel;
 import org.llm4j.api.LanguageModelFactory;
+import org.llm4j.huggingface.request.TextGenerationRequest;
+import org.llm4j.huggingface.request.TextGenerationResponse;
 
 public class HFLanguageModel implements LanguageModel {
 
@@ -24,6 +27,11 @@ public String process(String text) {
         return response.getGeneratedText();
     }
 
+    @Override
+    public String process(ChatHistory history) {
+        return null;
+    }
+
     public static final class Builder implements LanguageModelFactory {
         private Configuration config;
         private HFApiClient client;
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/HFRequestOptions.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/HFRequestOptions.java
@@ -1,4 +1,4 @@
-package org.llm4j.huggingface;
+package org.llm4j.huggingface.request;
 
 import org.apache.commons.configuration2.Configuration;
 
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/HFRequestParameters.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/HFRequestParameters.java
@@ -1,4 +1,4 @@
-package org.llm4j.huggingface;
+package org.llm4j.huggingface.request;
 
 import org.apache.commons.configuration2.Configuration;
 
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/TextEmbeddingRequest.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/TextEmbeddingRequest.java
@@ -1,4 +1,4 @@
-package org.llm4j.huggingface;
+package org.llm4j.huggingface.request;
 
 import org.apache.commons.configuration2.Configuration;
 
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/TextGenerationRequest.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/TextGenerationRequest.java
@@ -1,4 +1,4 @@
-package org.llm4j.huggingface;
+package org.llm4j.huggingface.request;
 
 import org.apache.commons.configuration2.Configuration;
 
diff --git a/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/TextGenerationResponse.java b/llm4j-huggingface/src/main/java/org/llm4j/huggingface/request/TextGenerationResponse.java
@@ -1,4 +1,4 @@
-package org.llm4j.huggingface;
+package org.llm4j.huggingface.request;
 
 public class TextGenerationResponse {
 
diff --git a/llm4j-palm/src/main/java/org/llm4j/palm/PaLMClient.java b/llm4j-palm/src/main/java/org/llm4j/palm/PaLMClient.java
@@ -13,8 +13,10 @@
 public class PaLMClient {
 
     private final TextServiceClient textClient;
+    private final DiscussServiceClient chatClient;
 
     PaLMClient(Builder builder) {
+        this.chatClient = builder.chatClient;
         this.textClient = builder.textClient;
     }
 
@@ -26,13 +28,23 @@ public TextCompletion generate(GenerateTextRequest request) {
         return returnedText;
     }
 
+    public Message chat(GenerateMessageRequest request) {
+        GenerateMessageResponse response = chatClient.generateMessage(request);
+
+        Message returnedText = response.getCandidatesList().get(0);
+        return returnedText;
+    }
+
+
     static class Builder {
         TextServiceClient textClient;
+        DiscussServiceClient chatClient;
 
         Builder withConfig(Configuration config) {
             HashMap<String, String> headers = new HashMap<>();
             headers.put("x-goog-api-key", config.getString("palm.apiKey"));
             this.textClient = buildTextServiceClient(headers);
+            this.chatClient = buildDiscussServiceClient(headers);
             return this;
         }
 
@@ -55,6 +67,25 @@ TextServiceClient buildTextServiceClient(HashMap<String, String> headers) {
             return client;
         }
 
+        DiscussServiceClient buildDiscussServiceClient(HashMap<String, String> headers) {
+            DiscussServiceClient client = null;
+            try {
+                TransportChannelProvider provider = InstantiatingGrpcChannelProvider.newBuilder()
+                        .setHeaderProvider(FixedHeaderProvider.create(headers))
+                        .build();
+
+                DiscussServiceSettings settings = DiscussServiceSettings.newBuilder()
+                        .setTransportChannelProvider(provider)
+                        .setCredentialsProvider(FixedCredentialsProvider.create(null))
+                        .build();
+
+                client = DiscussServiceClient.create(settings);
+            } catch (IOException e) {
+                System.err.println("Failed to create a PaLM Text client");
+            }
+            return client;
+        }
+
         public PaLMClient build() {
             return new PaLMClient(this);
         }
diff --git a/llm4j-palm/src/main/java/org/llm4j/palm/PaLMLanguageModel.java b/llm4j-palm/src/main/java/org/llm4j/palm/PaLMLanguageModel.java
@@ -1,13 +1,14 @@
 package org.llm4j.palm;
 
-import com.google.ai.generativelanguage.v1beta2.GenerateTextRequest;
-import com.google.ai.generativelanguage.v1beta2.GenerateTextResponse;
-import com.google.ai.generativelanguage.v1beta2.TextCompletion;
+import com.google.ai.generativelanguage.v1beta2.*;
 import org.apache.commons.configuration2.Configuration;
+import org.llm4j.api.ChatHistory;
 import org.llm4j.api.LanguageModel;
 import org.llm4j.api.LanguageModelFactory;
+import org.llm4j.palm.request.MessageRequestFactory;
+import org.llm4j.palm.request.TextRequestFactory;
 
-import com.google.ai.generativelanguage.v1beta2.TextPrompt;
+import java.util.Map;
 
 public class PaLMLanguageModel implements LanguageModel {
 
@@ -21,7 +22,7 @@ public PaLMLanguageModel(Builder builder) {
 
     @Override
     public String process(String text) {
-        GenerateTextRequest request = new GenerateTextRequestFactory()
+        GenerateTextRequest request = new TextRequestFactory()
                 .withInputs(text)
                 .withConfig(config)
                 .build();
@@ -31,6 +32,25 @@ public String process(String text) {
         return response.getOutput();
     }
 
+    @Override
+    public String process(ChatHistory history) {
+        MessageRequestFactory factory = new MessageRequestFactory();
+        factory.withContext(history.getContext());
+        for(Map.Entry<ChatHistory.Message, ChatHistory.Message> pair: history.getExampleList()) {
+            ChatHistory.Message left = pair.getKey();
+            ChatHistory.Message right = pair.getValue();
+            factory.withExample(left.toString(), right.toString());
+        }
+        for(ChatHistory.Message message: history.getMessageList()) {
+            factory.withMessage(message.getAuthor(), message.getContent());
+        }
+        GenerateMessageRequest request = factory.withConfig(config).build();
+
+        Message response = client.chat(request);
+
+        return response.toString();
+    }
+
     public static final class Builder implements LanguageModelFactory {
         private Configuration config;
         private PaLMClient client;
diff --git a/llm4j-palm/src/main/java/org/llm4j/palm/request/MessageRequestFactory.java b/llm4j-palm/src/main/java/org/llm4j/palm/request/MessageRequestFactory.java
@@ -0,0 +1,82 @@
+package org.llm4j.palm.request;
+
+import com.google.ai.generativelanguage.v1beta2.*;
+import org.apache.commons.configuration2.Configuration;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class MessageRequestFactory {
+
+    static String CHAT_MODEL_ID = "models/chat-bison-001";
+
+    /**
+     * which model to use to generate the result
+     */
+    private String modelId;
+    private PaLMRequestParameters parameters;
+    /**
+     * E.g. "Respond to all questions with a rhyming poem."
+     */
+    private String context;
+    private List<Message> messageList = new ArrayList<>();
+    private List<Example> exampleList = new ArrayList<>();
+
+    public MessageRequestFactory withMessage(String content) {
+        messageList.add(Message.newBuilder()
+                .setContent(content)
+                .build());
+        return this;
+    }
+
+    public MessageRequestFactory withMessage(String author, String content) {
+        Message.Builder builder = Message.newBuilder();
+        if(author!=null) builder.setAuthor(author);
+        builder.setContent(content);
+        messageList.add(builder.build());
+        return this;
+    }
+
+    public MessageRequestFactory withExample(String text1, String text2) {
+        exampleList.add(Example.newBuilder()
+                .setInput(Message.newBuilder()
+                        .setContent(text1)
+                        .build())
+                .setOutput(Message.newBuilder()
+                        .setContent(text2)
+                        .build())
+                .build());
+
+        return this;
+    }
+
+    public MessageRequestFactory withContext(String context) {
+        this.context = context;
+        return this;
+    }
+
+    public MessageRequestFactory withConfig(Configuration configs) {
+        this.modelId = configs.getString("modelId", CHAT_MODEL_ID);
+        this.parameters = PaLMRequestParameters.builder()
+                        .withConfig(configs)
+                        .build();
+        return this;
+    }
+
+    public GenerateMessageRequest build() {
+        MessagePrompt prompt = MessagePrompt.newBuilder()
+                .setContext(context)
+                .addAllMessages(messageList)
+                .addAllExamples(exampleList)
+                .build();
+        GenerateMessageRequest request = GenerateMessageRequest.newBuilder()
+                .setModel(modelId)
+                .setPrompt(prompt)
+                .setTemperature(parameters.temperature.floatValue())
+                .setCandidateCount(parameters.candidateCount)
+                .setTopK(parameters.topK)
+                .setTopP(parameters.topP.floatValue())
+                .build();
+        return request;
+    }
+}
diff --git a/llm4j-palm/src/main/java/org/llm4j/palm/request/PaLMRequestParameters.java b/llm4j-palm/src/main/java/org/llm4j/palm/request/PaLMRequestParameters.java
@@ -1,4 +1,4 @@
-package org.llm4j.palm;
+package org.llm4j.palm.request;
 
 import org.apache.commons.configuration2.Configuration;
 
diff --git a/llm4j-palm/src/main/java/org/llm4j/palm/request/TextRequestFactory.java b/llm4j-palm/src/main/java/org/llm4j/palm/request/TextRequestFactory.java
diff --git a/llm4j-palm/src/test/java/org/llm4j/palm/PaLMLanguageModelTest.java b/llm4j-palm/src/test/java/org/llm4j/palm/PaLMLanguageModelTest.java
diff --git a/llm4j-palm/src/test/resources/llm4j.properties b/llm4j-palm/src/test/resources/llm4j.properties

Original file line number	Diff line number	Diff line change
`@@ -3,5 +3,7 @@`
`3`	`3`	`public interface LanguageModel {`
`4`	`4`
`5`	`5`	`String process(String text);`
	`6`	`+`
	`7`	`+ String process(ChatHistory history);`
`6`	`8`	`}`
`7`	`9`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package org.llm4j.huggingface;`
	`1`	`+package org.llm4j.huggingface.request;`
`2`	`2`
`3`	`3`	`import org.apache.commons.configuration2.Configuration;`
`4`	`4`