| @@ -21,7 +21,9 @@ func GetAllLogs(c *gin.Context) { | |||||
| channel, _ := strconv.Atoi(c.Query("channel")) | channel, _ := strconv.Atoi(c.Query("channel")) | ||||
| group := c.Query("group") | group := c.Query("group") | ||||
| requestId := c.Query("request_id") | requestId := c.Query("request_id") | ||||
| logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId) | |||||
| chatId := c.Query("chat_id") | |||||
| upstreamId := c.Query("upstream_id") | |||||
| logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, chatId, upstreamId) | |||||
| if err != nil { | if err != nil { | ||||
| common.ApiError(c, err) | common.ApiError(c, err) | ||||
| return | return | ||||
| @@ -42,7 +44,9 @@ func GetUserLogs(c *gin.Context) { | |||||
| modelName := c.Query("model_name") | modelName := c.Query("model_name") | ||||
| group := c.Query("group") | group := c.Query("group") | ||||
| requestId := c.Query("request_id") | requestId := c.Query("request_id") | ||||
| logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId) | |||||
| chatId := c.Query("chat_id") | |||||
| upstreamId := c.Query("upstream_id") | |||||
| logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, chatId, upstreamId) | |||||
| if err != nil { | if err != nil { | ||||
| common.ApiError(c, err) | common.ApiError(c, err) | ||||
| return | return | ||||
| @@ -0,0 +1,94 @@ | |||||
| package controller | |||||
| import ( | |||||
| "encoding/json" | |||||
| "net/http" | |||||
| "net/http/httptest" | |||||
| "testing" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/QuantumNous/new-api/model" | |||||
| "github.com/glebarez/sqlite" | |||||
| "github.com/gin-gonic/gin" | |||||
| "github.com/stretchr/testify/require" | |||||
| "gorm.io/gorm" | |||||
| ) | |||||
| func setupControllerLogIdentityDB(t *testing.T) *gorm.DB { | |||||
| t.Helper() | |||||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||||
| require.NoError(t, err) | |||||
| origLogDB := model.LOG_DB | |||||
| origSQLite := common.UsingSQLite | |||||
| origMySQL := common.UsingMySQL | |||||
| origPostgreSQL := common.UsingPostgreSQL | |||||
| model.LOG_DB = db | |||||
| common.UsingSQLite = true | |||||
| common.UsingMySQL = false | |||||
| common.UsingPostgreSQL = false | |||||
| require.NoError(t, db.AutoMigrate(&model.Log{})) | |||||
| t.Cleanup(func() { | |||||
| model.LOG_DB = origLogDB | |||||
| common.UsingSQLite = origSQLite | |||||
| common.UsingMySQL = origMySQL | |||||
| common.UsingPostgreSQL = origPostgreSQL | |||||
| }) | |||||
| return db | |||||
| } | |||||
| func TestGetAllLogsPassesChatIDAndUpstreamIDFilters(t *testing.T) { | |||||
| db := setupControllerLogIdentityDB(t) | |||||
| require.NoError(t, db.Create(&model.Log{ | |||||
| UserId: 1, | |||||
| Username: "alice", | |||||
| CreatedAt: 1714465001, | |||||
| Type: model.LogTypeConsume, | |||||
| Content: "usage-a", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| RequestId: "req_a", | |||||
| ChatId: "chat_a", | |||||
| UpstreamId: "up_a", | |||||
| }).Error) | |||||
| require.NoError(t, db.Create(&model.Log{ | |||||
| UserId: 1, | |||||
| Username: "alice", | |||||
| CreatedAt: 1714465002, | |||||
| Type: model.LogTypeConsume, | |||||
| Content: "usage-b", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| RequestId: "req_b", | |||||
| ChatId: "chat_b", | |||||
| UpstreamId: "up_b", | |||||
| }).Error) | |||||
| gin.SetMode(gin.TestMode) | |||||
| router := gin.New() | |||||
| router.GET("/api/log", GetAllLogs) | |||||
| req := httptest.NewRequest(http.MethodGet, "/api/log?type=2&chat_id=chat_a&upstream_id=up_a", nil) | |||||
| w := httptest.NewRecorder() | |||||
| router.ServeHTTP(w, req) | |||||
| require.Equal(t, http.StatusOK, w.Code) | |||||
| var payload struct { | |||||
| Success bool `json:"success"` | |||||
| Data struct { | |||||
| Items []model.Log `json:"items"` | |||||
| } `json:"data"` | |||||
| } | |||||
| require.NoError(t, json.Unmarshal(w.Body.Bytes(), &payload)) | |||||
| require.True(t, payload.Success) | |||||
| require.Len(t, payload.Data.Items, 1) | |||||
| require.Equal(t, "chat_a", payload.Data.Items[0].ChatId) | |||||
| require.Equal(t, "up_a", payload.Data.Items[0].UpstreamId) | |||||
| } | |||||
| @@ -8,6 +8,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/logger" | "github.com/QuantumNous/new-api/logger" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||||
| "github.com/QuantumNous/new-api/types" | "github.com/QuantumNous/new-api/types" | ||||
| "github.com/gin-gonic/gin" | "github.com/gin-gonic/gin" | ||||
| @@ -36,6 +37,8 @@ type Log struct { | |||||
| Group string `json:"group" gorm:"index"` | Group string `json:"group" gorm:"index"` | ||||
| Ip string `json:"ip" gorm:"index;default:''"` | Ip string `json:"ip" gorm:"index;default:''"` | ||||
| RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"` | RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"` | ||||
| ChatId string `json:"chat_id,omitempty" gorm:"type:varchar(128);index:idx_logs_chat_id;default:''"` | |||||
| UpstreamId string `json:"upstream_id,omitempty" gorm:"type:varchar(128);index:idx_logs_upstream_id;default:''"` | |||||
| Other string `json:"other"` | Other string `json:"other"` | ||||
| } | } | ||||
| @@ -94,6 +97,8 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, | |||||
| logger.LogInfo(c, fmt.Sprintf("record error log: userId=%d, channelId=%d, modelName=%s, tokenName=%s, content=%s", userId, channelId, modelName, tokenName, content)) | logger.LogInfo(c, fmt.Sprintf("record error log: userId=%d, channelId=%d, modelName=%s, tokenName=%s, content=%s", userId, channelId, modelName, tokenName, content)) | ||||
| username := c.GetString("username") | username := c.GetString("username") | ||||
| requestId := c.GetString(common.RequestIdKey) | requestId := c.GetString(common.RequestIdKey) | ||||
| chatId := relaycommon.GetRelayChatID(c) | |||||
| upstreamId := relaycommon.GetRelayUpstreamID(c) | |||||
| otherStr := common.MapToJsonStr(other) | otherStr := common.MapToJsonStr(other) | ||||
| // 判断是否需要记录 IP | // 判断是否需要记录 IP | ||||
| needRecordIp := false | needRecordIp := false | ||||
| @@ -124,8 +129,10 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, | |||||
| } | } | ||||
| return "" | return "" | ||||
| }(), | }(), | ||||
| RequestId: requestId, | |||||
| Other: otherStr, | |||||
| RequestId: requestId, | |||||
| ChatId: chatId, | |||||
| UpstreamId: upstreamId, | |||||
| Other: otherStr, | |||||
| } | } | ||||
| err := LOG_DB.Create(log).Error | err := LOG_DB.Create(log).Error | ||||
| if err != nil { | if err != nil { | ||||
| @@ -155,6 +162,8 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) | |||||
| logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params))) | logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params))) | ||||
| username := c.GetString("username") | username := c.GetString("username") | ||||
| requestId := c.GetString(common.RequestIdKey) | requestId := c.GetString(common.RequestIdKey) | ||||
| chatId := relaycommon.GetRelayChatID(c) | |||||
| upstreamId := relaycommon.GetRelayUpstreamID(c) | |||||
| otherStr := common.MapToJsonStr(params.Other) | otherStr := common.MapToJsonStr(params.Other) | ||||
| // 判断是否需要记录 IP | // 判断是否需要记录 IP | ||||
| needRecordIp := false | needRecordIp := false | ||||
| @@ -185,8 +194,10 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) | |||||
| } | } | ||||
| return "" | return "" | ||||
| }(), | }(), | ||||
| RequestId: requestId, | |||||
| Other: otherStr, | |||||
| RequestId: requestId, | |||||
| ChatId: chatId, | |||||
| UpstreamId: upstreamId, | |||||
| Other: otherStr, | |||||
| } | } | ||||
| err := LOG_DB.Create(log).Error | err := LOG_DB.Create(log).Error | ||||
| if err != nil { | if err != nil { | ||||
| @@ -242,7 +253,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { | |||||
| } | } | ||||
| } | } | ||||
| func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string) (logs []*Log, total int64, err error) { | |||||
| func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, chatId string, upstreamId string) (logs []*Log, total int64, err error) { | |||||
| var tx *gorm.DB | var tx *gorm.DB | ||||
| if logType == LogTypeUnknown { | if logType == LogTypeUnknown { | ||||
| tx = LOG_DB | tx = LOG_DB | ||||
| @@ -262,6 +273,12 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName | |||||
| if requestId != "" { | if requestId != "" { | ||||
| tx = tx.Where("logs.request_id = ?", requestId) | tx = tx.Where("logs.request_id = ?", requestId) | ||||
| } | } | ||||
| if chatId != "" { | |||||
| tx = tx.Where("logs.chat_id = ?", chatId) | |||||
| } | |||||
| if upstreamId != "" { | |||||
| tx = tx.Where("logs.upstream_id = ?", upstreamId) | |||||
| } | |||||
| if startTimestamp != 0 { | if startTimestamp != 0 { | ||||
| tx = tx.Where("logs.created_at >= ?", startTimestamp) | tx = tx.Where("logs.created_at >= ?", startTimestamp) | ||||
| } | } | ||||
| @@ -328,7 +345,7 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName | |||||
| const logSearchCountLimit = 10000 | const logSearchCountLimit = 10000 | ||||
| func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string) (logs []*Log, total int64, err error) { | |||||
| func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, chatId string, upstreamId string) (logs []*Log, total int64, err error) { | |||||
| var tx *gorm.DB | var tx *gorm.DB | ||||
| if logType == LogTypeUnknown { | if logType == LogTypeUnknown { | ||||
| tx = LOG_DB.Where("logs.user_id = ?", userId) | tx = LOG_DB.Where("logs.user_id = ?", userId) | ||||
| @@ -349,6 +366,12 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int | |||||
| if requestId != "" { | if requestId != "" { | ||||
| tx = tx.Where("logs.request_id = ?", requestId) | tx = tx.Where("logs.request_id = ?", requestId) | ||||
| } | } | ||||
| if chatId != "" { | |||||
| tx = tx.Where("logs.chat_id = ?", chatId) | |||||
| } | |||||
| if upstreamId != "" { | |||||
| tx = tx.Where("logs.upstream_id = ?", upstreamId) | |||||
| } | |||||
| if startTimestamp != 0 { | if startTimestamp != 0 { | ||||
| tx = tx.Where("logs.created_at >= ?", startTimestamp) | tx = tx.Where("logs.created_at >= ?", startTimestamp) | ||||
| } | } | ||||
| @@ -0,0 +1,197 @@ | |||||
| package model | |||||
| import ( | |||||
| "net/http/httptest" | |||||
| "testing" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||||
| "github.com/gin-gonic/gin" | |||||
| "github.com/glebarez/sqlite" | |||||
| "github.com/stretchr/testify/require" | |||||
| "gorm.io/gorm" | |||||
| ) | |||||
| func setupLogIdentityDB(t *testing.T) *gorm.DB { | |||||
| t.Helper() | |||||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||||
| require.NoError(t, err) | |||||
| origDB := DB | |||||
| origLogDB := LOG_DB | |||||
| origSQLite := common.UsingSQLite | |||||
| origMySQL := common.UsingMySQL | |||||
| origPostgres := common.UsingPostgreSQL | |||||
| DB = db | |||||
| LOG_DB = db | |||||
| common.UsingSQLite = true | |||||
| common.UsingMySQL = false | |||||
| common.UsingPostgreSQL = false | |||||
| initCol() | |||||
| require.NoError(t, db.AutoMigrate(&Log{})) | |||||
| t.Cleanup(func() { | |||||
| DB = origDB | |||||
| LOG_DB = origLogDB | |||||
| common.UsingSQLite = origSQLite | |||||
| common.UsingMySQL = origMySQL | |||||
| common.UsingPostgreSQL = origPostgres | |||||
| initCol() | |||||
| }) | |||||
| return db | |||||
| } | |||||
| func TestLogPersistsChatIDAndUpstreamID(t *testing.T) { | |||||
| db := setupLogIdentityDB(t) | |||||
| entry := &Log{ | |||||
| UserId: 1, | |||||
| Username: "alice", | |||||
| CreatedAt: 1714464000, | |||||
| Type: LogTypeConsume, | |||||
| Content: "usage", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| ChatId: "chat_123", | |||||
| UpstreamId: "req_upstream_123", | |||||
| RequestId: "req_internal_123", | |||||
| } | |||||
| require.NoError(t, db.Create(entry).Error) | |||||
| var saved Log | |||||
| require.NoError(t, db.First(&saved).Error) | |||||
| require.Equal(t, "chat_123", saved.ChatId) | |||||
| require.Equal(t, "req_upstream_123", saved.UpstreamId) | |||||
| require.Equal(t, "req_internal_123", saved.RequestId) | |||||
| } | |||||
| func TestGetAllLogsFiltersByChatIDAndUpstreamID(t *testing.T) { | |||||
| db := setupLogIdentityDB(t) | |||||
| require.NoError(t, db.Create(&Log{ | |||||
| UserId: 1, | |||||
| Username: "alice", | |||||
| CreatedAt: 1714464001, | |||||
| Type: LogTypeConsume, | |||||
| Content: "usage-a", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| ChatId: "chat_a", | |||||
| UpstreamId: "up_a", | |||||
| RequestId: "req_a", | |||||
| }).Error) | |||||
| require.NoError(t, db.Create(&Log{ | |||||
| UserId: 2, | |||||
| Username: "bob", | |||||
| CreatedAt: 1714464002, | |||||
| Type: LogTypeConsume, | |||||
| Content: "usage-b", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| ChatId: "chat_b", | |||||
| UpstreamId: "up_b", | |||||
| RequestId: "req_b", | |||||
| }).Error) | |||||
| logs, total, err := GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 0, 20, 0, "", "", "chat_a", "") | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, int64(1), total) | |||||
| require.Len(t, logs, 1) | |||||
| require.Equal(t, "chat_a", logs[0].ChatId) | |||||
| logs, total, err = GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 0, 20, 0, "", "", "", "up_b") | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, int64(1), total) | |||||
| require.Len(t, logs, 1) | |||||
| require.Equal(t, "up_b", logs[0].UpstreamId) | |||||
| } | |||||
| func TestGetUserLogsFiltersByChatIDAndUpstreamID(t *testing.T) { | |||||
| db := setupLogIdentityDB(t) | |||||
| require.NoError(t, db.Create(&Log{ | |||||
| UserId: 9, | |||||
| Username: "owner", | |||||
| CreatedAt: 1714464010, | |||||
| Type: LogTypeConsume, | |||||
| Content: "usage-owner-a", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| ChatId: "chat-owner-a", | |||||
| UpstreamId: "up-owner-a", | |||||
| RequestId: "req-owner-a", | |||||
| }).Error) | |||||
| require.NoError(t, db.Create(&Log{ | |||||
| UserId: 9, | |||||
| Username: "owner", | |||||
| CreatedAt: 1714464011, | |||||
| Type: LogTypeConsume, | |||||
| Content: "usage-owner-b", | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| ChatId: "chat-owner-b", | |||||
| UpstreamId: "up-owner-b", | |||||
| RequestId: "req-owner-b", | |||||
| }).Error) | |||||
| logs, total, err := GetUserLogs(9, LogTypeConsume, 0, 0, "", "", 0, 20, "", "", "chat-owner-b", "") | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, int64(1), total) | |||||
| require.Len(t, logs, 1) | |||||
| require.Equal(t, "chat-owner-b", logs[0].ChatId) | |||||
| logs, total, err = GetUserLogs(9, LogTypeConsume, 0, 0, "", "", 0, 20, "", "", "", "up-owner-a") | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, int64(1), total) | |||||
| require.Len(t, logs, 1) | |||||
| require.Equal(t, "up-owner-a", logs[0].UpstreamId) | |||||
| } | |||||
| func TestRecordConsumeLogPersistsContextChatIDAndUpstreamID(t *testing.T) { | |||||
| db := setupLogIdentityDB(t) | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| c.Set("username", "alice") | |||||
| c.Set(common.RequestIdKey, "req_internal_ctx") | |||||
| relaycommon.SetRelayChatID(c, "chat_ctx") | |||||
| relaycommon.SetRelayUpstreamID(c, "up_ctx") | |||||
| RecordConsumeLog(c, 1, RecordConsumeLogParams{ | |||||
| ChannelId: 1, | |||||
| ModelName: "gpt-4o-mini", | |||||
| TokenName: "demo", | |||||
| Content: "usage", | |||||
| Quota: 10, | |||||
| }) | |||||
| var saved Log | |||||
| require.NoError(t, db.First(&saved).Error) | |||||
| require.Equal(t, "chat_ctx", saved.ChatId) | |||||
| require.Equal(t, "up_ctx", saved.UpstreamId) | |||||
| } | |||||
| func TestRecordErrorLogPersistsContextChatIDAndUpstreamID(t *testing.T) { | |||||
| db := setupLogIdentityDB(t) | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| c.Set("username", "alice") | |||||
| c.Set(common.RequestIdKey, "req_internal_err") | |||||
| relaycommon.SetRelayChatID(c, "chat_err") | |||||
| relaycommon.SetRelayUpstreamID(c, "up_err") | |||||
| RecordErrorLog(c, 1, 2, "gpt-4o-mini", "demo", "upstream failed", 3, 4, false, "default", map[string]interface{}{}) | |||||
| var saved Log | |||||
| require.NoError(t, db.First(&saved).Error) | |||||
| require.Equal(t, "chat_err", saved.ChatId) | |||||
| require.Equal(t, "up_err", saved.UpstreamId) | |||||
| } | |||||
| @@ -360,10 +360,14 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody | |||||
| targetHeader.Set(key, value) | targetHeader.Set(key, value) | ||||
| } | } | ||||
| targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type")) | targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type")) | ||||
| targetConn, _, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader) | |||||
| targetConn, resp, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader) | |||||
| if err != nil { | if err != nil { | ||||
| return nil, fmt.Errorf("dial failed to %s: %w", fullRequestURL, err) | return nil, fmt.Errorf("dial failed to %s: %w", fullRequestURL, err) | ||||
| } | } | ||||
| if resp != nil { | |||||
| common.CaptureUpstreamIDFromHTTPResponse(info, resp) | |||||
| common.SetRelayUpstreamID(c, info.UpstreamID) | |||||
| } | |||||
| // send request body | // send request body | ||||
| //all, err := io.ReadAll(requestBody) | //all, err := io.ReadAll(requestBody) | ||||
| //err = service.WssString(c, targetConn, string(all)) | //err = service.WssString(c, targetConn, string(all)) | ||||
| @@ -512,6 +516,8 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http | |||||
| if resp == nil { | if resp == nil { | ||||
| return nil, errors.New("resp is nil") | return nil, errors.New("resp is nil") | ||||
| } | } | ||||
| common.CaptureUpstreamIDFromHTTPResponse(info, resp) | |||||
| common.SetRelayUpstreamID(c, info.UpstreamID) | |||||
| _ = req.Body.Close() | _ = req.Body.Close() | ||||
| _ = c.Request.Body.Close() | _ = c.Request.Body.Close() | ||||
| @@ -87,6 +87,8 @@ type RelayInfo struct { | |||||
| TokenKey string | TokenKey string | ||||
| TokenGroup string | TokenGroup string | ||||
| UserId int | UserId int | ||||
| ChatID string | |||||
| UpstreamID string | |||||
| UsingGroup string // 使用的分组,当auto跨分组重试时,会变动 | UsingGroup string // 使用的分组,当auto跨分组重试时,会变动 | ||||
| UserGroup string // 用户所在分组 | UserGroup string // 用户所在分组 | ||||
| TokenUnlimited bool | TokenUnlimited bool | ||||
| @@ -452,6 +454,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { | |||||
| RequestId: reqId, | RequestId: reqId, | ||||
| UserId: common.GetContextKeyInt(c, constant.ContextKeyUserId), | UserId: common.GetContextKeyInt(c, constant.ContextKeyUserId), | ||||
| ChatID: GetRelayChatID(c), | |||||
| UsingGroup: common.GetContextKeyString(c, constant.ContextKeyUsingGroup), | UsingGroup: common.GetContextKeyString(c, constant.ContextKeyUsingGroup), | ||||
| UserGroup: common.GetContextKeyString(c, constant.ContextKeyUserGroup), | UserGroup: common.GetContextKeyString(c, constant.ContextKeyUserGroup), | ||||
| UserQuota: common.GetContextKeyInt(c, constant.ContextKeyUserQuota), | UserQuota: common.GetContextKeyInt(c, constant.ContextKeyUserQuota), | ||||
| @@ -0,0 +1,82 @@ | |||||
| package common | |||||
| import ( | |||||
| "net/http" | |||||
| "strings" | |||||
| commonpkg "github.com/QuantumNous/new-api/common" | |||||
| "github.com/gin-gonic/gin" | |||||
| ) | |||||
| const ( | |||||
| ContextKeyRelayChatID = "relay_chat_id" | |||||
| ContextKeyRelayUpstreamID = "relay_upstream_id" | |||||
| ) | |||||
| func SetRelayChatID(c *gin.Context, chatID string) { | |||||
| if c == nil { | |||||
| return | |||||
| } | |||||
| chatID = strings.TrimSpace(chatID) | |||||
| if chatID == "" { | |||||
| return | |||||
| } | |||||
| c.Set(ContextKeyRelayChatID, chatID) | |||||
| } | |||||
| func GetRelayChatID(c *gin.Context) string { | |||||
| if c == nil { | |||||
| return "" | |||||
| } | |||||
| return strings.TrimSpace(c.GetString(ContextKeyRelayChatID)) | |||||
| } | |||||
| func SetRelayUpstreamID(c *gin.Context, upstreamID string) { | |||||
| if c == nil { | |||||
| return | |||||
| } | |||||
| upstreamID = strings.TrimSpace(upstreamID) | |||||
| if upstreamID == "" { | |||||
| return | |||||
| } | |||||
| c.Set(ContextKeyRelayUpstreamID, upstreamID) | |||||
| } | |||||
| func GetRelayUpstreamID(c *gin.Context) string { | |||||
| if c == nil { | |||||
| return "" | |||||
| } | |||||
| return strings.TrimSpace(c.GetString(ContextKeyRelayUpstreamID)) | |||||
| } | |||||
| func ExtractTopLevelChatID(c *gin.Context) (string, error) { | |||||
| if c == nil { | |||||
| return "", nil | |||||
| } | |||||
| bodyStorage, err := commonpkg.GetBodyStorage(c) | |||||
| if err != nil { | |||||
| return "", err | |||||
| } | |||||
| bodyBytes, err := bodyStorage.Bytes() | |||||
| if err != nil || len(bodyBytes) == 0 { | |||||
| return "", err | |||||
| } | |||||
| var payload map[string]interface{} | |||||
| if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { | |||||
| return "", nil | |||||
| } | |||||
| return strings.TrimSpace(commonpkg.Interface2String(payload["chat_id"])), nil | |||||
| } | |||||
| func CaptureUpstreamIDFromHTTPResponse(info *RelayInfo, resp *http.Response) { | |||||
| if info == nil || resp == nil { | |||||
| return | |||||
| } | |||||
| for _, key := range []string{"x-request-id", "request-id"} { | |||||
| value := strings.TrimSpace(resp.Header.Get(key)) | |||||
| if value != "" { | |||||
| info.UpstreamID = value | |||||
| return | |||||
| } | |||||
| } | |||||
| } | |||||
| @@ -0,0 +1,89 @@ | |||||
| package common | |||||
| import ( | |||||
| "bytes" | |||||
| "net/http" | |||||
| "net/http/httptest" | |||||
| "testing" | |||||
| "github.com/gin-gonic/gin" | |||||
| "github.com/stretchr/testify/require" | |||||
| ) | |||||
| func TestExtractChatIDFromReusableBody(t *testing.T) { | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}`)) | |||||
| c.Request.Header.Set("Content-Type", "application/json") | |||||
| chatID, err := ExtractTopLevelChatID(c) | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, "chat_abc", chatID) | |||||
| } | |||||
| func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) { | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`)) | |||||
| c.Request.Header.Set("Content-Type", "application/json") | |||||
| chatID, err := ExtractTopLevelChatID(c) | |||||
| require.NoError(t, err) | |||||
| require.Empty(t, chatID) | |||||
| } | |||||
| func TestCaptureUpstreamIDFromHeaderPrefersXRequestID(t *testing.T) { | |||||
| info := &RelayInfo{} | |||||
| resp := &http.Response{ | |||||
| Header: http.Header{ | |||||
| "X-Request-Id": []string{"xreq_123"}, | |||||
| "Request-Id": []string{"req_456"}, | |||||
| }, | |||||
| } | |||||
| CaptureUpstreamIDFromHTTPResponse(info, resp) | |||||
| require.Equal(t, "xreq_123", info.UpstreamID) | |||||
| } | |||||
| func TestCaptureUpstreamIDFallsBackToRequestID(t *testing.T) { | |||||
| info := &RelayInfo{} | |||||
| resp := &http.Response{ | |||||
| Header: http.Header{ | |||||
| "Request-Id": []string{"req_456"}, | |||||
| }, | |||||
| } | |||||
| CaptureUpstreamIDFromHTTPResponse(info, resp) | |||||
| require.Equal(t, "req_456", info.UpstreamID) | |||||
| } | |||||
| func TestSetAndGetChatIDFromGinContext(t *testing.T) { | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| SetRelayChatID(c, "chat_ctx_1") | |||||
| require.Equal(t, "chat_ctx_1", GetRelayChatID(c)) | |||||
| } | |||||
| func TestCaptureUpstreamIDDoesNotOverwriteExistingValueWithEmptyHeader(t *testing.T) { | |||||
| info := &RelayInfo{UpstreamID: "existing_upstream"} | |||||
| resp := &http.Response{Header: http.Header{}} | |||||
| CaptureUpstreamIDFromHTTPResponse(info, resp) | |||||
| require.Equal(t, "existing_upstream", info.UpstreamID) | |||||
| } | |||||
| func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) { | |||||
| info := &RelayInfo{} | |||||
| resp := &http.Response{ | |||||
| Header: http.Header{ | |||||
| "X-Request-Id": []string{"ws_upstream_1"}, | |||||
| }, | |||||
| } | |||||
| CaptureUpstreamIDFromHTTPResponse(info, resp) | |||||
| require.Equal(t, "ws_upstream_1", info.UpstreamID) | |||||
| } | |||||
| @@ -10,6 +10,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/dto" | "github.com/QuantumNous/new-api/dto" | ||||
| "github.com/QuantumNous/new-api/logger" | "github.com/QuantumNous/new-api/logger" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||||
| relayconstant "github.com/QuantumNous/new-api/relay/constant" | relayconstant "github.com/QuantumNous/new-api/relay/constant" | ||||
| "github.com/QuantumNous/new-api/types" | "github.com/QuantumNous/new-api/types" | ||||
| @@ -17,6 +18,10 @@ import ( | |||||
| ) | ) | ||||
| func GetAndValidateRequest(c *gin.Context, format types.RelayFormat) (request dto.Request, err error) { | func GetAndValidateRequest(c *gin.Context, format types.RelayFormat) (request dto.Request, err error) { | ||||
| if chatID, extractErr := relaycommon.ExtractTopLevelChatID(c); extractErr == nil && chatID != "" { | |||||
| relaycommon.SetRelayChatID(c, chatID) | |||||
| } | |||||
| relayMode := relayconstant.Path2RelayMode(c.Request.URL.Path) | relayMode := relayconstant.Path2RelayMode(c.Request.URL.Path) | ||||
| switch format { | switch format { | ||||
| @@ -102,6 +102,24 @@ const LogsFilters = ({ | |||||
| size='small' | size='small' | ||||
| /> | /> | ||||
| <Form.Input | |||||
| field='chat_id' | |||||
| prefix={<IconSearch />} | |||||
| placeholder={t('Chat ID')} | |||||
| showClear | |||||
| pure | |||||
| size='small' | |||||
| /> | |||||
| <Form.Input | |||||
| field='upstream_id' | |||||
| prefix={<IconSearch />} | |||||
| placeholder={t('Upstream ID')} | |||||
| showClear | |||||
| pure | |||||
| size='small' | |||||
| /> | |||||
| {isAdminUser && ( | {isAdminUser && ( | ||||
| <> | <> | ||||
| <Form.Input | <Form.Input | ||||
| @@ -95,6 +95,8 @@ export const useLogsData = () => { | |||||
| channel: '', | channel: '', | ||||
| group: '', | group: '', | ||||
| request_id: '', | request_id: '', | ||||
| chat_id: '', | |||||
| upstream_id: '', | |||||
| dateRange: [ | dateRange: [ | ||||
| timestamp2string(getTodayStartTimestamp()), | timestamp2string(getTodayStartTimestamp()), | ||||
| timestamp2string(now.getTime() / 1000 + 3600), | timestamp2string(now.getTime() / 1000 + 3600), | ||||
| @@ -232,6 +234,8 @@ export const useLogsData = () => { | |||||
| channel: formValues.channel || '', | channel: formValues.channel || '', | ||||
| group: formValues.group || '', | group: formValues.group || '', | ||||
| request_id: formValues.request_id || '', | request_id: formValues.request_id || '', | ||||
| chat_id: formValues.chat_id || '', | |||||
| upstream_id: formValues.upstream_id || '', | |||||
| logType: formValues.logType ? parseInt(formValues.logType) : 0, | logType: formValues.logType ? parseInt(formValues.logType) : 0, | ||||
| }; | }; | ||||
| }; | }; | ||||
| @@ -356,6 +360,18 @@ export const useLogsData = () => { | |||||
| value: logs[i].request_id, | value: logs[i].request_id, | ||||
| }); | }); | ||||
| } | } | ||||
| if (logs[i].chat_id) { | |||||
| expandDataLocal.push({ | |||||
| key: t('Chat ID'), | |||||
| value: logs[i].chat_id, | |||||
| }); | |||||
| } | |||||
| if (logs[i].upstream_id) { | |||||
| expandDataLocal.push({ | |||||
| key: t('Upstream ID'), | |||||
| value: logs[i].upstream_id, | |||||
| }); | |||||
| } | |||||
| if (other?.ws || other?.audio) { | if (other?.ws || other?.audio) { | ||||
| expandDataLocal.push({ | expandDataLocal.push({ | ||||
| key: t('语音输入'), | key: t('语音输入'), | ||||
| @@ -647,6 +663,8 @@ export const useLogsData = () => { | |||||
| channel, | channel, | ||||
| group, | group, | ||||
| request_id, | request_id, | ||||
| chat_id, | |||||
| upstream_id, | |||||
| logType: formLogType, | logType: formLogType, | ||||
| } = getFormValues(); | } = getFormValues(); | ||||
| @@ -660,9 +678,9 @@ export const useLogsData = () => { | |||||
| let localStartTimestamp = Date.parse(start_timestamp) / 1000; | let localStartTimestamp = Date.parse(start_timestamp) / 1000; | ||||
| let localEndTimestamp = Date.parse(end_timestamp) / 1000; | let localEndTimestamp = Date.parse(end_timestamp) / 1000; | ||||
| if (isAdminUser) { | if (isAdminUser) { | ||||
| url = `/api/log/?p=${startIdx}&page_size=${pageSize}&type=${currentLogType}&username=${username}&token_name=${token_name}&model_name=${model_name}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&channel=${channel}&group=${group}&request_id=${request_id}`; | |||||
| url = `/api/log/?p=${startIdx}&page_size=${pageSize}&type=${currentLogType}&username=${username}&token_name=${token_name}&model_name=${model_name}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&channel=${channel}&group=${group}&request_id=${request_id}&chat_id=${chat_id}&upstream_id=${upstream_id}`; | |||||
| } else { | } else { | ||||
| url = `/api/log/self/?p=${startIdx}&page_size=${pageSize}&type=${currentLogType}&token_name=${token_name}&model_name=${model_name}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&group=${group}&request_id=${request_id}`; | |||||
| url = `/api/log/self/?p=${startIdx}&page_size=${pageSize}&type=${currentLogType}&token_name=${token_name}&model_name=${model_name}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&group=${group}&request_id=${request_id}&chat_id=${chat_id}&upstream_id=${upstream_id}`; | |||||
| } | } | ||||
| url = encodeURI(url); | url = encodeURI(url); | ||||
| const res = await API.get(url); | const res = await API.get(url); | ||||