AI observability upgrade (#1587)

Co-authored-by: Kent Dong <ch3cho@qq.com>

plugins/wasm-go/extensions/ai-cache/core.go

@@ -74,6 +74,9 @@ func processCacheHit(key string, response string, stream bool, ctx wrapper.HttpC
 
 	ctx.SetContext(CACHE_KEY_CONTEXT_KEY, nil)
 
+	ctx.SetUserAttribute("cache_status", "hit")
+	ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+
 	if stream {
 		proxywasm.SendHttpResponseWithDetail(200, "ai-cache.hit", [][2]string{{"content-type", "text/event-stream; charset=utf-8"}}, []byte(fmt.Sprintf(c.StreamResponseTemplate, escapedResponse)), -1)
 	} else {

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

@@ -8,14 +8,14 @@ replace github.com/alibaba/higress/plugins/wasm-go => ../..
 
 require (
 	github.com/alibaba/higress/plugins/wasm-go v1.4.2
-	github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f
+	github.com/google/uuid v1.6.0
+	github.com/higress-group/proxy-wasm-go-sdk v1.0.0
 	github.com/tidwall/gjson v1.17.3
 	github.com/tidwall/resp v0.1.1
 // github.com/weaviate/weaviate-go-client/v4 v4.15.1
 )
 
 require (
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 // indirect
 	github.com/magefile/mage v1.14.0 // indirect
 	github.com/stretchr/testify v1.9.0 // indirect

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

@@ -3,8 +3,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.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-20240711023527-ba358c48772f h1:ZIiIBRvIw62gA5MJhuwp1+2wWbqL9IGElQ499rUsYYg=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/higress-group/proxy-wasm-go-sdk v1.0.0 h1:BZRNf4R7jr9hwRivg/E29nkVaKEak5MWjBDhWjuHijU=
+github.com/higress-group/proxy-wasm-go-sdk v1.0.0/go.mod h1:iiSyFbo+rAtbtGt/bsefv8GU57h9CCLYGJA74/tF5/0=
 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=

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

@@ -128,9 +128,15 @@ func onHttpRequestBody(ctx wrapper.HttpContext, c config.PluginConfig, body []by
 func onHttpResponseHeaders(ctx wrapper.HttpContext, c config.PluginConfig, log wrapper.Log) types.Action {
 	skipCache := ctx.GetContext(SKIP_CACHE_HEADER)
 	if skipCache != nil {
+		ctx.SetUserAttribute("cache_status", "skip")
+		ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 		ctx.DontReadResponseBody()
 		return types.ActionContinue
 	}
+	if ctx.GetContext(CACHE_KEY_CONTEXT_KEY) != nil {
+		ctx.SetUserAttribute("cache_status", "miss")
+		ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+	}
 	contentType, _ := proxywasm.GetHttpResponseHeader("content-type")
 	if strings.Contains(contentType, "text/event-stream") {
 		ctx.SetContext(STREAM_CONTEXT_KEY, struct{}{})

plugins/wasm-go/extensions/ai-security-guard/README.md

@@ -31,6 +31,7 @@ description: 阿里云内容安全检测
 | `denyMessage` | string | optional | openai格式的流式/非流式响应 | 指定内容非法时的响应内容 |
 | `protocol` | string | optional | openai | 协议格式，非openai协议填`original` |
 | `riskLevelBar` | string | optional | high | 拦截风险等级，取值为 max, high, medium, low |
+| `timeout` | int | optional | 2000 | 调用内容安全服务时的超时时间 |
 
 补充说明一下 `denyMessage`，对非法请求的处理逻辑为：
 - 如果配置了 `denyMessage`，返回内容为 `denyMessage` 配置内容，格式为openai格式的流式/非流式响应

plugins/wasm-go/extensions/ai-security-guard/main.go

@@ -53,6 +53,7 @@ const (
 	DefaultStreamingResponseJsonPath = "choices.0.delta.content"
 	DefaultDenyCode                  = 200
 	DefaultDenyMessage               = "很抱歉，我无法回答您的问题"
+	DefaultTimeout                   = 2000
 
 	AliyunUserAgent = "CIPFrom/AIGateway"
 	LengthLimit     = 1800
@@ -100,6 +101,7 @@ type AISecurityConfig struct {
 	denyMessage                   string
 	protocolOriginal              bool
 	riskLevelBar                  string
+	timeout                       uint32
 	metrics                       map[string]proxywasm.MetricCounter
 }
 
@@ -225,6 +227,11 @@ func parseConfig(json gjson.Result, config *AISecurityConfig, log wrapper.Log) e
 	} else {
 		config.riskLevelBar = HighRisk
 	}
+	if obj := json.Get("timeout"); obj.Exists() {
+		config.timeout = uint32(obj.Int())
+	} else {
+		config.timeout = DefaultTimeout
+	}
 	config.client = wrapper.NewClusterClient(wrapper.FQDNCluster{
 		FQDN: serviceName,
 		Port: servicePort,
@@ -253,6 +260,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log
 
 func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []byte, log wrapper.Log) types.Action {
 	log.Debugf("checking request body...")
+	startTime := time.Now().UnixMilli()
 	content := gjson.GetBytes(body, config.requestContentJsonPath).String()
 	model := gjson.GetBytes(body, "model").String()
 	ctx.SetContext("requestModel", model)
@@ -279,6 +287,10 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
 		}
 		if riskLevelToInt(response.Data.RiskLevel) < riskLevelToInt(config.riskLevelBar) {
 			if contentIndex >= len(content) {
+				endTime := time.Now().UnixMilli()
+				ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
+				ctx.SetUserAttribute("safecheck_status", "request pass")
+				ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 				proxywasm.ResumeHttpRequest()
 			} else {
 				singleCall()
@@ -305,7 +317,9 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
 		}
 		ctx.DontReadResponseBody()
 		config.incrementCounter("ai_sec_request_deny", 1)
-		ctx.SetUserAttribute("safecheck_status", "request deny")
+		endTime := time.Now().UnixMilli()
+		ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
+		ctx.SetUserAttribute("safecheck_status", "reqeust deny")
 		if response.Data.Advice != nil {
 			ctx.SetUserAttribute("safecheck_riskLabel", response.Data.Result[0].Label)
 			ctx.SetUserAttribute("safecheck_riskWords", response.Data.Result[0].RiskWords)
@@ -345,7 +359,7 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
 			reqParams.Add(k, v)
 		}
 		reqParams.Add("Signature", signature)
-		err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil, callback)
+		err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil, callback, config.timeout)
 		if err != nil {
 			log.Errorf("failed call the safe check service: %v", err)
 			proxywasm.ResumeHttpRequest()
@@ -364,20 +378,6 @@ func convertHeaders(hs [][2]string) map[string][]string {
 	return ret
 }
 
-// headers: map[string][]string -> [][2]string
-func reconvertHeaders(hs map[string][]string) [][2]string {
-	var ret [][2]string
-	for k, vs := range hs {
-		for _, v := range vs {
-			ret = append(ret, [2]string{k, v})
-		}
-	}
-	sort.SliceStable(ret, func(i, j int) bool {
-		return ret[i][0] < ret[j][0]
-	})
-	return ret
-}
-
 func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log wrapper.Log) types.Action {
 	if !config.checkResponse {
 		log.Debugf("response checking is disabled")
@@ -401,6 +401,7 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log
 
 func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body []byte, log wrapper.Log) types.Action {
 	log.Debugf("checking response body...")
+	startTime := time.Now().UnixMilli()
 	hdsMap := ctx.GetContext("headers").(map[string][]string)
 	isStreamingResponse := strings.Contains(strings.Join(hdsMap["content-type"], ";"), "event-stream")
 	model := ctx.GetStringContext("requestModel", "unknown")
@@ -433,6 +434,10 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
 		}
 		if riskLevelToInt(response.Data.RiskLevel) < riskLevelToInt(config.riskLevelBar) {
 			if contentIndex >= len(content) {
+				endTime := time.Now().UnixMilli()
+				ctx.SetUserAttribute("safecheck_response_rt", endTime-startTime)
+				ctx.SetUserAttribute("safecheck_status", "response pass")
+				ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 				proxywasm.ResumeHttpResponse()
 			} else {
 				singleCall()
@@ -458,6 +463,8 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
 			proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
 		}
 		config.incrementCounter("ai_sec_response_deny", 1)
+		endTime := time.Now().UnixMilli()
+		ctx.SetUserAttribute("safecheck_response_rt", endTime-startTime)
 		ctx.SetUserAttribute("safecheck_status", "response deny")
 		if response.Data.Advice != nil {
 			ctx.SetUserAttribute("safecheck_riskLabel", response.Data.Result[0].Label)
@@ -498,7 +505,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
 			reqParams.Add(k, v)
 		}
 		reqParams.Add("Signature", signature)
-		err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil, callback)
+		err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil, callback, config.timeout)
 		if err != nil {
 			log.Errorf("failed call the safe check service: %v", err)
 			proxywasm.ResumeHttpResponse()

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

@@ -3,15 +3,13 @@ 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-20240711023527-ba358c48772f h1:ZIiIBRvIw62gA5MJhuwp1+2wWbqL9IGElQ499rUsYYg=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/higress-group/proxy-wasm-go-sdk v1.0.0 h1:BZRNf4R7jr9hwRivg/E29nkVaKEak5MWjBDhWjuHijU=
 github.com/higress-group/proxy-wasm-go-sdk v1.0.0/go.mod h1:iiSyFbo+rAtbtGt/bsefv8GU57h9CCLYGJA74/tF5/0=
 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/tidwall/gjson v1.14.3 h1:9jvXn7olKEHU1S9vwoMGliaT8jq1vJ7IH/n9zD9Dnlw=
-github.com/tidwall/gjson v1.14.3/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/gjson v1.17.3 h1:bwWLZU7icoKRG+C+0PNwIKC6FCJO/Q3p2pZvuP0jN94=
 github.com/tidwall/gjson v1.17.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=

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

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"strconv"
 	"strings"
 	"time"
 
@@ -28,14 +27,15 @@ func main() {
 }
 
 const (
-	// Trace span prefix
-	TracePrefix = "trace_span_tag."
 	// Context consts
 	StatisticsRequestStartTime = "ai-statistics-request-start-time"
 	StatisticsFirstTokenTime   = "ai-statistics-first-token-time"
 	CtxGeneralAtrribute        = "attributes"
 	CtxLogAtrribute            = "logAttributes"
 	CtxStreamingBodyBuffer     = "streamingBodyBuffer"
+	RouteName                  = "route"
+	ClusterName                = "cluster"
+	APIName                    = "api"
 
 	// Source Type
 	FixedValue            = "fixed_value"
@@ -46,12 +46,14 @@ const (
 	ResponseBody          = "response_body"
 
 	// Inner metric & log attributes name
-	Model                 = "model"
-	InputToken            = "input_token"
-	OutputToken           = "output_token"
-	LLMFirstTokenDuration = "llm_first_token_duration"
-	LLMServiceDuration    = "llm_service_duration"
-	LLMDurationCount      = "llm_duration_count"
+	Model                  = "model"
+	InputToken             = "input_token"
+	OutputToken            = "output_token"
+	LLMFirstTokenDuration  = "llm_first_token_duration"
+	LLMServiceDuration     = "llm_service_duration"
+	LLMDurationCount       = "llm_duration_count"
+	LLMStreamDurationCount = "llm_stream_duration_count"
+	ResponseType           = "response_type"
 
 	// Extract Rule
 	RuleFirst   = "first"
@@ -91,6 +93,19 @@ func getRouteName() (string, error) {
 	}
 }
 
+func getAPIName() (string, error) {
+	if raw, err := proxywasm.GetProperty([]string{"route_name"}); err != nil {
+		return "-", err
+	} else {
+		parts := strings.Split(string(raw), "@")
+		if len(parts) != 5 {
+			return "-", errors.New("not api type")
+		} else {
+			return strings.Join(parts[:3], "@"), nil
+		}
+	}
+}
+
 func getClusterName() (string, error) {
 	if raw, err := proxywasm.GetProperty([]string{"cluster_name"}); err != nil {
 		return "-", err
@@ -133,8 +148,15 @@ func parseConfig(configJson gjson.Result, config *AIStatisticsConfig, log wrappe
 }
 
 func onHttpRequestHeaders(ctx wrapper.HttpContext, config AIStatisticsConfig, log wrapper.Log) types.Action {
-	ctx.SetContext(CtxGeneralAtrribute, map[string]string{})
-	ctx.SetContext(CtxLogAtrribute, map[string]string{})
+	route, _ := getRouteName()
+	cluster, _ := getClusterName()
+	api, api_error := getAPIName()
+	if api_error == nil {
+		route = api
+	}
+	ctx.SetContext(RouteName, route)
+	ctx.SetContext(ClusterName, cluster)
+	ctx.SetUserAttribute(APIName, api)
 	ctx.SetContext(StatisticsRequestStartTime, time.Now().UnixMilli())
 
 	// Set user defined log & span attributes which type is fixed_value
@@ -149,6 +171,9 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config AIStatisticsConfig, lo
 func onHttpRequestBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body []byte, log wrapper.Log) types.Action {
 	// Set user defined log & span attributes.
 	setAttributeBySource(ctx, config, RequestBody, body, log)
+
+	// Write log
+	ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 	return types.ActionContinue
 }
 
@@ -177,6 +202,8 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat
 		ctx.SetContext(CtxStreamingBodyBuffer, streamingBodyBuffer)
 	}
 
+	ctx.SetUserAttribute(ResponseType, "stream")
+
 	// Get requestStartTime from http context
 	requestStartTime, ok := ctx.GetContext(StatisticsRequestStartTime).(int64)
 	if !ok {
@@ -188,28 +215,19 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat
 	if ctx.GetContext(StatisticsFirstTokenTime) == nil {
 		firstTokenTime := time.Now().UnixMilli()
 		ctx.SetContext(StatisticsFirstTokenTime, firstTokenTime)
-		attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-		attributes[LLMFirstTokenDuration] = fmt.Sprint(firstTokenTime - requestStartTime)
-		ctx.SetContext(CtxGeneralAtrribute, attributes)
+		ctx.SetUserAttribute(LLMFirstTokenDuration, firstTokenTime-requestStartTime)
 	}
 
 	// Set information about this request
-
 	if model, inputToken, outputToken, ok := getUsage(data); ok {
-		attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-		// Record Log Attributes
-		attributes[Model] = model
-		attributes[InputToken] = fmt.Sprint(inputToken)
-		attributes[OutputToken] = fmt.Sprint(outputToken)
-		// Set attributes to http context
-		ctx.SetContext(CtxGeneralAtrribute, attributes)
+		ctx.SetUserAttribute(Model, model)
+		ctx.SetUserAttribute(InputToken, inputToken)
+		ctx.SetUserAttribute(OutputToken, outputToken)
 	}
 	// If the end of the stream is reached, record metrics/logs/spans.
 	if endOfStream {
 		responseEndTime := time.Now().UnixMilli()
-		attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-		attributes[LLMServiceDuration] = fmt.Sprint(responseEndTime - requestStartTime)
-		ctx.SetContext(CtxGeneralAtrribute, attributes)
+		ctx.SetUserAttribute(LLMServiceDuration, responseEndTime-requestStartTime)
 
 		// Set user defined log & span attributes.
 		if config.shouldBufferStreamingBody {
@@ -220,11 +238,8 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat
 			setAttributeBySource(ctx, config, ResponseStreamingBody, streamingBodyBuffer, log)
 		}
 
-		// Write inner filter states which can be used by other plugins such as ai-token-ratelimit
-		writeFilterStates(ctx, log)
-
 		// Write log
-		writeLog(ctx, log)
+		ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 
 		// Write metrics
 		writeMetric(ctx, config, log)
@@ -233,33 +248,26 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat
 }
 
 func onHttpResponseBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body []byte, log wrapper.Log) types.Action {
-	// Get attributes from http context
-	attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-
 	// Get requestStartTime from http context
 	requestStartTime, _ := ctx.GetContext(StatisticsRequestStartTime).(int64)
 
 	responseEndTime := time.Now().UnixMilli()
-	attributes[LLMServiceDuration] = fmt.Sprint(responseEndTime - requestStartTime)
+	ctx.SetUserAttribute(LLMServiceDuration, responseEndTime-requestStartTime)
+
+	ctx.SetUserAttribute(ResponseType, "normal")
 
 	// Set information about this request
-	model, inputToken, outputToken, ok := getUsage(body)
-	if ok {
-		attributes[Model] = model
-		attributes[InputToken] = fmt.Sprint(inputToken)
-		attributes[OutputToken] = fmt.Sprint(outputToken)
-		// Update attributes
-		ctx.SetContext(CtxGeneralAtrribute, attributes)
+	if model, inputToken, outputToken, ok := getUsage(body); ok {
+		ctx.SetUserAttribute(Model, model)
+		ctx.SetUserAttribute(InputToken, inputToken)
+		ctx.SetUserAttribute(OutputToken, outputToken)
 	}
 
 	// Set user defined log & span attributes.
 	setAttributeBySource(ctx, config, ResponseBody, body, log)
 
-	// Write inner filter states which can be used by other plugins such as ai-token-ratelimit
-	writeFilterStates(ctx, log)
-
 	// Write log
-	writeLog(ctx, log)
+	ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 
 	// Write metrics
 	writeMetric(ctx, config, log)
@@ -294,57 +302,45 @@ func getUsage(data []byte) (model string, inputTokenUsage int64, outputTokenUsag
 
 // fetches the tracing span value from the specified source.
 func setAttributeBySource(ctx wrapper.HttpContext, config AIStatisticsConfig, source string, body []byte, log wrapper.Log) {
-	attributes, ok := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-	if !ok {
-		log.Error("failed to get attributes from http context")
-		return
-	}
 	for _, attribute := range config.attributes {
+		var key, value string
+		var err error
 		if source == attribute.ValueSource {
+			key = attribute.Key
 			switch source {
 			case FixedValue:
 				log.Debugf("[attribute] source type: %s, key: %s, value: %s", source, attribute.Key, attribute.Value)
-				attributes[attribute.Key] = attribute.Value
+				value = attribute.Value
 			case RequestHeader:
-				if value, err := proxywasm.GetHttpRequestHeader(attribute.Value); err == nil {
+				if value, err = proxywasm.GetHttpRequestHeader(attribute.Value); err == nil {
 					log.Debugf("[attribute] source type: %s, key: %s, value: %s", source, attribute.Key, value)
-					attributes[attribute.Key] = value
 				}
 			case RequestBody:
 				raw := gjson.GetBytes(body, attribute.Value).Raw
-				var value string
 				if len(raw) > 2 {
 					value = raw[1 : len(raw)-1]
 				}
 				log.Debugf("[attribute] source type: %s, key: %s, value: %s", source, attribute.Key, value)
-				attributes[attribute.Key] = value
 			case ResponseHeader:
-				if value, err := proxywasm.GetHttpResponseHeader(attribute.Value); err == nil {
+				if value, err = proxywasm.GetHttpResponseHeader(attribute.Value); err == nil {
 					log.Debugf("[log attribute] source type: %s, key: %s, value: %s", source, attribute.Key, value)
-					attributes[attribute.Key] = value
 				}
 			case ResponseStreamingBody:
-				value := extractStreamingBodyByJsonPath(body, attribute.Value, attribute.Rule, log)
+				value = extractStreamingBodyByJsonPath(body, attribute.Value, attribute.Rule, log)
 				log.Debugf("[log attribute] source type: %s, key: %s, value: %s", source, attribute.Key, value)
-				attributes[attribute.Key] = value
 			case ResponseBody:
-				value := gjson.GetBytes(body, attribute.Value).Raw
-				if len(value) > 2 && value[0] == '"' && value[len(value)-1] == '"' {
-					value = value[1 : len(value)-1]
-				}
+				value = gjson.GetBytes(body, attribute.Value).String()
 				log.Debugf("[log attribute] source type: %s, key: %s, value: %s", source, attribute.Key, value)
-				attributes[attribute.Key] = value
 			default:
 			}
-		}
-		if attribute.ApplyToLog {
-			setLogAttribute(ctx, attribute.Key, attributes[attribute.Key], log)
-		}
-		if attribute.ApplyToSpan {
-			setSpanAttribute(attribute.Key, attributes[attribute.Key], log)
+			if attribute.ApplyToLog {
+				ctx.SetUserAttribute(key, value)
+			}
+			if attribute.ApplyToSpan {
+				setSpanAttribute(key, value, log)
+			}
 		}
 	}
-	ctx.SetContext(CtxGeneralAtrribute, attributes)
 }
 
 func extractStreamingBodyByJsonPath(data []byte, jsonPath string, rule string, log wrapper.Log) string {
@@ -368,9 +364,9 @@ func extractStreamingBodyByJsonPath(data []byte, jsonPath string, rule string, l
 	} else if rule == RuleAppend {
 		// extract llm response
 		for _, chunk := range chunks {
-			raw := gjson.GetBytes(chunk, jsonPath).Raw
-			if len(raw) > 2 && raw[0] == '"' && raw[len(raw)-1] == '"' {
-				value += raw[1 : len(raw)-1]
+			jsonObj := gjson.GetBytes(chunk, jsonPath)
+			if jsonObj.Exists() {
+				value += jsonObj.String()
 			}
 		}
 	} else {
@@ -379,123 +375,49 @@ func extractStreamingBodyByJsonPath(data []byte, jsonPath string, rule string, l
 	return value
 }
 
-func setFilterState(key, value string, log wrapper.Log) {
-	if value != "" {
-		if e := proxywasm.SetProperty([]string{key}, []byte(fmt.Sprint(value))); e != nil {
-			log.Errorf("failed to set %s in filter state: %v", key, e)
-		}
-	} else {
-		log.Debugf("failed to write filter state [%s], because it's value is empty")
-	}
-}
-
 // Set the tracing span with value.
 func setSpanAttribute(key, value string, log wrapper.Log) {
 	if value != "" {
-		traceSpanTag := TracePrefix + key
+		traceSpanTag := wrapper.TraceSpanTagPrefix + key
 		if e := proxywasm.SetProperty([]string{traceSpanTag}, []byte(value)); e != nil {
-			log.Errorf("failed to set %s in filter state: %v", traceSpanTag, e)
+			log.Warnf("failed to set %s in filter state: %v", traceSpanTag, e)
 		}
 	} else {
 		log.Debugf("failed to write span attribute [%s], because it's value is empty")
 	}
 }
 
-// fetches the tracing span value from the specified source.
-func setLogAttribute(ctx wrapper.HttpContext, key string, value interface{}, log wrapper.Log) {
-	logAttributes, ok := ctx.GetContext(CtxLogAtrribute).(map[string]string)
-	if !ok {
-		log.Error("failed to get logAttributes from http context")
-		return
-	}
-	logAttributes[key] = fmt.Sprint(value)
-	ctx.SetContext(CtxLogAtrribute, logAttributes)
-}
-
-func writeFilterStates(ctx wrapper.HttpContext, log wrapper.Log) {
-	attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-	setFilterState(Model, attributes[Model], log)
-	setFilterState(InputToken, attributes[InputToken], log)
-	setFilterState(OutputToken, attributes[OutputToken], log)
-}
-
 func writeMetric(ctx wrapper.HttpContext, config AIStatisticsConfig, log wrapper.Log) {
-	attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-	route, _ := getRouteName()
-	cluster, _ := getClusterName()
-	model, ok := attributes["model"]
-	if !ok {
-		log.Errorf("Get model failed")
+	route := ctx.GetContext(RouteName).(string)
+	cluster := ctx.GetContext(ClusterName).(string)
+	// Generate usage metrics
+	var model string
+	var inputToken, outputToken int64
+	if ctx.GetUserAttribute(Model) == nil || ctx.GetUserAttribute(InputToken) == nil || ctx.GetUserAttribute(OutputToken) == nil {
+		log.Warnf("get usage information failed, skip metric record")
 		return
 	}
-	if inputToken, ok := attributes[InputToken]; ok {
-		inputTokenUint64, err := strconv.ParseUint(inputToken, 10, 0)
-		if err != nil || inputTokenUint64 == 0 {
-			log.Errorf("inputToken convert failed, value is %d, err msg is [%v]", inputTokenUint64, err)
-			return
-		}
-		config.incrementCounter(generateMetricName(route, cluster, model, InputToken), inputTokenUint64)
+	model = ctx.GetUserAttribute(Model).(string)
+	inputToken = ctx.GetUserAttribute(InputToken).(int64)
+	outputToken = ctx.GetUserAttribute(OutputToken).(int64)
+	if inputToken == 0 || outputToken == 0 {
+		log.Warnf("inputToken and outputToken cannot equal to 0, skip metric record")
+		return
 	}
-	if outputToken, ok := attributes[OutputToken]; ok {
-		outputTokenUint64, err := strconv.ParseUint(outputToken, 10, 0)
-		if err != nil || outputTokenUint64 == 0 {
-			log.Errorf("outputToken convert failed, value is %d, err msg is [%v]", outputTokenUint64, err)
-			return
-		}
-		config.incrementCounter(generateMetricName(route, cluster, model, OutputToken), outputTokenUint64)
-	}
-	if llmFirstTokenDuration, ok := attributes[LLMFirstTokenDuration]; ok {
-		llmFirstTokenDurationUint64, err := strconv.ParseUint(llmFirstTokenDuration, 10, 0)
-		if err != nil || llmFirstTokenDurationUint64 == 0 {
-			log.Errorf("llmFirstTokenDuration convert failed, value is %d, err msg is [%v]", llmFirstTokenDurationUint64, err)
-			return
-		}
-		config.incrementCounter(generateMetricName(route, cluster, model, LLMFirstTokenDuration), llmFirstTokenDurationUint64)
-	}
-	if llmServiceDuration, ok := attributes[LLMServiceDuration]; ok {
-		llmServiceDurationUint64, err := strconv.ParseUint(llmServiceDuration, 10, 0)
-		if err != nil || llmServiceDurationUint64 == 0 {
-			log.Errorf("llmServiceDuration convert failed, value is %d, err msg is [%v]", llmServiceDurationUint64, err)
-			return
-		}
-		config.incrementCounter(generateMetricName(route, cluster, model, LLMServiceDuration), llmServiceDurationUint64)
-	}
-	config.incrementCounter(generateMetricName(route, cluster, model, LLMDurationCount), 1)
-}
+	config.incrementCounter(generateMetricName(route, cluster, model, InputToken), uint64(inputToken))
+	config.incrementCounter(generateMetricName(route, cluster, model, OutputToken), uint64(outputToken))
 
-func writeLog(ctx wrapper.HttpContext, log wrapper.Log) {
-	attributes, _ := ctx.GetContext(CtxGeneralAtrribute).(map[string]string)
-	logAttributes, _ := ctx.GetContext(CtxLogAtrribute).(map[string]string)
-	// Set inner log fields
-	if attributes[Model] != "" {
-		logAttributes[Model] = attributes[Model]
+	// Generate duration metrics
+	var llmFirstTokenDuration, llmServiceDuration int64
+	// Is stream response
+	if ctx.GetUserAttribute(LLMFirstTokenDuration) != nil {
+		llmFirstTokenDuration = ctx.GetUserAttribute(LLMFirstTokenDuration).(int64)
+		config.incrementCounter(generateMetricName(route, cluster, model, LLMFirstTokenDuration), uint64(llmFirstTokenDuration))
+		config.incrementCounter(generateMetricName(route, cluster, model, LLMStreamDurationCount), 1)
 	}
-	if attributes[InputToken] != "" {
-		logAttributes[InputToken] = attributes[InputToken]
-	}
-	if attributes[OutputToken] != "" {
-		logAttributes[OutputToken] = attributes[OutputToken]
-	}
-	if attributes[LLMFirstTokenDuration] != "" {
-		logAttributes[LLMFirstTokenDuration] = attributes[LLMFirstTokenDuration]
-	}
-	if attributes[LLMServiceDuration] != "" {
-		logAttributes[LLMServiceDuration] = attributes[LLMServiceDuration]
-	}
-	// Traverse log fields
-	items := []string{}
-	for k, v := range logAttributes {
-		items = append(items, fmt.Sprintf(`"%s":"%s"`, k, v))
-	}
-	aiLogField := fmt.Sprintf(`{%s}`, strings.Join(items, ","))
-	// log.Infof("ai request json log: %s", aiLogField)
-	jsonMap := map[string]string{
-		"ai_log": aiLogField,
-	}
-	serialized, _ := json.Marshal(jsonMap)
-	jsonLogRaw := gjson.GetBytes(serialized, "ai_log").Raw
-	jsonLog := jsonLogRaw[1 : len(jsonLogRaw)-1]
-	if err := proxywasm.SetProperty([]string{"ai_log"}, []byte(jsonLog)); err != nil {
-		log.Errorf("failed to set ai_log in filter state: %v", err)
+	if ctx.GetUserAttribute(LLMServiceDuration) != nil {
+		llmServiceDuration = ctx.GetUserAttribute(LLMServiceDuration).(int64)
+		config.incrementCounter(generateMetricName(route, cluster, model, LLMServiceDuration), uint64(llmServiceDuration))
+		config.incrementCounter(generateMetricName(route, cluster, model, LLMDurationCount), 1)
 	}
 }

plugins/wasm-go/extensions/ai-token-ratelimit/go.sum

@@ -5,8 +5,7 @@ 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-20240711023527-ba358c48772f h1:ZIiIBRvIw62gA5MJhuwp1+2wWbqL9IGElQ499rUsYYg=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/higress-group/proxy-wasm-go-sdk v1.0.0 h1:BZRNf4R7jr9hwRivg/E29nkVaKEak5MWjBDhWjuHijU=
 github.com/higress-group/proxy-wasm-go-sdk v1.0.0/go.mod h1:iiSyFbo+rAtbtGt/bsefv8GU57h9CCLYGJA74/tF5/0=
 github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
 github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
@@ -14,8 +13,7 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 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/gjson v1.17.3 h1:bwWLZU7icoKRG+C+0PNwIKC6FCJO/Q3p2pZvuP0jN94=
 github.com/tidwall/gjson v1.17.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=

plugins/wasm-go/extensions/ai-token-ratelimit/main.go

@@ -15,6 +15,7 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"net/url"
@@ -61,9 +62,9 @@ const (
 	ConsumerHeader string = "x-mse-consumer" // LimitByConsumer从该request header获取consumer的名字
 	CookieHeader   string = "cookie"
 
-	RateLimitLimitHeader     string = "X-RateLimit-Limit"     // 限制的总请求数
-	RateLimitRemainingHeader string = "X-RateLimit-Remaining" // 剩余还可以发送的请求数
-	RateLimitResetHeader     string = "X-RateLimit-Reset"     // 限流重置时间（触发限流时返回）
+	RateLimitLimitHeader     string = "X-TokenRateLimit-Limit"     // 限制的总请求数
+	RateLimitRemainingHeader string = "X-TokenRateLimit-Remaining" // 剩余还可以发送的请求数
+	RateLimitResetHeader     string = "X-TokenRateLimit-Reset"     // 限流重置时间（触发限流时返回）
 )
 
 type LimitContext struct {
@@ -124,6 +125,8 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config ClusterKeyRateLimitCon
 		}
 		if context.remaining < 0 {
 			// 触发限流
+			ctx.SetUserAttribute("token_ratelimit_status", "limited")
+			ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
 			rejected(config, context)
 		} else {
 			proxywasm.ResumeHttpRequest()
@@ -137,39 +140,49 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config ClusterKeyRateLimitCon
 }
 
 func onHttpStreamingBody(ctx wrapper.HttpContext, config ClusterKeyRateLimitConfig, data []byte, endOfStream bool, log wrapper.Log) []byte {
-	if !endOfStream {
-		return data
+	var inputToken, outputToken int64
+	if inputToken, outputToken, ok := getUsage(data); ok {
+		ctx.SetContext("input_token", inputToken)
+		ctx.SetContext("output_token", outputToken)
 	}
-	inputTokenStr, err := proxywasm.GetProperty([]string{"filter_state", "wasm.input_token"})
-	if err != nil {
-		return data
+	if endOfStream {
+		if ctx.GetContext("input_token") == nil || ctx.GetContext("output_token") == nil {
+			return data
+		}
+		inputToken = ctx.GetContext("input_token").(int64)
+		outputToken = ctx.GetContext("output_token").(int64)
+		limitRedisContext, ok := ctx.GetContext(LimitRedisContextKey).(LimitRedisContext)
+		if !ok {
+			return data
+		}
+		keys := []interface{}{limitRedisContext.key}
+		args := []interface{}{limitRedisContext.count, limitRedisContext.window, inputToken + outputToken}
+		err := config.redisClient.Eval(ResponsePhaseFixedWindowScript, 1, keys, args, nil)
+		if err != nil {
+			log.Errorf("redis call failed: %v", err)
+		}
 	}
-	outputTokenStr, err := proxywasm.GetProperty([]string{"filter_state", "wasm.output_token"})
-	if err != nil {
-		return data
-	}
-	inputToken, err := strconv.Atoi(string(inputTokenStr))
-	if err != nil {
-		return data
-	}
-	outputToken, err := strconv.Atoi(string(outputTokenStr))
-	if err != nil {
-		return data
-	}
-	limitRedisContext, ok := ctx.GetContext(LimitRedisContextKey).(LimitRedisContext)
-	if !ok {
-		return data
-	}
-	keys := []interface{}{limitRedisContext.key}
-	args := []interface{}{limitRedisContext.count, limitRedisContext.window, inputToken + outputToken}
+	return data
+}
 
-	err = config.redisClient.Eval(ResponsePhaseFixedWindowScript, 1, keys, args, nil)
-	if err != nil {
-		log.Errorf("redis call failed: %v", err)
-		return data
-	} else {
-		return data
+func getUsage(data []byte) (inputTokenUsage int64, outputTokenUsage int64, ok bool) {
+	chunks := bytes.Split(bytes.TrimSpace(data), []byte("\n\n"))
+	for _, chunk := range chunks {
+		// the feature strings are used to identify the usage data, like:
+		// {"model":"gpt2","usage":{"prompt_tokens":1,"completion_tokens":1}}
+		if !bytes.Contains(chunk, []byte("prompt_tokens")) || !bytes.Contains(chunk, []byte("completion_tokens")) {
+			continue
+		}
+		inputTokenObj := gjson.GetBytes(chunk, "usage.prompt_tokens")
+		outputTokenObj := gjson.GetBytes(chunk, "usage.completion_tokens")
+		if inputTokenObj.Exists() && outputTokenObj.Exists() {
+			inputTokenUsage = inputTokenObj.Int()
+			outputTokenUsage = outputTokenObj.Int()
+			ok = true
+			return
+		}
 	}
+	return
 }
 
 func checkRequestAgainstLimitRule(ctx wrapper.HttpContext, ruleItems []LimitRuleItem, log wrapper.Log) (string, *LimitRuleItem, *LimitConfigItem) {

plugins/wasm-go/pkg/wrapper/plugin_wrapper.go

@@ -45,6 +45,8 @@ type HttpContext interface {
 	GetStringContext(key, defaultValue string) string
 	GetUserAttribute(key string) interface{}
 	SetUserAttribute(key string, value interface{})
+	SetUserAttributeMap(kvmap map[string]interface{})
+	GetUserAttributeMap() map[string]interface{}
 	// You can call this function to set custom log
 	WriteUserAttributeToLog() error
 	// You can call this function to set custom log with your specific key
@@ -403,6 +405,14 @@ func (ctx *CommonHttpCtx[PluginConfig]) GetUserAttribute(key string) interface{}
 	return ctx.userAttribute[key]
 }
 
+func (ctx *CommonHttpCtx[PluginConfig]) SetUserAttributeMap(kvmap map[string]interface{}) {
+	ctx.userAttribute = kvmap
+}
+
+func (ctx *CommonHttpCtx[PluginConfig]) GetUserAttributeMap() map[string]interface{} {
+	return ctx.userAttribute
+}
+
 func (ctx *CommonHttpCtx[PluginConfig]) WriteUserAttributeToLog() error {
 	return ctx.WriteUserAttributeToLogWithKey(CustomLogKey)
 }