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.
This commit is contained in:
CalciumIon
2024-12-23 20:48:10 +08:00
parent 58fac129d6
commit f3f1817aea
3 changed files with 13 additions and 6 deletions

View File

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

View File

@@ -12,6 +12,7 @@ import (
"one-api/service" "one-api/service"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/gin-gonic/gin" "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) SetupContextForSelectedChannel(c, channel, modelRequest.Model)
c.Next() c.Next()
} }

View File

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