| @@ -21,7 +21,9 @@ func GetAllLogs(c *gin.Context) { | |||
| channel, _ := strconv.Atoi(c.Query("channel")) | |||
| group := c.Query("group") | |||
| 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 { | |||
| common.ApiError(c, err) | |||
| return | |||
| @@ -42,7 +44,9 @@ func GetUserLogs(c *gin.Context) { | |||
| modelName := c.Query("model_name") | |||
| group := c.Query("group") | |||
| 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 { | |||
| common.ApiError(c, err) | |||
| 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/logger" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/gin-gonic/gin" | |||
| @@ -36,6 +37,8 @@ type Log struct { | |||
| Group string `json:"group" gorm:"index"` | |||
| Ip string `json:"ip" gorm:"index;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"` | |||
| } | |||
| @@ -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)) | |||
| username := c.GetString("username") | |||
| requestId := c.GetString(common.RequestIdKey) | |||
| chatId := relaycommon.GetRelayChatID(c) | |||
| upstreamId := relaycommon.GetRelayUpstreamID(c) | |||
| otherStr := common.MapToJsonStr(other) | |||
| // 判断是否需要记录 IP | |||
| needRecordIp := false | |||
| @@ -124,8 +129,10 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, | |||
| } | |||
| return "" | |||
| }(), | |||
| RequestId: requestId, | |||
| Other: otherStr, | |||
| RequestId: requestId, | |||
| ChatId: chatId, | |||
| UpstreamId: upstreamId, | |||
| Other: otherStr, | |||
| } | |||
| err := LOG_DB.Create(log).Error | |||
| 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))) | |||
| username := c.GetString("username") | |||
| requestId := c.GetString(common.RequestIdKey) | |||
| chatId := relaycommon.GetRelayChatID(c) | |||
| upstreamId := relaycommon.GetRelayUpstreamID(c) | |||
| otherStr := common.MapToJsonStr(params.Other) | |||
| // 判断是否需要记录 IP | |||
| needRecordIp := false | |||
| @@ -185,8 +194,10 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) | |||
| } | |||
| return "" | |||
| }(), | |||
| RequestId: requestId, | |||
| Other: otherStr, | |||
| RequestId: requestId, | |||
| ChatId: chatId, | |||
| UpstreamId: upstreamId, | |||
| Other: otherStr, | |||
| } | |||
| err := LOG_DB.Create(log).Error | |||
| 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 | |||
| if logType == LogTypeUnknown { | |||
| tx = LOG_DB | |||
| @@ -262,6 +273,12 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName | |||
| if 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 { | |||
| tx = tx.Where("logs.created_at >= ?", startTimestamp) | |||
| } | |||
| @@ -328,7 +345,7 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName | |||
| 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 | |||
| if logType == LogTypeUnknown { | |||
| tx = LOG_DB.Where("logs.user_id = ?", userId) | |||
| @@ -349,6 +366,12 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int | |||
| if 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 { | |||
| 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("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 { | |||
| 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 | |||
| //all, err := io.ReadAll(requestBody) | |||
| //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 { | |||
| return nil, errors.New("resp is nil") | |||
| } | |||
| common.CaptureUpstreamIDFromHTTPResponse(info, resp) | |||
| common.SetRelayUpstreamID(c, info.UpstreamID) | |||
| _ = req.Body.Close() | |||
| _ = c.Request.Body.Close() | |||
| @@ -87,6 +87,8 @@ type RelayInfo struct { | |||
| TokenKey string | |||
| TokenGroup string | |||
| UserId int | |||
| ChatID string | |||
| UpstreamID string | |||
| UsingGroup string // 使用的分组,当auto跨分组重试时,会变动 | |||
| UserGroup string // 用户所在分组 | |||
| TokenUnlimited bool | |||
| @@ -452,6 +454,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { | |||
| RequestId: reqId, | |||
| UserId: common.GetContextKeyInt(c, constant.ContextKeyUserId), | |||
| ChatID: GetRelayChatID(c), | |||
| UsingGroup: common.GetContextKeyString(c, constant.ContextKeyUsingGroup), | |||
| UserGroup: common.GetContextKeyString(c, constant.ContextKeyUserGroup), | |||
| 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/dto" | |||
| "github.com/QuantumNous/new-api/logger" | |||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||
| relayconstant "github.com/QuantumNous/new-api/relay/constant" | |||
| "github.com/QuantumNous/new-api/types" | |||
| @@ -17,6 +18,10 @@ import ( | |||
| ) | |||
| 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) | |||
| switch format { | |||
| @@ -102,6 +102,24 @@ const LogsFilters = ({ | |||
| 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 && ( | |||
| <> | |||
| <Form.Input | |||
| @@ -95,6 +95,8 @@ export const useLogsData = () => { | |||
| channel: '', | |||
| group: '', | |||
| request_id: '', | |||
| chat_id: '', | |||
| upstream_id: '', | |||
| dateRange: [ | |||
| timestamp2string(getTodayStartTimestamp()), | |||
| timestamp2string(now.getTime() / 1000 + 3600), | |||
| @@ -232,6 +234,8 @@ export const useLogsData = () => { | |||
| channel: formValues.channel || '', | |||
| group: formValues.group || '', | |||
| request_id: formValues.request_id || '', | |||
| chat_id: formValues.chat_id || '', | |||
| upstream_id: formValues.upstream_id || '', | |||
| logType: formValues.logType ? parseInt(formValues.logType) : 0, | |||
| }; | |||
| }; | |||
| @@ -356,6 +360,18 @@ export const useLogsData = () => { | |||
| 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) { | |||
| expandDataLocal.push({ | |||
| key: t('语音输入'), | |||
| @@ -647,6 +663,8 @@ export const useLogsData = () => { | |||
| channel, | |||
| group, | |||
| request_id, | |||
| chat_id, | |||
| upstream_id, | |||
| logType: formLogType, | |||
| } = getFormValues(); | |||
| @@ -660,9 +678,9 @@ export const useLogsData = () => { | |||
| let localStartTimestamp = Date.parse(start_timestamp) / 1000; | |||
| let localEndTimestamp = Date.parse(end_timestamp) / 1000; | |||
| 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 { | |||
| 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); | |||
| const res = await API.get(url); | |||