llmgw/internal/control/adminapi/guard.go

package adminapi

import (
	"context"
	"encoding/json"
	"fmt"
	"maps"
	"net/http"
	"slices"
	"strings"
	"time"

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

// GuardStore is where policy attachments are kept.
type GuardStore interface {
	GuardAttachments(ctx context.Context) ([]guard.Attachment, error)
	AttachGuard(ctx context.Context, a guard.Attachment, by string) error
	DetachGuard(ctx context.Context, a guard.Attachment) (bool, error)
}

// GuardAdmin is what the guard endpoints need: the library from server.yaml,
// the store, and a way to have the gate read attachments again at once.
type GuardAdmin struct {
	Library    guard.Config
	Store      GuardStore
	Invalidate func()
}

// SetGuard enables the guard endpoints.
func (a *API) SetGuard(g *GuardAdmin) { a.guard = g }

// GuardView is the library and where its policies are attached.
type GuardView struct {
	Policies    []string              `json:"policies"`
	Default     []string              `json:"default"`
	Attachments []GuardAttachmentView `json:"attachments"`
	// Details says what each policy does, in the order of Policies.
	Details []GuardPolicyView `json:"details"`
}

// GuardPolicyView is what one policy does, in brief.
type GuardPolicyView struct {
	Name string `json:"name"`
	// Mode is enforce, shadow or off; Fail is open or closed.
	Mode    string `json:"mode"`
	Fail    string `json:"fail"`
	NoStore bool   `json:"no_store,omitempty"`
	// Judge is the judge model, when one reads the text.
	Judge string `json:"judge,omitempty"`
	// Checks are what it looks for, one line each: "input: pii email → redact".
	Checks []string `json:"checks"`
}

// GuardHistory is what the guardrails did, from the usage record; the store
// has it, and a store that does not leaves the events endpoint unanswered.
type GuardHistory interface {
	GuardEvents(ctx context.Context, since time.Time, limit int) ([]store.GuardEvent, error)
	GuardCounts(ctx context.Context, since time.Time) ([]store.GuardCount, error)
}

// GuardEventsView is what the guardrails did since a moment: totals, and the
// newest requests.
type GuardEventsView struct {
	Since  time.Time        `json:"since"`
	Counts []GuardCountView `json:"counts"`
	Events []GuardEventView `json:"events"`
}

// GuardCountView is how many requests one policy did one thing to.
type GuardCountView struct {
	Project  string `json:"project,omitempty"`
	Action   string `json:"action"`
	Policy   string `json:"policy"`
	Category string `json:"category,omitempty"`
	Count    int64  `json:"count"`
}

// GuardEventView is one request the guardrails did something to.
type GuardEventView struct {
	At        time.Time `json:"at"`
	RequestID string    `json:"request_id"`
	Project   string    `json:"project,omitempty"`
	Model     string    `json:"model"`
	Status    int       `json:"status"`
	Action    string    `json:"action"`
	Policy    string    `json:"policy"`
	Category  string    `json:"category,omitempty"`
	GuardMs   int64     `json:"guard_ms"`
}

// policyView sums one policy up.
func policyView(name string, p guard.Policy) GuardPolicyView {
	v := GuardPolicyView{Name: name, Mode: or(p.Mode, "enforce"), Fail: or(p.Fail, "open"), NoStore: p.NoStore, Checks: []string{}}
	rules := func(dir string, r guard.Rules) {
		for _, kind := range slices.Sorted(maps.Keys(r.PII)) {
			v.Checks = append(v.Checks, fmt.Sprintf("%s: pii %s → %s", dir, kind, r.PII[kind]))
		}
		if r.Injection != "" {
			v.Checks = append(v.Checks, fmt.Sprintf("%s: prompt injection → %s", dir, r.Injection))
		}
		for _, d := range r.Deny {
			v.Checks = append(v.Checks, fmt.Sprintf("%s: pattern %s → %s", dir, d.Name, or(d.Action, "block")))
		}
		if r.MaxChars > 0 {
			v.Checks = append(v.Checks, fmt.Sprintf("%s: at most %d characters", dir, r.MaxChars))
		}
	}
	rules("input", p.Input)
	rules("output", p.Output)
	if j := p.Judge; j != nil {
		v.Judge = j.Model
		var dirs []string
		if j.Input != "" {
			dirs = append(dirs, "input → "+j.Input)
		}
		if j.Output != "" {
			dirs = append(dirs, "output → "+j.Output)
		}
		line := "judge " + j.Model + ": " + strings.Join(dirs, ", ")
		if len(j.Categories) > 0 {
			line += " (" + strings.Join(j.Categories, " ") + ")"
		}
		v.Checks = append(v.Checks, line)
	}
	if len(p.Tools.Allow) > 0 {
		v.Checks = append(v.Checks, "tools: only "+strings.Join(p.Tools.Allow, ", "))
	}
	if len(p.Tools.Deny) > 0 {
		v.Checks = append(v.Checks, "tools: never "+strings.Join(p.Tools.Deny, ", "))
	}
	if p.MaxOutputTokens > 0 {
		v.Checks = append(v.Checks, fmt.Sprintf("answers of at most %d tokens", p.MaxOutputTokens))
	}
	if p.SystemPrompt != "" {
		v.Checks = append(v.Checks, "a system prompt of its own")
	}
	return v
}

func or(s, dflt string) string {
	if s == "" {
		return dflt
	}
	return s
}

// GuardAttachmentView is one attachment, with the name people know its
// scope by.
type GuardAttachmentView struct {
	guard.Attachment
	Name string `json:"name,omitempty"`
}

// guardHistory is where what the guardrails did is read, when there is one.
func (a *API) guardHistory() (GuardHistory, bool) {
	if a.guard == nil {
		return nil, false
	}
	h, ok := a.guard.Store.(GuardHistory)
	return h, ok
}

// guardTarget reads where a policy is to be attached.
func (a *API) guardTarget(ctx context.Context, target string) (guard.Attachment, error) {
	kind, what, ok := strings.Cut(strings.TrimSpace(target), ":")
	what = strings.TrimSpace(what)
	if !ok || what == "" {
		return guard.Attachment{}, fmt.Errorf("%q is not org:, team:, project:, key: or user: and what", target)
	}
	switch kind {
	case "key", "user":
		return guard.Attachment{Scope: kind, ScopeID: what}, nil
	case "org", "team", "project":
		if a.quotas == nil {
			return guard.Attachment{Scope: kind, ScopeID: what}, nil
		}
		_, id, _, err := a.resolveScope(ctx, target)
		if err != nil {
			return guard.Attachment{}, err
		}
		return guard.Attachment{Scope: kind, ScopeID: id}, nil
	}
	return guard.Attachment{}, fmt.Errorf("%q: attach to org, team, project, key or user", kind)
}

// GuardAttachmentRequest attaches a policy to a part of the tenancy tree, a
// key or a person, or with DELETE takes it off again.
type GuardAttachmentRequest struct {
	// Target is org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT (or the ID
	// after the colon), key:KEY-ID or user:USER-ID.
	Target string `json:"target"`
	// Policy is a policy's name in server.yaml's guard.policies.
	Policy string `json:"policy"`
}

func (a *API) registerGuard(mux Mux) {
	mux.HandleFunc("GET /api/v1/guard", func(w http.ResponseWriter, r *http.Request) {
		if a.guard == nil {
			writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard: in server.yaml")
			return
		}
		att, err := a.guard.Store.GuardAttachments(r.Context())
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		v := GuardView{Default: a.guard.Library.Default, Attachments: []GuardAttachmentView{}}
		var names map[string]string
		if a.quotas != nil {
			names, _ = a.scopeNames(r.Context())
		}
		for _, at := range att {
			v.Attachments = append(v.Attachments, GuardAttachmentView{Attachment: at, Name: names[at.Scope+":"+at.ScopeID]})
		}
		for name := range a.guard.Library.Policies {
			v.Policies = append(v.Policies, name)
		}
		slices.Sort(v.Policies)
		for _, name := range v.Policies {
			v.Details = append(v.Details, policyView(name, a.guard.Library.Policies[name]))
		}
		writeJSON(w, http.StatusOK, v)
	})
	mux.HandleFunc("GET /api/v1/guard/events", func(w http.ResponseWriter, r *http.Request) {
		h, ok := a.guardHistory()
		if !ok {
			writeError(w, http.StatusNotImplemented, "no guard policies, or no store to record what they did")
			return
		}
		since := 24 * time.Hour
		if q := r.URL.Query().Get("since"); q != "" {
			d, err := parseSince(q)
			if err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			since = d
		}
		from := time.Now().Add(-since)
		counts, err := h.GuardCounts(r.Context(), from)
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		events, err := h.GuardEvents(r.Context(), from, 100)
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		var names map[string]string
		if a.quotas != nil {
			names, _ = a.scopeNames(r.Context())
		}
		project := func(id string) string {
			if n := names["project:"+id]; n != "" {
				return n
			}
			return id
		}
		v := GuardEventsView{Since: from, Counts: []GuardCountView{}, Events: []GuardEventView{}}
		for _, c := range counts {
			v.Counts = append(v.Counts, GuardCountView{Project: project(c.ProjectID), Action: c.Action,
				Policy: c.Policy, Category: c.Category, Count: c.Count})
		}
		slices.SortStableFunc(v.Counts, func(a, b GuardCountView) int { return int(b.Count - a.Count) })
		for _, e := range events {
			v.Events = append(v.Events, GuardEventView{At: e.At, RequestID: e.RequestID, Project: project(e.ProjectID),
				Model: e.Model, Status: e.Status, Action: e.Action, Policy: e.Policy, Category: e.Category, GuardMs: e.GuardMs})
		}
		writeJSON(w, http.StatusOK, v)
	})
	change := func(attach bool) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			if a.guard == nil {
				writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard: in server.yaml")
				return
			}
			var at GuardAttachmentRequest
			if err := json.NewDecoder(r.Body).Decode(&at); err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			att, err := a.guardTarget(r.Context(), at.Target)
			if err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			att.Policy = at.Policy
			if _, ok := a.guard.Library.Policies[att.Policy]; !ok && attach {
				writeError(w, http.StatusBadRequest, fmt.Sprintf("no policy %q in server.yaml's guard.policies", att.Policy))
				return
			}
			changed := true
			if attach {
				err = a.guard.Store.AttachGuard(r.Context(), att, actorOf(r))
			} else {
				changed, err = a.guard.Store.DetachGuard(r.Context(), att)
			}
			if err != nil {
				writeError(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !changed {
				writeError(w, http.StatusNotFound, "no such attachment")
				return
			}
			if a.guard.Invalidate != nil {
				a.guard.Invalidate()
			}
			writeJSON(w, http.StatusOK, att)
		}
	}
	mux.HandleFunc("POST /api/v1/guard/attachments", change(true))
	mux.HandleFunc("DELETE /api/v1/guard/attachments", change(false))
}