Просмотр исходного кода

feat: Add request start time context key and update middleware

- Introduced a new constant `ContextKeyRequestStartTime` to store the request start time in the context, enhancing request tracking.
- Updated the `Distribute` middleware to set the request start time in the context using the new constant.
- Modified the `GenRelayInfo` function to retrieve the request start time from the context, ensuring accurate timing information is used in relay operations.
CalciumIon 1 год назад
Родитель
Сommit
f3f1817aea
3 измененных файлов с 13 добавлено и 6 удалено
  1. 4 0
      constant/context_key.go
  2. 2 0
      middleware/distributor.go
  3. 7 6
      relay/common/relay_info.go

+ 4 - 0
constant/context_key.go

@@ -1 +1,5 @@
 package constant
+
+const (
+	ContextKeyRequestStartTime = "request_start_time"
+)

+ 2 - 0
middleware/distributor.go

@@ -12,6 +12,7 @@ import (
 	"one-api/service"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
 )
@@ -112,6 +113,7 @@ func Distribute() func(c *gin.Context) {
 				}
 			}
 		}
+		c.Set(constant.ContextKeyRequestStartTime, time.Now())
 		SetupContextForSelectedChannel(c, channel, modelRequest.Model)
 		c.Next()
 	}

+ 7 - 6
relay/common/relay_info.go

@@ -2,8 +2,9 @@ package common
 
 import (
 	"one-api/common"
+	"one-api/constant"
 	"one-api/dto"
-	"one-api/relay/constant"
+	relayconstant "one-api/relay/constant"
 	"strings"
 	"time"
 
@@ -66,13 +67,13 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 	userId := c.GetInt("id")
 	group := c.GetString("group")
 	tokenUnlimited := c.GetBool("token_unlimited_quota")
-	startTime := time.Now()
+	startTime := c.GetTime(constant.ContextKeyRequestStartTime)
 	// firstResponseTime = time.Now() - 1 second
 
-	apiType, _ := constant.ChannelType2APIType(channelType)
+	apiType, _ := relayconstant.ChannelType2APIType(channelType)
 
 	info := &RelayInfo{
-		RelayMode:         constant.Path2RelayMode(c.Request.URL.Path),
+		RelayMode:         relayconstant.Path2RelayMode(c.Request.URL.Path),
 		BaseUrl:           c.GetString("base_url"),
 		RequestURLPath:    c.Request.URL.String(),
 		ChannelType:       channelType,
@@ -158,10 +159,10 @@ func GenTaskRelayInfo(c *gin.Context) *TaskRelayInfo {
 	group := c.GetString("group")
 	startTime := time.Now()
 
-	apiType, _ := constant.ChannelType2APIType(channelType)
+	apiType, _ := relayconstant.ChannelType2APIType(channelType)
 
 	info := &TaskRelayInfo{
-		RelayMode:      constant.Path2RelayMode(c.Request.URL.Path),
+		RelayMode:      relayconstant.Path2RelayMode(c.Request.URL.Path),
 		BaseUrl:        c.GetString("base_url"),
 		RequestURLPath: c.Request.URL.String(),
 		ChannelType:    channelType,