File size: 2,772 Bytes
7def60a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
package functions_test
import (
. "github.com/mudler/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.ToJSONStructure("function", "arguments")
Expect(len(js.OneOf)).To(Equal(2))
fnName := js.OneOf[0].Properties["function"].(FunctionName)
fnArgs := js.OneOf[0].Properties["arguments"].(Argument)
Expect(fnName.Const).To(Equal("create_event"))
Expect(fnArgs.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
Expect(fnArgs.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
fnName = js.OneOf[1].Properties["function"].(FunctionName)
fnArgs = js.OneOf[1].Properties["arguments"].(Argument)
Expect(fnName.Const).To(Equal("search"))
Expect(fnArgs.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
// Test with custom keys
jsN := functions.ToJSONStructure("name", "arguments")
Expect(len(jsN.OneOf)).To(Equal(2))
fnName = jsN.OneOf[0].Properties["name"].(FunctionName)
fnArgs = jsN.OneOf[0].Properties["arguments"].(Argument)
Expect(fnName.Const).To(Equal("create_event"))
Expect(fnArgs.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
Expect(fnArgs.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
fnName = jsN.OneOf[1].Properties["name"].(FunctionName)
fnArgs = jsN.OneOf[1].Properties["arguments"].(Argument)
Expect(fnName.Const).To(Equal("search"))
Expect(fnArgs.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"))
})
})
})
|