diff --git a/controller/relay.go b/controller/relay.go index c97ab45b4ac4..a80d7de2ff99 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -15,6 +15,7 @@ import ( "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/monitor" "github.com/QuantumNous/new-api/relay" relaycommon "github.com/QuantumNous/new-api/relay/common" relayconstant "github.com/QuantumNous/new-api/relay/constant" @@ -221,6 +222,17 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { if newAPIError == nil { relayInfo.LastError = nil + monitor.RecordRelayRequest(&monitor.RelayMetricsData{ + ChannelId: relayInfo.ChannelId, + ChannelType: relayInfo.ChannelType, + Model: relayInfo.OriginModelName, + RelayMode: relayInfo.RelayMode, + StatusCode: 200, + StartTime: relayInfo.StartTime, + FirstTokenTime: relayInfo.FirstResponseTime, + IsStream: relayInfo.IsStream, + }) + monitor.RecordTokenRequest(relayInfo.UserId, relayInfo.TokenId) return } @@ -232,6 +244,22 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) { break } + monitor.RecordRelayRetry(relayInfo.ChannelId, relayInfo.ChannelType) + } + + // Record metrics for failed relay + if newAPIError != nil { + monitor.RecordRelayRequest(&monitor.RelayMetricsData{ + ChannelId: relayInfo.ChannelId, + ChannelType: relayInfo.ChannelType, + Model: relayInfo.OriginModelName, + RelayMode: relayInfo.RelayMode, + StatusCode: newAPIError.StatusCode, + StartTime: relayInfo.StartTime, + IsStream: relayInfo.IsStream, + ErrorType: string(newAPIError.GetErrorCode()), + }) + monitor.RecordTokenRequest(relayInfo.UserId, relayInfo.TokenId) } useChannel := c.GetStringSlice("use_channel") diff --git a/go.mod b/go.mod index a078b2091ad3..977b0f659738 100644 --- a/go.mod +++ b/go.mod @@ -35,6 +35,7 @@ require ( github.com/nicksnyder/go-i18n/v2 v2.6.1 github.com/pkg/errors v0.9.1 github.com/pquerna/otp v1.5.0 + github.com/prometheus/client_golang v1.22.0 github.com/samber/hot v0.11.0 github.com/samber/lo v1.52.0 github.com/shirou/gopsutil v3.21.11+incompatible @@ -115,7 +116,6 @@ require ( github.com/ncruces/go-strftime v0.1.9 // indirect github.com/pelletier/go-toml/v2 v2.2.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect - github.com/prometheus/client_golang v1.22.0 // indirect github.com/prometheus/client_model v0.6.1 // indirect github.com/prometheus/common v0.62.0 // indirect github.com/prometheus/procfs v0.15.1 // indirect diff --git a/go.sum b/go.sum index 8b687906336c..a0a712a1860c 100644 --- a/go.sum +++ b/go.sum @@ -186,6 +186,8 @@ github.com/kr/pty v1.1.8/go.mod h1:O1sed60cT9XZ5uDucP5qwvh+TE3NnUj51EiZO/lmSfw= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII= github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= diff --git a/main.go b/main.go index dbbf44a1826b..7322fa40cb55 100644 --- a/main.go +++ b/main.go @@ -18,6 +18,7 @@ import ( "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/monitor" "github.com/QuantumNous/new-api/oauth" "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/router" @@ -165,6 +166,7 @@ func main() { // This will cause SSE not to work!!! //server.Use(gzip.Gzip(gzip.DefaultCompression)) server.Use(middleware.RequestId()) + server.Use(monitor.PrometheusMiddleware()) server.Use(middleware.PoweredBy()) server.Use(middleware.I18n()) middleware.SetUpLogger(server) @@ -252,6 +254,12 @@ func InitResources() error { // 加载环境变量 common.InitEnv() + // Initialize Prometheus metrics + if os.Getenv("PROMETHEUS_ENABLED") != "false" { + monitor.InitMetrics() + go monitor.StartMetricsServer() + } + logger.SetupLogger() // Initialize model settings @@ -291,6 +299,11 @@ func InitResources() error { return err } + // Register Prometheus metrics hook for Redis + if common.RedisEnabled { + monitor.RegisterRedisHook(common.RDB) + } + // 启动系统监控 common.StartSystemMonitor() diff --git a/model/main.go b/model/main.go index f37cb667cd43..ff08914df0fb 100644 --- a/model/main.go +++ b/model/main.go @@ -10,6 +10,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/monitor" "github.com/glebarez/sqlite" "gorm.io/driver/mysql" @@ -195,6 +196,10 @@ func InitDB() (err error) { sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000)) sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60))) + // Register Prometheus DB metrics + monitor.RegisterDBCallbacks(DB) + monitor.StartDBPoolCollector(sqlDB) + if !common.IsMasterNode { return nil } diff --git a/monitor/db.go b/monitor/db.go new file mode 100644 index 000000000000..474c926b3676 --- /dev/null +++ b/monitor/db.go @@ -0,0 +1,100 @@ +package monitor + +import ( + "database/sql" + "time" + + "gorm.io/gorm" +) + +const dbStartTimeKey = "prom_start_time" + +// RegisterDBCallbacks adds Prometheus instrumentation callbacks to a GORM DB. +func RegisterDBCallbacks(db *gorm.DB) { + if db == nil || !Enabled() { + return + } + + registerCallback(db, "create") + registerCallback(db, "query") + registerCallback(db, "update") + registerCallback(db, "delete") + registerCallback(db, "raw") +} + +func registerCallback(db *gorm.DB, operation string) { + beforeName := "prometheus:before_" + operation + afterName := "prometheus:after_" + operation + + before := func(db *gorm.DB) { + db.Set(dbStartTimeKey, time.Now()) + } + after := func(db *gorm.DB) { + recordDBMetrics(db, operation) + } + + switch operation { + case "create": + _ = db.Callback().Create().Before("*").Register(beforeName, before) + _ = db.Callback().Create().After("*").Register(afterName, after) + case "query": + _ = db.Callback().Query().Before("*").Register(beforeName, before) + _ = db.Callback().Query().After("*").Register(afterName, after) + case "update": + _ = db.Callback().Update().Before("*").Register(beforeName, before) + _ = db.Callback().Update().After("*").Register(afterName, after) + case "delete": + _ = db.Callback().Delete().Before("*").Register(beforeName, before) + _ = db.Callback().Delete().After("*").Register(afterName, after) + case "raw": + _ = db.Callback().Raw().Before("*").Register(beforeName, before) + _ = db.Callback().Raw().After("*").Register(afterName, after) + } +} + +func recordDBMetrics(db *gorm.DB, operation string) { + if !Enabled() { + return + } + status := "success" + if db.Error != nil { + status = "error" + } + DBQueriesTotal.WithLabelValues(operation, status).Inc() + + if v, ok := db.Get(dbStartTimeKey); ok { + if start, ok := v.(time.Time); ok { + DBQueryDuration.WithLabelValues(operation).Observe(time.Since(start).Seconds()) + } + } +} + +// StartDBPoolCollector starts a background goroutine that periodically collects DB pool stats. +func StartDBPoolCollector(sqlDB *sql.DB) { + if sqlDB == nil || !Enabled() { + return + } + + go func() { + var lastWaitCount int64 + var lastWaitDuration time.Duration + + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + for range ticker.C { + stats := sqlDB.Stats() + DBOpenConnections.WithLabelValues("open").Set(float64(stats.OpenConnections)) + DBOpenConnections.WithLabelValues("in_use").Set(float64(stats.InUse)) + DBOpenConnections.WithLabelValues("idle").Set(float64(stats.Idle)) + + if delta := stats.WaitCount - lastWaitCount; delta > 0 { + DBWaitTotal.Add(float64(delta)) + lastWaitCount = stats.WaitCount + } + if delta := stats.WaitDuration - lastWaitDuration; delta > 0 { + DBWaitDuration.Add(delta.Seconds()) + lastWaitDuration = stats.WaitDuration + } + } + }() +} diff --git a/monitor/metrics.go b/monitor/metrics.go new file mode 100644 index 000000000000..a8961ed02564 --- /dev/null +++ b/monitor/metrics.go @@ -0,0 +1,235 @@ +package monitor + +import ( + "os" + "sync" + + "github.com/prometheus/client_golang/prometheus" +) + +const namespace = "newapi" + +var ( + initOnce sync.Once + enabled bool + tier3Enabled bool +) + +// ---- Tier 1: System-Level (Low Cardinality) ---- + +var ( + HTTPRequestsTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "http_requests_total", + Help: "Total number of HTTP requests.", + }, + []string{"method", "path", "status_code"}, + ) + HTTPRequestDuration = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Name: "http_request_duration_seconds", + Help: "HTTP request latency in seconds.", + Buckets: prometheus.DefBuckets, + }, + []string{"method", "path", "status_code"}, + ) + HTTPActiveConnections = prometheus.NewGauge( + prometheus.GaugeOpts{ + Namespace: namespace, + Name: "http_active_connections", + Help: "Number of currently active HTTP connections.", + }, + ) + + RedisCommandsTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "redis_commands_total", + Help: "Total number of Redis commands executed.", + }, + []string{"command", "status"}, + ) + RedisCommandDuration = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Name: "redis_command_duration_seconds", + Help: "Redis command latency in seconds.", + Buckets: []float64{.0005, .001, .005, .01, .025, .05, .1, .25, .5, 1}, + }, + []string{"command"}, + ) + + DBQueriesTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "db_queries_total", + Help: "Total number of database queries.", + }, + []string{"operation", "status"}, + ) + DBQueryDuration = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Name: "db_query_duration_seconds", + Help: "Database query latency in seconds.", + Buckets: []float64{.001, .005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5}, + }, + []string{"operation"}, + ) + DBOpenConnections = prometheus.NewGaugeVec( + prometheus.GaugeOpts{ + Namespace: namespace, + Name: "db_open_connections", + Help: "Number of database connections by state.", + }, + []string{"state"}, + ) + DBWaitTotal = prometheus.NewCounter( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "db_wait_total", + Help: "Total number of waits for a database connection.", + }, + ) + DBWaitDuration = prometheus.NewCounter( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "db_wait_duration_seconds_total", + Help: "Total time spent waiting for a database connection in seconds.", + }, + ) +) + +// ---- Tier 2: Relay/Channel-Level (Moderate Cardinality) ---- + +var ( + RelayRequestsTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "relay_requests_total", + Help: "Total relay requests per channel.", + }, + []string{"channel_id", "channel_type", "model", "status_code", "relay_mode"}, + ) + RelayRequestDuration = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Name: "relay_request_duration_seconds", + Help: "End-to-end relay request latency in seconds.", + Buckets: []float64{.1, .25, .5, 1, 2.5, 5, 10, 30, 60, 120}, + }, + []string{"channel_id", "channel_type", "model", "relay_mode"}, + ) + RelayFirstTokenDuration = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Name: "relay_first_token_seconds", + Help: "Time to first token for streaming requests in seconds.", + Buckets: []float64{.05, .1, .25, .5, 1, 2.5, 5, 10, 30}, + }, + []string{"channel_id", "channel_type", "model"}, + ) + RelayTokensUsedTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "relay_tokens_used_total", + Help: "Total tokens consumed via relay.", + }, + []string{"channel_id", "channel_type", "model", "direction"}, + ) + RelayErrorsTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "relay_errors_total", + Help: "Total relay errors by type.", + }, + []string{"channel_id", "channel_type", "error_type"}, + ) + RelayRetriesTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "relay_retries_total", + Help: "Total relay retries.", + }, + []string{"channel_id", "channel_type"}, + ) +) + +// ---- Tier 3: User/Token-Level (High Cardinality, Counters Only) ---- + +var ( + TokenRequestsTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "token_requests_total", + Help: "Total requests per user per token.", + }, + []string{"user_id", "token_id"}, + ) + TokenTokensUsedTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "token_tokens_used_total", + Help: "Total tokens consumed per user per token.", + }, + []string{"user_id", "token_id", "model", "direction"}, + ) + TokenQuotaConsumedTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "token_quota_consumed_total", + Help: "Total quota units consumed per user per token.", + }, + []string{"user_id", "token_id"}, + ) +) + +// InitMetrics registers all Prometheus metrics. Safe to call multiple times. +func InitMetrics() { + initOnce.Do(func() { + // Tier 1 + prometheus.MustRegister( + HTTPRequestsTotal, + HTTPRequestDuration, + HTTPActiveConnections, + RedisCommandsTotal, + RedisCommandDuration, + DBQueriesTotal, + DBQueryDuration, + DBOpenConnections, + DBWaitTotal, + DBWaitDuration, + ) + // Tier 2 + prometheus.MustRegister( + RelayRequestsTotal, + RelayRequestDuration, + RelayFirstTokenDuration, + RelayTokensUsedTotal, + RelayErrorsTotal, + RelayRetriesTotal, + ) + // Tier 3 (high cardinality — opt-in via PROMETHEUS_TIER3_ENABLED) + if os.Getenv("PROMETHEUS_TIER3_ENABLED") == "true" { + prometheus.MustRegister( + TokenRequestsTotal, + TokenTokensUsedTotal, + TokenQuotaConsumedTotal, + ) + tier3Enabled = true + } + enabled = true + }) +} + +// Enabled returns whether Prometheus metrics are initialized. +func Enabled() bool { + return enabled +} + +// Tier3Enabled returns whether high-cardinality user/token metrics are enabled. +func Tier3Enabled() bool { + return tier3Enabled +} diff --git a/monitor/middleware.go b/monitor/middleware.go new file mode 100644 index 000000000000..8a4a69cce67d --- /dev/null +++ b/monitor/middleware.go @@ -0,0 +1,38 @@ +package monitor + +import ( + "strconv" + "time" + + "github.com/gin-gonic/gin" +) + +// PrometheusMiddleware returns a Gin middleware that records HTTP request metrics. +func PrometheusMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if !Enabled() { + c.Next() + return + } + + HTTPActiveConnections.Inc() + start := time.Now() + + defer func() { + HTTPActiveConnections.Dec() + + path := c.FullPath() + if path == "" { + path = "unknown" + } + method := c.Request.Method + statusCode := strconv.Itoa(c.Writer.Status()) + duration := time.Since(start).Seconds() + + HTTPRequestsTotal.WithLabelValues(method, path, statusCode).Inc() + HTTPRequestDuration.WithLabelValues(method, path, statusCode).Observe(duration) + }() + + c.Next() + } +} diff --git a/monitor/redis_hook.go b/monitor/redis_hook.go new file mode 100644 index 000000000000..1bd75237eb1c --- /dev/null +++ b/monitor/redis_hook.go @@ -0,0 +1,67 @@ +package monitor + +import ( + "context" + "strings" + "time" + + "github.com/go-redis/redis/v8" +) + +type redisMetricsHookKey struct{} + +type redisMetricsHook struct{} + +var _ redis.Hook = (*redisMetricsHook)(nil) + +func (h *redisMetricsHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) { + return context.WithValue(ctx, redisMetricsHookKey{}, time.Now()), nil +} + +func (h *redisMetricsHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error { + if !Enabled() { + return nil + } + command := strings.ToUpper(cmd.Name()) + status := "success" + if cmd.Err() != nil && cmd.Err() != redis.Nil { + status = "error" + } + RedisCommandsTotal.WithLabelValues(command, status).Inc() + + if start, ok := ctx.Value(redisMetricsHookKey{}).(time.Time); ok { + RedisCommandDuration.WithLabelValues(command).Observe(time.Since(start).Seconds()) + } + return nil +} + +func (h *redisMetricsHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) { + return context.WithValue(ctx, redisMetricsHookKey{}, time.Now()), nil +} + +func (h *redisMetricsHook) AfterProcessPipeline(ctx context.Context, cmds []redis.Cmder) error { + if !Enabled() { + return nil + } + start, _ := ctx.Value(redisMetricsHookKey{}).(time.Time) + for _, cmd := range cmds { + command := strings.ToUpper(cmd.Name()) + status := "success" + if cmd.Err() != nil && cmd.Err() != redis.Nil { + status = "error" + } + RedisCommandsTotal.WithLabelValues(command, status).Inc() + } + if !start.IsZero() { + RedisCommandDuration.WithLabelValues("PIPELINE").Observe(time.Since(start).Seconds()) + } + return nil +} + +// RegisterRedisHook adds the Prometheus metrics hook to a Redis client. +func RegisterRedisHook(rdb *redis.Client) { + if rdb == nil || !Enabled() { + return + } + rdb.AddHook(&redisMetricsHook{}) +} diff --git a/monitor/relay.go b/monitor/relay.go new file mode 100644 index 000000000000..5b16e0d01da3 --- /dev/null +++ b/monitor/relay.go @@ -0,0 +1,70 @@ +package monitor + +import ( + "strconv" + "time" +) + +// RelayMetricsData holds the data needed to record relay metrics. +type RelayMetricsData struct { + ChannelId int + ChannelType int + Model string + RelayMode int + StatusCode int + StartTime time.Time + FirstTokenTime time.Time + IsStream bool + ErrorType string +} + +// RecordRelayRequest records relay request count, duration, first-token time, and errors. +func RecordRelayRequest(d *RelayMetricsData) { + if !Enabled() || d == nil { + return + } + + chID := strconv.Itoa(d.ChannelId) + chType := strconv.Itoa(d.ChannelType) + model := d.Model + mode := strconv.Itoa(d.RelayMode) + status := strconv.Itoa(d.StatusCode) + + RelayRequestsTotal.WithLabelValues(chID, chType, model, status, mode).Inc() + + duration := time.Since(d.StartTime).Seconds() + RelayRequestDuration.WithLabelValues(chID, chType, model, mode).Observe(duration) + + if d.IsStream && !d.FirstTokenTime.IsZero() { + ttft := d.FirstTokenTime.Sub(d.StartTime).Seconds() + RelayFirstTokenDuration.WithLabelValues(chID, chType, model).Observe(ttft) + } + + if d.ErrorType != "" { + RelayErrorsTotal.WithLabelValues(chID, chType, d.ErrorType).Inc() + } +} + +// RecordRelayRetry records a relay retry event. +func RecordRelayRetry(channelId, channelType int) { + if !Enabled() { + return + } + RelayRetriesTotal.WithLabelValues(strconv.Itoa(channelId), strconv.Itoa(channelType)).Inc() +} + +// RecordRelayTokens records token usage for a relay request. +func RecordRelayTokens(channelId, channelType int, model string, inputTokens, outputTokens int) { + if !Enabled() { + return + } + chID := strconv.Itoa(channelId) + chType := strconv.Itoa(channelType) + + if inputTokens > 0 { + RelayTokensUsedTotal.WithLabelValues(chID, chType, model, "input").Add(float64(inputTokens)) + } + if outputTokens > 0 { + RelayTokensUsedTotal.WithLabelValues(chID, chType, model, "output").Add(float64(outputTokens)) + } +} diff --git a/monitor/server.go b/monitor/server.go new file mode 100644 index 000000000000..637c532f60ec --- /dev/null +++ b/monitor/server.go @@ -0,0 +1,35 @@ +package monitor + +import ( + "fmt" + "log" + "net/http" + "os" + "strconv" + + "github.com/prometheus/client_golang/prometheus/promhttp" +) + +// StartMetricsServer starts a dedicated HTTP server for Prometheus metrics. +// It blocks, so call it in a goroutine. +func StartMetricsServer() { + if !Enabled() { + return + } + + port := 9090 + if p := os.Getenv("PROMETHEUS_PORT"); p != "" { + if v, err := strconv.Atoi(p); err == nil && v > 0 { + port = v + } + } + + mux := http.NewServeMux() + mux.Handle("/metrics", promhttp.Handler()) + + addr := fmt.Sprintf(":%d", port) + log.Printf("Prometheus metrics server listening on %s", addr) + if err := http.ListenAndServe(addr, mux); err != nil { + log.Printf("Prometheus metrics server error: %v", err) + } +} diff --git a/monitor/token.go b/monitor/token.go new file mode 100644 index 000000000000..db72f9be1f78 --- /dev/null +++ b/monitor/token.go @@ -0,0 +1,32 @@ +package monitor + +import "strconv" + +// RecordTokenRequest increments the per-user per-token request counter. +func RecordTokenRequest(userId, tokenId int) { + if !Tier3Enabled() { + return + } + uid := strconv.Itoa(userId) + tid := strconv.Itoa(tokenId) + TokenRequestsTotal.WithLabelValues(uid, tid).Inc() +} + +// RecordTokenUsage records token and quota usage for a specific user/token. +func RecordTokenUsage(userId, tokenId int, model string, inputTokens, outputTokens, quota int) { + if !Tier3Enabled() { + return + } + uid := strconv.Itoa(userId) + tid := strconv.Itoa(tokenId) + + if inputTokens > 0 { + TokenTokensUsedTotal.WithLabelValues(uid, tid, model, "input").Add(float64(inputTokens)) + } + if outputTokens > 0 { + TokenTokensUsedTotal.WithLabelValues(uid, tid, model, "output").Add(float64(outputTokens)) + } + if quota > 0 { + TokenQuotaConsumedTotal.WithLabelValues(uid, tid).Add(float64(quota)) + } +} diff --git a/service/text_quota.go b/service/text_quota.go index 8caee8f28799..db4cbbcedf0c 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -10,6 +10,7 @@ import ( "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/monitor" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/QuantumNous/new-api/types" @@ -303,6 +304,10 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us adminRejectReason := common.GetContextKeyString(ctx, constant.ContextKeyAdminRejectReason) summary := calculateTextQuotaSummary(ctx, relayInfo, usage) + // Record Prometheus metrics for token usage + monitor.RecordRelayTokens(relayInfo.ChannelId, relayInfo.ChannelType, summary.ModelName, summary.PromptTokens, summary.CompletionTokens) + monitor.RecordTokenUsage(relayInfo.UserId, relayInfo.TokenId, summary.ModelName, summary.PromptTokens, summary.CompletionTokens, summary.Quota) + if summary.WebSearchCallCount > 0 { extraContent = append(extraContent, fmt.Sprintf("Web Search 调用 %d 次，调用花费 %s", summary.WebSearchCallCount, decimal.NewFromFloat(summary.WebSearchPrice).Mul(decimal.NewFromInt(int64(summary.WebSearchCallCount))).Div(decimal.NewFromInt(1000)).Mul(decimal.NewFromFloat(summary.GroupRatio)).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).String())) }