Merge pull request HamaWhiteGG#14 from HamaWhiteGG/dev

Support Prompt templates of ChatModels
starksm64 · Jun 20, 2023 · a5af5fb · a5af5fb
2 parents cdcb379 + cf2dd71
commit a5af5fb
Show file tree

Hide file tree

Showing 19 changed files with 560 additions and 29 deletions.
diff --git a/langchain-core/src/main/java/com/hw/langchain/base/language/BaseLanguageModel.java b/langchain-core/src/main/java/com/hw/langchain/base/language/BaseLanguageModel.java
@@ -25,7 +25,8 @@
 import java.util.List;
 
 /**
- * BaseLanguageModel
+ * BaseLanguageModel is an interface for interacting with a language model.
+ *
  * @author HamaWhite
  */
 public interface BaseLanguageModel {
@@ -35,11 +36,21 @@ public interface BaseLanguageModel {
      */
     LLMResult generatePrompt(List<PromptValue> prompts, List<String> stop);
 
+    /**
+     * Predict text from text.
+     */
+    String predict(String text);
+
     /**
      * Predict text from text.
      */
     String predict(String text, List<String> stop);
 
+    /**
+     * Predict message from messages.
+     */
+    BaseMessage predictMessages(List<BaseMessage> messages);
+
     /**
      * Predict message from messages.
      */

diff --git a/langchain-core/src/main/java/com/hw/langchain/chat/models/base/BaseChatModel.java b/langchain-core/src/main/java/com/hw/langchain/chat/models/base/BaseChatModel.java
@@ -90,6 +90,11 @@ public BaseMessage call(List<BaseMessage> messages, List<String> stop) {
         }
     }
 
+    @Override
+    public String predict(String text) {
+        return predict(text, null);
+    }
+
     @Override
     public String predict(String text, List<String> stop) {
         List<String> copyStop = stop != null ? List.copyOf(stop) : null;
@@ -99,6 +104,11 @@ public String predict(String text, List<String> stop) {
         return result.getContent();
     }
 
+    @Override
+    public BaseMessage predictMessages(List<BaseMessage> messages) {
+        return predictMessages(messages, null);
+    }
+
     @Override
     public BaseMessage predictMessages(List<BaseMessage> messages, List<String> stop) {
         List<String> copyStop = stop != null ? List.copyOf(stop) : null;

diff --git a/langchain-core/src/main/java/com/hw/langchain/chat/models/openai/OpenAI.java b/langchain-core/src/main/java/com/hw/langchain/chat/models/openai/OpenAI.java
@@ -61,7 +61,7 @@ public static BaseMessage convertOpenAiToLangChain(Message message) {
                 return new SystemMessage(content);
             }
             default -> {
-                return new ChatMessage(role.getValue(), content);
+                return new ChatMessage(content, role.getValue());
             }
         }
     }

diff --git a/langchain-core/src/main/java/com/hw/langchain/llms/base/BaseLLM.java b/langchain-core/src/main/java/com/hw/langchain/llms/base/BaseLLM.java
@@ -19,6 +19,7 @@
 package com.hw.langchain.llms.base;
 
 import com.hw.langchain.base.language.BaseLanguageModel;
+import com.hw.langchain.schema.BaseMessage;
 import com.hw.langchain.schema.LLMResult;
 import com.hw.langchain.schema.PromptValue;
 
@@ -68,4 +69,24 @@ public LLMResult generatePrompt(List<PromptValue> prompts, List<String> stop) {
                 .toList();
         return generate(promptStrings, stop);
     }
+
+    @Override
+    public String predict(String text) {
+        return predict(text, null);
+    }
+
+    @Override
+    public String predict(String text, List<String> stop) {
+        return null;
+    }
+
+    @Override
+    public BaseMessage predictMessages(List<BaseMessage> messages) {
+        return predictMessages(messages, null);
+    }
+
+    @Override
+    public BaseMessage predictMessages(List<BaseMessage> messages, List<String> stop) {
+        return null;
+    }
 }
diff --git a/langchain-core/src/main/java/com/hw/langchain/llms/openai/BaseOpenAI.java b/langchain-core/src/main/java/com/hw/langchain/llms/openai/BaseOpenAI.java
@@ -19,7 +19,6 @@
 package com.hw.langchain.llms.openai;
 
 import com.hw.langchain.llms.base.BaseLLM;
-import com.hw.langchain.schema.BaseMessage;
 import com.hw.langchain.schema.Generation;
 import com.hw.langchain.schema.LLMResult;
 import com.hw.openai.OpenAiClient;
