@@ -81,11 +81,11 @@ type Response struct {
8181}
8282
8383type Data struct {
84- RiskLevel string `json:"RiskLevel"`
85- AttackLevel string `json:"AttackLevel,omitempty"`
86- Result []Result `json:"Result,omitempty"`
87- Advice []Advice `json:"Advice,omitempty"`
88- Detail []Detail `json:"Detail,omitempty"`
84+ RiskLevel string `json:"RiskLevel"`
85+ AttackLevel string `json:"AttackLevel,omitempty"`
86+ Result []Result `json:"Result,omitempty"`
87+ Advice []Advice `json:"Advice,omitempty"`
88+ Detail []Detail `json:"Detail,omitempty"`
8989}
9090
9191type Result struct {
@@ -123,6 +123,7 @@ type AISecurityConfig struct {
123123 denyCode int64
124124 denyMessage string
125125 protocolOriginal bool
126+ riskLevelBar string
126127 contentModerationLevelBar string
127128 promptAttackLevelBar string
128129 sensitiveDataLevelBar string
@@ -224,7 +225,7 @@ func isRiskLevelAcceptable(action string, data Data, config AISecurityConfig) bo
224225 }
225226 return true
226227 } else {
227- return levelToInt (data .RiskLevel ) < levelToInt (config .contentModerationLevelBar )
228+ return levelToInt (data .RiskLevel ) < levelToInt (config .riskLevelBar )
228229 }
229230}
230231
@@ -281,8 +282,17 @@ func parseConfig(json gjson.Result, config *AISecurityConfig) error {
281282 if config .ak == "" || config .sk == "" {
282283 return errors .New ("invalid AK/SK config" )
283284 }
285+ if obj := json .Get ("riskLevelBar" ); obj .Exists () {
286+ config .riskLevelBar = obj .String ()
287+ } else {
288+ config .riskLevelBar = HighRisk
289+ }
284290 config .token = json .Get ("securityToken" ).String ()
285- config .action = json .Get ("action" ).String ()
291+ if obj := json .Get ("action" ); obj .Exists () {
292+ config .action = json .Get ("action" ).String ()
293+ } else {
294+ config .action = "TextModerationPlus"
295+ }
286296 config .checkRequest = json .Get ("checkRequest" ).Bool ()
287297 config .checkResponse = json .Get ("checkResponse" ).Bool ()
288298 config .protocolOriginal = json .Get ("protocol" ).String () == "original"
0 commit comments