add plugin: ai-statistics (#1011)

plugins/wasm-go/extensions/ai-statistics/main.go

@@ -0,0 +1,143 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+
+	"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"
+)
+
+func main() {
+	wrapper.SetCtx(
+		"ai-statistics",
+		wrapper.ParseConfigBy(parseConfig),
+		wrapper.ProcessResponseHeadersBy(onHttpResponseHeaders),
+		wrapper.ProcessStreamingResponseBodyBy(onHttpStreamingBody),
+		wrapper.ProcessResponseBodyBy(onHttpResponseBody),
+	)
+}
+
+type AIStatisticsConfig struct {
+	enable  bool
+	metrics map[string]proxywasm.MetricCounter
+}
+
+func (config *AIStatisticsConfig) incrementCounter(metricName string, inc uint64, log wrapper.Log) {
+	counter, ok := config.metrics[metricName]
+	if !ok {
+		counter = proxywasm.DefineCounterMetric(metricName)
+		config.metrics[metricName] = counter
+	}
+	counter.Increment(inc)
+}
+
+func parseConfig(json gjson.Result, config *AIStatisticsConfig, log wrapper.Log) error {
+	config.enable = json.Get("enable").Bool()
+	config.metrics = make(map[string]proxywasm.MetricCounter)
+	return nil
+}
+
+func onHttpResponseHeaders(ctx wrapper.HttpContext, config AIStatisticsConfig, log wrapper.Log) types.Action {
+	if !config.enable {
+		ctx.DontReadResponseBody()
+		return types.ActionContinue
+	}
+	contentType, _ := proxywasm.GetHttpResponseHeader("content-type")
+	if !strings.Contains(contentType, "text/event-stream") {
+		ctx.BufferResponseBody()
+	}
+	return types.ActionContinue
+}
+
+func getLastChunk(data []byte) []byte {
+	chunks := strings.Split(strings.TrimSpace(string(data)), "\n\n")
+	length := len(chunks)
+	if length < 2 {
+		return data
+	}
+	// ai-proxy append extra usage chunk
+	return []byte(chunks[length-1])
+}
+
+func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, data []byte, endOfStream bool, log wrapper.Log) []byte {
+	lastChunk := getLastChunk(data)
+	modelObj := gjson.GetBytes(lastChunk, "model")
+	inputTokenObj := gjson.GetBytes(lastChunk, "usage.prompt_tokens")
+	outputTokenObj := gjson.GetBytes(lastChunk, "usage.completion_tokens")
+	if modelObj.Exists() && inputTokenObj.Exists() && outputTokenObj.Exists() {
+		ctx.SetContext("model", modelObj.String())
+		ctx.SetContext("input_token", inputTokenObj.Int())
+		ctx.SetContext("output_token", outputTokenObj.Int())
+	}
+
+	if endOfStream {
+		var route, cluster string
+		if raw, err := proxywasm.GetProperty([]string{"route_name"}); err == nil {
+			route = string(raw)
+		}
+		if raw, err := proxywasm.GetProperty([]string{"cluster_name"}); err == nil {
+			cluster = string(raw)
+		}
+		model, ok := ctx.GetContext("model").(string)
+		if !ok {
+			log.Error("Get model failed!")
+			return data
+		}
+		inputToken, ok := ctx.GetContext("input_token").(int64)
+		if !ok {
+			log.Error("Get input_token failed!")
+			return data
+		}
+		outputToken, ok := ctx.GetContext("output_token").(int64)
+		if !ok {
+			log.Error("Get output_token failed!")
+			return data
+		}
+		config.incrementCounter("route."+route+".upstream."+cluster+".model."+model+".input_token", uint64(inputToken), log)
+		config.incrementCounter("route."+route+".upstream."+cluster+".model."+model+".output_token", uint64(outputToken), log)
+		proxywasm.SetProperty([]string{"model"}, []byte(model))
+		proxywasm.SetProperty([]string{"input_token"}, []byte(fmt.Sprint(inputToken)))
+		proxywasm.SetProperty([]string{"output_token"}, []byte(fmt.Sprint(outputToken)))
+	}
+
+	return data
+}
+
+func onHttpResponseBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body []byte, log wrapper.Log) types.Action {
+	modeObj := gjson.GetBytes(body, "model")
+	inputTokenObj := gjson.GetBytes(body, "usage.prompt_tokens")
+	outputTokenObj := gjson.GetBytes(body, "usage.completion_tokens")
+	if !modeObj.Exists() {
+		log.Error("Get model failed")
+		return types.ActionContinue
+	}
+	if !inputTokenObj.Exists() {
+		log.Error("Get input_token failed")
+		return types.ActionContinue
+	}
+	if !outputTokenObj.Exists() {
+		log.Error("Get output_token failed")
+		return types.ActionContinue
+	}
+	model := modeObj.String()
+	inputToken := inputTokenObj.Int()
+	outputToken := outputTokenObj.Int()
+	var route, cluster string
+	if raw, err := proxywasm.GetProperty([]string{"route_name"}); err == nil {
+		route = string(raw)
+	}
+	if raw, err := proxywasm.GetProperty([]string{"cluster_name"}); err == nil {
+		cluster = string(raw)
+	}
+	config.incrementCounter("route."+route+".upstream."+cluster+".model."+model+".input_token", uint64(inputToken), log)
+	config.incrementCounter("route."+route+".upstream."+cluster+".model."+model+".output_token", uint64(outputToken), log)
+
+	proxywasm.SetProperty([]string{"model"}, []byte(model))
+	proxywasm.SetProperty([]string{"input_token"}, []byte(fmt.Sprint(inputToken)))
+	proxywasm.SetProperty([]string{"output_token"}, []byte(fmt.Sprint(outputToken)))
+
+	return types.ActionContinue
+}