feat: Enhance sensitive word detection with detailed logging

This commit is contained in:
1808837298@qq.com
2025-02-21 16:57:30 +08:00
parent 2d42145b66
commit 9cc6385b0c
4 changed files with 28 additions and 32 deletions

View File

@@ -13,6 +13,7 @@ import (
"one-api/relay/helper" "one-api/relay/helper"
"one-api/service" "one-api/service"
"one-api/setting" "one-api/setting"
"strings"
) )
func getAndValidAudioRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.AudioRequest, error) { func getAndValidAudioRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.AudioRequest, error) {
@@ -27,8 +28,9 @@ func getAndValidAudioRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.
return nil, errors.New("model is required") return nil, errors.New("model is required")
} }
if setting.ShouldCheckPromptSensitive() { if setting.ShouldCheckPromptSensitive() {
err := service.CheckSensitiveInput(audioRequest.Input) words, err := service.CheckSensitiveInput(audioRequest.Input)
if err != nil { if err != nil {
common.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ",")))
return nil, err return nil, err
} }
} }

View File

@@ -61,8 +61,9 @@ func getAndValidImageRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.
// return service.OpenAIErrorWrapper(errors.New("n must be between 1 and 10"), "invalid_field_value", http.StatusBadRequest) // return service.OpenAIErrorWrapper(errors.New("n must be between 1 and 10"), "invalid_field_value", http.StatusBadRequest)
//} //}
if setting.ShouldCheckPromptSensitive() { if setting.ShouldCheckPromptSensitive() {
err := service.CheckSensitiveInput(imageRequest.Prompt) words, err := service.CheckSensitiveInput(imageRequest.Prompt)
if err != nil { if err != nil {
common.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ",")))
return nil, err return nil, err
} }
} }

View File

@@ -78,8 +78,9 @@ func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
} }
if setting.ShouldCheckPromptSensitive() { if setting.ShouldCheckPromptSensitive() {
err = checkRequestSensitive(textRequest, relayInfo) words, err := checkRequestSensitive(textRequest, relayInfo)
if err != nil { if err != nil {
common.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
return service.OpenAIErrorWrapperLocal(err, "sensitive_words_detected", http.StatusBadRequest) return service.OpenAIErrorWrapperLocal(err, "sensitive_words_detected", http.StatusBadRequest)
} }
} }
@@ -219,19 +220,20 @@ func getPromptTokens(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.Re
return promptTokens, err return promptTokens, err
} }
func checkRequestSensitive(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) error { func checkRequestSensitive(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) ([]string, error) {
var err error var err error
var words []string
switch info.RelayMode { switch info.RelayMode {
case relayconstant.RelayModeChatCompletions: case relayconstant.RelayModeChatCompletions:
err = service.CheckSensitiveMessages(textRequest.Messages) words, err = service.CheckSensitiveMessages(textRequest.Messages)
case relayconstant.RelayModeCompletions: case relayconstant.RelayModeCompletions:
err = service.CheckSensitiveInput(textRequest.Prompt) words, err = service.CheckSensitiveInput(textRequest.Prompt)
case relayconstant.RelayModeModerations: case relayconstant.RelayModeModerations:
err = service.CheckSensitiveInput(textRequest.Input) words, err = service.CheckSensitiveInput(textRequest.Input)
case relayconstant.RelayModeEmbeddings: case relayconstant.RelayModeEmbeddings:
err = service.CheckSensitiveInput(textRequest.Input) words, err = service.CheckSensitiveInput(textRequest.Input)
} }
return err return words, err
} }
// 预扣费并返回用户剩余配额 // 预扣费并返回用户剩余配额

View File

@@ -8,39 +8,30 @@ import (
"strings" "strings"
) )
func CheckSensitiveMessages(messages []dto.Message) error { func CheckSensitiveMessages(messages []dto.Message) ([]string, error) {
for _, message := range messages { for _, message := range messages {
if len(message.Content) > 0 { arrayContent := message.ParseContent()
if message.IsStringContent() { for _, m := range arrayContent {
stringContent := message.StringContent() if m.Type == "image_url" {
if ok, words := SensitiveWordContains(stringContent); ok { // TODO: check image url
return errors.New("sensitive words: " + strings.Join(words, ",")) } else {
} if ok, words := SensitiveWordContains(m.Text); ok {
} return words, errors.New("sensitive words detected")
} 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 return nil, nil
} }
func CheckSensitiveText(text string) error { func CheckSensitiveText(text string) ([]string, error) {
if ok, words := SensitiveWordContains(text); ok { if ok, words := SensitiveWordContains(text); ok {
return errors.New("sensitive words: " + strings.Join(words, ",")) return words, errors.New("sensitive words detected")
} }
return nil return nil, nil
} }
func CheckSensitiveInput(input any) error { func CheckSensitiveInput(input any) ([]string, error) {
switch v := input.(type) { switch v := input.(type) {
case string: case string:
return CheckSensitiveText(v) return CheckSensitiveText(v)
@@ -60,7 +51,7 @@ func SensitiveWordContains(text string) (bool, []string) {
return false, nil return false, nil
} }
checkText := strings.ToLower(text) checkText := strings.ToLower(text)
return AcSearch(checkText, setting.SensitiveWords, false) return AcSearch(checkText, setting.SensitiveWords, true)
} }
// SensitiveWordReplace 敏感词替换,返回是否包含敏感词和替换后的文本 // SensitiveWordReplace 敏感词替换,返回是否包含敏感词和替换后的文本