jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-06-03 09:37:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

Optimize the method for judging streaming responses (#1017)

Browse Source
This commit is contained in:
澄潭
2024-05-30 17:50:28 +08:00
committed by GitHub
parent 26ec0d3d55
commit 38ddc49360
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
plugins/wasm-go/extensions/ai-cache/main.go
View File

@@ -30,6 +30,7 @@ func main() {
wrapper.ParseConfigBy(parseConfig), wrapper.ParseConfigBy(parseConfig),
wrapper.ProcessRequestHeadersBy(onHttpRequestHeaders), wrapper.ProcessRequestHeadersBy(onHttpRequestHeaders),
wrapper.ProcessRequestBodyBy(onHttpRequestBody), wrapper.ProcessRequestBodyBy(onHttpRequestBody),
wrapper.ProcessResponseHeadersBy(onHttpResponseHeaders),
wrapper.ProcessStreamingResponseBodyBy(onHttpResponseBody), wrapper.ProcessStreamingResponseBodyBy(onHttpResponseBody),
) )
} }
@@ -181,12 +182,6 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config PluginConfig, log wrap
ctx.DontReadRequestBody() ctx.DontReadRequestBody()
return types.ActionContinue return types.ActionContinue
} }
// compatiable with qwen
x_dashscope_sse, _ := proxywasm.GetHttpRequestHeader("X-DashScope-SSE")
accept, _ := proxywasm.GetHttpRequestHeader("Accept")
if x_dashscope_sse == "enable" || strings.Contains(accept, "text/event-stream") {
ctx.SetContext(StreamContextKey, struct{}{})
}
proxywasm.RemoveHttpRequestHeader("Accept-Encoding") proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
// The request has a body and requires delaying the header transmission until a cache miss occurs, // The request has a body and requires delaying the header transmission until a cache miss occurs,
// at which point the header should be sent. // at which point the header should be sent.
@@ -274,6 +269,14 @@ func processSSEMessage(ctx wrapper.HttpContext, config PluginConfig, sseMessage
return "" return ""
} }
func onHttpResponseHeaders(ctx wrapper.HttpContext, config PluginConfig, log wrapper.Log) types.Action {
contentType, _ := proxywasm.GetHttpResponseHeader("content-type")
if strings.Contains(contentType, "text/event-stream") {
ctx.SetContext(StreamContextKey, struct{}{})
}
return types.ActionContinue
}
func onHttpResponseBody(ctx wrapper.HttpContext, config PluginConfig, chunk []byte, isLastChunk bool, log wrapper.Log) []byte { func onHttpResponseBody(ctx wrapper.HttpContext, config PluginConfig, chunk []byte, isLastChunk bool, log wrapper.Log) []byte {
if ctx.GetContext(ToolCallsContextKey) != nil { if ctx.GetContext(ToolCallsContextKey) != nil {
// we should not cache tool call result // we should not cache tool call result