Skip to content

Commit 5467339

Browse files
daixijundaixijun
committed
feat(ai-statistics): add request_model to context
Signed-off-by: Xijun Dai <[email protected]>
1 parent 5ee8781 commit 5467339

File tree

3 files changed

+22
-29
lines changed

3 files changed

+22
-29
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ toolchain go1.24.4
66

77
require (
88
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250611100342-5654e89a7a80
9-
github.com/higress-group/wasm-go v1.0.2-0.20250729071413-2478fd585950
9+
github.com/higress-group/wasm-go v1.0.2-0.20250804091745-dbabcf425bb2
1010
github.com/tidwall/gjson v1.18.0
1111
)
1212

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250611100342-5654e89a7a80 h1
66
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250611100342-5654e89a7a80/go.mod h1:tRI2LfMudSkKHhyv1uex3BWzcice2s/l8Ah8axporfA=
77
github.com/higress-group/wasm-go v1.0.2-0.20250729071413-2478fd585950 h1:X4a+wzGEuLkCcAX2XiDf/vcVOIdZWxtEo0YkT+F/mcM=
88
github.com/higress-group/wasm-go v1.0.2-0.20250729071413-2478fd585950/go.mod h1:9k7L730huS/q4V5iH9WLDgf5ZUHEtfhM/uXcegKDG/M=
9+
github.com/higress-group/wasm-go v1.0.2-0.20250804091745-dbabcf425bb2 h1:qYUeTPlWhVlpDrfISm0EArwZQC1Jl/uiD3l1h9PFjuA=
10+
github.com/higress-group/wasm-go v1.0.2-0.20250804091745-dbabcf425bb2/go.mod h1:9k7L730huS/q4V5iH9WLDgf5ZUHEtfhM/uXcegKDG/M=
911
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
1012
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
1113
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=

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

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,6 @@ const (
5959
LLMDurationCount = "llm_duration_count"
6060
LLMStreamDurationCount = "llm_stream_duration_count"
6161
ResponseType = "response_type"
62-
ChatID = "chat_id"
6362
ChatRound = "chat_round"
6463

6564
// Inner span attributes
@@ -192,6 +191,9 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config AIStatisticsConfig) ty
192191
if consumer, _ := proxywasm.GetHttpRequestHeader(ConsumerKey); consumer != "" {
193192
ctx.SetContext(ConsumerKey, consumer)
194193
}
194+
if requestModel, _ := proxywasm.GetHttpRequestHeader("x-higress-llm-model"); requestModel != "" {
195+
ctx.SetContext(tokenusage.CtxKeyRequestModel, requestModel)
196+
}
195197

196198
ctx.SetRequestBodyBufferLimit(defaultMaxBodyBytes)
197199

@@ -214,14 +216,15 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body
214216
requestModel = model.String()
215217
} else {
216218
requestPath := ctx.GetStringContext(RequestPath, "")
217-
if strings.Contains(requestPath, "generateContent") || strings.Contains(requestPath, "streamGenerateContent") { // Google Gemini GenerateContent
218-
reg := regexp.MustCompile(`^.*/(?P<api_version>[^/]+)/models/(?P<model>[^:]+):\w+Content$`)
219+
if strings.Contains(requestPath, "generateContent") || strings.Contains(requestPath, "streamGenerateContent") || strings.Contains(requestPath, "countTokens") || strings.Contains(requestPath, "predictLongRunning") { // Google Gemini GenerateContent
220+
reg := regexp.MustCompile(`^.*/(?P<api_version>[^/]+)/models/(?P<model>[^:]+):(\w+Content|countTokens|predictLongRunning)$`)
219221
matches := reg.FindStringSubmatch(requestPath)
220222
if len(matches) == 3 {
221223
requestModel = matches[2]
222224
}
223225
}
224226
}
227+
ctx.SetContext(tokenusage.CtxKeyRequestModel, requestModel)
225228
setSpanAttribute(ArmsRequestModel, requestModel)
226229
// Set the number of conversation rounds
227230

@@ -242,7 +245,7 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body
242245
ctx.SetUserAttribute(ChatRound, userPromptCount)
243246

244247
// Write log
245-
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
248+
_ = ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
246249
return types.ActionContinue
247250
}
248251

