From 74f7c955c9a7dc1ef14d36e5082fc1e535bdcbdf Mon Sep 17 00:00:00 2001
From: liyunfei <2716499331@qq.com>
Date: Sun, 15 Jun 2025 14:38:17 +0800
Subject: [PATCH] add token metrics

---
 metrics/metrics.go  | 50 +++++++++++++++++++++++++++++++++++++++++++++
 relay/relay-text.go | 14 ++++++++++++-
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/metrics/metrics.go b/metrics/metrics.go
index e0e484117ca8..6b5e4bb8b6fc 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -21,6 +21,11 @@ func RegisterMetrics(registry prometheus.Registerer) {
 	registry.MustRegister(relayRequestE2ESuccessCounter)
 	registry.MustRegister(relayRequestE2EFailedCounter)
 	registry.MustRegister(relayRequestE2EDurationObsever)
+	// token metrics
+	registry.MustRegister(inputTokensCounter)
+	registry.MustRegister(outputTokensCounter)
+	registry.MustRegister(cacheHitTokensCounter)
+	registry.MustRegister(inferenceTokensCounter)
 }
 
 var (
@@ -84,6 +89,34 @@ var (
 		},
 		[]string{"channel", "model", "group", "token_key", "token_name"},
 	)
+	// Token metrics
+	inputTokensCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Subsystem: Namespace,
+			Name:      "input_tokens_total",
+			Help:      "Total number of input tokens processed",
+		}, []string{"channel", "model", "group", "user_id"})
+
+	outputTokensCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Subsystem: Namespace,
+			Name:      "output_tokens_total",
+			Help:      "Total number of output tokens generated",
+		}, []string{"channel", "model", "group", "user_id"})
+
+	cacheHitTokensCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Subsystem: Namespace,
+			Name:      "cache_hit_tokens_total",
+			Help:      "Total number of tokens served from cache",
+		}, []string{"channel", "model", "group", "user_id"})
+
+	inferenceTokensCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Subsystem: Namespace,
+			Name:      "inference_tokens_total",
+			Help:      "Total number of tokens processed during inference",
+		}, []string{"channel", "model", "group", "user_id"})
 )
 
 func IncrementRelayRequestTotalCounter(channel, tag, baseURL, model, group string, add float64) {
@@ -121,3 +154,20 @@ func IncrementRelayRequestE2EFailedCounter(channel, model, group, code, tokenKey
 func ObserveRelayRequestE2EDuration(channel, model, group, tokenKey, tokenName string, duration float64) {
 	relayRequestE2EDurationObsever.WithLabelValues(channel, model, group, tokenKey, tokenName).Observe(duration)
 }
+
+// Token metrics functions
+func IncrementInputTokens(channel, model, group, userId string, add float64) {
+	inputTokensCounter.WithLabelValues(channel, model, group, userId).Add(add)
+}
+
+func IncrementOutputTokens(channel, model, group, userId string, add float64) {
+	outputTokensCounter.WithLabelValues(channel, model, group, userId).Add(add)
+}
+
+func IncrementCacheHitTokens(channel, model, group, userId string, add float64) {
+	cacheHitTokensCounter.WithLabelValues(channel, model, group, userId).Add(add)
+}
+
+func IncrementInferenceTokens(channel, model, group, userId string, add float64) {
+	inferenceTokensCounter.WithLabelValues(channel, model, group, userId).Add(add)
+}
diff --git a/relay/relay-text.go b/relay/relay-text.go
index ce1e07182f84..cc86e4797a79 100644
--- a/relay/relay-text.go
+++ b/relay/relay-text.go
@@ -138,6 +138,9 @@ func TextHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest *d
 		c.Set("prompt_tokens", promptTokens)
 	}
 
+	// Record input tokens metric
+	metrics.IncrementInputTokens(strconv.Itoa(relayInfo.ChannelId), textRequest.Model, relayInfo.Group, strconv.Itoa(relayInfo.UserId), float64(promptTokens))
+
 	priceData, err := helper.ModelPriceHelper(c, relayInfo, promptTokens, int(textRequest.MaxTokens))
 	if err != nil {
 		funcErr = service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError)
@@ -291,6 +294,7 @@ func TextHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest *d
 	} else {
 		postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "")
 	}
+
 	return nil
 }
 
@@ -405,7 +409,6 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 	useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
 	promptTokens := usage.PromptTokens
 	cacheTokens := usage.PromptTokensDetails.CachedTokens
-	//
 
 	completionTokens := usage.CompletionTokens
 	thinkingTokens := usage.CompletionTokenDetails.ReasoningTokens
@@ -479,6 +482,15 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 		logContent += ", " + extraContent
 	}
 
+	// Record token metrics
+	metrics.IncrementOutputTokens(strconv.Itoa(relayInfo.ChannelId), modelName, relayInfo.Group, strconv.Itoa(relayInfo.UserId), float64(completionTokens))
+
+	if cacheTokens > 0 {
+		metrics.IncrementCacheHitTokens(strconv.Itoa(relayInfo.ChannelId), modelName, relayInfo.Group, strconv.Itoa(relayInfo.UserId), float64(cacheTokens))
+	}
+
+	metrics.IncrementInferenceTokens(strconv.Itoa(relayInfo.ChannelId), modelName, relayInfo.Group, strconv.Itoa(relayInfo.UserId), float64(thinkingTokens))
+
 	other := service.GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, cacheTokens, cacheRatio, modelPrice)
 	model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, promptTokens, completionTokens, thinkingTokens, logModel,
 		tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other)