forked from mudler/LocalAI
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor(template): isolate and add tests (mudler#2069)
* refactor(template): isolate and add tests Signed-off-by: Ettore Di Giacinto <[email protected]> --------- Signed-off-by: Ettore Di Giacinto <[email protected]> Signed-off-by: Dave <[email protected]> Co-authored-by: Dave <[email protected]>
- Loading branch information
1 parent
852316c
commit 27ec848
Showing
6 changed files
with
218 additions
and
95 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| }) | ||
| }) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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") | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters