From 133a30b8d5980cc30197bf9cf15c0b53e9d35dc5 Mon Sep 17 00:00:00 2001 From: Se7en Date: Wed, 22 Jan 2025 11:28:37 +0800 Subject: [PATCH] fix: stream response buffer issue (#1703) --- plugins/wasm-go/extensions/ai-proxy/main.go | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go index 88b85949b..220243f98 100644 --- a/plugins/wasm-go/extensions/ai-proxy/main.go +++ b/plugins/wasm-go/extensions/ai-proxy/main.go @@ -95,10 +95,9 @@ func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConf // Disable the route re-calculation since the plugin may modify some headers related to the chosen route. ctx.DisableReroute() - _, needHandleStreamingBody := activeProvider.(provider.StreamingResponseBodyHandler) - if needHandleStreamingBody { - proxywasm.RemoveHttpRequestHeader("Accept-Encoding") - } + // Always remove the Accept-Encoding header to prevent the LLM from sending compressed responses, + // allowing plugins to inspect or modify the response correctly + proxywasm.RemoveHttpRequestHeader("Accept-Encoding") if handler, ok := activeProvider.(provider.RequestHeadersHandler); ok { // Set the apiToken for the current request. @@ -200,8 +199,11 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, pluginConfig config.PluginCo util.ReplaceResponseHeaders(headers) checkStream(ctx, log) + _, needHandleBody := activeProvider.(provider.TransformResponseBodyHandler) _, needHandleStreamingBody := activeProvider.(provider.StreamingResponseBodyHandler) - if !needHandleStreamingBody { + if !needHandleBody && !needHandleStreamingBody { + ctx.DontReadResponseBody() + } else if !needHandleStreamingBody { ctx.BufferResponseBody() }