@@ -271,14 +274,8 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat
271274
}
272275

273276
ctx.SetUserAttribute(ResponseType, "stream")
274-
if chatID := wrapper.GetValueFromBody(data, []string{
275-
"id",
276-
"response.id",
277-
"responseId", // Gemini generateContent
278-
"message.id", // anthropic messages
279-
}); chatID != nil {
280-
ctx.SetUserAttribute(ChatID, chatID.String())
281-
}
277+
278+
tokenusage.ExtractChatId(ctx, data)
282279

283280
// Get requestStartTime from http context
284281
requestStartTime, ok := ctx.GetContext(StatisticsRequestStartTime).(int64)
@@ -319,7 +316,7 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat
319316
}
320317

321318
// Write log
322-
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
319+
_ = ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
323320

324321
// Write metrics
325322
writeMetric(ctx, config)
@@ -335,14 +332,8 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body
335332
ctx.SetUserAttribute(LLMServiceDuration, responseEndTime-requestStartTime)
336333

337334
ctx.SetUserAttribute(ResponseType, "normal")
338-
if chatID := wrapper.GetValueFromBody(body, []string{
339-
"id",
340-
"response.id",
341-
"responseId", // Gemini generateContent
342-
"message.id", // anthropic messages
343-
}); chatID != nil {
344-
ctx.SetUserAttribute(ChatID, chatID.String())
345-
}
335+
336+
tokenusage.ExtractChatId(ctx, body)
346337

347338
// Set information about this request
348339
if !config.disableOpenaiUsage {
@@ -359,7 +350,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body
359350
setAttributeBySource(ctx, config, ResponseBody, body)
360351

361352
// Write log
362-
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
353+
_ = ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
363354

364355
// Write metrics
365356
writeMetric(ctx, config)
@@ -372,7 +363,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body
372363
func setAttributeBySource(ctx wrapper.HttpContext, config AIStatisticsConfig, source string, body []byte) {
373364
for _, attribute := range config.attributes {
374365
var key string
375-
var value interface{}
366+
var value any
376367
if source == attribute.ValueSource {
377368
key = attribute.Key
378369
switch source {
@@ -418,9 +409,9 @@ func setAttributeBySource(ctx wrapper.HttpContext, config AIStatisticsConfig, so
418409
}
419410
}
420411

421-
func extractStreamingBodyByJsonPath(data []byte, jsonPath string, rule string) interface{} {
412+
func extractStreamingBodyByJsonPath(data []byte, jsonPath string, rule string) any {
422413
chunks := bytes.Split(bytes.TrimSpace(wrapper.UnifySSEChunk(data)), []byte("\n\n"))
423-
var value interface{}
414+
var value any
424415
if rule == RuleFirst {
425416
for _, chunk := range chunks {
426417
jsonObj := gjson.GetBytes(chunk, jsonPath)
@@ -453,10 +444,10 @@ func extractStreamingBodyByJsonPath(data []byte, jsonPath string, rule string) i
453444
}
454445

455446
// Set the tracing span with value.
456-
func setSpanAttribute(key string, value interface{}) {
447+
func setSpanAttribute(key string, value any) {
457448
if value != "" {
458449
traceSpanTag := wrapper.TraceSpanTagPrefix + key
459-
if e := proxywasm.SetProperty([]string{traceSpanTag}, []byte(fmt.Sprint(value))); e != nil {
450+
if e := proxywasm.SetProperty([]string{traceSpanTag}, fmt.Append(nil, value)); e != nil {
460451
log.Warnf("failed to set %s in filter state: %v", traceSpanTag, e)
461452
}
462453
} else {
@@ -532,7 +523,7 @@ func writeMetric(ctx wrapper.HttpContext, config AIStatisticsConfig) {
532523
}
533524
}
534525

535-
func convertToUInt(val interface{}) (uint64, bool) {
526+
func convertToUInt(val any) (uint64, bool) {
536527
switch v := val.(type) {
537528
case float32:
538529
return uint64(v), true

0 commit comments

Comments
 (0)