LocalAI/pkg/functions/functions_test.go
Ettore Di Giacinto efa32a2677
feat(grammar): support models with specific construct (#2291)
When enabling grammar with functions, it might be useful to
allow more flexibility to support models that are fine-tuned against returning
function calls of the form of { "name": "function_name", "arguments" {...} }
rather then { "function": "function_name", "arguments": {..} }.

This might call out to a more generic approach later on, but for the moment being we can easily support both
as we have just to specific different types.

If needed we can expand on this later on

Signed-off-by: mudler <mudler@localai.io>
2024-05-12 01:13:22 +02:00

72 lines
2.4 KiB
Go

package functions_test
import (
. "github.com/go-skynet/LocalAI/pkg/functions"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var _ = Describe("LocalAI grammar functions", func() {
Describe("ToJSONStructure()", func() {
It("converts a list of functions to a JSON structure that can be parsed to a grammar", func() {
var functions Functions = []Function{
{
Name: "create_event",
Parameters: map[string]interface{}{
"properties": map[string]interface{}{
"event_name": map[string]interface{}{
"type": "string",
},
"event_date": map[string]interface{}{
"type": "string",
},
},
},
},
{
Name: "search",
Parameters: map[string]interface{}{
"properties": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
},
},
},
},
}
js := functions.ToJSONFunctionStructure()
Expect(len(js.OneOf)).To(Equal(2))
Expect(js.OneOf[0].Properties.Function.Const).To(Equal("create_event"))
Expect(js.OneOf[0].Properties.Arguments.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
Expect(js.OneOf[0].Properties.Arguments.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
Expect(js.OneOf[1].Properties.Function.Const).To(Equal("search"))
Expect(js.OneOf[1].Properties.Arguments.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
jsN := functions.ToJSONNameStructure()
Expect(len(jsN.OneOf)).To(Equal(2))
Expect(jsN.OneOf[0].Properties.Function.Const).To(Equal("create_event"))
Expect(jsN.OneOf[0].Properties.Arguments.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
Expect(jsN.OneOf[0].Properties.Arguments.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
Expect(jsN.OneOf[1].Properties.Function.Const).To(Equal("search"))
Expect(jsN.OneOf[1].Properties.Arguments.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
})
})
Context("Select()", func() {
It("selects one of the functions and returns a list containing only the selected one", func() {
var functions Functions = []Function{
{
Name: "create_event",
},
{
Name: "search",
},
}
functions = functions.Select("create_event")
Expect(len(functions)).To(Equal(1))
Expect(functions[0].Name).To(Equal("create_event"))
})
})
})