llmgw/internal/guard/output.go

package guard

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"time"
)

// The tools a request offered, kept from Input until its answer is judged:
// a call's arguments are checked against the parameters declared for it.
type toolMemo struct {
	mu   sync.Mutex
	byID map[string]toolSet
}

type toolSet struct {
	schemas map[string]json.RawMessage
	// answer is the schema response_format asks the answer to fit.
	answer json.RawMessage
	// user is the request's last user turn, for a judge reading the answer.
	user string
	at   time.Time
}

func (m *toolMemo) put(id string, tools []tool, answer json.RawMessage, user string, now time.Time) {
	if id == "" || len(tools) == 0 && answer == nil && user == "" {
		return
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.byID == nil {
		m.byID = map[string]toolSet{}
	}
	for k, v := range m.byID {
		if now.Sub(v.at) > 10*time.Minute {
			delete(m.byID, k)
		}
	}
	set := toolSet{schemas: map[string]json.RawMessage{}, answer: answer, user: user, at: now}
	for _, t := range tools {
		set.schemas[t.name] = t.parameters
	}
	m.byID[id] = set
}

func (m *toolMemo) get(id string) toolSet {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.byID[id]
}

func outputWatched(p Policy) bool {
	o, t := p.Output, p.Tools
	return len(o.PII) > 0 || o.Injection != "" || len(o.Deny) > 0 || o.MaxChars > 0 || p.ValidateStructured ||
		p.Judge != nil && p.Judge.Output != "" ||
		len(t.Allow) > 0 || len(t.Deny) > 0 || t.MaxCallsPerTurn > 0 || t.ValidateArguments
}

// Output is the filter for one answer, or nil when the policy in force
// watches no answers.
func (g *Gate) Output(ctx context.Context, s Subject) OutputFilter {
	p, names := g.Policy(s)
	if !p.active() || !outputWatched(p) {
		return nil
	}
	memo := g.tools.get(s.RequestID)
	return &outFilter{ctx: ctx, g: g, p: p, names: names, shadow: p.Mode == "shadow",
		schemas: memo.schemas, answer: memo.answer, user: memo.user, choices: map[int]*choiceState{}}
}

// outFilter judges one answer.
type outFilter struct {
	ctx     context.Context
	g       *Gate
	p       Policy
	names   string
	shadow  bool
	schemas map[string]json.RawMessage
	answer  json.RawMessage
	full    strings.Builder // the whole answer, when it must fit a schema
	user    string          // the request's last user turn, for the judge
	sofar   strings.Builder // the answer so far, for the judge
	judged  int             // how much of it the judge has read
	v       Verdict
	took    time.Duration
	stopped bool

	held     []string // events not yet released
	sent     bool     // whether a chunk this filter made has gone out
	template map[string]json.RawMessage
	choices  map[int]*choiceState
	total    int
}

type choiceState struct {
	pending strings.Builder // content held, not yet judged
	// carried says pending begins with text whose events were already
	// replaced: it reaches the client only in a chunk this filter makes.
	carried bool
	calls   map[int]*call
}

type call struct {
	name string
	args strings.Builder
}

func (f *outFilter) Verdict() Verdict {
	v := f.v
	v.Took = f.took
	return v
}

func (f *outFilter) record(fs []finding) finding {
	best, found := strictest(fs)
	if found && best.action != Allow && rank(best.action) >= rank(f.v.Action) {
		f.v = Verdict{Action: best.action, Stage: Output, Policy: f.names, Category: best.category,
			Rule: best.rule, Shadow: f.shadow}
	}
	return best
}

// judgeText judges answer text, returning the decisive finding and the text
// to send in its place.
func (f *outFilter) judgeText(text string) (finding, string) {
	start := time.Now()
	defer func() { f.took += time.Since(start) }()
	f.total += len(text)
	fs, red := f.g.judge(text, f.p.Output)
	if f.p.Output.MaxChars > 0 && f.total > f.p.Output.MaxChars {
		fs = append(fs, finding{action: Block, category: "size", rule: "max_chars"})
	}
	best := f.record(fs)
	if f.shadow || len(red) == 0 {
		return best, text
	}
	f.g.countRedactions(red)
	return best, newPlaceholders().rewrite(text, red)
}

// judgeCalls checks an answer's tool calls: allowed names, how many, and
// arguments that fit what the request declared.
func (f *outFilter) judgeCalls(calls []call) finding {
	t := f.p.Tools
	var fs []finding
	if t.MaxCallsPerTurn > 0 && len(calls) > t.MaxCallsPerTurn {
		fs = append(fs, finding{action: Block, category: "tool", rule: "max_calls_per_turn"})
	}
	for _, c := range calls {
		if !toolAllowed(t, c.name) {
			fs = append(fs, finding{action: Block, category: "tool", rule: "tool:" + c.name})
			continue
		}
		if t.ValidateArguments {
			var v any
			if err := json.Unmarshal([]byte(c.args.String()), &v); err != nil {
				fs = append(fs, finding{action: Block, category: "tool", rule: "arguments:" + c.name})
				continue
			}
			if schema, ok := f.schemas[c.name]; ok {
				if err := validate(schema, v); err != nil {
					fs = append(fs, finding{action: Block, category: "tool", rule: "arguments:" + c.name})
				}
			}
		}
		// What a tool is told to do is answer text too.
		if fd, _ := f.judgeText(c.args.String()); fd.action == Block {
			fs = append(fs, fd)
		}
	}
	return f.record(fs)
}

// judgeEvery is how much new answer text a streamed answer's judge waits
// for before reading it again.
const judgeEvery = 600

// askJudge has the judge read the answer so far, when the policy asks and
// enough has arrived since it last did, or at the end.
func (f *outFilter) askJudge(text string, final bool) finding {
	jp := f.p.Judge
	if jp == nil || jp.Output == "" {
		return finding{}
	}
	f.sofar.WriteString(text)
	if !final && f.sofar.Len()-f.judged < judgeEvery || f.sofar.Len() == f.judged {
		return finding{}
	}
	f.judged = f.sofar.Len()
	start := time.Now()
	res, err := f.g.judger.ask(f.ctx, jp, f.user, f.sofar.String())
	f.took += time.Since(start)
	if err != nil {
		f.g.log.Warn("the judge gave no verdict on an answer", "judge", jp.Name(), "err", err)
		f.g.noJudger(jp, err)
	}
	jf := judgeFinding(f.p, jp.Output, res, err)
	if jf.failOpen {
		f.g.m.failOpen.Inc()
	}
	return f.record([]finding{jf})
}

// judgeStructured checks a whole answer against the schema the request
// asked it to fit.
func (f *outFilter) judgeStructured(text string, enforce bool) finding {
	if !f.p.ValidateStructured || f.answer == nil {
		return finding{}
	}
	var v any
	err := json.Unmarshal([]byte(text), &v)
	if err == nil {
		err = validate(f.answer, v)
	}
	if err == nil {
		return finding{}
	}
	a := Block
	if !enforce {
		a = Flag
	}
	return f.record([]finding{{action: a, category: "structured", rule: "response_format"}})
}

// --- buffered ---

func (f *outFilter) Buffered(body []byte) []byte {
	var top map[string]json.RawMessage
	if json.Unmarshal(body, &top) != nil {
		return body
	}
	var choices []map[string]json.RawMessage
	if json.Unmarshal(top["choices"], &choices) != nil {
		return body
	}
	blocked := false
	changed := false
	for _, ch := range choices {
		var msg map[string]json.RawMessage
		if json.Unmarshal(ch["message"], &msg) != nil {
			if t, ok := str(ch["text"]); ok { // a completion
				fd, out := f.judgeText(t)
				blocked = blocked || fd.action == Block
				if out != t {
					ch["text"] = rawString(out)
					changed = true
				}
			}
			continue
		}
		if t, ok := str(msg["content"]); ok {
			fd, out := f.judgeText(t)
			blocked = blocked || fd.action == Block
			if f.judgeStructured(t, true).action == Block {
				blocked = true
			}
			if f.askJudge(t, true).action == Block {
				blocked = true
			}
			if out != t {
				msg["content"] = rawString(out)
				ch["message"], _ = json.Marshal(msg)
				changed = true
			}
		}
		var tcs []struct {
			Function struct {
				Name      string `json:"name"`
				Arguments string `json:"arguments"`
			} `json:"function"`
		}
		if json.Unmarshal(msg["tool_calls"], &tcs) == nil && len(tcs) > 0 {
			calls := make([]call, len(tcs))
			for i, tc := range tcs {
				calls[i].name = tc.Function.Name
				calls[i].args.WriteString(tc.Function.Arguments)
			}
			if f.judgeCalls(calls).action == Block {
				blocked = true
			}
		}
	}
	if f.shadow {
		return body
	}
	if blocked {
		top["choices"] = json.RawMessage(`[{"index":0,"message":{"role":"assistant","content":""},"finish_reason":"content_filter"}]`)
		out, _ := json.Marshal(top)
		return out
	}
	if changed {
		top["choices"], _ = json.Marshal(choices)
		out, _ := json.Marshal(top)
		return out
	}
	return body
}

// --- streamed ---

type chunk struct {
	Choices []struct {
		Index int `json:"index"`
		Delta struct {
			Content   *string `json:"content"`
			ToolCalls []struct {
				Index    int `json:"index"`
				Function struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"delta"`
		Text         *string `json:"text"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
}

// releaseAt is how much held text is judged at once when no sentence ends.
const releaseAt = 200

func (f *outFilter) OnEvent(data string) ([]string, bool) {
	if f.stopped {
		return nil, true
	}
	var c chunk
	if json.Unmarshal([]byte(data), &c) != nil {
		return []string{data}, false
	}
	if f.template == nil {
		_ = json.Unmarshal([]byte(data), &f.template)
	}
	finished := false
	for _, ch := range c.Choices {
		st := f.choices[ch.Index]
		if st == nil {
			st = &choiceState{calls: map[int]*call{}}
			f.choices[ch.Index] = st
		}
		if ch.Delta.Content != nil {
			st.pending.WriteString(*ch.Delta.Content)
		}
		if ch.Text != nil {
			st.pending.WriteString(*ch.Text)
		}
		for _, tc := range ch.Delta.ToolCalls {
			cl := st.calls[tc.Index]
			if cl == nil {
				cl = &call{}
				st.calls[tc.Index] = cl
			}
			if tc.Function.Name != "" {
				cl.name = tc.Function.Name
			}
			cl.args.WriteString(tc.Function.Arguments)
		}
		if ch.FinishReason != nil && *ch.FinishReason != "" {
			finished = true
		}
	}
	if f.shadow {
		// Passed on at once; judged as it goes, to be recorded.
		if finished || f.ready() {
			f.flush(finished)
		}
		return []string{data}, false
	}
	f.held = append(f.held, data)
	if finished || f.ready() {
		return f.release(finished)
	}
	return nil, false
}

// ready reports whether held text has reached a place to judge it: a
// finished sentence, a line, or releaseAt characters — never while a tool
// call is still arriving, nor when the whole answer is to be held.
func (f *outFilter) ready() bool {
	if f.p.Stream == "buffer_all" {
		return false
	}
	for _, st := range f.choices {
		if len(st.calls) > 0 {
			return false
		}
	}
	for _, st := range f.choices {
		p := st.pending.String()
		if len(p) >= releaseAt || strings.Contains(p, "\n") ||
			strings.Contains(p, ". ") || strings.Contains(p, "! ") || strings.Contains(p, "? ") {
			return true
		}
	}
	return false
}

// flush judges what each choice holds — up to its last whitespace, keeping
// a word that may not be whole, unless final — and with final its tool calls.
// It returns whether anything blocks, each choice's text to send, and
// whether that differs from what the held events carry.
func (f *outFilter) flush(final bool) (bool, map[int]string, bool) {
	blocked, differs := false, false
	out := map[int]string{}
	for i, st := range f.choices {
		t := st.pending.String()
		head, tail := t, ""
		if !final {
			cut := strings.LastIndexAny(t, " \n\t") + 1
			head, tail = t[:cut], t[cut:]
		}
		st.pending.Reset()
		st.pending.WriteString(tail)
		if tail != "" || st.carried {
			differs = true
		}
		st.carried = tail != ""
		if f.answer != nil {
			f.full.WriteString(head)
		}
		if i == 0 && f.askJudge(head, final).action == Block {
			blocked = true
		}
		if head != "" {
			fd, r := f.judgeText(head)
			blocked = blocked || fd.action == Block
			if r != head {
				differs = true
			}
			out[i] = r
		}
		if final && i == 0 && f.answer != nil {
			// Held whole, it can still be stopped; released a sentence at
			// a time, it is flagged.
			if f.judgeStructured(f.full.String(), f.p.Stream == "buffer_all").action == Block {
				blocked = true
			}
		}
		if final && len(st.calls) > 0 {
			calls := make([]call, 0, len(st.calls))
			for j := 0; j < len(st.calls); j++ {
				if c := st.calls[j]; c != nil {
					calls = append(calls, *c)
				}
			}
			if f.judgeCalls(calls).action == Block {
				blocked = true
			}
			st.calls = map[int]*call{}
		}
	}
	return blocked, out, differs
}

func (f *outFilter) release(final bool) ([]string, bool) {
	blocked, texts, differs := f.flush(final)
	held := f.held
	f.held = nil
	if blocked {
		f.stopped = true
		return []string{f.synth(map[int]string{0: ""}, "content_filter")}, true
	}
	if !differs {
		return held, false
	}
	// Rewritten, or cut short of a word still arriving: sent as one chunk
	// per choice in place of the held ones, and what else they carried —
	// tool calls, a finish — after.
	var out []string
	if len(texts) > 0 {
		out = append(out, f.synth(texts, ""))
	}
	for _, h := range held {
		var c chunk
		if json.Unmarshal([]byte(h), &c) != nil {
			continue
		}
		for _, ch := range c.Choices {
			if len(ch.Delta.ToolCalls) > 0 || ch.FinishReason != nil && *ch.FinishReason != "" {
				out = append(out, withoutContent(h))
				break
			}
		}
	}
	return out, false
}

// synth is a chunk like the stream's own, carrying text per choice, or a
// finish reason.
func (f *outFilter) synth(texts map[int]string, finish string) string {
	top := map[string]json.RawMessage{}
	for k, v := range f.template {
		if k != "choices" && k != "usage" {
			top[k] = v
		}
	}
	var choices []map[string]any
	for i, t := range texts {
		delta := map[string]any{"content": t}
		if !f.sent {
			delta["role"] = "assistant"
		}
		c := map[string]any{"index": i, "delta": delta, "finish_reason": nil}
		if finish != "" {
			c["delta"] = map[string]any{}
			c["finish_reason"] = finish
		}
		choices = append(choices, c)
	}
	f.sent = true
	top["choices"], _ = json.Marshal(choices)
	b, _ := json.Marshal(top)
	return string(b)
}

// withoutContent is a chunk with its text taken out, sent after a chunk
// that already carried it rewritten.
func withoutContent(data string) string {
	var top map[string]json.RawMessage
	if json.Unmarshal([]byte(data), &top) != nil {
		return data
	}
	var choices []map[string]json.RawMessage
	if json.Unmarshal(top["choices"], &choices) != nil {
		return data
	}
	for _, c := range choices {
		var d map[string]json.RawMessage
		if json.Unmarshal(c["delta"], &d) == nil {
			delete(d, "content")
			c["delta"], _ = json.Marshal(d)
		}
		delete(c, "text")
	}
	top["choices"], _ = json.Marshal(choices)
	b, _ := json.Marshal(top)
	return string(b)
}

func (f *outFilter) OnEnd() []string {
	switch {
	case f.stopped:
		return nil
	case f.shadow:
		f.flush(true)
		return nil
	}
	out, _ := f.release(true)
	return out
}