llmgw/internal/bench/eval/schema.go

package eval

import (
	"encoding/json"
	"fmt"
	"math"
	"reflect"
	"regexp"
	"slices"
	"sort"
	"strings"
	"unicode/utf8"
)

// A JSON Schema validator for what a suite asks of an answer's shape: types,
// properties, required ones, items, enums, bounds and patterns. Not the whole
// standard — no $ref, no oneOf — and a schema using what is not here is
// refused when the suite is read rather than passed by an answer it never
// checked.

type schema struct {
	types                []string
	properties           map[string]*schema
	required             []string
	additional           *schema // nil: anything; noneAllowed: nothing
	items                *schema
	enum                 []any
	constant             any
	hasConst             bool
	minimum, maximum     *float64
	exclMin, exclMax     *float64
	minLength, maxLength *int
	minItems, maxItems   *int
	pattern              *regexp.Regexp
	anyOf, allOf         []*schema
	never                bool
}

// noneAllowed is additionalProperties: false.
var noneAllowed = &schema{never: true}

// annotations are keywords that describe rather than constrain.
var annotations = map[string]bool{"$schema": true, "$id": true, "title": true, "description": true,
	"default": true, "examples": true, "format": true, "$comment": true}

func compileSchema(v any) (*schema, error) {
	// Through JSON, so a schema written in YAML and one in JSON are the same
	// thing by the time they are read: numbers as float64, maps as map[string]any.
	b, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	var raw any
	if err := json.Unmarshal(b, &raw); err != nil {
		return nil, err
	}
	return compile(raw, "")
}

func compile(raw any, at string) (*schema, error) {
	if b, ok := raw.(bool); ok {
		if b {
			return &schema{}, nil
		}
		return noneAllowed, nil
	}
	m, ok := raw.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("%s is not a schema object", where(at))
	}
	s := &schema{}
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := m[k]
		var err error
		switch k {
		case "type":
			switch t := v.(type) {
			case string:
				s.types = []string{t}
			case []any:
				for _, x := range t {
					if str, ok := x.(string); ok {
						s.types = append(s.types, str)
					}
				}
			}
			for _, t := range s.types {
				if !slices.Contains([]string{"object", "array", "string", "number", "integer", "boolean", "null"}, t) {
					return nil, fmt.Errorf("%s: %q is not a JSON type", where(at), t)
				}
			}
		case "properties":
			props, ok := v.(map[string]any)
			if !ok {
				return nil, fmt.Errorf("%s: properties is not an object", where(at))
			}
			s.properties = map[string]*schema{}
			for name, p := range props {
				if s.properties[name], err = compile(p, at+"."+name); err != nil {
					return nil, err
				}
			}
		case "required":
			list, _ := v.([]any)
			for _, x := range list {
				if str, ok := x.(string); ok {
					s.required = append(s.required, str)
				}
			}
		case "additionalProperties":
			if s.additional, err = compile(v, at+".*"); err != nil {
				return nil, err
			}
		case "items":
			if s.items, err = compile(v, at+"[]"); err != nil {
				return nil, err
			}
		case "enum":
			s.enum, _ = v.([]any)
		case "const":
			s.constant, s.hasConst = v, true
		case "minimum", "maximum", "exclusiveMinimum", "exclusiveMaximum":
			f, ok := v.(float64)
			if !ok {
				return nil, fmt.Errorf("%s: %s is not a number", where(at), k)
			}
			*map[string]**float64{"minimum": &s.minimum, "maximum": &s.maximum,
				"exclusiveMinimum": &s.exclMin, "exclusiveMaximum": &s.exclMax}[k] = &f
		case "minLength", "maxLength", "minItems", "maxItems":
			f, ok := v.(float64)
			if !ok || f < 0 || f != math.Trunc(f) {
				return nil, fmt.Errorf("%s: %s is not a whole number", where(at), k)
			}
			n := int(f)
			*map[string]**int{"minLength": &s.minLength, "maxLength": &s.maxLength,
				"minItems": &s.minItems, "maxItems": &s.maxItems}[k] = &n
		case "pattern":
			str, _ := v.(string)
			if s.pattern, err = regexp.Compile(str); err != nil {
				return nil, fmt.Errorf("%s: pattern: %w", where(at), err)
			}
		case "anyOf", "allOf":
			list, ok := v.([]any)
			if !ok || len(list) == 0 {
				return nil, fmt.Errorf("%s: %s is not a list of schemas", where(at), k)
			}
			for i, x := range list {
				sub, err := compile(x, fmt.Sprintf("%s(%s %d)", at, k, i+1))
				if err != nil {
					return nil, err
				}
				if k == "anyOf" {
					s.anyOf = append(s.anyOf, sub)
				} else {
					s.allOf = append(s.allOf, sub)
				}
			}
		default:
			if !annotations[k] {
				return nil, fmt.Errorf("%s: %s is not supported here (type, properties, required, "+
					"additionalProperties, items, enum, const, minimum, maximum, exclusiveMinimum, exclusiveMaximum, "+
					"minLength, maxLength, minItems, maxItems, pattern, anyOf and allOf are)", where(at), k)
			}
		}
	}
	return s, nil
}

