From 5a5af4ecbf72a960a01d0b9de8a5914573fd1aa0 Mon Sep 17 00:00:00 2001
From: rinfx <yucheng.lxr@alibaba-inc.com>
Date: Tue, 11 Mar 2025 09:32:11 +0800
Subject: [PATCH] support default value (#1873)

---
 .../extensions/ai-statistics/README.md        |  1 +
 .../extensions/ai-statistics/README_EN.md     |  1 +
 .../wasm-go/extensions/ai-statistics/main.go  | 24 ++++++++++++-------
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/plugins/wasm-go/extensions/ai-statistics/README.md b/plugins/wasm-go/extensions/ai-statistics/README.md
index fad5ebfc4..ad71061ff 100644
--- a/plugins/wasm-go/extensions/ai-statistics/README.md
+++ b/plugins/wasm-go/extensions/ai-statistics/README.md
@@ -31,6 +31,7 @@ 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`|
 | `apply_to_log`      | bool | 非必填  | false  | 是否将提取的信息记录在日志中 |
 | `apply_to_span`      | bool | 非必填  | false  | 是否将提取的信息记录在链路追踪span中 |
diff --git a/plugins/wasm-go/extensions/ai-statistics/README_EN.md b/plugins/wasm-go/extensions/ai-statistics/README_EN.md
index 37e58186d..3b5dd9d3c 100644
--- a/plugins/wasm-go/extensions/ai-statistics/README_EN.md
+++ b/plugins/wasm-go/extensions/ai-statistics/README_EN.md
@@ -31,6 +31,7 @@ Attribute Configuration instructions:
 | `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 |
+| `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 |
diff --git a/plugins/wasm-go/extensions/ai-statistics/main.go b/plugins/wasm-go/extensions/ai-statistics/main.go
index 1516c3ca5..f8d315ec8 100644
--- a/plugins/wasm-go/extensions/ai-statistics/main.go
+++ b/plugins/wasm-go/extensions/ai-statistics/main.go
@@ -64,12 +64,13 @@ const (
 
 // TracingSpan is the tracing span configuration.
 type Attribute struct {
-	Key         string `json:"key"`
-	ValueSource string `json:"value_source"`
-	Value       string `json:"value"`
-	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"`
+	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"`
 }
 
 type AIStatisticsConfig struct {
@@ -294,10 +295,14 @@ func getUsage(data []byte) (model string, inputTokenUsage int64, outputTokenUsag
 			continue
 		}
 		modelObj := gjson.GetBytes(chunk, "model")
+		if modelObj.Exists() {
+			model = modelObj.String()
+		} else {
+			model = "unknown"
+		}
 		inputTokenObj := gjson.GetBytes(chunk, "usage.prompt_tokens")
 		outputTokenObj := gjson.GetBytes(chunk, "usage.completion_tokens")
-		if modelObj.Exists() && inputTokenObj.Exists() && outputTokenObj.Exists() {
-			model = modelObj.String()
+		if inputTokenObj.Exists() && outputTokenObj.Exists() {
 			inputTokenUsage = inputTokenObj.Int()
 			outputTokenUsage = outputTokenObj.Int()
 			ok = true
@@ -329,6 +334,9 @@ func setAttributeBySource(ctx wrapper.HttpContext, config AIStatisticsConfig, so
 				value = gjson.GetBytes(body, attribute.Value).Value()
 			default:
 			}
+			if (value == nil || value == "") && attribute.DefaultValue != "" {
+				value = attribute.DefaultValue
+			}
 			log.Debugf("[attribute] source type: %s, key: %s, value: %+v", source, key, value)
 			if attribute.ApplyToLog {
 				ctx.SetUserAttribute(key, value)