feat: custom sensitive words

Author: Kerwin

service/src/chatgpt/index.ts

@@ -90,9 +90,6 @@ async function chatReplyProcess(options: RequestOptions) {
   const model = isNotEmptyString(config.apiModel) ? config.apiModel : 'gpt-3.5-turbo'
   const { message, lastContext, process, systemMessage, temperature, top_p } = options
 
-  if ((config.auditConfig?.enabled ?? false) && !await auditText(config.auditConfig, message))
-    return sendResponse({ type: 'Fail', message: '含有敏感词 | Contains sensitive words' })
-
   try {
     const timeoutMs = (await getCacheConfig()).timeoutMs
     let options: SendMessageOptions = { timeoutMs }
@@ -135,11 +132,19 @@ export function initAuditService(audit: AuditConfig) {
   auditService = new Service(audit.options)
 }
 
-async function auditText(audit: AuditConfig, text: string): Promise<boolean> {
-  if (!auditService)
-    initAuditService(audit)
-
-  return await auditService.audit(text)
+async function containsSensitiveWords(audit: AuditConfig, text: string): Promise<boolean> {
+  if (audit.customizeEnabled && isNotEmptyString(audit.sensitiveWords)) {
+    const textLower = text.toLowerCase()
+    const notSafe = audit.sensitiveWords.split('\n').filter(d => textLower.includes(d.trim().toLowerCase())).length > 0
+    if (notSafe)
+      return true
+  }
+  if (audit.enabled) {
+    if (!auditService)
+      initAuditService(audit)
+    return await auditService.containsSensitiveWords(text)
+  }
+  return false
 }
 let cachedBanlance: number | undefined
 let cacheExpiration = 0
@@ -193,7 +198,7 @@ async function fetchBalance() {
 
     // 计算剩余额度
     cachedBanlance = totalAmount - totalUsage
-    cacheExpiration = now + 10 * 60 * 1000
+    cacheExpiration = now + 60 * 60 * 1000
 
     return Promise.resolve(cachedBanlance.toFixed(3))
   }
@@ -254,4 +259,4 @@ initApi()
 
 export type { ChatContext, ChatMessage }
 
-export { chatReplyProcess, chatConfig, currentModel, auditText }
+export { chatReplyProcess, chatConfig, currentModel, containsSensitiveWords }

service/src/index.ts

@@ -3,7 +3,7 @@ import jwt from 'jsonwebtoken'
 import * as dotenv from 'dotenv'
 import type { RequestProps } from './types'
 import type { ChatContext, ChatMessage } from './chatgpt'
-import { auditText, chatConfig, chatReplyProcess, currentModel, initApi, initAuditService } from './chatgpt'
+import { chatConfig, chatReplyProcess, containsSensitiveWords, currentModel, initApi, initAuditService } from './chatgpt'
 import { auth } from './middleware/auth'
 import { clearConfigCache, getCacheConfig, getOriginConfig } from './storage/config'
 import type { AuditConfig, ChatInfo, ChatOptions, Config, MailConfig, SiteConfig, UsageResponse, UserInfo } from './storage/model'
@@ -277,6 +277,16 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
   let result
   let message: ChatInfo
   try {
+    const config = await getCacheConfig()
+    if (config.auditConfig.enabled || config.auditConfig.customizeEnabled) {
+      const userId = req.headers.userId.toString()
+      const user = await getUserById(userId)
+      if (user.email.toLowerCase() !== process.env.ROOT_USER && await containsSensitiveWords(config.auditConfig, prompt)) {
+        res.send({ status: 'Fail', message: '含有敏感词 | Contains sensitive words', data: null })
+        return
+      }
+    }
+
     message = regenerate
       ? await getChat(roomId, uuid)
       : await insertChat(uuid, prompt, roomId, options as ChatOptions)
@@ -602,7 +612,8 @@ router.post('/setting-audit', rootAuth, async (req, res) => {
     thisConfig.auditConfig = config
     const result = await updateConfig(thisConfig)
     clearConfigCache()
-    initAuditService(config)
+    if (config.enabled)
+      initAuditService(config)
     res.send({ status: 'Success', message: '操作成功 | Successfully', data: result.auditConfig })
   }
   catch (error) {
@@ -614,10 +625,12 @@ router.post('/audit-test', rootAuth, async (req, res) => {
   try {
     const { audit, text } = req.body as { audit: AuditConfig; text: string }
     const config = await getCacheConfig()
-    initAuditService(audit)
-    const result = await auditText(audit, text)
-    initAuditService(config.auditConfig)
-    res.send({ status: 'Success', message: !result ? '含敏感词 | Contains sensitive words' : '不含敏感词 | Does not contain sensitive words.', data: null })
+    if (audit.enabled)
+      initAuditService(audit)
+    const result = await containsSensitiveWords(audit, text)
+    if (audit.enabled)
+      initAuditService(config.auditConfig)
+    res.send({ status: 'Success', message: result ? '含敏感词 | Contains sensitive words' : '不含敏感词 | Does not contain sensitive words.', data: null })
   }
   catch (error) {
     res.send({ status: 'Fail', message: error.message, data: null })

service/src/storage/model.ts

@@ -162,6 +162,8 @@ export class AuditConfig {
     public provider: TextAuditServiceProvider,
     public options: TextAuditServiceOptions,
     public textType: TextAudioType,
+    public customizeEnabled: boolean,
+    public sensitiveWords: string,
   ) { }
 }
 

service/src/utils/textAudit.ts

@@ -7,7 +7,7 @@ export interface TextAuditServiceOptions {
 }
 
 export interface TextAuditService {
-  audit(text: string): Promise<boolean>
+  containsSensitiveWords(text: string): Promise<boolean>
 }
 
 /**
@@ -19,23 +19,23 @@ export class BaiduTextAuditService implements TextAuditService {
 
   constructor(private options: TextAuditServiceOptions) { }
 
-  async audit(text: string): Promise<boolean> {
+  async containsSensitiveWords(text: string): Promise<boolean> {
     if (!await this.refreshAccessToken())
-      return
+      throw new Error('Access Token Error')
     const url = `https://aip.baidubce.com/rest/2.0/solution/v1/text_censor/v2/user_defined?access_token=${this.accessToken}`
     let headers: {
       'Content-Type': 'application/x-www-form-urlencoded'
       'Accept': 'application/json'
     }
     const response = await fetch(url, { headers, method: 'POST', body: `text=${encodeURIComponent(text)}` })
-    const data = await response.json() as { conclusionType: number; data: any }
+    const data = await response.json() as { conclusionType: number; data: any; error_msg: string }
 
     if (data.error_msg)
       throw new Error(data.error_msg)
 
     // 审核结果类型，可取值1、2、3、4，分别代表1：合规，2：不合规，3：疑似，4：审核失败
     if (data.conclusionType === 1)
-      return true
+      return false
 
     // https://ai.baidu.com/ai-doc/ANTIPORN/Nk3h6xbb2#%E7%BB%86%E5%88%86%E6%A0%87%E7%AD%BE%E5%AF%B9%E7%85%A7%E8%A1%A8
 
@@ -46,9 +46,9 @@ export class BaiduTextAuditService implements TextAuditService {
     const str = JSON.stringify(data)
     for (const l of this.options.label.split(',')) {
       if (str.indexOf(l))
-        return false
+        return true
     }
-    return true
+    return false
   }
 
   async refreshAccessToken() {

src/components/common/Setting/Audit.vue

@@ -64,6 +64,27 @@ onMounted(() => {
   <NSpin :show="loading">
     <div class="p-4 space-y-5 min-h-[200px]">
       <div class="space-y-6">
+        <div class="flex items-center space-x-4">
+          <span class="flex-shrink-0 w-[100px]">{{ $t('setting.auditCustomizeEnabled') }}</span>
+          <div class="flex-1">
+            <NSwitch
+              :round="false" :value="config && config.customizeEnabled"
+              @update:value="(val) => { if (config) config.customizeEnabled = val }"
+            />
+          </div>
+        </div>
+        <div v-if="config && config.customizeEnabled" class="flex items-center space-x-4">
+          <span class="flex-shrink-0 w-[100px]">{{ $t('setting.auditCustomizeWords') }}</span>
+          <div class="flex-1">
+            <NInput
+              :value="config && config.sensitiveWords"
+              placeholder="One word per line"
+              type="textarea"
+              :autosize="{ minRows: 1, maxRows: 4 }"
+              @input="(val) => { if (config) config.sensitiveWords = val }"
+            />
+          </div>
+        </div>
         <div class="flex items-center space-x-4">
           <span class="flex-shrink-0 w-[100px]">{{ $t('setting.auditEnabled') }}</span>
           <div class="flex-1">
@@ -117,7 +138,7 @@ onMounted(() => {
             <a target="_blank" href="https://ai.baidu.com/ai-doc/ANTIPORN/Nk3h6xbb2#%E7%BB%86%E5%88%86%E6%A0%87%E7%AD%BE%E5%AF%B9%E7%85%A7%E8%A1%A8">{{ $t('setting.auditBaiduLabelLink') }}</a>
           </p>
         </div>
-        <div v-if="config && config.enabled" class="flex items-center space-x-4">
+        <div v-if="config && (config.enabled || config.customizeEnabled)" class="flex items-center space-x-4">
           <span class="flex-shrink-0 w-[100px]">{{ $t('setting.auditTest') }}</span>
           <div class="flex-1">
             <NInput

src/components/common/Setting/index.vue

@@ -40,7 +40,7 @@ const show = computed({
 </script>
 
 <template>
-  <NModal v-model:show="show" :auto-focus="false" preset="card" style="width: 95%; max-width: 640px">
+  <NModal v-model:show="show" :auto-focus="false" preset="card" style="width: 95%; max-width: 1024px">
     <div>
       <NTabs v-model:value="active" type="line" animated>
         <NTabPane name="General" tab="General">

src/components/common/Setting/model.ts

@@ -50,4 +50,6 @@ export class AuditConfig {
   provider?: TextAuditServiceProvider
   options?: TextAuditServiceOptions
   textType?: TextAudioType
+  customizeEnabled?: boolean
+  sensitiveWords?: string
 }

src/locales/en-US.ts

@@ -103,14 +103,16 @@ export default {
     loginSalt: 'Login Salt',
     loginSaltTip: 'Changes will invalidate all logged in',
     monthlyUsage: 'Monthly Usage',
-    auditEnabled: 'Audit Enabled',
+    auditEnabled: 'Third Party',
     auditProvider: 'Provider',
     auditApiKey: 'Api Key',
     auditApiSecret: 'Api Secret',
     auditTest: 'Test Text',
     auditBaiduLabel: 'Label',
     auditBaiduLabelTip: 'English comma separated, If empty, only politics.',
     auditBaiduLabelLink: 'Goto Label Detail',
+    auditCustomizeEnabled: 'Customize',
+    auditCustomizeWords: 'Sensitive Words',
   },
   store: {
     siderButton: 'Prompt Store',

src/locales/ko-KR.ts

@@ -101,14 +101,16 @@ export default {
     loginSalt: '로그인 정보 암호화 Salt',
     loginSaltTip: '변경하면 모든사용자의 로그인이 풀립니다.',
     monthlyUsage: '월간 사용량',
-    auditEnabled: '승인 상태',
+    auditEnabled: '타사',
     auditProvider: '공급자',
     auditApiKey: 'Api Key',
     auditApiSecret: 'Api Secret',
     auditTest: '테스트 텍스트',
     auditBaiduLabel: 'Label',
     auditBaiduLabelTip: '영어 쉼표로 구분, If empty, only politics.',
     auditBaiduLabelLink: '레이블 세부 정보로 이동',
+    auditCustomizeEnabled: '맞춤화하다',
+    auditCustomizeWords: '단어 맞춤설정',
   },
   store: {
     siderButton: '프롬프트 스토어',

src/locales/zh-CN.ts

@@ -103,14 +103,16 @@ export default {
     loginSalt: '登录混淆盐',
     loginSaltTip: '变更会导致所有已登录失效',
     monthlyUsage: '本月使用量',
-    auditEnabled: '审核状态',
+    auditEnabled: '第三方',
     auditProvider: '提供商',
     auditApiKey: 'Api Key',
     auditApiSecret: 'Api Secret',
     auditTest: '测试文本',
     auditBaiduLabel: '二级分类',
     auditBaiduLabelTip: '英文逗号分隔, 如果空, 仅政治',
     auditBaiduLabelLink: '查看细分类型',
+    auditCustomizeEnabled: '自定义',
+    auditCustomizeWords: '敏感词',
   },
   store: {
     siderButton: '提示词商店',