From 352ef053d5d6577b9ecbc9ca608acd3c8a19bd19 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Sun, 17 Sep 2023 21:07:00 +0800
Subject: [PATCH] =?UTF-8?q?=E5=90=88=E5=B9=B6=E6=9C=80=E6=96=B0=E4=BB=A3?=
 =?UTF-8?q?=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 controller/log.go        | 2 ++
 controller/relay-mj.go   | 9 +++++----
 controller/relay-text.go | 1 -
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/controller/log.go b/controller/log.go
index c43c5483..3e9721a3 100644
--- a/controller/log.go
+++ b/controller/log.go
@@ -2,6 +2,7 @@ package controller
 
 import (
 	"github.com/gin-gonic/gin"
+	"net/http"
 	"one-api/common"
 	"one-api/model"
 	"strconv"
@@ -123,6 +124,7 @@ func GetLogsStat(c *gin.Context) {
 	tokenName := c.Query("token_name")
 	username := c.Query("username")
 	modelName := c.Query("model_name")
+	channel, _ := strconv.Atoi(c.Query("channel"))
 	stat := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel)
 	//tokenNum := model.SumUsedToken(logType, startTimestamp, endTimestamp, modelName, username, "")
 	c.JSON(http.StatusOK, gin.H{
diff --git a/controller/relay-mj.go b/controller/relay-mj.go
index 211480f7..adcf027c 100644
--- a/controller/relay-mj.go
+++ b/controller/relay-mj.go
@@ -2,6 +2,7 @@ package controller
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -145,7 +146,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 	userId := c.GetInt("id")
 	consumeQuota := c.GetBool("consume_quota")
 	group := c.GetString("group")
-
+	channelId := c.GetInt("channel_id")
 	var midjRequest MidjourneyRequest
 	if consumeQuota {
 		err := common.UnmarshalBodyReusable(c, &midjRequest)
@@ -308,7 +309,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 	}
 	var midjResponse MidjourneyResponse
 
-	defer func() {
+	defer func(ctx context.Context) {
 		if consumeQuota {
 			err := model.PostConsumeTokenQuota(tokenId, quota)
 			if err != nil {
@@ -321,13 +322,13 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 			if quota != 0 {
 				tokenName := c.GetString("token_name")
 				logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
-				model.RecordConsumeLog(userId, 0, 0, imageModel, tokenName, quota, logContent, tokenId)
+				model.RecordConsumeLog(ctx, userId, channelId, 0, 0, imageModel, tokenName, quota, logContent, tokenId)
 				model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
 				channelId := c.GetInt("channel_id")
 				model.UpdateChannelUsedQuota(channelId, quota)
 			}
 		}
-	}()
+	}(c.Request.Context())
 
 	//if consumeQuota {
 	//
diff --git a/controller/relay-text.go b/controller/relay-text.go
index 1081c2b5..b5c84f9b 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -372,7 +372,6 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 
 	var textResponse TextResponse
 	tokenName := c.GetString("token_name")
-	channelId := c.GetInt("channel_id")
 
 	defer func(ctx context.Context) {
 		// c.Writer.Flush()