llmgw/internal/guard/schema.go

package guard

import (
	"encoding/json"
	"fmt"
	"math"
)

// validate checks a value against the part of JSON Schema tool parameters
// are written in: type, properties, required, enum, items, and
// additionalProperties false. Enough to stop a model calling a tool with
// arguments the tool never declared; not a full validator.
func validate(schema json.RawMessage, value any) error {
	if len(schema) == 0 {
		return nil
	}
	var s struct {
		Type                 any                        `json:"type"`
		Properties           map[string]json.RawMessage `json:"properties"`
		Required             []string                   `json:"required"`
		Enum                 []any                      `json:"enum"`
		Items                json.RawMessage            `json:"items"`
		AdditionalProperties *bool                      `json:"additionalProperties"`
	}
	if json.Unmarshal(schema, &s) != nil {
		return nil // a schema this cannot read is the engine's to apply
	}
	if t, ok := s.Type.(string); ok && !isType(value, t) {
		return fmt.Errorf("is not %s", t)
	}
	if list, ok := s.Type.([]any); ok {
		match := false
		for _, t := range list {
			if ts, ok := t.(string); ok && isType(value, ts) {
				match = true
			}
		}
		if !match {
			return fmt.Errorf("is none of the types allowed")
		}
	}
	if len(s.Enum) > 0 {
		b, _ := json.Marshal(value)
		found := false
		for _, e := range s.Enum {
			eb, _ := json.Marshal(e)
			if string(eb) == string(b) {
				found = true
			}
		}
		if !found {
			return fmt.Errorf("is not one of the values allowed")
		}
	}
	if obj, ok := value.(map[string]any); ok {
		for _, r := range s.Required {
			if _, ok := obj[r]; !ok {
				return fmt.Errorf("lacks %q", r)
			}
		}
		for k, v := range obj {
			sub, declared := s.Properties[k]
			if !declared {
				if s.AdditionalProperties != nil && !*s.AdditionalProperties {
					return fmt.Errorf("has %q, which is not declared", k)
				}
				continue
			}
			if err := validate(sub, v); err != nil {
				return fmt.Errorf("%s %w", k, err)
			}
		}
	}
	if arr, ok := value.([]any); ok && len(s.Items) > 0 {
		for i, v := range arr {
			if err := validate(s.Items, v); err != nil {
				return fmt.Errorf("[%d] %w", i, err)
			}
		}
	}
	return nil
}

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 true
}