feat(functions): don't use yaml.MapSlice

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
mudler · May 19, 2024 · 48bc0c8 · 48bc0c8
1 parent 73566a2
commit 48bc0c8
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/pkg/functions/parse.go b/pkg/functions/parse.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/go-skynet/LocalAI/pkg/utils"
 	"github.com/rs/zerolog/log"
-	"gopkg.in/yaml.v2"
 )
 
 // FunctionsConfig is the configuration for the tool/function call.
@@ -44,14 +43,19 @@ type FunctionsConfig struct {
 	GrammarPrefix string `yaml:"grammar_prefix"`
 
 	// ReplaceResults allow to replace strings in the results before parsing them
-	ReplaceResults yaml.MapSlice `yaml:"replace_results"`
+	ReplaceResults []ReplaceResult `yaml:"replace_results"`
 
 	// FunctionName enable the LLM to return { "name": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }
 	// instead of { "function": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }.
 	// This might be useful for certain models trained with the function name as the first token.
 	FunctionName bool `yaml:"return_name_in_function_response"`
 }
 
+type ReplaceResult struct {
+	Key   string `yaml:"key"`
+	Value string `yaml:"value"`
+}
+
 type FuncCallResults struct {
 	Name      string
 	Arguments string
@@ -61,7 +65,7 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 	log.Debug().Msgf("LLM result: %s", llmresult)
 
 	for _, item := range functionConfig.ReplaceResults {
-		k, v := item.Key.(string), item.Value.(string)
+		k, v := item.Key, item.Value
 		log.Debug().Msgf("Replacing %s with %s", k, v)
 		re := regexp.MustCompile(k)
 		llmresult = re.ReplaceAllString(llmresult, v)

diff --git a/pkg/functions/parse_test.go b/pkg/functions/parse_test.go
@@ -4,7 +4,6 @@ import (
 	. "github.com/go-skynet/LocalAI/pkg/functions"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
-	"gopkg.in/yaml.v2"
 )
 
 var _ = Describe("LocalAI function parse tests", func() {
@@ -121,7 +120,7 @@ Some text before the JSON
 Some text after the JSON
 `
 
-			functionConfig.ReplaceResults = yaml.MapSlice{
+			functionConfig.ReplaceResults = []ReplaceResult{
 				{Key: `(?s)^[^{\[]*`, Value: ""},
 				{Key: `(?s)[^}\]]*$`, Value: ""},
 			}
@@ -138,7 +137,7 @@ Some text before the JSON
 [{"function": "add", "arguments": {"x": 5, "y": 3}}, {"function": "subtract", "arguments": {"x": 10, "y": 7}}]
 Some text after the JSON
 `
-			functionConfig.ReplaceResults = yaml.MapSlice{
+			functionConfig.ReplaceResults = []ReplaceResult{
 				{Key: `(?s)^[^{\[]*`, Value: ""},
 				{Key: `(?s)[^}\]]*$`, Value: ""},
 			}
@@ -164,7 +163,7 @@ Some text after the JSON
 			// Regex to match non-JSON characters after the JSON structure
 			//reAfter := regexp.MustCompile(`(?s)(?<=\}|\]).*$`)
 
-			functionConfig.ReplaceResults = yaml.MapSlice{
+			functionConfig.ReplaceResults = []ReplaceResult{
 				{Key: `(?s)^[^{\[]*`, Value: ""},
 				{Key: `(?s)[^}\]]*$`, Value: ""},
 				// Regex pattern to match single quotes around keys and values
@@ -197,7 +196,7 @@ Some text after the JSON
 			// Regex to match non-JSON characters after the JSON structure
 			//reAfter := regexp.MustCompile(`(?s)(?<=\}|\]).*$`)
 
-			functionConfig.ReplaceResults = yaml.MapSlice{
+			functionConfig.ReplaceResults = []ReplaceResult{
 				{Key: `(?s)^[^{\[]*`, Value: ""},
 				{Key: `(?s)[^}\]]*$`, Value: ""},
 				// Regex pattern to match single quotes around keys and values