func where(at string) string {
	if at == "" {
		return "the schema"
	}
	return "the schema at " + strings.TrimPrefix(at, ".")
}

// validate says what is wrong with v, or "".
func (s *schema) validate(v any, at string) string {
	if s.never {
		return fmt.Sprintf("%s is not allowed", place(at))
	}
	if len(s.types) > 0 && !slices.ContainsFunc(s.types, func(t string) bool { return isType(v, t) }) {
		return fmt.Sprintf("%s is %s, not %s", place(at), typeOf(v), strings.Join(s.types, " or "))
	}
	if s.hasConst && !equal(v, s.constant) {
		return fmt.Sprintf("%s is %s, not %s", place(at), short(v), short(s.constant))
	}
	if s.enum != nil && !slices.ContainsFunc(s.enum, func(e any) bool { return equal(v, e) }) {
		return fmt.Sprintf("%s is %s, not one of %s", place(at), short(v), short(s.enum))
	}
	switch x := v.(type) {
	case float64:
		switch {
		case s.minimum != nil && x < *s.minimum:
			return fmt.Sprintf("%s is %v, below %v", place(at), x, *s.minimum)
		case s.maximum != nil && x > *s.maximum:
			return fmt.Sprintf("%s is %v, above %v", place(at), x, *s.maximum)
		case s.exclMin != nil && x <= *s.exclMin:
			return fmt.Sprintf("%s is %v, not above %v", place(at), x, *s.exclMin)
		case s.exclMax != nil && x >= *s.exclMax:
			return fmt.Sprintf("%s is %v, not below %v", place(at), x, *s.exclMax)
		}
	case string:
		n := utf8.RuneCountInString(x)
		switch {
		case s.minLength != nil && n < *s.minLength:
			return fmt.Sprintf("%s is %d characters, fewer than %d", place(at), n, *s.minLength)
		case s.maxLength != nil && n > *s.maxLength:
			return fmt.Sprintf("%s is %d characters, more than %d", place(at), n, *s.maxLength)
		case s.pattern != nil && !s.pattern.MatchString(x):
			return fmt.Sprintf("%s, %s, does not match %s", place(at), short(x), s.pattern)
		}
	case []any:
		switch {
		case s.minItems != nil && len(x) < *s.minItems:
			return fmt.Sprintf("%s has %d items, fewer than %d", place(at), len(x), *s.minItems)
		case s.maxItems != nil && len(x) > *s.maxItems:
			return fmt.Sprintf("%s has %d items, more than %d", place(at), len(x), *s.maxItems)
		}
		if s.items != nil {
			for i, e := range x {
				if why := s.items.validate(e, fmt.Sprintf("%s[%d]", at, i)); why != "" {
					return why
				}
			}
		}
	case map[string]any:
		for _, r := range s.required {
			if _, ok := x[r]; !ok {
				return fmt.Sprintf("%s has no %s", place(at), r)
			}
		}
		keys := make([]string, 0, len(x))
		for k := range x {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			if p, ok := s.properties[k]; ok {
				if why := p.validate(x[k], at+"."+k); why != "" {
					return why
				}
			} else if s.additional != nil {
				if why := s.additional.validate(x[k], at+"."+k); why != "" {
					return why
				}
			}
		}
	}
	for _, sub := range s.allOf {
		if why := sub.validate(v, at); why != "" {
			return why
		}
	}
	if len(s.anyOf) > 0 {
		var first string
		for _, sub := range s.anyOf {
			why := sub.validate(v, at)
			if why == "" {
				return ""
			}
			if first == "" {
				first = why
			}
		}
		return first + " (nor any other choice)"
	}
	return ""
}

func place(at string) string {
	if at == "" {
		return "the answer"
	}
	return strings.TrimPrefix(at, ".")
}

func isType(v any, t string) bool {
	switch t {
	case "object":
		_, ok := v.(map[string]any)
		return ok
	case "array":
		_, ok := v.([]any)
		return ok
	case "string":
		_, ok := v.(string)
		return ok
	case "number":
		_, ok := v.(float64)
		return ok
	case "integer":
		f, ok := v.(float64)
		return ok && f == math.Trunc(f)
	case "boolean":
		_, ok := v.(bool)
		return ok
	case "null":
		return v == nil
	}
	return false
}

func typeOf(v any) string {
	switch x := v.(type) {
	case map[string]any:
		return "an object"
	case []any:
		return "a list"
	case string:
		return "a string"
	case float64:
		if x == math.Trunc(x) {
			return "an integer"
		}
		return "a number"
	case bool:
		return "a boolean"
	case nil:
		return "null"
	}
	return fmt.Sprintf("%T", v)
}

func equal(a, b any) bool { return reflect.DeepEqual(a, b) }

func short(v any) string {
	b, _ := json.Marshal(v)
	s := string(b)
	if len(s) > 60 {
		s = s[:57] + "..."
	}
	return s
}