ai-gateway-go/internal/contentpolicy/output.go

package contentpolicy

import (
	"bytes"
	"encoding/json"
	"strings"
)

// OutputRedact 对一条模型响应 JSON(完整响应或 SSE data 载荷)应用全部
// redact 策略的规则:提取 choices[].message.content / choices[].delta.content
// 文本并替换敏感信息。与输入侧不同,输出 JSON 是给客户端消费的展示数据,
// 重编码可接受。返回替换后的字节与是否发生替换。
func (e *Engine) OutputRedact(payload []byte) ([]byte, bool) {
	if e == nil || len(payload) == 0 {
		return payload, false
	}
	rules := e.outputRules()
	if len(rules) == 0 {
		return payload, false
	}
	var document any
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.UseNumber()
	if decoder.Decode(&document) != nil {
		return payload, false
	}
	changed := false
	redactTree(&document, &changed, rules)
	if !changed {
		return payload, false
	}
	encoded, err := json.Marshal(document)
	if err != nil {
		return payload, false
	}
	return encoded, true
}

// outputRules 返回全部启用策略的 redact 规则(输出侧不区分端点)。
func (e *Engine) outputRules() []compiledRule {
	current := e.current.Load()
	if current == nil {
		return nil
	}
	var rules []compiledRule
	for _, policy := range current.policies {
		if policy.Action == "redact" {
			rules = append(rules, policy.rules...)
		}
	}
	return rules
}

// redactTree 递归遍历 JSON,仅对 content 文本应用规则。
func redactTree(value *any, changed *bool, rules []compiledRule) {
	switch current := (*value).(type) {
	case map[string]any:
		for key, child := range current {
			local := child
			if strings.ToLower(key) == "content" {
				if text, ok := local.(string); ok {
					next := text
					for _, rule := range rules {
						if rule.expression.MatchString(next) {
							next = rule.expression.ReplaceAllString(next, rule.replacement)
						}
					}
					if next != text {
						*changed = true
						local = next
					}
				}
			}
			redactTree(&local, changed, rules)
			current[key] = local
		}
	case []any:
		for index, child := range current {
			local := child
			redactTree(&local, changed, rules)
			current[index] = local
		}
	}
}