llmgw/internal/guard/service.go

package guard

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"
)

// A guardrail that is someone else's service: Azure Content Safety, Lakera,
// NeMo Guardrails, a company's own classifier. A site may be required to use
// one; this asks it beside the patterns and the judge, and its answer counts
// like theirs. Told only the text, never who sent it, and nothing of the
// call is kept but its verdict.

// Service is one outside guardrail, asked over HTTP.
type Service struct {
	// Name says which service decided, in the verdict's rule.
	Name string `koanf:"name"`
	// URL is POSTed the body.
	URL string `koanf:"url"`
	// HeaderFiles are headers whose values are read from files each time,
	// so a rotated key needs no restart: {Authorization: /etc/…/lakera}.
	HeaderFiles map[string]string `koanf:"header_files"`
	// Body is the request, a template given .Text; {{json .Text}} writes it
	// as a JSON string. Empty is {"input": <text>}.
	Body string `koanf:"body"`
	// Blocked is where in the answer the service says the text is unsafe:
	// a dotted path, numbers for list items (results.0.flagged). True, a
	// number above zero, a non-empty list or string other than "false" is
	// unsafe.
	Blocked string `koanf:"blocked"`
	// Reason is where it says why, for the verdict; optional.
	Reason string `koanf:"reason"`
	// On: input, output or both. Empty is input.
	On string `koanf:"on"`
	// Action on unsafe: block or flag. Empty is block.
	Action string `koanf:"action"`
	// Timeout bounds one call. Empty is 3s; past it the policy's fail says.
	Timeout time.Duration `koanf:"timeout"`
}

func (s Service) validate() error {
	if s.Name == "" {
		return fmt.Errorf("services: one has no name")
	}
	u, err := url.Parse(s.URL)
	if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
		return fmt.Errorf("services.%s: url %q is not an http(s) address", s.Name, s.URL)
	}
	if s.Blocked == "" {
		return fmt.Errorf("services.%s: blocked says where the answer says unsafe, and is empty", s.Name)
	}
	switch s.On {
	case "", "input", "output", "both":
	default:
		return fmt.Errorf("services.%s: on is %q: input, output or both", s.Name, s.On)
	}
	switch s.Action {
	case "", "block", "flag":
	default:
		return fmt.Errorf("services.%s: action is %q: block or flag", s.Name, s.Action)
	}
	if _, err := s.template(); err != nil {
		return fmt.Errorf("services.%s: body: %w", s.Name, err)
	}
	return nil
}

func (s Service) onInput() bool  { return s.On == "" || s.On == "input" || s.On == "both" }
func (s Service) onOutput() bool { return s.On == "output" || s.On == "both" }

func (s Service) action() Action {
	if s.Action == "flag" {
		return Flag
	}
	return Block
}

func (s Service) template() (*template.Template, error) {
	body := s.Body
	if body == "" {
		body = `{"input": {{json .Text}}}`
	}
	return template.New(s.Name).Funcs(template.FuncMap{"json": func(v any) (string, error) {
		b, err := json.Marshal(v)
		return string(b), err
	}}).Parse(body)
}

// serviceClient asks every service; each call is bounded by its own timeout.
var serviceClient = &http.Client{}

// ask sends the service the text and says whether it found it unsafe, and why.
func (s Service) ask(ctx context.Context, text string) (unsafe bool, reason string, err error) {
	t, err := s.template()
	if err != nil {
		return false, "", err
	}
	var body bytes.Buffer
	if err := t.Execute(&body, struct{ Text string }{text}); err != nil {
		return false, "", err
	}
	timeout := s.Timeout
	if timeout <= 0 {
		timeout = 3 * time.Second
	}
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.URL, &body)
	if err != nil {
		return false, "", err
	}
	req.Header.Set("Content-Type", "application/json")
	for h, file := range s.HeaderFiles {
		v, err := os.ReadFile(file)
		if err != nil {
			return false, "", fmt.Errorf("header %s: %w", h, err)
		}
		req.Header.Set(h, strings.TrimSpace(string(v)))
	}
	resp, err := serviceClient.Do(req)
	if err != nil {
		return false, "", err
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return false, "", err
	}
	if resp.StatusCode/100 != 2 {
		return false, "", fmt.Errorf("answered %d", resp.StatusCode)
	}
	var doc any
	if err := json.Unmarshal(b, &doc); err != nil {
		return false, "", fmt.Errorf("answer is not JSON: %w", err)
	}
	v, ok := at(doc, s.Blocked)
	if !ok || v == nil {
		// null is no answer: a service that cannot say is not saying safe.
		return false, "", fmt.Errorf("the answer has no %s", s.Blocked)
	}
	if s.Reason != "" {
		if r, ok := at(doc, s.Reason); ok {
			reason = fmt.Sprint(r)
		}
	}
	return truthy(v), reason, nil
}

// at is the value at a dotted path in decoded JSON.
func at(doc any, path string) (any, bool) {
	cur := doc
	for _, k := range strings.Split(path, ".") {
		switch c := cur.(type) {
		case map[string]any:
			v, ok := c[k]
			if !ok {
				return nil, false
			}
			cur = v
		case []any:
			i, err := strconv.Atoi(k)
			if err != nil || i < 0 || i >= len(c) {
				return nil, false
			}
			cur = c[i]
		default:
			return nil, false
		}
	}
	return cur, true
}

func truthy(v any) bool {
	switch x := v.(type) {
	case bool:
		return x
	case float64:
		return x > 0
	case string:
		return x != "" && !strings.EqualFold(x, "false") && x != "0"
	case []any:
		return len(x) > 0
	case map[string]any:
		return len(x) > 0
	}
	return false
}

// askServices asks each service of a direction, and says what the policy
// makes of their answers: a finding per service that found the text unsafe,
// or could not be asked under fail: closed. failed is true when a service
// gave no answer, for the caller to refuse under fail: closed.
func (g *Gate) askServices(ctx context.Context, p Policy, input bool, text string) (fs []finding, failed error) {
	for _, s := range p.Services {
		if input && !s.onInput() || !input && !s.onOutput() || text == "" {
			continue
		}
		unsafe, reason, err := s.ask(ctx, text)
		if err != nil {
			g.log.Warn("a guardrail service gave no verdict", "service", s.Name, "err", err)
			if p.Fail == "closed" {
				failed = fmt.Errorf("the guardrail service %s could not be asked: %w", s.Name, err)
				continue
			}
			g.m.failOpen.Inc()
			continue
		}
		if unsafe {
			rule := "service:" + s.Name
			if reason != "" {
				rule += ":" + reason
			}
			fs = append(fs, finding{action: s.action(), category: "service", rule: rule})
		}
	}
	return fs, failed
}