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

2026-02-06 23:21:08 +08:00 · 2025-06-25 09:28:14 +08:00
parent f4d3fec228
commit 76ada0b844
5 changed files with 78 additions and 32 deletions
--- a/plugins/wasm-go/extensions/ai-statistics/README.md
+++ b/plugins/wasm-go/extensions/ai-statistics/README.md
@@ -29,22 +29,24 @@ Attribute 配置说明:

 | 名称             | 数据类型  | 填写要求 | 默认值 | 描述                     |
 |----------------|-------|-----|-----|------------------------|
-| `key`         | string | 必填  | -   | attrribute 名称           |
-| `value_source` | string | 必填  | -   | attrribute 取值来源，可选值为 `fixed_value`, `request_header`, `request_body`, `response_header`, `response_body`, `response_streaming_body`             |
-| `value`      | string | 必填  | -   | attrribute 取值 key value/path |
-| `default_value`      | string | 非必填  | -   | attrribute 默认值 |
-| `rule`      | string | 非必填  | -   | 从流式响应中提取 attrribute 的规则，可选值为 `first`, `replace`, `append`|
+| `key`         | string | 必填  | -   | attribute 名称           |
+| `value_source` | string | 必填  | -   | attribute 取值来源，可选值为 `fixed_value`, `request_header`, `request_body`, `response_header`, `response_body`, `response_streaming_body`             |
+| `value`      | string | 必填  | -   | attribute 取值 key value/path |
+| `default_value`      | string | 非必填  | -   | attribute 默认值 |
+| `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_body` ：attrribute 值通过请求 body 获取，value 配置格式为 gjson 的 jsonpath
- `response_header` ：attrribute 值通过 http 响应头获取，value 配置为header key
- `response_body` ：attrribute 值通过响应 body 获取，value 配置格式为 gjson 的 jsonpath
- `response_streaming_body` ：attrribute 值通过流式响应 body 获取，value 配置格式为 gjson 的 jsonpath
+- `request_header` ： attribute 值通过 http 请求头获取，value 配置为 header key
+- `request_body` ：attribute 值通过请求 body 获取，value 配置格式为 gjson 的 jsonpath
+- `response_header` ：attribute 值通过 http 响应头获取，value 配置为header key
+- `response_body` ：attribute 值通过响应 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)%"}'
+```
+
 ### 空配置
 #### 监控

--- a/plugins/wasm-go/extensions/ai-statistics/README_EN.md
+++ b/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           |
-| `value_source` | string | required  | -   | attrribute 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 |
+| `key`         | string | required  | -   | attribute key           |
+| `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 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_body`: The attrribute is obtained through the http request body
- `response_header`: The attrribute is obtained through the http response header
- `response_body`: The attrribute is obtained through the http response body
- `response_streaming_body`: The attrribute is obtained through the http streaming response body
+- `request_header`: The attribute is obtained through the http request header
+- `request_body`: The attribute is obtained through the http request body
+- `response_header`: The attribute is obtained through the http response header
+- `response_body`: The attribute is obtained through the http 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

--- a/plugins/wasm-go/extensions/ai-statistics/main.go
+++ b/plugins/wasm-go/extensions/ai-statistics/main.go
@@ -75,13 +75,15 @@ const (

 // TracingSpan is the tracing span configuration.
 type Attribute struct {
-	Key          string `json:"key"`
-	ValueSource  string `json:"value_source"`
-	Value        string `json:"value"`
-	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"`
+	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 {
-				ctx.SetUserAttribute(key, value)
+				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)
 			}
 		}
@@ -481,10 +493,10 @@ func writeMetric(ctx wrapper.HttpContext, config AIStatisticsConfig, log wrapper
 		log.Warnf("ClusterName typd assert failed, skip metric record")
 		return
 	}
-	
+
 	if config.disableOpenaiUsage {
 		return
-	} 
+	}

 	if ctx.GetUserAttribute(Model) == nil || ctx.GetUserAttribute(InputToken) == nil || ctx.GetUserAttribute(OutputToken) == nil {
 		log.Warnf("get usage information failed, skip metric record")