2023-07-02 09:13:51 +00:00
|
|
|
package grammar
|
|
|
|
|
|
|
|
// a golang port of https://github.com/ggerganov/llama.cpp/pull/1887
|
|
|
|
|
|
|
|
import (
|
|
|
|
"encoding/json"
|
|
|
|
"fmt"
|
|
|
|
"regexp"
|
|
|
|
"sort"
|
|
|
|
"strings"
|
|
|
|
)
|
|
|
|
|
|
|
|
var (
|
|
|
|
SPACE_RULE = `" "?`
|
|
|
|
|
|
|
|
PRIMITIVE_RULES = map[string]string{
|
|
|
|
"boolean": `("true" | "false") space`,
|
|
|
|
"number": `[0-9]+ space`, // TODO complete
|
|
|
|
"string": `"\"" [ \t!#-\[\]-~]* "\"" space`, // TODO complete
|
|
|
|
"null": `"null" space`,
|
|
|
|
}
|
|
|
|
|
|
|
|
INVALID_RULE_CHARS_RE = regexp.MustCompile(`[^a-zA-Z0-9-]+`)
|
|
|
|
GRAMMAR_LITERAL_ESCAPE_RE = regexp.MustCompile(`[\r\n"]`)
|
|
|
|
GRAMMAR_LITERAL_ESCAPES = map[string]string{
|
|
|
|
"\r": `\r`,
|
|
|
|
"\n": `\n`,
|
|
|
|
`"`: `\"`,
|
|
|
|
}
|
|
|
|
)
|
|
|
|
|
|
|
|
type JSONSchemaConverter struct {
|
|
|
|
propOrder map[string]int
|
|
|
|
rules map[string]string
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewJSONSchemaConverter(propOrder string) *JSONSchemaConverter {
|
|
|
|
propOrderSlice := strings.Split(propOrder, ",")
|
|
|
|
propOrderMap := make(map[string]int)
|
|
|
|
for idx, name := range propOrderSlice {
|
|
|
|
propOrderMap[name] = idx
|
|
|
|
}
|
|
|
|
|
|
|
|
rules := make(map[string]string)
|
|
|
|
rules["space"] = SPACE_RULE
|
|
|
|
|
|
|
|
return &JSONSchemaConverter{
|
|
|
|
propOrder: propOrderMap,
|
|
|
|
rules: rules,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (sc *JSONSchemaConverter) formatLiteral(literal interface{}) string {
|
|
|
|
escaped := GRAMMAR_LITERAL_ESCAPE_RE.ReplaceAllStringFunc(jsonString(literal), func(match string) string {
|
|
|
|
return GRAMMAR_LITERAL_ESCAPES[match]
|
|
|
|
})
|
|
|
|
return fmt.Sprintf(`"%s"`, escaped)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (sc *JSONSchemaConverter) addRule(name, rule string) string {
|
|
|
|
escName := INVALID_RULE_CHARS_RE.ReplaceAllString(name, "-")
|
|
|
|
key := escName
|
|
|
|
if existingRule, ok := sc.rules[escName]; ok && existingRule != rule {
|
|
|
|
i := 0
|
|
|
|
for {
|
|
|
|
key = fmt.Sprintf("%s%d", escName, i)
|
|
|
|
if _, ok := sc.rules[key]; !ok {
|
|
|
|
break
|
|
|
|
}
|
|
|
|
i++
|
|
|
|
}
|
|
|
|
}
|
|
|
|
sc.rules[key] = rule
|
|
|
|
return key
|
|
|
|
}
|
|
|
|
|
|
|
|
func (sc *JSONSchemaConverter) formatGrammar() string {
|
|
|
|
var lines []string
|
|
|
|
for name, rule := range sc.rules {
|
|
|
|
lines = append(lines, fmt.Sprintf("%s ::= %s", name, rule))
|
|
|
|
}
|
|
|
|
return strings.Join(lines, "\n")
|
|
|
|
}
|
|
|
|
|
|
|
|
func (sc *JSONSchemaConverter) visit(schema map[string]interface{}, name string) string {
|
|
|
|
st, existType := schema["type"]
|
|
|
|
var schemaType string
|
|
|
|
if existType {
|
|
|
|
schemaType = st.(string)
|
|
|
|
}
|
|
|
|
ruleName := name
|
|
|
|
if name == "" {
|
|
|
|
ruleName = "root"
|
|
|
|
}
|
|
|
|
_, oneOfExists := schema["oneOf"]
|
|
|
|
_, anyOfExists := schema["anyOf"]
|
|
|
|
if oneOfExists || anyOfExists {
|
|
|
|
var alternatives []string
|
|
|
|
oneOfSchemas, oneOfExists := schema["oneOf"].([]interface{})
|
|
|
|
anyOfSchemas, anyOfExists := schema["anyOf"].([]interface{})
|
|
|
|
|
|
|
|
if oneOfExists {
|
|
|
|
for i, altSchema := range oneOfSchemas {
|
|
|
|
alternative := sc.visit(altSchema.(map[string]interface{}), fmt.Sprintf("%s-%d", ruleName, i))
|
|
|
|
alternatives = append(alternatives, alternative)
|
|
|
|
}
|
|
|
|
} else if anyOfExists {
|
|
|
|
for i, altSchema := range anyOfSchemas {
|
|
|
|
alternative := sc.visit(altSchema.(map[string]interface{}), fmt.Sprintf("%s-%d", ruleName, i))
|
|
|
|
alternatives = append(alternatives, alternative)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
rule := strings.Join(alternatives, " | ")
|
|
|
|
return sc.addRule(ruleName, rule)
|
|
|
|
} else if constVal, exists := schema["const"]; exists {
|
|
|
|
return sc.addRule(ruleName, sc.formatLiteral(constVal))
|
|
|
|
} else if enumVals, exists := schema["enum"].([]interface{}); exists {
|
|
|
|
var enumRules []string
|
|
|
|
for _, enumVal := range enumVals {
|
|
|
|
enumRule := sc.formatLiteral(enumVal)
|
|
|
|
enumRules = append(enumRules, enumRule)
|
|
|
|
}
|
|
|
|
rule := strings.Join(enumRules, " | ")
|
|
|
|
return sc.addRule(ruleName, rule)
|
|
|
|
} else if properties, exists := schema["properties"].(map[string]interface{}); schemaType == "object" && exists {
|
|
|
|
propOrder := sc.propOrder
|
|
|
|
var propPairs []struct {
|
|
|
|
propName string
|
|
|
|
propSchema map[string]interface{}
|
|
|
|
}
|
|
|
|
|
|
|
|
for propName, propSchema := range properties {
|
|
|
|
propPairs = append(propPairs, struct {
|
|
|
|
propName string
|
|
|
|
propSchema map[string]interface{}
|
|
|
|
}{propName: propName, propSchema: propSchema.(map[string]interface{})})
|
|
|
|
}
|
|
|
|
|
|
|
|
sort.Slice(propPairs, func(i, j int) bool {
|
|
|
|
iOrder := propOrder[propPairs[i].propName]
|
|
|
|
jOrder := propOrder[propPairs[j].propName]
|
|
|
|
if iOrder != 0 && jOrder != 0 {
|
|
|
|
return iOrder < jOrder
|
|
|
|
}
|
|
|
|
return propPairs[i].propName < propPairs[j].propName
|
|
|
|
})
|
|
|
|
|
|
|
|
var rule strings.Builder
|
|
|
|
rule.WriteString(`"{" space`)
|
|
|
|
|
|
|
|
for i, propPair := range propPairs {
|
|
|
|
propName := propPair.propName
|
|
|
|
propSchema := propPair.propSchema
|
|
|
|
propRuleName := sc.visit(propSchema, fmt.Sprintf("%s-%s", ruleName, propName))
|
|
|
|
|
|
|
|
if i > 0 {
|
|
|
|
rule.WriteString(` "," space`)
|
|
|
|
}
|
|
|
|
|
|
|
|
rule.WriteString(fmt.Sprintf(` %s space ":" space %s`, sc.formatLiteral(propName), propRuleName))
|
|
|
|
}
|
|
|
|
|
|
|
|
rule.WriteString(` "}" space`)
|
|
|
|
return sc.addRule(ruleName, rule.String())
|
|
|
|
} else if items, exists := schema["items"].(map[string]interface{}); schemaType == "array" && exists {
|
|
|
|
itemRuleName := sc.visit(items, fmt.Sprintf("%s-item", ruleName))
|
|
|
|
rule := fmt.Sprintf(`"[" space (%s ("," space %s)*)? "]" space`, itemRuleName, itemRuleName)
|
|
|
|
return sc.addRule(ruleName, rule)
|
|
|
|
} else {
|
|
|
|
primitiveRule, exists := PRIMITIVE_RULES[schemaType]
|
|
|
|
if !exists {
|
|
|
|
panic(fmt.Sprintf("Unrecognized schema: %v", schema))
|
|
|
|
}
|
|
|
|
return sc.addRule(schemaType, primitiveRule)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (sc *JSONSchemaConverter) Grammar(schema map[string]interface{}) string {
|
|
|
|
sc.visit(schema, "")
|
|
|
|
return sc.formatGrammar()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (sc *JSONSchemaConverter) GrammarFromBytes(b []byte) string {
|
|
|
|
var schema map[string]interface{}
|
|
|
|
_ = json.Unmarshal(b, &schema)
|
|
|
|
return sc.Grammar(schema)
|
|
|
|
}
|
|
|
|
|
|
|
|
func jsonString(v interface{}) string {
|
|
|
|
b, _ := json.Marshal(v)
|
|
|
|
return string(b)
|
|
|
|
}
|
|
|
|
|
|
|
|
type FunctionName struct {
|
|
|
|
Const string `json:"const"`
|
|
|
|
}
|
|
|
|
|
|
|
|
type Properties struct {
|
|
|
|
Function FunctionName `json:"function"`
|
|
|
|
Arguments Argument `json:"arguments"`
|
|
|
|
}
|
|
|
|
|
|
|
|
type Argument struct {
|
|
|
|
Type string `json:"type"`
|
|
|
|
Properties map[string]interface{} `json:"properties"`
|
|
|
|
}
|
|
|
|
|
|
|
|
type Item struct {
|
|
|
|
Type string `json:"type"`
|
|
|
|
Properties Properties `json:"properties"`
|
|
|
|
}
|
|
|
|
|
2023-07-06 16:05:52 +00:00
|
|
|
type JSONFunctionStructure struct {
|
2023-07-02 09:13:51 +00:00
|
|
|
OneOf []Item `json:"oneOf,omitempty"`
|
|
|
|
AnyOf []Item `json:"anyOf,omitempty"`
|
|
|
|
}
|
|
|
|
|
2023-07-06 16:05:52 +00:00
|
|
|
func (j JSONFunctionStructure) Grammar(propOrder string) string {
|
2023-07-02 09:13:51 +00:00
|
|
|
dat, _ := json.Marshal(j)
|
|
|
|
return NewJSONSchemaConverter(propOrder).GrammarFromBytes(dat)
|
|
|
|
}
|