mirror of
https://github.com/alibaba/higress.git
synced 2026-02-24 20:50:51 +08:00
254 lines
8.5 KiB
Go
254 lines
8.5 KiB
Go
// File generated by hgctl. Modify as required.
|
|
// See: https://higress.io/zh-cn/docs/user/wasm-go#2-%E7%BC%96%E5%86%99-maingo-%E6%96%87%E4%BB%B6
|
|
|
|
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"net/url"
|
|
"strings"
|
|
|
|
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/config"
|
|
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/provider"
|
|
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/util"
|
|
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
|
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
|
|
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
const (
|
|
pluginName = "ai-proxy"
|
|
|
|
ctxKeyApiName = "apiName"
|
|
|
|
defaultMaxBodyBytes uint32 = 10 * 1024 * 1024
|
|
)
|
|
|
|
func main() {
|
|
wrapper.SetCtx(
|
|
pluginName,
|
|
wrapper.ParseOverrideConfigBy(parseGlobalConfig, parseOverrideRuleConfig),
|
|
wrapper.ProcessRequestHeadersBy(onHttpRequestHeader),
|
|
wrapper.ProcessRequestBodyBy(onHttpRequestBody),
|
|
wrapper.ProcessResponseHeadersBy(onHttpResponseHeaders),
|
|
wrapper.ProcessStreamingResponseBodyBy(onStreamingResponseBody),
|
|
wrapper.ProcessResponseBodyBy(onHttpResponseBody),
|
|
)
|
|
}
|
|
|
|
func parseGlobalConfig(json gjson.Result, pluginConfig *config.PluginConfig, log wrapper.Log) error {
|
|
//log.Debugf("loading global config: %s", json.String())
|
|
|
|
pluginConfig.FromJson(json)
|
|
if err := pluginConfig.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if err := pluginConfig.Complete(); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func parseOverrideRuleConfig(json gjson.Result, global config.PluginConfig, pluginConfig *config.PluginConfig, log wrapper.Log) error {
|
|
//log.Debugf("loading override rule config: %s", json.String())
|
|
|
|
*pluginConfig = global
|
|
|
|
pluginConfig.FromJson(json)
|
|
if err := pluginConfig.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if err := pluginConfig.Complete(); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConfig, log wrapper.Log) types.Action {
|
|
activeProvider := pluginConfig.GetProvider()
|
|
|
|
if activeProvider == nil {
|
|
log.Debugf("[onHttpRequestHeader] no active provider, skip processing")
|
|
ctx.DontReadRequestBody()
|
|
return types.ActionContinue
|
|
}
|
|
|
|
log.Debugf("[onHttpRequestHeader] provider=%s", activeProvider.GetProviderType())
|
|
|
|
rawPath := ctx.Path()
|
|
path, _ := url.Parse(rawPath)
|
|
apiName := getOpenAiApiName(path.Path)
|
|
providerConfig := pluginConfig.GetProviderConfig()
|
|
if apiName == "" && !providerConfig.IsOriginal() {
|
|
log.Debugf("[onHttpRequestHeader] unsupported path: %s", path.Path)
|
|
_ = util.SendResponse(404, "ai-proxy.unknown_api", util.MimeTypeTextPlain, "API not found: "+path.Path)
|
|
return types.ActionContinue
|
|
}
|
|
ctx.SetContext(ctxKeyApiName, apiName)
|
|
|
|
if handler, ok := activeProvider.(provider.RequestHeadersHandler); ok {
|
|
// Disable the route re-calculation since the plugin may modify some headers related to the chosen route.
|
|
ctx.DisableReroute()
|
|
hasRequestBody := wrapper.HasRequestBody()
|
|
action, err := handler.OnRequestHeaders(ctx, apiName, log)
|
|
if err == nil {
|
|
if hasRequestBody {
|
|
ctx.SetRequestBodyBufferLimit(defaultMaxBodyBytes)
|
|
// Always return types.HeaderStopIteration to support fallback routing,
|
|
// as long as onHttpRequestBody can be called.
|
|
return types.HeaderStopIteration
|
|
}
|
|
return action
|
|
}
|
|
_ = util.SendResponse(500, "ai-proxy.proc_req_headers_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to process request headers: %v", err))
|
|
return types.ActionContinue
|
|
}
|
|
|
|
return types.ActionContinue
|
|
}
|
|
|
|
func onHttpRequestBody(ctx wrapper.HttpContext, pluginConfig config.PluginConfig, body []byte, log wrapper.Log) types.Action {
|
|
activeProvider := pluginConfig.GetProvider()
|
|
|
|
if activeProvider == nil {
|
|
log.Debugf("[onHttpRequestBody] no active provider, skip processing")
|
|
return types.ActionContinue
|
|
}
|
|
|
|
log.Debugf("[onHttpRequestBody] provider=%s", activeProvider.GetProviderType())
|
|
|
|
if handler, ok := activeProvider.(provider.RequestBodyHandler); ok {
|
|
apiName, _ := ctx.GetContext(ctxKeyApiName).(provider.ApiName)
|
|
|
|
newBody, settingErr := pluginConfig.GetProviderConfig().ReplaceByCustomSettings(body)
|
|
if settingErr != nil {
|
|
_ = util.SendResponse(500, "ai-proxy.proc_req_body_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to rewrite request body by custom settings: %v", settingErr))
|
|
return types.ActionContinue
|
|
}
|
|
|
|
log.Debugf("[onHttpRequestBody] newBody=%s", newBody)
|
|
body = newBody
|
|
action, err := handler.OnRequestBody(ctx, apiName, body, log)
|
|
if err == nil {
|
|
return action
|
|
}
|
|
_ = util.SendResponse(500, "ai-proxy.proc_req_body_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to process request body: %v", err))
|
|
}
|
|
return types.ActionContinue
|
|
}
|
|
|
|
func onHttpResponseHeaders(ctx wrapper.HttpContext, pluginConfig config.PluginConfig, log wrapper.Log) types.Action {
|
|
if !wrapper.IsResponseFromUpstream() {
|
|
// Response is not coming from the upstream. Let it pass through.
|
|
ctx.DontReadResponseBody()
|
|
return types.ActionContinue
|
|
}
|
|
|
|
activeProvider := pluginConfig.GetProvider()
|
|
|
|
if activeProvider == nil {
|
|
log.Debugf("[onHttpResponseHeaders] no active provider, skip processing")
|
|
ctx.DontReadResponseBody()
|
|
return types.ActionContinue
|
|
}
|
|
|
|
log.Debugf("[onHttpResponseHeaders] provider=%s", activeProvider.GetProviderType())
|
|
|
|
status, err := proxywasm.GetHttpResponseHeader(":status")
|
|
if err != nil || status != "200" {
|
|
if err != nil {
|
|
log.Errorf("unable to load :status header from response: %v", err)
|
|
}
|
|
ctx.DontReadResponseBody()
|
|
return types.ActionContinue
|
|
}
|
|
|
|
if handler, ok := activeProvider.(provider.ResponseHeadersHandler); ok {
|
|
apiName, _ := ctx.GetContext(ctxKeyApiName).(provider.ApiName)
|
|
action, err := handler.OnResponseHeaders(ctx, apiName, log)
|
|
if err == nil {
|
|
checkStream(&ctx, &log)
|
|
return action
|
|
}
|
|
_ = util.SendResponse(500, "ai-proxy.proc_resp_headers_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to process response headers: %v", err))
|
|
return types.ActionContinue
|
|
}
|
|
|
|
checkStream(&ctx, &log)
|
|
_, needHandleBody := activeProvider.(provider.ResponseBodyHandler)
|
|
_, needHandleStreamingBody := activeProvider.(provider.StreamingResponseBodyHandler)
|
|
if !needHandleBody && !needHandleStreamingBody {
|
|
ctx.DontReadResponseBody()
|
|
} else if !needHandleStreamingBody {
|
|
ctx.BufferResponseBody()
|
|
}
|
|
|
|
return types.ActionContinue
|
|
}
|
|
|
|
func onStreamingResponseBody(ctx wrapper.HttpContext, pluginConfig config.PluginConfig, chunk []byte, isLastChunk bool, log wrapper.Log) []byte {
|
|
activeProvider := pluginConfig.GetProvider()
|
|
|
|
if activeProvider == nil {
|
|
log.Debugf("[onStreamingResponseBody] no active provider, skip processing")
|
|
return chunk
|
|
}
|
|
|
|
log.Debugf("[onStreamingResponseBody] provider=%s", activeProvider.GetProviderType())
|
|
log.Debugf("isLastChunk=%v chunk: %s", isLastChunk, string(chunk))
|
|
|
|
if handler, ok := activeProvider.(provider.StreamingResponseBodyHandler); ok {
|
|
apiName, _ := ctx.GetContext(ctxKeyApiName).(provider.ApiName)
|
|
modifiedChunk, err := handler.OnStreamingResponseBody(ctx, apiName, chunk, isLastChunk, log)
|
|
if err == nil && modifiedChunk != nil {
|
|
return modifiedChunk
|
|
}
|
|
return chunk
|
|
}
|
|
return chunk
|
|
}
|
|
|
|
func onHttpResponseBody(ctx wrapper.HttpContext, pluginConfig config.PluginConfig, body []byte, log wrapper.Log) types.Action {
|
|
activeProvider := pluginConfig.GetProvider()
|
|
|
|
if activeProvider == nil {
|
|
log.Debugf("[onHttpResponseBody] no active provider, skip processing")
|
|
return types.ActionContinue
|
|
}
|
|
|
|
log.Debugf("[onHttpResponseBody] provider=%s", activeProvider.GetProviderType())
|
|
//log.Debugf("response body: %s", string(body))
|
|
|
|
if handler, ok := activeProvider.(provider.ResponseBodyHandler); ok {
|
|
apiName, _ := ctx.GetContext(ctxKeyApiName).(provider.ApiName)
|
|
action, err := handler.OnResponseBody(ctx, apiName, body, log)
|
|
if err == nil {
|
|
return action
|
|
}
|
|
_ = util.SendResponse(500, "ai-proxy.proc_resp_body_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to process response body: %v", err))
|
|
return types.ActionContinue
|
|
}
|
|
return types.ActionContinue
|
|
}
|
|
|
|
func getOpenAiApiName(path string) provider.ApiName {
|
|
if strings.HasSuffix(path, "/v1/chat/completions") {
|
|
return provider.ApiNameChatCompletion
|
|
}
|
|
if strings.HasSuffix(path, "/v1/embeddings") {
|
|
return provider.ApiNameEmbeddings
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func checkStream(ctx *wrapper.HttpContext, log *wrapper.Log) {
|
|
contentType, err := proxywasm.GetHttpResponseHeader("Content-Type")
|
|
if err != nil || !strings.HasPrefix(contentType, "text/event-stream") {
|
|
if err != nil {
|
|
log.Errorf("unable to load content-type header from response: %v", err)
|
|
}
|
|
(*ctx).BufferResponseBody()
|
|
}
|
|
}
|