refactor(template): isolate and add tests (#2069)

* refactor(template): isolate and add tests

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Signed-off-by: Dave <dave@gray101.com>
Co-authored-by: Dave <dave@gray101.com>

pkg/templates/cache.go

@@ -0,0 +1,103 @@
+package templates
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"path/filepath"
+	"sync"
+	"text/template"
+
+	"github.com/go-skynet/LocalAI/pkg/utils"
+
+	"github.com/Masterminds/sprig/v3"
+)
+
+// Keep this in sync with config.TemplateConfig. Is there a more idiomatic way to accomplish this in go?
+// Technically, order doesn't _really_ matter, but the count must stay in sync, see tests/integration/reflect_test.go
+type TemplateType int
+
+type TemplateCache struct {
+	mu            sync.Mutex
+	templatesPath string
+	templates     map[TemplateType]map[string]*template.Template
+}
+
+func NewTemplateCache(templatesPath string) *TemplateCache {
+	tc := &TemplateCache{
+		templatesPath: templatesPath,
+		templates:     make(map[TemplateType]map[string]*template.Template),
+	}
+	return tc
+}
+
+func (tc *TemplateCache) initializeTemplateMapKey(tt TemplateType) {
+	if _, ok := tc.templates[tt]; !ok {
+		tc.templates[tt] = make(map[string]*template.Template)
+	}
+}
+
+func (tc *TemplateCache) EvaluateTemplate(templateType TemplateType, templateName string, in interface{}) (string, error) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+
+	tc.initializeTemplateMapKey(templateType)
+	m, ok := tc.templates[templateType][templateName]
+	if !ok {
+		// return "", fmt.Errorf("template not loaded: %s", templateName)
+		loadErr := tc.loadTemplateIfExists(templateType, templateName)
+		if loadErr != nil {
+			return "", loadErr
+		}
+		m = tc.templates[templateType][templateName] // ok is not important since we check m on the next line, and wealready checked
+	}
+	if m == nil {
+		return "", fmt.Errorf("failed loading a template for %s", templateName)
+	}
+
+	var buf bytes.Buffer
+
+	if err := m.Execute(&buf, in); err != nil {
+		return "", err
+	}
+	return buf.String(), nil
+}
+
+func (tc *TemplateCache) loadTemplateIfExists(templateType TemplateType, templateName string) error {
+
+	// Check if the template was already loaded
+	if _, ok := tc.templates[templateType][templateName]; ok {
+		return nil
+	}
+
+	// Check if the model path exists
+	// skip any error here - we run anyway if a template does not exist
+	modelTemplateFile := fmt.Sprintf("%s.tmpl", templateName)
+
+	dat := ""
+	file := filepath.Join(tc.templatesPath, modelTemplateFile)
+
+	// Security check
+	if err := utils.VerifyPath(modelTemplateFile, tc.templatesPath); err != nil {
+		return fmt.Errorf("template file outside path: %s", file)
+	}
+
+	if utils.ExistsInPath(tc.templatesPath, modelTemplateFile) {
+		d, err := os.ReadFile(file)
+		if err != nil {
+			return err
+		}
+		dat = string(d)
+	} else {
+		dat = templateName
+	}
+
+	// Parse the template
+	tmpl, err := template.New("prompt").Funcs(sprig.FuncMap()).Parse(dat)
+	if err != nil {
+		return err
+	}
+	tc.templates[templateType][templateName] = tmpl
+
+	return nil
+}

pkg/templates/cache_test.go

@@ -0,0 +1,73 @@
+package templates_test
+
+import (
+	"os"
+	"path/filepath"
+
+	"github.com/go-skynet/LocalAI/pkg/templates" // Update with your module path
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("TemplateCache", func() {
+	var (
+		templateCache *templates.TemplateCache
+		tempDir       string
+	)
+
+	BeforeEach(func() {
+		var err error
+		tempDir, err = os.MkdirTemp("", "templates")
+		Expect(err).NotTo(HaveOccurred())
+
+		// Writing example template files
+		err = os.WriteFile(filepath.Join(tempDir, "example.tmpl"), []byte("Hello, {{.Name}}!"), 0644)
+		Expect(err).NotTo(HaveOccurred())
+		err = os.WriteFile(filepath.Join(tempDir, "empty.tmpl"), []byte(""), 0644)
+		Expect(err).NotTo(HaveOccurred())
+
+		templateCache = templates.NewTemplateCache(tempDir)
+	})
+
+	AfterEach(func() {
+		os.RemoveAll(tempDir) // Clean up
+	})
+
+	Describe("EvaluateTemplate", func() {
+		Context("when template is loaded successfully", func() {
+			It("should evaluate the template correctly", func() {
+				result, err := templateCache.EvaluateTemplate(1, "example", map[string]string{"Name": "Gopher"})
+				Expect(err).NotTo(HaveOccurred())
+				Expect(result).To(Equal("Hello, Gopher!"))
+			})
+		})
+
+		Context("when template isn't a file", func() {
+			It("should parse from string", func() {
+				result, err := templateCache.EvaluateTemplate(1, "{{.Name}}", map[string]string{"Name": "Gopher"})
+				Expect(err).ToNot(HaveOccurred())
+				Expect(result).To(Equal("Gopher"))
+			})
+		})
+
+		Context("when template is empty", func() {
+			It("should return an empty string", func() {
+				result, err := templateCache.EvaluateTemplate(1, "empty", nil)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(result).To(Equal(""))
+			})
+		})
+	})
+
+	Describe("concurrency", func() {
+		It("should handle multiple concurrent accesses", func(done Done) {
+			go func() {
+				_, _ = templateCache.EvaluateTemplate(1, "example", map[string]string{"Name": "Gopher"})
+			}()
+			go func() {
+				_, _ = templateCache.EvaluateTemplate(1, "example", map[string]string{"Name": "Gopher"})
+			}()
+			close(done)
+		}, 0.1) // timeout in seconds
+	})
+})

pkg/templates/utils_suite_test.go

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