llmgw/internal/guard/guard.go

// Package guard is what the gateway asks about a request's text and its
// answer's: whether to let it through, flag it, rewrite it or stop it.
//
// The router calls a Guard at two points. Input, once the request is
// authorised, limited and has its stored conversation spliced in, and before
// it is dispatched: the whole upstream body. Output, for the answer: a filter
// over the response as it is relayed, buffered or streamed. A router without
// a Guard calls nothing and pays nothing.
//
// Nothing here may log, record or return the text it judges. Verdicts carry
// policy, rule and category names, never what matched.
package guard

import (
	"context"
	"time"
)

// Action is what a verdict does to a request or a response.
type Action string

const (
	// Allow lets it through untouched.
	Allow Action = ""
	// Flag lets it through, and says so: in a header and in the usage row.
	Flag Action = "flagged"
	// Redact lets it through rewritten.
	Redact Action = "redacted"
	// Block stops it.
	Block Action = "blocked"
)

// Stage is where a verdict was reached.
type Stage string

const (
	Input  Stage = "input"
	Output Stage = "output"
)

// Subject is who is asking, for which model: what a policy is chosen by.
type Subject struct {
	RequestID string
	KeyID     string
	UserID    string
	ProjectID string
	TeamID    string
	OrgID     string
	// Model is the canonical model that will serve.
	Model string
	// Path is the API path asked: /v1/chat/completions, /v1/completions, ….
	Path string
	// Stream says whether the answer is streamed.
	Stream bool
	// Tools and Vision are what the model's catalog entry says it can do.
	Tools, Vision bool
	// Restored are the values the request's redaction took out, by
	// placeholder, which restore puts back in the answer: outside services
	// reading the answer are not told them.
	Restored map[string]string
}

// Verdict is a judgement on one stage of one request.
type Verdict struct {
	Action   Action
	Stage    Stage
	Policy   string
	Rule     string
	Category string
	// Message is what a blocked caller is told. It names the policy, never
	// the text that tripped it.
	Message string
	// Body replaces the request's upstream body when the input was
	// rewritten — redacted, or given a pinned prompt or an output limit —
	// whatever the Action; nil leaves it as it was.
	Body []byte
	// Took is how long the judgement took.
	Took time.Duration
	// Shadow says the policy only records: the action is what it would have
	// done, and nothing is done.
	Shadow bool
	// Restore maps each placeholder the input was given to the value it
	// stands for, when the policy puts them back in the answer. Held for
	// the request's life only: never stored, logged or traced.
	Restore map[string]string
}

// Stricter reports whether v outranks w: block over redact over flag over
// allow. A response judged at several points carries its strictest verdict.
func (v Verdict) Stricter(w Verdict) bool { return rank(v.Action) > rank(w.Action) }

func rank(a Action) int {
	switch a {
	case Block:
		return 3
	case Redact:
		return 2
	case Flag:
		return 1
	}
	return 0
}

// Guard judges requests and their answers.
type Guard interface {
	// Input judges the body about to be sent upstream. An error means the
	// guard could not judge; the router refuses the request rather than
	// serve it unjudged — a guard that wants to fail open says Allow.
	Input(ctx context.Context, s Subject, body []byte) (Verdict, error)
	// Output is the filter for one attempt's answer, or nil when nothing
	// watches it.
	Output(ctx context.Context, s Subject) OutputFilter
	// Directives are what the policy in force asks of the gateway itself,
	// known before anything is stored or sent anywhere.
	Directives(s Subject) Directives
}

// Directives are what a policy asks of the gateway beyond judging text.
type Directives struct {
	// NoStore keeps nothing of the request: no stored conversation or
	// memory note, no prefix routing, nothing sent outside the cluster.
	NoStore bool
}

// OutputFilter watches one answer as it is relayed.
type OutputFilter interface {
	// Buffered sees a whole non-streamed response body before anything is
	// sent, and returns the body to send.
	Buffered(body []byte) []byte
	// OnEvent sees each streamed event's data — a chunk's JSON — and
	// returns the events to send in its place: itself to pass it on,
	// nothing to hold it back, more than one to release what it held. stop
	// ends the stream after those are sent.
	OnEvent(data string) (emit []string, stop bool)
	// OnEnd is called when the upstream stream ends, and returns what is
	// still held.
	OnEnd() []string
	// Verdict is the strictest judgement so far.
	Verdict() Verdict
}