From c4186f13c33836244a5991f31175ccd072a5fcaf Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Tue, 14 May 2024 00:32:32 +0200
Subject: [PATCH] feat(functions): support models with no grammar and no regex
 (#2315)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 pkg/functions/parse.go      | 41 ++++++++++++++++++++++---------------
 pkg/functions/parse_test.go | 28 +++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/pkg/functions/parse.go b/pkg/functions/parse.go
index 13ac11e5..304f336f 100644
--- a/pkg/functions/parse.go
+++ b/pkg/functions/parse.go
@@ -41,24 +41,33 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 	// if no grammar is used, we have to extract function and arguments from the result
 	if !useGrammars {
 		// the response is a string that we have to parse
-
-		// 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
-		// TODO: optimize this and pre-compile it
-		var respRegex = regexp.MustCompile(functionConfig.ResponseRegex)
-		match := respRegex.FindStringSubmatch(llmresult)
 		result := make(map[string]string)
-		for i, name := range respRegex.SubexpNames() {
-			if i != 0 && name != "" && len(match) > i {
-				result[name] = match[i]
-			}
-		}
 
-		// TODO: open point about multiple results and/or mixed with chat messages
-		// This is not handled as for now, we only expect one function call per response
-		functionName := result[functionNameKey]
-		if functionName == "" {
-			return results
+		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
+			// TODO: optimize this and pre-compile it
+			var respRegex = regexp.MustCompile(functionConfig.ResponseRegex)
+			match := respRegex.FindStringSubmatch(llmresult)
+			for i, name := range respRegex.SubexpNames() {
+				if i != 0 && name != "" && len(match) > i {
+					result[name] = match[i]
+				}
+			}
+
+			// TODO: open point about multiple results and/or mixed with chat messages
+			// This is not handled as for now, we only expect one function call per response
+			functionName := result[functionNameKey]
+			if functionName == "" {
+				return results
+			}
+		} else {
+			// We expect the result to be a JSON object with a function name and arguments
+			err := json.Unmarshal([]byte(llmresult), &result)
+			if err != nil {
+				log.Error().Err(err).Str("llmresult", llmresult).Msg("unable to unmarshal llm result")
+				return results
+			}
 		}
 
 		return append(results, FuncCallResults{Name: result[functionNameKey], Arguments: result["arguments"]})
diff --git a/pkg/functions/parse_test.go b/pkg/functions/parse_test.go
index 5168a7d1..4c2208ed 100644
--- a/pkg/functions/parse_test.go
+++ b/pkg/functions/parse_test.go
@@ -82,4 +82,32 @@ var _ = Describe("LocalAI function parse tests", func() {
 			Expect(results[1].Arguments).To(Equal(`{"x":10,"y":7}`))
 		})
 	})
+
+	Context("without grammars and without regex", func() {
+		It("should parse the function name and arguments correctly with the name key", func() {
+			input := `{"name": "add", "arguments": {"x": 5, "y": 3}}`
+			functionConfig.ParallelCalls = false
+			functionConfig.NoGrammar = false
+			functionConfig.ResponseRegex = ""
+			functionConfig.FunctionName = true
+
+			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":3}`))
+		})
+
+		It("should parse the function name and arguments correctly with the function key", func() {
+			input := `{"function": "add", "arguments": {"x": 5, "y": 3}}`
+			functionConfig.ParallelCalls = false
+			functionConfig.NoGrammar = false
+			functionConfig.ResponseRegex = ""
+			functionConfig.FunctionName = false
+
+			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":3}`))
+		})
+	})
 })