llmgw/internal/guard/policy.go

package guard

import (
	"fmt"
	"regexp"
	"slices"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard/detect"
)

// Config is the gateway's policy library, from server.yaml's guard: section.
//
//	guard:
//	  default: [baseline]
//	  policies:
//	    baseline:
//	      input:  {pii: {secret: block, card: redact}, injection: flag}
//	      output: {pii: {secret: block}}
//	    finance:
//	      mode: enforce
//	      fail: closed
//	      input: {pii: {iban: redact, tckn: redact}}
//
// Policies are attached to parts of the tenancy tree in the store; Default
// applies to every request.
type Config struct {
	Default  []string          `koanf:"default"`
	Policies map[string]Policy `koanf:"policies"`
}

// Policy is one named set of rules.
type Policy struct {
	// Mode: enforce acts, shadow only records what it would have done, off
	// does nothing. Empty is enforce.
	Mode string `koanf:"mode"`
	// Fail: what happens when a policy cannot judge — a judge model down or
	// over its latency budget. open serves, closed refuses. Empty is open.
	Fail string `koanf:"fail"`
	// NoStore keeps nothing of the request: no stored conversation, no
	// memory notes, no prefix routing, no fallback outside the cluster.
	NoStore bool `koanf:"no_store"`

	Input  Rules `koanf:"input"`
	Output Rules `koanf:"output"`
	Tools  Tools `koanf:"tools"`

	// MaxOutputTokens is put in a request that asks for no limit; one
	// asking for more is lowered to it.
	MaxOutputTokens int `koanf:"max_output_tokens"`
	// SystemPrompt is put before every request's messages.
	SystemPrompt string `koanf:"system_prompt"`
	// DropClientSystem removes the caller's own system and developer
	// messages, so only the pinned prompt instructs the model.
	DropClientSystem bool `koanf:"drop_client_system"`
	// RequireCapabilities refuses a request offering tools, or sending
	// images, to a model whose catalog entry does not say it takes them.
	RequireCapabilities bool `koanf:"require_capabilities"`
	// ValidateStructured checks an answer against the JSON schema the
	// request asked for in response_format: one that does not fit is
	// stopped — or, streamed a sentence at a time, flagged, having been
	// sent already.
	ValidateStructured bool `koanf:"validate_structured"`
	// Judge has a judge model on the cluster read the text too.
	Judge *JudgePolicy `koanf:"judge"`
	// Stream: sentence scans a streamed answer a sentence at a time,
	// releasing each once it passes; buffer_all holds the whole answer
	// until it is judged. Empty is sentence.
	Stream string `koanf:"stream"`
}

// Rules are what one direction of text is checked for.
type Rules struct {
	// PII maps a kind (email, phone, iban, card, tckn, ip, secret) to what is
	// done when it is found: flag, redact or block.
	PII map[string]string `koanf:"pii"`
	// Injection is what is done on signs of prompt injection: flag or block.
	Injection string `koanf:"injection"`
	// Deny are patterns, matched against the text as a model reads it.
	Deny []Pattern `koanf:"deny"`
	// MaxChars refuses text longer than this; zero is no limit.
	MaxChars int `koanf:"max_chars"`
}

// Pattern is a named regular expression and what its match does.
type Pattern struct {
	Name   string `koanf:"name"`
	Regex  string `koanf:"regex"`
	Action string `koanf:"action"`
}

// Tools are rules for tool definitions and calls.
type Tools struct {
	// Allow, when set, are the only tool names a request may offer and a
	// model may call.
	Allow []string `koanf:"allow"`
	// Deny are tool names never offered or called.
	Deny []string `koanf:"deny"`
	// OnDenied is strip (the tool is removed from the request) or block.
	// Empty is strip.
	OnDenied string `koanf:"on_denied"`
	// MaxCallsPerTurn stops an answer calling more tools than this.
	MaxCallsPerTurn int `koanf:"max_calls_per_turn"`
	// ValidateArguments checks a call's arguments against the parameters
	// schema the request declared for the tool.
	ValidateArguments bool `koanf:"validate_arguments"`
}

var actions = map[string]Action{"flag": Flag, "redact": Redact, "block": Block}

// Validate checks the library.
func (c Config) Validate() error {
	for _, d := range c.Default {
		if _, ok := c.Policies[d]; !ok {
			return fmt.Errorf("guard.default names %q, which guard.policies does not define", d)
		}
	}
	for name, p := range c.Policies {
		if err := p.validate(); err != nil {
			return fmt.Errorf("guard.policies.%s: %w", name, err)
		}
	}
	return nil
}

