From 5a6d120a5665db8c70b9071d5f2d78ad2b285a4d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 20 May 2024 08:31:06 +0200 Subject: [PATCH] feat(functions): don't use yaml.MapSlice (#2354) Signed-off-by: Ettore Di Giacinto --- pkg/functions/parse.go | 10 +++++++--- pkg/functions/parse_test.go | 9 ++++----- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/pkg/functions/parse.go b/pkg/functions/parse.go index 3d08ef56..b551a40e 100644 --- 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,7 +43,7 @@ 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" } }. @@ -52,6 +51,11 @@ type FunctionsConfig struct { 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 index 58c2c2bb..2485d70b 100644 --- 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