feat: add grammar and functions call support

pkg/grammar/json_schema.go

@@ -0,0 +1,222 @@
+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"`
+}
+
+type JSONStructure struct {
+	OneOf []Item `json:"oneOf,omitempty"`
+	AnyOf []Item `json:"anyOf,omitempty"`
+}
+
+func (j JSONStructure) Grammar(propOrder string) string {
+	dat, _ := json.Marshal(j)
+	return NewJSONSchemaConverter(propOrder).GrammarFromBytes(dat)
+}