feat(functions): allow to use JSONRegexMatch unconditionally (#2349)

* feat(functions): allow to use JSONRegexMatch unconditionally Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * feat(functions): make json_regex_match a list Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
mudler · May 19, 2024 · 73566a2 · 73566a2
1 parent 8ccd5ab
commit 73566a2
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 13 deletions.
diff --git a/pkg/functions/parse.go b/pkg/functions/parse.go
@@ -37,7 +37,7 @@ type FunctionsConfig struct {
 	ResponseRegex string `yaml:"response_regex"`
 
 	// JSONRegexMatch is a regex to extract the JSON object from the response
-	JSONRegexMatch string `yaml:"json_regex_match"`
+	JSONRegexMatch []string `yaml:"json_regex_match"`
 
 	// GrammarPrefix is the suffix to append to the grammar when being generated
 	// This is useful when models prepend a tag before returning JSON
@@ -124,6 +124,18 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 	// the response is a string that we have to parse
 	result := make(map[string]string)
 
+	if len(functionConfig.JSONRegexMatch) != 0 {
+		for _, r := range functionConfig.JSONRegexMatch {
+			// We use a regex to extract the JSON object from the response
+			var respRegex = regexp.MustCompile(r)
+			match := respRegex.FindStringSubmatch(llmresult)
+			if len(match) >= 2 {
+				llmresult = match[1]
+				break
+			}
+		}
+	}
+
 	if functionConfig.ResponseRegex != "" {
 		// We use named regexes here to extract the function name and arguments
 		// obviously, this expects the LLM to be stable and return correctly formatted JSON
@@ -143,16 +155,6 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 			return results
 		}
 		results = append(results, FuncCallResults{Name: result[functionNameKey], Arguments: result["arguments"]})
-	} else if functionConfig.JSONRegexMatch != "" {
-
-		// We use a regex to extract the JSON object from the response
-		var respRegex = regexp.MustCompile(functionConfig.JSONRegexMatch)
-		match := respRegex.FindStringSubmatch(llmresult)
-		if len(match) < 2 {
-			return results
-		}
-
-		results, _ = returnResult(match[1])
 	} else {
 		results, _ = returnResult(llmresult)
 	}

diff --git a/pkg/functions/parse_test.go b/pkg/functions/parse_test.go
@@ -91,7 +91,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 {"function": "add", "arguments": {"x": 5, "y": 3}}
 </tool_call>`
 
-			functionConfig.JSONRegexMatch = `(?s)<tool_call>(.*?)</tool_call>`
+			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(1))
@@ -104,7 +104,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 {"function": "add", "arguments": {"x": 5, "y": 3}}
 </tool_call>`
 
-			functionConfig.JSONRegexMatch = `(?s)(.*?)</tool_call>`
+			functionConfig.JSONRegexMatch = []string{`(?s)(.*?)</tool_call>`}
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(1))
@@ -157,6 +157,41 @@ Some text before the JSON
 {'function': '"add"', 'arguments': {'x': 5, 'z': '"v"', 'y': 'v"value"'}}
 Some text after the JSON
 `
+			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
+
+			// Regex to match non-JSON characters before the JSON structure
+			//reBefore := regexp.MustCompile(`(?s)^.*?(?=\{|\[)`)
+			// Regex to match non-JSON characters after the JSON structure
+			//reAfter := regexp.MustCompile(`(?s)(?<=\}|\]).*$`)
+
+			functionConfig.ReplaceResults = yaml.MapSlice{
+				{Key: `(?s)^[^{\[]*`, Value: ""},
+				{Key: `(?s)[^}\]]*$`, Value: ""},
+				// Regex pattern to match single quotes around keys and values
+				// Step 1: Replace single quotes around keys and values with double quotes
+				{Key: `'([^']*?)'`, Value: `_DQUOTE_${1}_DQUOTE_`},
+				// Step 2: Replace double quotes inside values with placeholders
+				{Key: `\\"`, Value: `__TEMP_QUOTE__`},
+				{Key: `"`, Value: `\"`},
+				{Key: `\'`, Value: `'`},
+				{Key: `_DQUOTE_`, Value: `"`},
+				{Key: `__TEMP_QUOTE__`, Value: `"`},
+			}
+
+			results := ParseFunctionCall(input, functionConfig)
+			Expect(results).To(HaveLen(1))
+			Expect(results[0].Name).To(Equal("\"add\""))
+			Expect(results[0].Arguments).To(Equal(`{"x":5,"y":"v\"value\"","z":"\"v\""}`))
+		})
+
+		It("should convert single-quoted key-value pairs to double-quoted and escape double quotes within values", func() {
+			input := `
+Some text before the JSON
+<tool_call>{'function': '"add"', 'arguments': {'x': 5, 'z': '"v"', 'y': 'v"value"'}}</tool_call>
+Some text after the JSON
+`
+			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
+
 			// Regex to match non-JSON characters before the JSON structure
 			//reBefore := regexp.MustCompile(`(?s)^.*?(?=\{|\[)`)
 			// Regex to match non-JSON characters after the JSON structure