🚀 feature: Add request duration

controller/relay-utils.go

@@ -14,6 +14,7 @@ import (
 	"one-api/types"
 	"reflect"
 	"strconv"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/go-playground/validator/v10"
@@ -129,28 +130,6 @@ func shouldEnableChannel(err error, openAIErr *types.OpenAIError) bool {
 	return true
 }
 
-func postConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuota int, userId int, channelId int, modelRatio float64, groupRatio float64, modelName string, tokenName string) {
-	// quotaDelta is remaining quota to be consumed
-	err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
-	if err != nil {
-		common.SysError("error consuming token remain quota: " + err.Error())
-	}
-	err = model.CacheUpdateUserQuota(userId)
-	if err != nil {
-		common.SysError("error update user quota cache: " + err.Error())
-	}
-	// totalQuota is total quota consumed
-	if totalQuota != 0 {
-		logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
-		model.RecordConsumeLog(ctx, userId, channelId, totalQuota, 0, modelName, tokenName, totalQuota, logContent)
-		model.UpdateUserUsedQuotaAndRequestCount(userId, totalQuota)
-		model.UpdateChannelUsedQuota(channelId, totalQuota)
-	}
-	if totalQuota <= 0 {
-		common.LogError(ctx, fmt.Sprintf("totalQuota consumed is %d, something is wrong", totalQuota))
-	}
-}
-
 func parseModelMapping(modelMapping string) (map[string]string, error) {
 	if modelMapping == "" || modelMapping == "{}" {
 		return nil, nil
@@ -270,8 +249,17 @@ func (q *QuotaInfo) completedQuotaConsumption(usage *types.Usage, tokenName stri
 		return errors.New("error consuming token remain quota: " + err.Error())
 	}
 	if quota != 0 {
+		requestTime := 0
+		requestStartTimeValue := ctx.Value("requestStartTime")
+		if requestStartTimeValue != nil {
+			requestStartTime, ok := requestStartTimeValue.(time.Time)
+			if ok {
+				requestTime = int(time.Since(requestStartTime).Milliseconds())
+			}
+		}
+
 		logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", q.modelRatio, q.groupRatio)
-		model.RecordConsumeLog(ctx, q.userId, q.channelId, promptTokens, completionTokens, q.modelName, tokenName, quota, logContent)
+		model.RecordConsumeLog(ctx, q.userId, q.channelId, promptTokens, completionTokens, q.modelName, tokenName, quota, logContent, requestTime)
 		model.UpdateUserUsedQuotaAndRequestCount(q.userId, quota)
 		model.UpdateChannelUsedQuota(q.channelId, quota)
 	}

middleware/request-id.go

@@ -2,8 +2,10 @@ package middleware
 
 import (
 	"context"
-	"github.com/gin-gonic/gin"
 	"one-api/common"
+	"time"
+
+	"github.com/gin-gonic/gin"
 )
 
 func RequestId() func(c *gin.Context) {
@@ -11,6 +13,7 @@ func RequestId() func(c *gin.Context) {
 		id := common.GetTimeString() + common.GetRandomString(8)
 		c.Set(common.RequestIdKey, id)
 		ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id)
+		ctx = context.WithValue(ctx, "requestStartTime", time.Now())
 		c.Request = c.Request.WithContext(ctx)
 		c.Header(common.RequestIdKey, id)
 		c.Next()

model/log.go

@@ -21,6 +21,7 @@ type Log struct {
 	PromptTokens     int    `json:"prompt_tokens" gorm:"default:0"`
 	CompletionTokens int    `json:"completion_tokens" gorm:"default:0"`
 	ChannelId        int    `json:"channel" gorm:"index"`
+	RequestTime      int    `json:"request_time" gorm:"default:0"`
 }
 
 type LogStatistic struct {
@@ -57,7 +58,7 @@ func RecordLog(userId int, logType int, content string) {
 	}
 }
 
-func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string) {
+func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string, requestTime int) {
 	common.LogInfo(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
 	if !common.LogConsumeEnabled {
 		return
@@ -74,6 +75,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 		ModelName:        modelName,
 		Quota:            quota,
 		ChannelId:        channelId,
+		RequestTime:      requestTime,
 	}
 	err := DB.Create(log).Error
 	if err != nil {

web/src/views/Log/component/TableHead.js

@@ -11,6 +11,7 @@ const LogTableHead = ({ userIsAdmin }) => {
         <TableCell>令牌</TableCell>
         <TableCell>类型</TableCell>
         <TableCell>模型</TableCell>
+        <TableCell>耗时</TableCell>
         <TableCell>提示</TableCell>
         <TableCell>补全</TableCell>
         <TableCell>额度</TableCell>

web/src/views/Log/component/TableRow.js

@@ -25,7 +25,25 @@ function renderType(type) {
   }
 }
 
+function requestTimeLabelOptions(request_time) {
+  let color = 'error';
+  if (request_time === 0) {
+    color = 'default';
+  } else if (request_time <= 1000) {
+    color = 'success';
+  } else if (request_time <= 3000) {
+    color = 'primary';
+  } else if (request_time <= 5000) {
+    color = 'secondary';
+  }
+
+  return color;
+}
+
 export default function LogTableRow({ item, userIsAdmin }) {
+  let request_time = item.request_time / 1000;
+  request_time = request_time.toFixed(2) + ' 秒';
+
   return (
     <>
       <TableRow tabIndex={item.id}>
@@ -54,6 +72,10 @@ export default function LogTableRow({ item, userIsAdmin }) {
             </Label>
           )}
         </TableCell>
+        <TableCell>
+          {' '}
+          <Label color={requestTimeLabelOptions(item.request_time)}> {item.request_time == 0 ? '无' : request_time} </Label>
+        </TableCell>
         <TableCell>{item.prompt_tokens || ''}</TableCell>
         <TableCell>{item.completion_tokens || ''}</TableCell>
         <TableCell>{item.quota ? renderQuota(item.quota, 6) : ''}</TableCell>