func (p Policy) validate() error {
	switch p.Mode {
	case "", "enforce", "shadow", "off":
	default:
		return fmt.Errorf("mode is %q: enforce, shadow or off", p.Mode)
	}
	switch p.Fail {
	case "", "open", "closed":
	default:
		return fmt.Errorf("fail is %q: open or closed", p.Fail)
	}
	switch p.Stream {
	case "", "sentence", "buffer_all":
	default:
		return fmt.Errorf("stream is %q: sentence or buffer_all", p.Stream)
	}
	if p.Judge != nil {
		if err := p.Judge.validate(); err != nil {
			return err
		}
	}
	switch p.Tools.OnDenied {
	case "", "strip", "block":
	default:
		return fmt.Errorf("tools.on_denied is %q: strip or block", p.Tools.OnDenied)
	}
	for dir, r := range map[string]Rules{"input": p.Input, "output": p.Output} {
		for kind, a := range r.PII {
			if !slices.Contains(detect.PIIKinds, kind) {
				return fmt.Errorf("%s.pii: %q is not one of %s", dir, kind, strings.Join(detect.PIIKinds, ", "))
			}
			if _, ok := actions[a]; !ok {
				return fmt.Errorf("%s.pii.%s is %q: flag, redact or block", dir, kind, a)
			}
		}
		switch r.Injection {
		case "", "flag", "block":
		default:
			return fmt.Errorf("%s.injection is %q: flag or block", dir, r.Injection)
		}
		for i, d := range r.Deny {
			if d.Name == "" {
				return fmt.Errorf("%s.deny[%d] has no name", dir, i)
			}
			if _, err := regexp.Compile(d.Regex); err != nil {
				return fmt.Errorf("%s.deny.%s: %w", dir, d.Name, err)
			}
			if _, ok := actions[d.Action]; !ok {
				return fmt.Errorf("%s.deny.%s action is %q: flag, redact or block", dir, d.Name, d.Action)
			}
		}
	}
	return nil
}

// Merge combines the policies that apply to one request into the one in
// force, by the rule limits narrow by: nothing below can loosen what is set
// above. The strictest mode and failure, the strictest action for each kind,
// every denied pattern and tool, only the tools every allow list allows, the
// smallest limits; any no_store; the pinned prompts in order.
func Merge(names []string, ps []Policy) Policy {
	var out Policy
	out.Input.PII, out.Output.PII = map[string]string{}, map[string]string{}
	var allowSets [][]string
	var prompts []string
	for _, p := range ps {
		// Unwritten is enforce, not the weakest: a policy that says nothing
		// of its mode, merged with one in shadow, must still act.
		mode := p.Mode
		if mode == "" {
			mode = "enforce"
		}
		out.Mode = stricterOf(out.Mode, mode, []string{"", "off", "shadow", "enforce"})
		out.Fail = stricterOf(out.Fail, p.Fail, []string{"", "open", "closed"})
		out.Stream = stricterOf(out.Stream, p.Stream, []string{"", "sentence", "buffer_all"})
		out.NoStore = out.NoStore || p.NoStore
		out.DropClientSystem = out.DropClientSystem || p.DropClientSystem
		out.RequireCapabilities = out.RequireCapabilities || p.RequireCapabilities
		out.ValidateStructured = out.ValidateStructured || p.ValidateStructured
		out.Judge = mergeJudge(out.Judge, p.Judge)
		out.MaxOutputTokens = minSet(out.MaxOutputTokens, p.MaxOutputTokens)
		if p.SystemPrompt != "" {
			prompts = append(prompts, p.SystemPrompt)
		}
		mergeRules(&out.Input, p.Input)
		mergeRules(&out.Output, p.Output)
		t := p.Tools
		if len(t.Allow) > 0 {
			allowSets = append(allowSets, t.Allow)
		}
		out.Tools.Deny = union(out.Tools.Deny, t.Deny)
		out.Tools.OnDenied = stricterOf(out.Tools.OnDenied, t.OnDenied, []string{"", "strip", "block"})
		out.Tools.MaxCallsPerTurn = minSet(out.Tools.MaxCallsPerTurn, t.MaxCallsPerTurn)
		out.Tools.ValidateArguments = out.Tools.ValidateArguments || t.ValidateArguments
	}
	for i, s := range allowSets {
		if i == 0 {
			out.Tools.Allow = slices.Clone(s)
			continue
		}
		out.Tools.Allow = slices.DeleteFunc(out.Tools.Allow, func(n string) bool { return !slices.Contains(s, n) })
		if len(out.Tools.Allow) == 0 {
			// Allow lists that share nothing allow nothing — said by a
			// name no tool has, not by an empty list, which allows all.
			out.Tools.Allow = []string{"\x00none"}
		}
	}
	out.SystemPrompt = strings.Join(prompts, "\n\n")
	return out
}

func mergeRules(dst *Rules, r Rules) {
	for k, a := range r.PII {
		if rank(actions[a]) > rank(actions[dst.PII[k]]) {
			dst.PII[k] = a
		}
	}
	dst.Injection = stricterOf(dst.Injection, r.Injection, []string{"", "flag", "block"})
	for _, d := range r.Deny {
		if !slices.ContainsFunc(dst.Deny, func(e Pattern) bool { return e.Name == d.Name && e.Regex == d.Regex }) {
			dst.Deny = append(dst.Deny, d)
		}
	}
	dst.MaxChars = minSet(dst.MaxChars, r.MaxChars)
}

func stricterOf(a, b string, order []string) string {
	if slices.Index(order, b) > slices.Index(order, a) {
		return b
	}
	return a
}

func minSet(a, b int) int {
	switch {
	case a == 0:
		return b
	case b == 0:
		return a
	}
	return min(a, b)
}

func union(a, b []string) []string {
	for _, s := range b {
		if !slices.Contains(a, s) {
			a = append(a, s)
		}
	}
	return a
}

// active reports whether a merged policy does anything.
func (p Policy) active() bool { return p.Mode != "" && p.Mode != "off" }