refactor: 重构敏感词

service/sensitive.go

@@ -1,13 +1,60 @@
 package service
 
 import (
-	"bytes"
+	"errors"
 	"fmt"
-	"github.com/anknown/ahocorasick"
+	"one-api/common"
 	"one-api/constant"
+	"one-api/dto"
 	"strings"
 )
 
+func CheckSensitiveMessages(messages []dto.Message) error {
+	for _, message := range messages {
+		if len(message.Content) > 0 {
+			if message.IsStringContent() {
+				stringContent := message.StringContent()
+				if ok, words := SensitiveWordContains(stringContent); ok {
+					return errors.New("sensitive words: " + strings.Join(words, ","))
+				}
+			}
+		} else {
+			arrayContent := message.ParseContent()
+			for _, m := range arrayContent {
+				if m.Type == "image_url" {
+					// TODO: check image url
+				} else {
+					if ok, words := SensitiveWordContains(m.Text); ok {
+						return errors.New("sensitive words: " + strings.Join(words, ","))
+					}
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func CheckSensitiveText(text string) error {
+	if ok, words := SensitiveWordContains(text); ok {
+		return errors.New("sensitive words: " + strings.Join(words, ","))
+	}
+	return nil
+}
+
+func CheckSensitiveInput(input any) error {
+	switch v := input.(type) {
+	case string:
+		return CheckSensitiveText(v)
+	case []string:
+		text := ""
+		for _, s := range v {
+			text += s
+		}
+		return CheckSensitiveText(text)
+	}
+	return CheckSensitiveText(fmt.Sprintf("%v", input))
+}
+
 // SensitiveWordContains 是否包含敏感词，返回是否包含敏感词和敏感词列表
 func SensitiveWordContains(text string) (bool, []string) {
 	if len(constant.SensitiveWords) == 0 {
@@ -15,7 +62,7 @@ func SensitiveWordContains(text string) (bool, []string) {
 	}
 	checkText := strings.ToLower(text)
 	// 构建一个AC自动机
-	m := initAc()
+	m := common.InitAc()
 	hits := m.MultiPatternSearch([]rune(checkText), false)
 	if len(hits) > 0 {
 		words := make([]string, 0)
@@ -33,7 +80,7 @@ func SensitiveWordReplace(text string, returnImmediately bool) (bool, []string,
 		return false, nil, text
 	}
 	checkText := strings.ToLower(text)
-	m := initAc()
+	m := common.InitAc()
 	hits := m.MultiPatternSearch([]rune(checkText), returnImmediately)
 	if len(hits) > 0 {
 		words := make([]string, 0)
@@ -47,25 +94,3 @@ func SensitiveWordReplace(text string, returnImmediately bool) (bool, []string,
 	}
 	return false, nil, text
 }
-
-func initAc() *goahocorasick.Machine {
-	m := new(goahocorasick.Machine)
-	dict := readRunes()
-	if err := m.Build(dict); err != nil {
-		fmt.Println(err)
-		return nil
-	}
-	return m
-}
-
-func readRunes() [][]rune {
-	var dict [][]rune
-
-	for _, word := range constant.SensitiveWords {
-		word = strings.ToLower(word)
-		l := bytes.TrimSpace([]byte(word))
-		dict = append(dict, bytes.Runes(l))
-	}
-
-	return dict
-}