llmgw/internal/router/guard.go

package router

import (
	"context"
	"encoding/json"
	"net/http"
	"strconv"
	"time"

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

// The headers a judged request is answered with: what was done, and by
// which policy. Set before anything is written, so on a streamed answer they
// carry the input's verdict, and the output's only on a buffered one.
const (
	guardHeader       = "x-llmgw-guard"
	guardPolicyHeader = "x-llmgw-guard-policy"
)

// policyViolation is the error code a blocked request is refused with.
const policyViolation = "content_policy_violation"

func (p *parsedRequest) guardSubject(model string) guard.Subject {
	return guard.Subject{
		RequestID: p.RequestID, KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID,
		TeamID: p.TeamID, OrgID: p.OrgID, Model: model, Path: p.Path, Stream: p.Stream,
	}
}

// subject is a request as a guard sees it, with what the model serving it
// can do.
func (r *Router) subject(req *parsedRequest, model string) guard.Subject {
	s := req.guardSubject(model)
	if _, entry := r.resolve(model); entry.Name != "" {
		s.Tools, s.Vision = entry.Capabilities.Tools, entry.Capabilities.Vision
	}
	return s
}

func setGuardHeaders(h http.Header, v guard.Verdict) {
	if v.Action == guard.Allow || v.Shadow {
		return
	}
	h.Set(guardHeader, string(v.Action))
	if v.Policy != "" {
		h.Set(guardPolicyHeader, v.Policy)
	}
}

// callerSubject is who a request comes from, for a request that is not an
// inference: a memory note.
func callerSubject(req *http.Request) guard.Subject {
	p, _ := auth.FromContext(req.Context())
	return guard.Subject{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID,
		OrgID: p.OrgID, Path: req.URL.Path}
}

// noStore reports whether the policy in force for a request keeps nothing.
func (r *Router) noStore(s guard.Subject) bool {
	return r.guard != nil && r.guard.Directives(s).NoStore
}

// errNoStore is a request asking the gateway to keep something under a
// policy that keeps nothing.
const errNoStoreMsg = "the gateway's policy for this caller keeps nothing: no stored conversations or memory notes, and no model outside the cluster"

// guardInput judges the request about to be dispatched. It answers the
// client itself and returns false when the request goes no further.
func (r *Router) guardInput(ctx context.Context, w http.ResponseWriter, req *parsedRequest, start time.Time) bool {
	if req.ContentType != "" {
		// Audio: no text for a guard to read on the way in.
		return true
	}
	model, _ := r.resolve(req.Model)
	t := time.Now()
	v, err := r.guard.Input(ctx, r.subject(req, model), req.UpstreamBody)
	if v.Took == 0 {
		v.Took = time.Since(t)
	}
	v.Stage = guard.Input
	if err != nil {
		// Not judged is not judged fine: a guard that means to fail open
		// says Allow instead.
		r.refuseGuarded(w, req, guard.Verdict{Action: guard.Block, Stage: guard.Input, Took: v.Took},
			http.StatusServiceUnavailable, "server_error", "guard_unavailable",
			"the request could not be checked against this gateway's policies; try again", start)
		return false
	}
	req.guardIn = v
	r.m.guard.observe(v, guard.Input, r.m.tenant.tenantOf(req))
	if v.Shadow {
		// Recorded, not done: the client sees nothing of it.
		return true
	}
	setGuardHeaders(w.Header(), v)
	if v.Action != guard.Block && v.Body != nil {
		r.useGuardBody(req, v.Body)
	}
	switch v.Action {
	case guard.Block:
		msg := v.Message
		if msg == "" {
			msg = "this request was refused by policy " + strconv.Quote(v.Policy)
		}
		r.refuseGuarded(w, req, v, http.StatusBadRequest, "invalid_request_error", policyViolation, msg, start)
		return false
	}
	return true
}

// useGuardBody sends the body a guard rewrote, and routes by it, so the
// prefix tree never holds what was taken out.
func (r *Router) useGuardBody(req *parsedRequest, body []byte) {
	req.UpstreamBody = body
	var b inferenceBody
	if json.Unmarshal(body, &b) == nil {
		req.PromptPrefix = promptPrefix(&b)
		if n := max(b.MaxTokens, b.MaxCompletionTokens); n > 0 {
			req.MaxTokens = n
		}
	}
}

// refuseGuarded answers a request a guard stopped. Nothing is charged —
// nothing was generated — but the row is kept, with the verdict, so a team
// running into a policy shows as one.
func (r *Router) refuseGuarded(w http.ResponseWriter, req *parsedRequest, v guard.Verdict,
	status int, kind, code, msg string, start time.Time,
) {
	writeErrorCode(w, status, kind, code, msg)
	r.m.observeRequest(r.metricModel(req.Model), "", status, time.Since(start), nil, 0)
	r.m.tenant.observe(req, status, 0, 0, 0)
	if r.acct != nil && req.KeyID != "" {
		u := store.Usage{
			At: time.Now(), KeyID: req.KeyID, UserID: req.UserID, ProjectID: req.ProjectID,
			Model: req.Model, ServedModel: req.Model, Status: status,
			LatencyMs: time.Since(start).Milliseconds(), RequestID: req.RequestID, QueueMs: req.QueueMs,
			EndUser: req.EndUser, Tags: req.Tags,
		}
		guardFields(&u, v, guard.Verdict{})
		// Settled at nothing: the estimate the limit check charged is
		// given back.
		r.acct.Record(req.reservation, u)
	}
	r.accessLog(req, dispatchResult{guard: v}, status, time.Since(start))
}

// observeOutput counts what the output filter decided for an answer that
// was judged, and a stream it stopped.
func (r *Router) observeOutput(req *parsedRequest, res dispatchResult, judged bool) {
	if !judged {
		return
	}
	r.m.guard.observe(res.guard, guard.Output, r.m.tenant.tenantOf(req))
	if res.outcome.Stopped {
		r.m.guard.stops.Inc()
	}
}

// outputFilter attaches the guard's filter for one attempt to the proxy
// request, and returns it so its verdict can be read afterwards.
func (r *Router) outputFilter(ctx context.Context, req *parsedRequest, model string, preq *proxy.Request) guard.OutputFilter {
	if r.guard == nil || req.internal {
		return nil
	}
	f := r.guard.Output(ctx, r.subject(req, model))
	if f == nil {
		return nil
	}
	if req.Stream {
		preq.StreamFilter = f
	}
	preq.Filter = func(body []byte, h http.Header) []byte {
		out := f.Buffered(body)
		v := f.Verdict()
		if v.Stricter(req.guardIn) {
			setGuardHeaders(h, v)
		}
		return out
	}
	return f
}

// guardFields records the stricter of the input's and output's verdicts.
func guardFields(u *store.Usage, in, out guard.Verdict) {
	v := in
	if out.Stricter(in) {
		v = out
	}
	u.GuardAction, u.GuardPolicy, u.GuardCategory = string(v.Action), v.Policy, v.Category
	if v.Shadow && v.Action != guard.Allow {
		u.GuardAction = "shadow-" + u.GuardAction
	}
	u.GuardMs = (in.Took + out.Took).Milliseconds()
}