@@ -149,16 +148,6 @@ public class BaseOpenAI extends BaseLLM {
      */
     protected Set<String> disallowedSpecial;
 
-    @Override
-    public String predict(String text, List<String> stop) {
-        return null;
-    }
-
-    @Override
-    public BaseMessage predictMessages(List<BaseMessage> messages, List<String> stop) {
-        return null;
-    }
-
     @Override
     public String llmType() {
         return "openai";

diff --git a/langchain-core/src/main/java/com/hw/langchain/llms/openai/OpenAIChat.java b/langchain-core/src/main/java/com/hw/langchain/llms/openai/OpenAIChat.java
@@ -19,7 +19,6 @@
 package com.hw.langchain.llms.openai;
 
 import com.hw.langchain.llms.base.BaseLLM;
-import com.hw.langchain.schema.BaseMessage;
 import com.hw.langchain.schema.Generation;
 import com.hw.langchain.schema.LLMResult;
 import com.hw.langchain.utils.Utils;
@@ -150,16 +149,6 @@ public OpenAIChat init() {
         return this;
     }
 
-    @Override
-    public String predict(String text, List<String> stop) {
-        return null;
-    }
-
-    @Override
-    public BaseMessage predictMessages(List<BaseMessage> messages, List<String> stop) {
-        return null;
-    }
-
     @Override
     public String llmType() {
         return "openai-chat";

diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/base/BasePromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/base/BasePromptTemplate.java
@@ -22,6 +22,7 @@
 import com.hw.langchain.schema.PromptValue;
 
 import lombok.Data;
+import lombok.NoArgsConstructor;
 
 import java.util.HashMap;
 import java.util.List;
@@ -32,6 +33,7 @@
  * @author HamaWhite
  */
 @Data
+@NoArgsConstructor
 public abstract class BasePromptTemplate {
 
     /**

diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/chat/AIMessagePromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/chat/AIMessagePromptTemplate.java
@@ -0,0 +1,40 @@
+/*
+ * 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.langchain.prompts.chat;
+
+import com.hw.langchain.schema.AIMessage;
+import com.hw.langchain.schema.BaseMessage;
+
+import java.util.Map;
+
+/**
+ * @author HamaWhite
+ */
+public class AIMessagePromptTemplate extends BaseStringMessagePromptTemplate {
+
+    public static AIMessagePromptTemplate fromTemplate(String template) {
+        return BaseStringMessagePromptTemplate.fromTemplate(AIMessagePromptTemplate.class, template);
+    }
+
+    @Override
+    public BaseMessage format(Map<String, Object> kwargs) {
+        String text = prompt.format(kwargs);
+        return new AIMessage(text);
+    }
+}
diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/chat/BaseChatPromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/chat/BaseChatPromptTemplate.java
@@ -0,0 +1,52 @@
+/*
+ * 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.langchain.prompts.chat;
+
+import com.hw.langchain.prompts.base.BasePromptTemplate;
+import com.hw.langchain.schema.BaseMessage;
+import com.hw.langchain.schema.PromptValue;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * @author HamaWhite
+ */
+public abstract class BaseChatPromptTemplate extends BasePromptTemplate {
+
+    public BaseChatPromptTemplate(List<String> inputVariables) {
+        super(inputVariables);
+    }
+
+    @Override
+    public String format(Map<String, Object> kwargs) {
+        return formatPrompt(kwargs).toString();
+    }
+
+    public PromptValue formatPrompt(Map<String, Object> kwargs) {
+        List<BaseMessage> messages = formatMessages(kwargs);
+        return new ChatPromptValue(messages);
+    }
+
+    /**
+     * Format kwargs into a list of messages.
+     */
+    public abstract List<BaseMessage> formatMessages(Map<String, Object> kwargs);
+
+}
diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/chat/BaseMessagePromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/chat/BaseMessagePromptTemplate.java
@@ -0,0 +1,46 @@
+/*
+ * 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.langchain.prompts.chat;
+
+import com.hw.langchain.schema.BaseMessage;
+
+import java.io.Serializable;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * @author HamaWhite
+ */
+public abstract class BaseMessagePromptTemplate implements Serializable {
+
+    /**
+     * To messages.
+     *
+     * @param kwargs keyword arguments
+     * @return a list of BaseMessage
+     */
+    public abstract List<BaseMessage> formatMessages(Map<String, Object> kwargs);
+
+    /**
+     * Input variables for this prompt template.
+     *
+     * @return a list of input variables
+     */
+    public abstract List<String> inputVariables();
+}
diff --git a/...ain-core/src/main/java/com/hw/langchain/prompts/chat/BaseStringMessagePromptTemplate.java b/...ain-core/src/main/java/com/hw/langchain/prompts/chat/BaseStringMessagePromptTemplate.java
@@ -0,0 +1,65 @@
+/*
+ * 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.langchain.prompts.chat;
+
+import com.hw.langchain.exception.LangChainException;
+import com.hw.langchain.prompts.base.StringPromptTemplate;
+import com.hw.langchain.prompts.prompt.PromptTemplate;
+import com.hw.langchain.schema.BaseMessage;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * @author HamaWhite
+ */
+public abstract class BaseStringMessagePromptTemplate extends BaseMessagePromptTemplate {
+
+    protected StringPromptTemplate prompt;
+
+    public static <T extends BaseStringMessagePromptTemplate> T fromTemplate(Class<T> cls, String template) {
+        StringPromptTemplate prompt = PromptTemplate.fromTemplate(template);
+        try {
+            T instance = cls.getDeclaredConstructor().newInstance();
+            instance.setPrompt(prompt);
+            return instance;
+        } catch (Exception e) {
+            throw new LangChainException("Failed to create instance of BaseStringMessagePromptTemplate", e);
+        }
+    }
+
+    /**
+     * To a BaseMessage.
+     */
+    public abstract BaseMessage format(Map<String, Object> kwargs);
+
+    @Override
+    public List<BaseMessage> formatMessages(Map<String, Object> kwargs) {
+        return List.of(this.format(kwargs));
+    }
+
+    @Override
+    public List<String> inputVariables() {
+        return prompt.getInputVariables();
+    }
+
+    public void setPrompt(StringPromptTemplate prompt) {
+        this.prompt = prompt;
+    }
+}