add trace_span_key & as_seperate_log_field configuration for ai-statistics (#2488)

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

@@ -29,22 +29,24 @@ Attribute 配置说明:
 
 | 名称             | 数据类型  | 填写要求 | 默认值 | 描述                     |
 |----------------|-------|-----|-----|------------------------|
-| `key`         | string | 必填  | -   | attrribute 名称           |
+| `key`         | string | 必填  | -   | attribute 名称           |
-| `value_source` | string | 必填  | -   | attrribute 取值来源，可选值为 `fixed_value`, `request_header`, `request_body`, `response_header`, `response_body`, `response_streaming_body`             |
+| `value_source` | string | 必填  | -   | attribute 取值来源，可选值为 `fixed_value`, `request_header`, `request_body`, `response_header`, `response_body`, `response_streaming_body`             |
-| `value`      | string | 必填  | -   | attrribute 取值 key value/path |
+| `value`      | string | 必填  | -   | attribute 取值 key value/path |
-| `default_value`      | string | 非必填  | -   | attrribute 默认值 |
+| `default_value`      | string | 非必填  | -   | attribute 默认值 |
-| `rule`      | string | 非必填  | -   | 从流式响应中提取 attrribute 的规则，可选值为 `first`, `replace`, `append`|
+| `rule`      | string | 非必填  | -   | 从流式响应中提取 attribute 的规则，可选值为 `first`, `replace`, `append`|
 | `apply_to_log`      | bool | 非必填  | false  | 是否将提取的信息记录在日志中 |
 | `apply_to_span`      | bool | 非必填  | false  | 是否将提取的信息记录在链路追踪span中 |
+| `trace_span_key`      | string | 非必填  | -  | 链路追踪attribute key，默认会使用`key`的设置 |
+| `as_separate_log_field`      | bool | 非必填  | false  | 记录日志时是否作为单独的字段，日志字段名使用`key`的设置 |
 
 `value_source` 的各种取值含义如下：
 
 - `fixed_value`：固定值
-- `request_header` ： attrribute 值通过 http 请求头获取，value 配置为 header key
+- `request_header` ： attribute 值通过 http 请求头获取，value 配置为 header key
-- `request_body` ：attrribute 值通过请求 body 获取，value 配置格式为 gjson 的 jsonpath
+- `request_body` ：attribute 值通过请求 body 获取，value 配置格式为 gjson 的 jsonpath
-- `response_header` ：attrribute 值通过 http 响应头获取，value 配置为header key
+- `response_header` ：attribute 值通过 http 响应头获取，value 配置为header key
-- `response_body` ：attrribute 值通过响应 body 获取，value 配置格式为 gjson 的 jsonpath
+- `response_body` ：attribute 值通过响应 body 获取，value 配置格式为 gjson 的 jsonpath
-- `response_streaming_body` ：attrribute 值通过流式响应 body 获取，value 配置格式为 gjson 的 jsonpath
+- `response_streaming_body` ：attribute 值通过流式响应 body 获取，value 配置格式为 gjson 的 jsonpath
 
 
 当 `value_source` 为 `response_streaming_body` 时，应当配置 `rule`，用于指定如何从流式body中获取指定值，取值含义如下：
@@ -60,6 +62,21 @@ Attribute 配置说明:
 '{"ai_log":"%FILTER_STATE(wasm.ai_log:PLAIN)%"}'
 ```
 
+如果字段设置了 `as_separate_log_field`，例如：
+```yaml
+attributes:
+  - key: consumer
+    value_source: request_header
+    value: x-mse-consumer
+    apply_to_log: true
+    as_separate_log_field: true
+```
+
+那么要在日志中打印，需要额外设置log_format：
+```
+'{"consumer":"%FILTER_STATE(wasm.consumer:PLAIN)%"}'
+```
+
 ### 空配置
 #### 监控
 

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

@@ -29,22 +29,24 @@ Attribute Configuration instructions:
 
 | Name             | Type  | Required | Default | Description |
 |----------------|-------|-----|-----|------------------------|
-| `key`         | string | required  | -   | attrribute key           |
+| `key`         | string | required  | -   | attribute key           |
-| `value_source` | string | required  | -   | attrribute value source, optional values ​​are `fixed_value`, `request_header`, `request_body`, `response_header`, `response_body`, `response_streaming_body`             |
+| `value_source` | string | required  | -   | attribute value source, optional values ​​are `fixed_value`, `request_header`, `request_body`, `response_header`, `response_body`, `response_streaming_body`             |
-| `value`      | string | required  | -   | how to get attrribute value |
+| `value`      | string | required  | -   | how to get attribute value |
 | `default_value`      | string | optional  | -   | default value for attribute |
 | `rule`      | string | optional  | -   | Rule to extract attribute from streaming response, optional values ​​are `first`, `replace`, `append`|
 | `apply_to_log`      | bool | optional  | false  | Whether to record the extracted information in the log |
 | `apply_to_span`      | bool | optional  | false  | Whether to record the extracted information in the link tracking span |
+| `trace_span_key`      | string | optional  | -  | span attribute key, default is the value of `key` |
+| `as_separate_log_field`      | bool | optional  | false  | Whether to use a separate log field, the field name is equal to the value of `key` |
 
 The meanings of various values for `value_source` ​​are as follows:
 
 - `fixed_value`: fixed value
-- `request_header`: The attrribute is obtained through the http request header
+- `request_header`: The attribute is obtained through the http request header
-- `request_body`: The attrribute is obtained through the http request body
+- `request_body`: The attribute is obtained through the http request body
-- `response_header`: The attrribute is obtained through the http response header
+- `response_header`: The attribute is obtained through the http response header
-- `response_body`: The attrribute is obtained through the http response body
+- `response_body`: The attribute is obtained through the http response body
-- `response_streaming_body`: The attrribute is obtained through the http streaming response body
+- `response_streaming_body`: The attribute is obtained through the http streaming response body
 
 
 When `value_source` is `response_streaming_body`, `rule` should be configured to specify how to obtain the specified value from the streaming body. The meaning of the value is as follows:
@@ -60,6 +62,21 @@ If you want to record ai-statistic related statistical values in the gateway acc
 '{"ai_log":"%FILTER_STATE(wasm.ai_log:PLAIN)%"}'
 ```
 
+If the field is set with `as_separate_log_field`, for example:
+```yaml
+attributes:
+  - key: consumer
+    value_source: request_header
+    value: x-mse-consumer
+    apply_to_log: true
+    as_separate_log_field: true
+```
+
+Then to print in the log, you need to set log_format additionally:
+```
+'{"consumer":"%FILTER_STATE(wasm.consumer:PLAIN)%"}'
+```
+
 ### Empty
 #### Metric
 

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

@@ -78,10 +78,12 @@ type Attribute struct {
 	Key                string `json:"key"`
 	ValueSource        string `json:"value_source"`
 	Value              string `json:"value"`
+	TraceSpanKey       string `json:"trace_span_key,omitempty"`
 	DefaultValue       string `json:"default_value,omitempty"`
 	Rule               string `json:"rule,omitempty"`
 	ApplyToLog         bool   `json:"apply_to_log,omitempty"`
 	ApplyToSpan        bool   `json:"apply_to_span,omitempty"`
+	AsSeparateLogField bool   `json:"as_separate_log_field,omitempty"`
 }
 
 type AIStatisticsConfig struct {
@@ -406,13 +408,23 @@ func setAttributeBySource(ctx wrapper.HttpContext, config AIStatisticsConfig, so
 			}
 			log.Debugf("[attribute] source type: %s, key: %s, value: %+v", source, key, value)
 			if attribute.ApplyToLog {
+				if attribute.AsSeparateLogField {
+					marshalledJsonStr := wrapper.MarshalStr(fmt.Sprint(value))
+					if err := proxywasm.SetProperty([]string{key}, []byte(marshalledJsonStr)); err != nil {
+						log.Warnf("failed to set %s in filter state, raw is %s, err is %v", key, marshalledJsonStr, err)
+					}
+				} else {
 					ctx.SetUserAttribute(key, value)
 				}
+			}
 			// for metrics
 			if key == Model || key == InputToken || key == OutputToken {
 				ctx.SetContext(key, value)
 			}
 			if attribute.ApplyToSpan {
+				if attribute.TraceSpanKey != "" {
+					key = attribute.TraceSpanKey
+				}
 				setSpanAttribute(key, value, log)
 			}
 		}

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

@@ -566,7 +566,7 @@ func (ctx *CommonHttpCtx[PluginConfig]) WriteUserAttributeToLogWithKey(key strin
 	newAttributeMap := map[string]interface{}{}
 	if string(preMarshalledJsonLogStr) != "" {
 		// e.g. {"field1":"value1","field2":"value2"}
-		preJsonLogStr := unmarshalStr(fmt.Sprintf(`"%s"`, string(preMarshalledJsonLogStr)))
+		preJsonLogStr := UnmarshalStr(fmt.Sprintf(`"%s"`, string(preMarshalledJsonLogStr)))
 		err := json.Unmarshal([]byte(preJsonLogStr), &newAttributeMap)
 		if err != nil {
 			ctx.plugin.vm.log.Warnf("Unmarshal failed, will overwrite %s, pre value is: %s", key, string(preMarshalledJsonLogStr))
@@ -580,7 +580,7 @@ func (ctx *CommonHttpCtx[PluginConfig]) WriteUserAttributeToLogWithKey(key strin
 	// e.g. {"field1":"value1","field2":2,"field3":"value3"}
 	jsonStr, _ := json.Marshal(newAttributeMap)
 	// e.g. {\"field1\":\"value1\",\"field2\":2,\"field3\":\"value3\"}
-	marshalledJsonStr := marshalStr(string(jsonStr))
+	marshalledJsonStr := MarshalStr(string(jsonStr))
 	if err := proxywasm.SetProperty([]string{key}, []byte(marshalledJsonStr)); err != nil {
 		ctx.plugin.vm.log.Warnf("failed to set %s in filter state, raw is %s, err is %v", key, marshalledJsonStr, err)
 		return err

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

@@ -7,7 +7,7 @@ import (
 	"github.com/tidwall/gjson"
 )
 
-func unmarshalStr(marshalledJsonStr string) string {
+func UnmarshalStr(marshalledJsonStr string) string {
 	// e.g. "{\"field1\":\"value1\",\"field2\":\"value2\"}"
 	var jsonStr string
 	err := json.Unmarshal([]byte(marshalledJsonStr), &jsonStr)
@@ -19,7 +19,7 @@ func unmarshalStr(marshalledJsonStr string) string {
 	return jsonStr
 }
 
-func marshalStr(raw string) string {
+func MarshalStr(raw string) string {
 	// e.g. {"field1":"value1","field2":"value2"}
 	helper := map[string]string{
 		"placeholder": raw,