Add tests, enhance output in case of configuration error

Signed-off-by: Ettore Di Giacinto <[email protected]>

core/backend/backend_suite_test.go

@@ -0,0 +1,13 @@
+package backend_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestBackend(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Backend test suite")
+}

core/backend/llm.go

@@ -9,6 +9,8 @@ import (
 	"sync"
 	"unicode/utf8"
 
+	"github.com/rs/zerolog/log"
+
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/schema"
 
@@ -181,7 +183,11 @@ func Finetune(config config.BackendConfig, input, prediction string) string {
 		mu.Lock()
 		reg, ok := cutstrings[c]
 		if !ok {
-			cutstrings[c] = regexp.MustCompile(c)
+			r, err := regexp.Compile(c)
+			if err != nil {
+				log.Fatal().Err(err).Msg("failed to compile regex")
+			}
+			cutstrings[c] = r
 			reg = cutstrings[c]
 		}
 		mu.Unlock()
@@ -190,12 +196,16 @@ func Finetune(config config.BackendConfig, input, prediction string) string {
 
 	// extract results from the response which can be for instance inside XML tags
 	var predResult string
-	for _, r := range config.ExtactRegex {
+	for _, r := range config.ExtractRegex {
 		mu.Lock()
 		reg, ok := cutstrings[r]
 		if !ok {
-			cutstrings[r] = regexp.MustCompile(r)
-			reg = cutstrings[r]
+			regex, err := regexp.Compile(r)
+			if err != nil {
+				log.Fatal().Err(err).Msg("failed to compile regex")
+			}
+			cutstrings[r] = regex
+			reg = regex
 		}
 		mu.Unlock()
 		predResult += reg.FindString(prediction)

core/backend/llm_test.go

@@ -0,0 +1,107 @@
+package backend_test
+
+import (
+	. "github.com/mudler/LocalAI/core/backend"
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/schema"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("LLM tests", func() {
+	var (
+		testConfig config.BackendConfig
+		input      string
+		prediction string
+		result     string
+	)
+
+	BeforeEach(func() {
+		testConfig = config.BackendConfig{
+			PredictionOptions: schema.PredictionOptions{
+				Echo: false,
+			},
+			LLMConfig: config.LLMConfig{
+				Cutstrings:   []string{`<.*?>`},                  // Example regex for removing XML tags
+				ExtractRegex: []string{`<result>(.*?)</result>`}, // Example regex to extract from tags
+				TrimSpace:    []string{" ", "\n"},
+				TrimSuffix:   []string{".", "!"},
+			},
+		}
+	})
+
+	Context("when echo is enabled", func() {
+		BeforeEach(func() {
+			testConfig.Echo = true
+			input = "Hello"
+			prediction = "World"
+		})
+
+		It("should prepend input to prediction", func() {
+			result = Finetune(testConfig, input, prediction)
+			Expect(result).To(Equal("HelloWorld"))
+		})
+	})
+
+	Context("when echo is disabled", func() {
+		BeforeEach(func() {
+			testConfig.Echo = false
+			input = "Hello"
+			prediction = "World"
+		})
+
+		It("should not modify the prediction with input", func() {
+			result = Finetune(testConfig, input, prediction)
+			Expect(result).To(Equal("World"))
+		})
+	})
+
+	Context("when cutstrings regex is applied", func() {
+		BeforeEach(func() {
+			input = ""
+			prediction = "<div>Hello</div> World"
+		})
+
+		It("should remove substrings matching cutstrings regex", func() {
+			result = Finetune(testConfig, input, prediction)
+			Expect(result).To(Equal("Hello World"))
+		})
+	})
+
+	Context("when extract regex is applied", func() {
+		BeforeEach(func() {
+			input = ""
+			prediction = "<response><result>42</result></response>"
+		})
+
+		It("should extract substrings matching the extract regex", func() {
+			result = Finetune(testConfig, input, prediction)
+			Expect(result).To(Equal("42"))
+		})
+	})
+
+	Context("when trimming spaces", func() {
+		BeforeEach(func() {
+			input = ""
+			prediction = "   Hello World   "
+		})
+
+		It("should trim spaces from the prediction", func() {
+			result = Finetune(testConfig, input, prediction)
+			Expect(result).To(Equal("Hello World"))
+		})
+	})
+
+	Context("when trimming suffixes", func() {
+		BeforeEach(func() {
+			input = ""
+			prediction = "Hello World."
+		})
+
+		It("should trim suffixes from the prediction", func() {
+			result = Finetune(testConfig, input, prediction)
+			Expect(result).To(Equal("Hello World"))
+		})
+	})
+})

core/config/backend_config.go

@@ -126,7 +126,7 @@ type LLMConfig struct {
 	Grammar         string   `yaml:"grammar"`
 	StopWords       []string `yaml:"stopwords"`
 	Cutstrings      []string `yaml:"cutstrings"`
-	ExtactRegex     []string `yaml:"extract_regex"`
+	ExtractRegex    []string `yaml:"extract_regex"`
 	TrimSpace       []string `yaml:"trimspace"`
 	TrimSuffix      []string `yaml:"trimsuffix"`