Skip to content

Commit

Permalink
feat(functions): parse broken JSON when we parse the raw results, use…
Browse files Browse the repository at this point in the history 
… dynamic rules for grammar keys (mudler#2912)

* feat(functions): enhance parsing with broken JSON when we parse the raw results

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

* breaking: make function name by default

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

* feat(grammar): dynamically generate grammars with mutating keys

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

* refactor: simplify condition

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

* Update docs

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

---------

Signed-off-by: Ettore Di Giacinto <[email protected]>
  • Loading branch information
@mudler
mudler authored Jul 18, 2024
1 parent 35d5557 commit bf9dd1d
Show file tree
Hide file tree
Showing 9 changed files with 279 additions and 202 deletions.
12 changes: 2 additions & 10 deletions core/http/endpoints/openai/chat.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,18 +225,10 @@ func ChatEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, startup
}

// Update input grammar
// Handle if we should return "name" instead of "functions"
if config.FunctionsConfig.FunctionName {
jsStruct := funcs.ToJSONNameStructure()
config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
} else {
jsStruct := funcs.ToJSONFunctionStructure()
config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
}
jsStruct := funcs.ToJSONStructure(config.FunctionsConfig.FunctionNameKey, config.FunctionsConfig.FunctionNameKey)
config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
case input.JSONFunctionGrammarObject != nil:
config.Grammar = input.JSONFunctionGrammarObject.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
case input.JSONFunctionGrammarObjectName != nil:
config.Grammar = input.JSONFunctionGrammarObjectName.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
default:
// Force picking one of the functions by the request
if config.FunctionToCall() != "" {
Expand Down
3 changes: 1 addition & 2 deletions core/schema/openai.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,8 +179,7 @@ type OpenAIRequest struct {
// A grammar to constrain the LLM output
Grammar string `json:"grammar" yaml:"grammar"`

JSONFunctionGrammarObject *functions.JSONFunctionStructureFunction `json:"grammar_json_functions" yaml:"grammar_json_functions"`
JSONFunctionGrammarObjectName *functions.JSONFunctionStructureName `json:"grammar_json_name" yaml:"grammar_json_name"`
JSONFunctionGrammarObject *functions.JSONFunctionStructure `json:"grammar_json_functions" yaml:"grammar_json_functions"`

Backend string `json:"backend" yaml:"backend"`

Expand Down
3 changes: 2 additions & 1 deletion docs/content/docs/advanced/advanced-usage.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,8 @@ function:
replace_function_results: [] # Placeholder to replace function call results with arbitrary strings or patterns.
replace_llm_results: [] # Replace language model results with arbitrary strings or patterns.
capture_llm_results: [] # Capture language model results as text result, among JSON, in function calls. For instance, if a model returns a block for "thinking" and a block for "response", this will allow you to capture the thinking block.
return_name_in_function_response: false # Some models might prefer to use "name" rather then "function" when returning JSON data. This will allow to use "name" as a key in the JSON response.
function_name_key: "name"
function_arguments_key: "arguments"
# Feature gating flags to enable experimental or optional features.
feature_flags: {}
Expand Down
76 changes: 25 additions & 51 deletions pkg/functions/functions.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,11 @@ import (
"github.com/rs/zerolog/log"
)

const (
defaultFunctionNameKey = "name"
defaultFunctionArgumentsKey = "arguments"
)

type Function struct {
Name string `json:"name"`
Description string `json:"description"`
Expand All @@ -19,50 +24,18 @@ type Tool struct {
}
type Tools []Tool

// ToJSONFunctionStructure converts a list of functions to a JSON structure that can be parsed to a grammar
// This allows the LLM to return a response of the type: { "function": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }
func (f Functions) ToJSONFunctionStructure() JSONFunctionStructureFunction {
js := JSONFunctionStructureFunction{}
for _, function := range f {
// t := function.Parameters["type"]
//tt := t.(string)

properties := function.Parameters["properties"]
defs := function.Parameters["$defs"]
dat, _ := json.Marshal(properties)
dat2, _ := json.Marshal(defs)
prop := map[string]interface{}{}
defsD := map[string]interface{}{}

err := json.Unmarshal(dat, &prop)
if err != nil {
log.Error().Err(err).Msg("error unmarshalling dat")
}
err = json.Unmarshal(dat2, &defsD)
if err != nil {
log.Error().Err(err).Msg("error unmarshalling dat2")
}
if js.Defs == nil {
js.Defs = defsD
}
js.OneOf = append(js.OneOf, ItemFunction{
Type: "object",
Properties: FunctionProperties{
Function: FunctionName{Const: function.Name},
Arguments: Argument{
Type: "object",
Properties: prop,
},
},
})
}
return js
}

// ToJSONNameStructure converts a list of functions to a JSON structure that can be parsed to a grammar
// This allows the LLM to return a response of the type: { "name": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }
func (f Functions) ToJSONNameStructure() JSONFunctionStructureName {
js := JSONFunctionStructureName{}
func (f Functions) ToJSONStructure(name, args string) JSONFunctionStructure {
nameKey := defaultFunctionNameKey
argsKey := defaultFunctionArgumentsKey
if name != "" {
nameKey = name
}
if args != "" {
argsKey = args
}
js := JSONFunctionStructure{}
for _, function := range f {
// t := function.Parameters["type"]
//tt := t.(string)
Expand All @@ -85,15 +58,16 @@ func (f Functions) ToJSONNameStructure() JSONFunctionStructureName {
if js.Defs == nil {
js.Defs = defsD
}
js.OneOf = append(js.OneOf, ItemName{
Type: "object",
Properties: NameProperties{
Function: FunctionName{Const: function.Name},
Arguments: Argument{
Type: "object",
Properties: prop,
},
},

property := map[string]interface{}{}
property[nameKey] = FunctionName{Const: function.Name}
property[argsKey] = Argument{
Type: "object",
Properties: prop,
}
js.OneOf = append(js.OneOf, Item{
Type: "object",
Properties: property,
})
}
return js
Expand Down
38 changes: 26 additions & 12 deletions pkg/functions/functions_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,21 +35,35 @@ var _ = Describe("LocalAI grammar functions", func() {
},
}

js := functions.ToJSONFunctionStructure()
js := functions.ToJSONStructure("function", "arguments")
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"))
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"))

jsN := functions.ToJSONNameStructure()
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))
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"))

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() {
Expand Down
44 changes: 8 additions & 36 deletions pkg/functions/grammar_json_schema.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -331,6 +331,7 @@ func (sc *JSONSchemaConverter) resolveReference(ref string, rootSchema map[strin

return def
}

func (sc *JSONSchemaConverter) Grammar(schema map[string]interface{}, options ...func(*GrammarOption)) string {
sc.addRule("freestring", PRIMITIVE_RULES["freestring"])
sc.visit(schema, "", schema)
Expand All @@ -352,52 +353,23 @@ type FunctionName struct {
Const string `json:"const"`
}

type FunctionProperties struct {
Function FunctionName `json:"function"`
Arguments Argument `json:"arguments"`
}

type NameProperties struct {
Function FunctionName `json:"name"`
Arguments Argument `json:"arguments"`
}

type Argument struct {
Type string `json:"type"`
Properties map[string]interface{} `json:"properties"`
}

type ItemName struct {
Type string `json:"type"`
Properties NameProperties `json:"properties"`
}

type ItemFunction struct {
Type string `json:"type"`
Properties FunctionProperties `json:"properties"`
}

type JSONFunctionStructureName struct {
OneOf []ItemName `json:"oneOf,omitempty"`
AnyOf []ItemName `json:"anyOf,omitempty"`
Defs map[string]interface{} `json:"$defs,omitempty"`
}

func (j JSONFunctionStructureName) Grammar(options ...func(*GrammarOption)) string {
grammarOpts := &GrammarOption{}
grammarOpts.Apply(options...)

dat, _ := json.Marshal(j)
return NewJSONSchemaConverter(grammarOpts.PropOrder).GrammarFromBytes(dat, options...)
type Item struct {
Type string `json:"type"`
Properties map[string]interface{} `json:"properties"`
}

type JSONFunctionStructureFunction struct {
OneOf []ItemFunction `json:"oneOf,omitempty"`
AnyOf []ItemFunction `json:"anyOf,omitempty"`
type JSONFunctionStructure struct {
OneOf []Item `json:"oneOf,omitempty"`
AnyOf []Item `json:"anyOf,omitempty"`
Defs map[string]interface{} `json:"$defs,omitempty"`
}

func (j JSONFunctionStructureFunction) Grammar(options ...func(*GrammarOption)) string {
func (j JSONFunctionStructure) Grammar(options ...func(*GrammarOption)) string {
grammarOpts := &GrammarOption{}
grammarOpts.Apply(options...)

Expand Down
Loading

0 comments on commit bf9dd1d

Please sign in to comment.