add StructuredOutputParser test cases

myteam23 · Jul 3, 2023 · 2376dd9 · 2376dd9
1 parent d81c3af
commit 2376dd9
Show file tree

Hide file tree

Showing 14 changed files with 353 additions and 38 deletions.
diff --git a/langchain-core/src/main/java/com/hw/langchain/agents/chat/base/ChatAgent.java b/langchain-core/src/main/java/com/hw/langchain/agents/chat/base/ChatAgent.java
@@ -21,6 +21,7 @@
 import com.hw.langchain.agents.agent.Agent;
 import com.hw.langchain.agents.agent.AgentOutputParser;
 import com.hw.langchain.agents.chat.output.parser.ChatOutputParser;
+import com.hw.langchain.agents.initialize.Initialize;
 import com.hw.langchain.base.language.BaseLanguageModel;
 import com.hw.langchain.chains.llm.LLMChain;
 import com.hw.langchain.prompts.base.BasePromptTemplate;
@@ -38,6 +39,7 @@
 
 import static com.hw.langchain.agents.chat.prompt.Prompt.*;
 import static com.hw.langchain.agents.utils.Utils.validateToolsSingleInput;
+import static com.hw.langchain.prompts.utils.FormatUtils.formatTemplate;
 
 /**
  * @author HamaWhite
@@ -92,10 +94,7 @@ public static BasePromptTemplate createPrompt(List<BaseTool> tools, String syste
         String toolStrings =
                 String.join("\n", tools.stream().map(tool -> tool.getName() + ": " + tool.getDescription()).toList());
 
-        formatInstructions = formatInstructions.replace("{tool_names}", toolNames);
-        // In Python format() method, the curly braces '{{}}' are used to represent the output '{}'.
-        formatInstructions = formatInstructions.replace("{{{{", "{{").replace("}}}}", "}}");
-
+        formatInstructions = formatTemplate(formatInstructions, Map.of("tool_names", toolNames));
         String template =
                 String.join("\n\n", systemMessagePrefix, toolStrings, formatInstructions, systemMessageSuffix);
 
@@ -110,6 +109,7 @@ public static BasePromptTemplate createPrompt(List<BaseTool> tools, String syste
 
     /**
      * Construct an agent from an LLM and tools.
+     * This method will be called by the {@link Initialize#initializeAgent} using MethodUtils.invokeStaticMethod.
      */
     public static Agent fromLLMAndTools(BaseLanguageModel llm, List<BaseTool> tools, Map<String, Object> kwargs) {
         return fromLLMAndTools(llm, tools, null, SYSTEM_MESSAGE_PREFIX, SYSTEM_MESSAGE_SUFFIX, HUMAN_MESSAGE,

diff --git a/langchain-core/src/main/java/com/hw/langchain/agents/initialize/Initialize.java b/langchain-core/src/main/java/com/hw/langchain/agents/initialize/Initialize.java
@@ -42,6 +42,9 @@
  */
 public class Initialize {
 
+    private Initialize() {
+    }
+
     public static AgentExecutor initializeAgent(List<BaseTool> tools, BaseLanguageModel llm, AgentType agent) {
         return initializeAgent(tools, llm, agent, null, Map.of(), Map.of());
     }

diff --git a/langchain-core/src/main/java/com/hw/langchain/agents/mrkl/base/ZeroShotAgent.java b/langchain-core/src/main/java/com/hw/langchain/agents/mrkl/base/ZeroShotAgent.java
@@ -31,6 +31,7 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.hw.langchain.agents.mrkl.prompt.Prompt.*;
+import static com.hw.langchain.prompts.utils.FormatUtils.formatTemplate;
 
 /**
  * Agent for the MRKL chain.
@@ -63,10 +64,7 @@ public static PromptTemplate createPrompt(List<BaseTool> tools, String prefix, S
                 String.join("\n", tools.stream().map(tool -> tool.getName() + ": " + tool.getDescription()).toList());
         String toolNames = String.join(", ", tools.stream().map(BaseTool::getName).toList());
 
-        formatInstructions = formatInstructions.replace("{tool_names}", toolNames);
-        // In Python format() method, the curly braces '{{}}' are used to represent the output '{}'.
-        formatInstructions = formatInstructions.replace("{{{{", "{{").replace("}}}}", "}}");
-
+        formatInstructions = formatTemplate(formatInstructions, Map.of("tool_names", toolNames));
         String template = String.join("\n\n", prefix, toolStrings, formatInstructions, suffix);
 
         if (inputVariables == null) {

diff --git a/langchain-core/src/main/java/com/hw/langchain/chains/sql/database/prompt/Prompt.java b/langchain-core/src/main/java/com/hw/langchain/chains/sql/database/prompt/Prompt.java
@@ -68,8 +68,8 @@ public class Prompt {
                     Relevant Table Names:
                     """;
 
-    public static PromptTemplate DECIDER_PROMPT = new PromptTemplate(List.of("query", "table_names"),
-            _DECIDER_TEMPLATE,
+    public static PromptTemplate DECIDER_PROMPT = new PromptTemplate(_DECIDER_TEMPLATE,
+            List.of("query", "table_names"),
             new CommaSeparatedListOutputParser());
 
     private static String _mysql_prompt =

diff --git a/...core/src/main/java/com/hw/langchain/output/parsers/structured/StructuredOutputParser.java b/...core/src/main/java/com/hw/langchain/output/parsers/structured/StructuredOutputParser.java
@@ -22,16 +22,18 @@
 import com.hw.langchain.schema.BaseOutputParser;
 
 import java.util.List;
+import java.util.Map;
 
 import static com.hw.langchain.output.parsers.FormatInstructions.STRUCTURED_FORMAT_INSTRUCTIONS;
 import static com.hw.langchain.output.parsers.json.Json.parseAndCheckJsonMarkdown;
+import static com.hw.langchain.prompts.utils.FormatUtils.formatTemplate;
 
 /**
  * @author HamaWhite
  */
 public class StructuredOutputParser extends BaseOutputParser<JsonNode> {
 
-    private static final String LINE_TEMPLATE = "\t\"%s\": %s  // %s";
+    private static final String LINE_TEMPLATE = "\t\"{name}\": {type}  // {description}";
 
     private final List<ResponseSchema> responseSchemas;
 
@@ -44,7 +46,10 @@ public static StructuredOutputParser fromResponseSchemas(List<ResponseSchema> re
     }
 
     private String getSubString(ResponseSchema schema) {
-        return String.format(LINE_TEMPLATE, schema.getName(), schema.getType(), schema.getDescription());
+        Map<String, Object> kwargs = Map.of("name", schema.getName(),
+                "description", schema.getDescription(),
+                "type", schema.getType());
+        return formatTemplate(LINE_TEMPLATE, kwargs);
     }
 
     @Override
@@ -58,6 +63,6 @@ public JsonNode parse(String text) {
     @Override
     public String getFormatInstructions() {
         var schemaStr = String.join("\n", responseSchemas.stream().map(this::getSubString).toList());
-        return String.format(STRUCTURED_FORMAT_INSTRUCTIONS, schemaStr);
+        return formatTemplate(STRUCTURED_FORMAT_INSTRUCTIONS, Map.of("format", schemaStr));
     }
 }
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
@@ -27,9 +27,11 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * Base class for all prompt templates, returning a prompt.
+ *
  * @author HamaWhite
  */
 @Data
@@ -48,11 +50,16 @@ public abstract class BasePromptTemplate {
 
     protected Map<String, Object> partialVariables = new HashMap<>();
 
-    public BasePromptTemplate(List<String> inputVariables) {
+    protected BasePromptTemplate(List<String> inputVariables) {
+        this.inputVariables = inputVariables;
+    }
+
+    protected BasePromptTemplate(List<String> inputVariables, Map<String, Object> partialVariables) {
         this.inputVariables = inputVariables;
+        this.partialVariables = partialVariables;
     }
 
-    public BasePromptTemplate(List<String> inputVariables, BaseOutputParser<?> outputParser) {
+    protected BasePromptTemplate(List<String> inputVariables, BaseOutputParser<?> outputParser) {
         this.inputVariables = inputVariables;
         this.outputParser = outputParser;
     }
@@ -62,10 +69,32 @@ public BasePromptTemplate(List<String> inputVariables, BaseOutputParser<?> outpu
      */
     public abstract PromptValue formatPrompt(Map<String, Object> kwargs);
 
+    /**
+     * Merge the partial variables and user variables into a single map.
+     *
+     * @param kwargs Additional user variables provided.
+     * @return Merged map containing partial variables and user variables.
+     */
+    public Map<String, Object> mergePartialAndUserVariables(Map<String, Object> kwargs) {
+        var mergedVariables = new HashMap<String, Object>(partialVariables.size() + kwargs.size());
+        // Add partial variables
+        partialVariables.forEach((key, value) -> {
+            if (value instanceof String stringValue) {
+                mergedVariables.put(key, stringValue);
+            } else if (value instanceof Supplier<?> supplier) {
+                mergedVariables.put(key, supplier.get());
+            }
+        });
+        // Add user variables
+        mergedVariables.putAll(kwargs);
+        return mergedVariables;
+    }
+
     /**
      * Format the prompt with the inputs.
+     *
      * @param kwargs Any arguments to be passed to the prompt template.
-     * @return  A formatted string.
+     * @return A formatted string.
      */
     public abstract String format(Map<String, Object> kwargs);
 

diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/base/StringPromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/base/StringPromptTemplate.java
@@ -29,17 +29,22 @@
 
 /**
  * String prompt should expose the format method, returning a prompt.
+ *
  * @author HamaWhite
  */
 @Data
 @NoArgsConstructor
 public abstract class StringPromptTemplate extends BasePromptTemplate {
 
-    public StringPromptTemplate(List<String> inputVariables) {
+    protected StringPromptTemplate(List<String> inputVariables) {
         super(inputVariables);
     }
 
-    public StringPromptTemplate(List<String> inputVariables, BaseOutputParser outputParser) {
+    protected StringPromptTemplate(List<String> inputVariables, Map<String, Object> partialVariables) {
+        super(inputVariables, partialVariables);
+    }
+
+    protected StringPromptTemplate(List<String> inputVariables, BaseOutputParser outputParser) {
         super(inputVariables, outputParser);
     }
 

diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/base/StringPromptValue.java b/langchain-core/src/main/java/com/hw/langchain/prompts/base/StringPromptValue.java
@@ -30,7 +30,7 @@
  */
 public class StringPromptValue implements PromptValue {
 
-    private String text;
+    private final String text;
 
     public StringPromptValue(String text) {
         this.text = 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
@@ -34,6 +34,10 @@ public BaseChatPromptTemplate(List<String> inputVariables) {
         super(inputVariables);
     }
 
+    public BaseChatPromptTemplate(List<String> inputVariables, Map<String, Object> partialVariables) {
+        super(inputVariables, partialVariables);
+    }
+
     @Override
     public String format(Map<String, Object> kwargs) {
         return formatPrompt(kwargs).toString();

diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/chat/ChatPromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/chat/ChatPromptTemplate.java
@@ -33,7 +33,11 @@ public class ChatPromptTemplate extends BaseChatPromptTemplate {
     private final List<?> messages;
 
     public ChatPromptTemplate(List<String> inputVariables, List<?> messages) {
-        super(inputVariables);
+        this(inputVariables, messages, new HashMap<>());
+    }
+
+    public ChatPromptTemplate(List<String> inputVariables, List<?> messages, Map<String, Object> partialVariables) {
+        super(inputVariables, partialVariables);
         this.messages = messages;
 
         validateInputVariables();
@@ -46,10 +50,13 @@ private void validateInputVariables() {
                 inputVars.addAll(promptTemplate.inputVariables());
             }
         }
+        if (partialVariables != null) {
+            inputVars.removeAll(partialVariables.keySet());
+        }
         if (inputVariables != null) {
             if (!inputVars.equals(new HashSet<>(inputVariables))) {
                 throw new IllegalArgumentException(String
-                        .format("Got mismatched input_variables. Expected: %s. Got: %s", inputVars, inputVariables));
+                        .format("Got mismatched inputVariables. Expected: %s. Got: %s", inputVars, inputVariables));
             }
         } else {
             inputVariables = List.copyOf(inputVars);
@@ -58,6 +65,7 @@ private void validateInputVariables() {
 
     @Override
     public List<BaseMessage> formatMessages(Map<String, Object> kwargs) {
+        kwargs = mergePartialAndUserVariables(kwargs);
         List<BaseMessage> result = new ArrayList<>();
         for (var messageTemplate : messages) {
             if (messageTemplate instanceof BaseMessage baseMessage) {

diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/prompt/PromptTemplate.java b/langchain-core/src/main/java/com/hw/langchain/prompts/prompt/PromptTemplate.java
@@ -21,16 +21,17 @@
 import com.hw.langchain.prompts.base.StringPromptTemplate;
 import com.hw.langchain.schema.BaseOutputParser;
 
-import org.apache.commons.text.StringSubstitutor;
-
 import lombok.Data;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
+import static com.hw.langchain.prompts.utils.FormatUtils.findVariables;
+import static com.hw.langchain.prompts.utils.FormatUtils.formatTemplate;
+
 /**
  * Schema to represent a prompt for an LLM.
+ *
  * @author HamaWhite
  */
 @Data
@@ -51,29 +52,24 @@ public PromptTemplate(List<String> inputVariables, String template) {
         this.template = template;
     }
 
-    public PromptTemplate(List<String> inputVariables, String template, BaseOutputParser<?> outputParser) {
+    public PromptTemplate(String template, List<String> inputVariables, Map<String, Object> partialVariables) {
+        super(inputVariables, partialVariables);
+        this.template = template;
+    }
+
+    public PromptTemplate(String template, List<String> inputVariables, BaseOutputParser<?> outputParser) {
         super(inputVariables, outputParser);
         this.template = template;
     }
 
     @Override
     public String format(Map<String, Object> kwargs) {
-        String text = StringSubstitutor.replace(template, kwargs, "{", "}");
-        // In Python format() method, the curly braces '{{}}' are used to represent the output '{}'.
-        return text.replace("{{", "{").replace("}}", "}");
+        kwargs = mergePartialAndUserVariables(kwargs);
+        return formatTemplate(template, kwargs);
     }
 
     public static PromptTemplate fromTemplate(String template) {
-        List<String> variableNames = new ArrayList<>();
-        StringSubstitutor substitutor = new StringSubstitutor(variable -> {
-            if (!variable.startsWith("{") && !variable.endsWith("}")) {
-                variableNames.add(variable);
-            }
-            return null;
-        });
-        substitutor.setVariablePrefix("{");
-        substitutor.setVariableSuffix("}");
-        substitutor.replace(template);
+        List<String> variableNames = findVariables(template);
         return new PromptTemplate(variableNames, template);
     }
 }
diff --git a/langchain-core/src/main/java/com/hw/langchain/prompts/utils/FormatUtils.java b/langchain-core/src/main/java/com/hw/langchain/prompts/utils/FormatUtils.java
@@ -0,0 +1,71 @@
+/*
+ * 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.utils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * @author HamaWhite
+ */
+public class FormatUtils {
+
+    private static final Pattern PATTERN = Pattern.compile("\\{([^{}]+)}(?!})");
+
+    private FormatUtils() {
+    }
+
+    /**
+     * Formats the given template string by replacing variables with corresponding values.
+     *
+     * @param template the template string to format
+     * @param kwargs   a map of variables and their corresponding values
+     * @return the formatted string
+     */
+    public static String formatTemplate(String template, Map<String, Object> kwargs) {
+        String result = template;
+        for (Map.Entry<String, Object> entry : kwargs.entrySet()) {
+            String placeholder = "{" + entry.getKey() + "}";
+            String value = entry.getValue().toString();
+            result = result.replace(placeholder, value);
+        }
+
+        // In Python format() method, the curly braces '{{}}' are used to represent the output '{}'.
+        return result.replace("{{", "{").replace("}}", "}");
+    }
+
+    /**
+     * Finds all variables enclosed in curly braces '{' and '}' in the input string,
+     * excluding variables enclosed in double curly braces '{{' or more.
+     *
+     * @param input the input string to search for variables
+     * @return a list of variables found in the input string
+     */
+    public static List<String> findVariables(String input) {
+        List<String> variables = new ArrayList<>();
+        Matcher matcher = PATTERN.matcher(input);
+        while (matcher.find()) {
+            variables.add(matcher.group(1));
+        }
+        return variables;
+    }
+}