add plugin: ai-statistics (#1011)

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

@@ -0,0 +1,2 @@
+main.wasm
+config.yaml

plugins/wasm-go/extensions/ai-statistics/README.md

@@ -0,0 +1,44 @@
+# 介绍
+提供AI可观测基础能力，其后需接ai-proxy插件，如果不接ai-proxy插件的话，则只支持openai协议。
+
+# 配置说明
+
+| 名称         | 数据类型   | 填写要求 | 默认值 | 描述               |
+|------------|--------|------|-----|------------------|
+| `enable` | bool | 必填   | -   | 是否开启ai统计功能 |
+
+开启后 metrics 示例：
+```
+route_upstream_model_input_token{ai_route="openai",ai_cluster="qwen",ai_model="qwen-max"} 21
+route_upstream_model_output_token{ai_route="openai",ai_cluster="qwen",ai_model="qwen-max"} 17
+```
+
+日志示例：
+
+```json
+{
+    "model": "qwen-max",
+    "input_token": "21",
+    "output_token": "17",
+    "authority": "dashscope.aliyuncs.com",
+    "bytes_received": "336",
+    "bytes_sent": "1675",
+    "duration": "1590",
+    "istio_policy_status": "-",
+    "method": "POST",
+    "path": "/v1/chat/completions",
+    "protocol": "HTTP/1.1",
+    "request_id": "5895f5a9-e4e3-425b-98db-6c6a926195b7",
+    "requested_server_name": "-",
+    "response_code": "200",
+    "response_flags": "-",
+    "route_name": "openai",
+    "start_time": "2024-06-18T09:37:14.078Z",
+    "trace_id": "-",
+    "upstream_cluster": "qwen",
+    "upstream_service_time": "496",
+    "upstream_transport_failure_reason": "-",
+    "user_agent": "PostmanRuntime/7.37.3",
+    "x_forwarded_for": "-"
+}
+```

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

@@ -0,0 +1,21 @@
+module ai-statistics
+
+go 1.18
+
+require (
+	github.com/alibaba/higress/plugins/wasm-go v1.3.6-0.20240522012622-fc6a6aad8906
+	github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240327114451-d6b7174a84fc
+	github.com/tidwall/gjson v1.14.3
+)
+
+require github.com/tetratelabs/wazero v1.7.1 // indirect
+
+require (
+	github.com/google/uuid v1.3.0 // indirect
+	github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 // indirect
+	github.com/magefile/mage v1.14.0 // indirect
+	github.com/tidwall/match v1.1.1 // indirect
+	github.com/tidwall/pretty v1.2.0 // indirect
+	github.com/tidwall/resp v0.1.1 // indirect
+	github.com/wasilibs/go-re2 v1.5.3
+)

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

@@ -0,0 +1,26 @@
+github.com/alibaba/higress/plugins/wasm-go v1.3.6-0.20240522012622-fc6a6aad8906 h1:RhEmB+ApLKsClZD7joTC4ifmsVgOVz4pFLdPR3xhNaE=
+github.com/alibaba/higress/plugins/wasm-go v1.3.6-0.20240522012622-fc6a6aad8906/go.mod h1:10jQXKsYFUF7djs+Oy7t82f4dbie9pISfP9FJwpPLuk=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 h1:IHDghbGQ2DTIXHBHxWfqCYQW1fKjyJ/I7W1pMyUDeEA=
+github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520/go.mod h1:Nz8ORLaFiLWotg6GeKlJMhv8cci8mM43uEnLA5t8iew=
+github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240327114451-d6b7174a84fc h1:t2AT8zb6N/59Y78lyRWedVoVWHNRSCBh0oWCC+bluTQ=
+github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240327114451-d6b7174a84fc/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
+github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/tetratelabs/wazero v1.7.1 h1:QtSfd6KLc41DIMpDYlJdoMc6k7QTN246DM2+n2Y/Dx8=
+github.com/tetratelabs/wazero v1.7.1/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=
+github.com/tidwall/gjson v1.14.3 h1:9jvXn7olKEHU1S9vwoMGliaT8jq1vJ7IH/n9zD9Dnlw=
+github.com/tidwall/gjson v1.14.3/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
+github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
+github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
+github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
+github.com/tidwall/resp v0.1.1 h1:Ly20wkhqKTmDUPlyM1S7pWo5kk0tDu8OoC/vFArXmwE=
+github.com/tidwall/resp v0.1.1/go.mod h1:3/FrruOBAxPTPtundW0VXgmsQ4ZBA0Aw714lVYgwFa0=
+github.com/wasilibs/go-re2 v1.5.3 h1:wiuTcgDZdLhu8NG8oqF5sF5Q3yIU14lPAvXqeYzDK3g=
+github.com/wasilibs/go-re2 v1.5.3/go.mod h1:PzpVPsBdFC7vM8QJbbEnOeTmwA0DGE783d/Gex8eCV8=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

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
+}