jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-05-26 21:57:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(ai-security-guard): replace denyMessage with structured DenyResponseBody (#3642)

Browse Source 
Co-authored-by: rinfx <yucheng.lxr@alibaba-inc.com>
This commit is contained in:
JianweiWang
2026-04-01 19:38:01 +08:00
committed by GitHub
parent 89587c1c9b
commit 1c9e981bf2
10 changed files with 820 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
plugins/wasm-go/extensions/ai-security-guard/lvwang/common/text/openai.go
View File

@@ -65,13 +65,19 @@ func HandleTextGenerationStreamingResponseBody(ctx wrapper.HttpContext, config c
return
}
if !cfg.IsRiskLevelAcceptable(config.Action, response.Data, config, consumer) {
denyMessage := cfg.DefaultDenyMessage
if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
denyMessage = "\n" + response.Data.Advice[0].Answer
} else if config.DenyMessage != "" {
denyMessage = config.DenyMessage
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
if err != nil {
log.Errorf("failed to build deny response body: %v", err)
endStream := ctx.GetContext("end_of_stream_received").(bool) && ctx.BufferQueueSize() == 0
proxywasm.InjectEncodedDataToFilterChain(bytes.Join(bufferQueue, []byte("")), endStream)
bufferQueue = [][]byte{}
if !endStream {
ctx.SetContext("during_call", false)
singleCall()
}
return
}
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
randomID := utils.GenerateRandomChatID()
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
proxywasm.InjectEncodedDataToFilterChain(jsonData, true)
@@ -199,21 +205,22 @@ func HandleTextGenerationResponseBody(ctx wrapper.HttpContext, config cfg.AISecu
}
return
}
denyMessage := cfg.DefaultDenyMessage
if config.DenyMessage != "" {
denyMessage = config.DenyMessage
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
denyMessage = response.Data.Advice[0].Answer
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
if err != nil {
log.Errorf("failed to build deny response body: %v", err)
proxywasm.ResumeHttpResponse()
return
}
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
if config.ProtocolOriginal {
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
} else if isStreamingResponse {
randomID := utils.GenerateRandomChatID()
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
} else {
randomID := utils.GenerateRandomChatID()
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
}