diff --git a/.github/workflows/docker-image-amd64.yml b/.github/workflows/docker-image-amd64.yml index 36236df25fa4..7a30c5b2144f 100644 --- a/.github/workflows/docker-image-amd64.yml +++ b/.github/workflows/docker-image-amd64.yml @@ -1,54 +1,40 @@ -name: Publish Docker image (amd64) +name: Build NewAPI Docker Image on: push: tags: - - '*' + - "main" workflow_dispatch: inputs: - name: - description: 'reason' - required: false + git_branch: + description: 'Git branch' + required: true + default: 'main' + remote_repo: + description: 'Remote repository' + required: true + default: 'https://github.com/Furion-cn/new-api.git' + jobs: - push_to_registries: - name: Push Docker image to multiple registries - runs-on: ubuntu-latest - permissions: - packages: write - contents: read + build-dev: + runs-on: ubuntu-22.04 steps: - - name: Check out the repo - uses: actions/checkout@v3 - - - name: Save version info - run: | - git describe --tags > VERSION - - - name: Log in to Docker Hub - uses: docker/login-action@v2 - with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} + - name: Checkout repository + uses: actions/checkout@v4 - - name: Log in to the Container registry + - name: Login to Tencent Cloud Registry uses: docker/login-action@v2 with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Extract metadata (tags, labels) for Docker - id: meta - uses: docker/metadata-action@v4 - with: - images: | - calciumion/new-api - ghcr.io/${{ github.repository }} + registry: furion-sh.tencentcloudcr.com + username: ${{ secrets.TCR_USERNAME }} + password: ${{ secrets.TCR_PASSWORD }} - - name: Build and push Docker images - uses: docker/build-push-action@v3 - with: - context: . - push: true - tags: ${{ steps.meta.outputs.tags }} - labels: ${{ steps.meta.outputs.labels }} \ No newline at end of file + - name: Build and Push SGLang Image + run: | + DATETIME=$(date '+%Y%m%d_%H%M') + ORG=$(echo "${{ inputs.remote_repo }}" | awk -F'/' '{print $(NF-1)}' | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9]/-/g') + BRANCH=$(echo "${{ inputs.git_branch }}" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9]/-/g') + VERSION="${DATETIME}_${ORG}_${BRANCH}" + git clone ${{ inputs.remote_repo }} && cd new-api + docker build . -t furion-sh.tencentcloudcr.com/furion/new-api:${VERSION} + docker push furion-sh.tencentcloudcr.com/furion/new-api:${VERSION} \ No newline at end of file diff --git a/.gitignore b/.gitignore index 28106a2406fc..6e3b5b7ab521 100644 --- a/.gitignore +++ b/.gitignore @@ -9,4 +9,13 @@ logs web/dist .env one-api -.DS_Store \ No newline at end of file +.DS_Store +new-api +script/ +.env_2 +*_test.go +out.log +out.log_2 +.env_3 +*.prof +.env* \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 214ceaa3d5f2..f9349cb71e09 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM oven/bun:latest AS builder +FROM swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/oven/bun:latest AS builder WORKDIR /build COPY web/package.json . @@ -7,24 +7,28 @@ COPY ./web . COPY ./VERSION . RUN DISABLE_ESLINT_PLUGIN='true' VITE_REACT_APP_VERSION=$(cat VERSION) bun run build -FROM golang:alpine AS builder2 +FROM swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/golang:alpine3.21 AS builder2 ENV GO111MODULE=on \ CGO_ENABLED=0 \ - GOOS=linux + GOOS=linux \ + GOPROXY=https://goproxy.cn WORKDIR /build ADD go.mod go.sum ./ +RUN go mod tidy RUN go mod download COPY . . COPY --from=builder /build/dist ./web/dist +RUN go mod tidy RUN go build -ldflags "-s -w -X 'one-api/common.Version=$(cat VERSION)'" -o one-api -FROM alpine +FROM swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/library/alpine:latest -RUN apk update \ +RUN sed -i 's/dl-cdn.alpinelinux.org/mirrors.aliyun.com/g' /etc/apk/repositories \ + && apk update \ && apk upgrade \ && apk add --no-cache ca-certificates tzdata ffmpeg \ && update-ca-certificates @@ -32,4 +36,4 @@ RUN apk update \ COPY --from=builder2 /build/one-api / EXPOSE 3000 WORKDIR /data -ENTRYPOINT ["/one-api"] +ENTRYPOINT ["/one-api"] \ No newline at end of file diff --git a/bin/groups.sql b/bin/groups.sql new file mode 100644 index 000000000000..9e0cd2c15a01 --- /dev/null +++ b/bin/groups.sql @@ -0,0 +1,7 @@ +CREATE TABLE IF NOT EXISTS `groups` ( + `id` int(11) NOT NULL AUTO_INCREMENT, + `name` varchar(255) NOT NULL DEFAULT '', + `ratio` int(11) NOT NULL DEFAULT 0, + PRIMARY KEY (`id`), + UNIQUE KEY `idx_name` (`name`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci; \ No newline at end of file diff --git a/bin/log_requestid.sql b/bin/log_requestid.sql new file mode 100644 index 000000000000..5a594f7e4be9 --- /dev/null +++ b/bin/log_requestid.sql @@ -0,0 +1 @@ +ALTER TABLE logs ADD COLUMN request_id VARCHAR(255) DEFAULT '' COMMENT '请求ID'; diff --git a/common/constants.go b/common/constants.go index bcab24fc0a36..e01593b0cd54 100644 --- a/common/constants.go +++ b/common/constants.go @@ -26,12 +26,16 @@ var DataExportEnabled = true var DataExportInterval = 5 // unit: minute var DataExportDefaultTime = "hour" // unit: minute var DefaultCollapseSidebar = false // default value of collapse sidebar +var MockResponseEnabled = false // whether to enable mock response for test traffic // Any options with "Secret", "Token" in its key won't be return by GetOptions var SessionSecret = uuid.New().String() var CryptoSecret = uuid.New().String() +// 需要先导入 model 包 +var Groups = map[string]int{} +var GroupRWMutex sync.RWMutex var OptionMap map[string]string var OptionMapRWMutex sync.RWMutex @@ -234,6 +238,8 @@ const ( ChannelTypeMokaAI = 44 ChannelTypeVolcEngine = 45 ChannelTypeBaiduV2 = 46 + ChannelTypeXai = 47 + ChannelTypeDoubaoOffline = 100 ChannelTypeDummy // this one is only for count, do not add any channel after this ) @@ -286,4 +292,58 @@ var ChannelBaseURLs = []string{ "https://api.moka.ai", //44 "https://ark.cn-beijing.volces.com", //45 "https://qianfan.baidubce.com", //46 + "", //47 + "https://api.x.ai", //48 + "", //49 + "", //50 + "", //51 + "", //52 + "", //53 + "", //54 + "", //55 + "", //56 + "", //57 + "", //58 + "", //59 + "", //60 + "", //61 + "", //62 + "", //63 + "", //64 + "", //65 + "", //66 + "", //67 + "", //68 + "", //69 + "", //70 + "", //71 + "", //72 + "", //73 + "", //74 + "", //75 + "", //76 + "", //77 + "", //78 + "", //79 + "", //80 + "", //81 + "", //82 + "", //83 + "", //84 + "", //85 + "", //86 + "", //87 + "", //88 + "", //89 + "", //90 + "", //91 + "", //92 + "", //93 + "", //94 + "", //95 + "", //96 + "", //97 + "", //98 + "", //99 + "https://ark.cn-beijing.volces.com", //100 - 豆包离线 } diff --git a/common/init.go b/common/init.go index 694e603ef067..22083dd0fca1 100644 --- a/common/init.go +++ b/common/init.go @@ -66,4 +66,8 @@ func LoadEnv() { } } } + + // Initialize mock response feature + MockResponseEnabled = os.Getenv("MOCK_RESPONSE_ENABLED") == "true" + SysLog(fmt.Sprintf("MockResponseEnabled: %v", MockResponseEnabled)) } diff --git a/common/logger.go b/common/logger.go index 86d15fa4db7a..735a3f8e38c6 100644 --- a/common/logger.go +++ b/common/logger.go @@ -4,23 +4,53 @@ import ( "context" "encoding/json" "fmt" - "github.com/bytedance/gopkg/util/gopool" - "github.com/gin-gonic/gin" "io" "log" "os" "path/filepath" + "runtime" + "sort" + "strings" "sync" "time" + + "one-api/metrics" + + "github.com/bytedance/gopkg/util/gopool" + "github.com/gin-gonic/gin" ) +// 是否开启透传日志打印 +var LogPassthroughEnabled = false + +// 日志打印采样比例（0-100之间的整数，表示百分比） +var LogSampleRatio = 100 + const ( loggerINFO = "INFO" loggerWarn = "WARN" loggerError = "ERR" + // 10GB in bytes + maxLogFileSize = 20 * 1024 * 1024 * 1024 + // 保留最近的1个日志文件 + maxLogFiles = 1 + // 日志计数上限 + maxLogCount = 1000000 ) -const maxLogCount = 1000000 +// 错误类型常量 +const ( + ErrorTypeOther = "other" + ErrorTypeParameter = "parameter_error" + ErrorTypeNoCandidates = "no_candidates" + ErrorTypeRequestFailed = "request_failed" + ErrorTypeBadGateway = "bad_gateway" + ErrorTypeResponseFailed = "response_failed" + ErrorTypeConnectionTimeout = "connection_timeout" + ErrorTypeTokenUnavailable = "token_unavailable" + ErrorTypeBadRequest = "bad_request" + ErrorTypeNoAvailableChannel = "no_available_channel" +) var logCount int var setupLogLock sync.Mutex @@ -37,24 +67,126 @@ func SetupLogger() { setupLogLock.Unlock() setupLogWorking = false }() + + // 创建日志目录 + if _, err := os.Stat(*LogDir); os.IsNotExist(err) { + if err := os.MkdirAll(*LogDir, 0755); err != nil { + log.Fatal("failed to create log directory") + } + } + + // 检查并清理旧的日志文件 + cleanOldLogs() + logPath := filepath.Join(*LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102150405"))) fd, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) if err != nil { log.Fatal("failed to open log file") } - gin.DefaultWriter = io.MultiWriter(os.Stdout, fd) - gin.DefaultErrorWriter = io.MultiWriter(os.Stderr, fd) + + // 创建一个自定义的 writer，用于检查文件大小 + writer := &logWriter{ + file: fd, + filepath: logPath, + size: 0, + } + + gin.DefaultWriter = io.MultiWriter(os.Stdout, writer) + gin.DefaultErrorWriter = io.MultiWriter(os.Stderr, writer) } } +// logWriter 是一个自定义的 writer，用于跟踪文件大小 +type logWriter struct { + file *os.File + filepath string + size int64 + mu sync.Mutex +} + +func (w *logWriter) Write(p []byte) (n int, err error) { + w.mu.Lock() + defer w.mu.Unlock() + + // 写入数据 + n, err = w.file.Write(p) + if err != nil { + return n, err + } + + // 更新文件大小 + w.size += int64(n) + + // 检查文件大小是否超过限制 + if w.size >= maxLogFileSize { + // 关闭当前文件 + w.file.Close() + + // 清理旧日志并创建新文件 + cleanOldLogs() + + // 创建新的日志文件 + logPath := filepath.Join(*LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102150405"))) + fd, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + return n, fmt.Errorf("failed to create new log file: %v", err) + } + + // 更新 writer 状态 + w.file = fd + w.filepath = logPath + w.size = 0 + + // 更新 gin 的 writer + gin.DefaultWriter = io.MultiWriter(os.Stdout, w) + gin.DefaultErrorWriter = io.MultiWriter(os.Stderr, w) + } + + return n, nil +} + +// cleanOldLogs 清理旧的日志文件，只保留最近的几个文件 +func cleanOldLogs() { + files, err := filepath.Glob(filepath.Join(*LogDir, "oneapi-*.log")) + if err != nil { + log.Printf("failed to list log files: %v", err) + return + } + + // 按修改时间排序 + sort.Slice(files, func(i, j int) bool { + fi, _ := os.Stat(files[i]) + fj, _ := os.Stat(files[j]) + return fi.ModTime().After(fj.ModTime()) + }) + + // 删除旧文件 + for i := maxLogFiles; i < len(files); i++ { + if err := os.Remove(files[i]); err != nil { + log.Printf("failed to remove old log file %s: %v", files[i], err) + } + } +} + +func getCallerInfo() string { + _, file, line, ok := runtime.Caller(3) // 增加调用栈深度到3，跳过日志函数本身 + if !ok { + return "unknown:0" + } + // 返回完整路径 + return fmt.Sprintf("%s:%d", file, line) +} + func SysLog(s string) { t := time.Now() - _, _ = fmt.Fprintf(gin.DefaultWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s) + caller := getCallerInfo() + _, _ = fmt.Fprintf(gin.DefaultWriter, "[SYS] %v | %s | %s \n", t.Format("2006/01/02 - 15:04:05"), caller, s) } func SysError(s string) { t := time.Now() - _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s) + caller := getCallerInfo() + _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s | %s \n", t.Format("2006/01/02 - 15:04:05"), caller, s) } func LogInfo(ctx context.Context, msg string) { @@ -69,14 +201,99 @@ func LogError(ctx context.Context, msg string) { logHelper(ctx, loggerError, msg) } +// 获取错误类型 +func getErrorType(msg string) (string, string) { + // 提取错误码（如果有） + errorCode := "unknown" + if strings.Contains(msg, "status code:") { + parts := strings.Split(msg, "status code:") + if len(parts) > 1 { + errorCode = strings.TrimSpace(parts[1]) + } + } + + // 根据错误消息内容判断错误类型 + switch { + case strings.Contains(msg, "One or more parameter"): + return ErrorTypeParameter, errorCode + case strings.Contains(msg, "No candidates"): + return ErrorTypeNoCandidates, errorCode + case strings.Contains(msg, "do request failed"): + return ErrorTypeRequestFailed, errorCode + case strings.Contains(msg, "status code: 502"): + return ErrorTypeBadGateway, errorCode + case strings.Contains(msg, "doResponse failed"): + return ErrorTypeResponseFailed, errorCode + case strings.Contains(msg, "write: connection timed out"): + return ErrorTypeConnectionTimeout, errorCode + case strings.Contains(msg, "该令牌状态不可用"): + return ErrorTypeTokenUnavailable, errorCode + case strings.Contains(msg, "bad response status code 400"): + return ErrorTypeBadRequest, errorCode + case strings.Contains(msg, "无可用渠道"): + return ErrorTypeNoAvailableChannel, errorCode + default: + return ErrorTypeOther, errorCode + } +} + func logHelper(ctx context.Context, level string, msg string) { + // 获取请求ID + var requestId string + if id := ctx.Value(RequestIdKey); id != nil { + requestId = id.(string) + } + + // 如果有请求ID，则检查是否需要打印日志 + if requestId != "" { + // 从上下文中获取哈希值 + if ginCtx, ok := ctx.Value("gin_context").(*gin.Context); ok { + hashValue := ginCtx.GetInt("hash_value") + if hashValue > int(LogSampleRatio) { + return + } + } + } + writer := gin.DefaultErrorWriter if level == loggerINFO { writer = gin.DefaultWriter } - id := ctx.Value(RequestIdKey) now := time.Now() - _, _ = fmt.Fprintf(writer, "[%s] %v | %s | %s \n", level, now.Format("2006/01/02 - 15:04:05"), id, msg) + caller := getCallerInfo() + _, _ = fmt.Fprintf(writer, "[%s] %v | %s | %s | %s \n", level, now.Format("2006/01/02 - 15:04:05"), requestId, caller, msg) + + // 如果是错误日志，增加错误计数 + if level == loggerError { + errorType, errorCode := getErrorType(msg) + // 从上下文中获取相关信息 + channel := "unknown" + channelName := "unknown" + model := "unknown" + group := "unknown" + tokenName := "unknown" + + if ginCtx, ok := ctx.Value("gin_context").(*gin.Context); ok { + if ch := ginCtx.GetString("channel"); ch != "" { + channel = ch + } + if chName := ginCtx.GetString("channel_name"); chName != "" { + channelName = chName + } + if m := ginCtx.GetString("model"); m != "" { + model = m + } + if g := ginCtx.GetString("group"); g != "" { + group = g + } + if tn := ginCtx.GetString("token_name"); tn != "" { + tokenName = tn + } + } + + metrics.IncrementErrorLog(channel, channelName, errorCode, errorType, model, group, tokenName, 1.0) + } + logCount++ // we don't need accurate count, so no lock here if logCount > maxLogCount && !setupLogWorking { logCount = 0 @@ -89,7 +306,8 @@ func logHelper(ctx context.Context, level string, msg string) { func FatalLog(v ...any) { t := time.Now() - _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v) + caller := getCallerInfo() + _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %s | %v \n", t.Format("2006/01/02 - 15:04:05"), caller, v) os.Exit(1) } diff --git a/common/origin.go b/common/origin.go new file mode 100644 index 000000000000..b48d996495c5 --- /dev/null +++ b/common/origin.go @@ -0,0 +1,27 @@ +package common + +import ( + "strconv" + + "github.com/gin-gonic/gin" +) + +// GetOriginUserId 获取原始用户ID，如果请求头中存在则使用请求头中的值 +func GetOriginUserId(c *gin.Context, defaultUserId int) int { + if originUserId := c.GetHeader("X-Origin-User-ID"); originUserId != "" { + if userId, err := strconv.Atoi(originUserId); err == nil { + return userId + } + } + return defaultUserId +} + +// GetOriginChannelId 获取原始渠道ID，如果请求头中存在则使用请求头中的值 +func GetOriginChannelId(c *gin.Context, defaultChannelId int) int { + if originChannelId := c.GetHeader("X-Origin-Channel-ID"); originChannelId != "" { + if channelId, err := strconv.Atoi(originChannelId); err == nil { + return channelId + } + } + return defaultChannelId +} diff --git a/common/time.go b/common/time.go new file mode 100644 index 000000000000..1ed23bbca7c4 --- /dev/null +++ b/common/time.go @@ -0,0 +1,89 @@ +package common + +import ( + "time" +) + +var ( + // BeijingLocation 北京时区 + // 使用 "Asia/Shanghai" 作为时区标识符，如果加载失败则打印错误日志 + BeijingLocation = func() *time.Location { + loc, err := time.LoadLocation("Asia/Shanghai") + if err != nil { + // 打印错误日志 + println("Failed to load Asia/Shanghai timezone:", err.Error()) + return time.UTC + } + return loc + }() +) + +// GetBeijingTime 获取当前北京时区的时间 +func GetBeijingTime() time.Time { + // 直接使用北京时区创建时间 + return time.Now().In(BeijingLocation) +} + +// GetBeijingTimestamp 获取当前北京时区的时间戳 +func GetBeijingTimestamp() int64 { + // 获取当前时间 + now := time.Now() + + // 获取系统时区信息 + zone, offset := now.Zone() + + // 检查是否是北京时区（CST 且 UTC+8） + if zone != "CST" || offset != 8*3600 { + // 如果不是北京时区，转换为北京时区 + now = now.In(BeijingLocation) + } + + // 返回北京时区的时间戳 + return now.Unix() +} + +// GetBeijingTimeFromTimestamp 从时间戳获取北京时区的时间 +func GetBeijingTimeFromTimestamp(timestamp int64) time.Time { + // 直接使用北京时区创建时间 + return time.Unix(timestamp, 0).In(BeijingLocation) +} + +// GetBeijingTimeString 获取北京时区的格式化时间字符串 +func GetBeijingTimeString() string { + return GetBeijingTime().Format("2006-01-02 15:04:05") +} + +// GetBeijingDate 获取北京时区的日期（年月日） +func GetBeijingDate() (year int, month time.Month, day int) { + beijingTime := GetBeijingTime() + return beijingTime.Year(), beijingTime.Month(), beijingTime.Day() +} + +// GetBeijingHour 获取北京时区的小时 +func GetBeijingHour() int { + return GetBeijingTime().Hour() +} + +// GetBeijingTimeFromString 从字符串解析北京时区的时间 +func GetBeijingTimeFromString(timeStr string) (time.Time, error) { + // 直接解析为北京时区的时间 + t, err := time.ParseInLocation("2006-01-02 15:04:05", timeStr, BeijingLocation) + if err != nil { + return time.Time{}, err + } + return t, nil +} + +// PrintTimeInfo 打印当前时区和时间信息 +func PrintTimeInfo() { + now := time.Now() + zone, offset := now.Zone() + beijingTime := GetBeijingTime() + beijingTimestamp := GetBeijingTimestamp() + + println("系统时区:", zone) + println("时区偏移:", offset/3600, "小时") + println("系统时间:", now.Format("2006-01-02 15:04:05")) + println("北京时间:", beijingTime.Format("2006-01-02 15:04:05")) + println("日志存储的时间戳:", beijingTimestamp) +} diff --git a/controller/channel-test.go b/controller/channel-test.go index 02a30593dc34..236673a5449a 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -159,7 +159,7 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr milliseconds := tok.Sub(tik).Milliseconds() consumedTime := float64(milliseconds) / 1000.0 other := service.GenerateTextOtherInfo(c, info, priceData.ModelRatio, priceData.GroupRatio, priceData.CompletionRatio, 0, 0.0, priceData.ModelPrice) - model.RecordConsumeLog(c, 1, channel.Id, usage.PromptTokens, usage.CompletionTokens, info.OriginModelName, "模型测试", + model.RecordConsumeLog(c, 1, channel.Id, usage.PromptTokens, usage.CompletionTokens, 0, info.OriginModelName, "模型测试", quota, "模型测试", 0, quota, int(consumedTime), false, info.Group, other) common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody))) return nil, nil diff --git a/controller/channel.go b/controller/channel.go index f3ec6b3a39c1..b2365b1b0b73 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -49,9 +49,15 @@ func GetAllChannels(c *gin.Context) { if pageSize < 0 { pageSize = common.ItemsPerPage } + + // 获取用户信息 + userRole := c.GetInt("role") + username := c.GetString("username") + channelData := make([]*model.Channel, 0) idSort, _ := strconv.ParseBool(c.Query("id_sort")) enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode")) + if enableTagMode { tags, err := model.GetPaginatedTags(p*pageSize, pageSize) if err != nil { @@ -65,7 +71,23 @@ func GetAllChannels(c *gin.Context) { if tag != nil && *tag != "" { tagChannel, err := model.GetChannelsByTag(*tag, idSort) if err == nil { - channelData = append(channelData, tagChannel...) + // 过滤非管理员可见的渠道 + if userRole < 100 { + filteredChannels := make([]*model.Channel, 0) + for _, channel := range tagChannel { + // 检查渠道的分组是否包含用户名 + groups := strings.Split(channel.Group, ",") + for _, group := range groups { + if strings.Contains(group, username) { + filteredChannels = append(filteredChannels, channel) + break + } + } + } + channelData = append(channelData, filteredChannels...) + } else { + channelData = append(channelData, tagChannel...) + } } } } @@ -78,8 +100,26 @@ func GetAllChannels(c *gin.Context) { }) return } - channelData = channels + + // 过滤非管理员可见的渠道 + if userRole < 100 { + filteredChannels := make([]*model.Channel, 0) + for _, channel := range channels { + // 检查渠道的分组是否包含用户名 + groups := strings.Split(channel.Group, ",") + for _, group := range groups { + if strings.Contains(group, username) { + filteredChannels = append(filteredChannels, channel) + break + } + } + } + channelData = filteredChannels + } else { + channelData = channels + } } + c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -171,6 +211,11 @@ func SearchChannels(c *gin.Context) { modelKeyword := c.Query("model") idSort, _ := strconv.ParseBool(c.Query("id_sort")) enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode")) + + // 获取用户信息 + userRole := c.GetInt("role") + username := c.GetString("username") + channelData := make([]*model.Channel, 0) if enableTagMode { tags, err := model.SearchTags(keyword, group, modelKeyword, idSort) @@ -185,7 +230,23 @@ func SearchChannels(c *gin.Context) { if tag != nil && *tag != "" { tagChannel, err := model.GetChannelsByTag(*tag, idSort) if err == nil { - channelData = append(channelData, tagChannel...) + // 过滤非管理员可见的渠道 + if userRole < 100 { + filteredChannels := make([]*model.Channel, 0) + for _, channel := range tagChannel { + // 检查渠道的分组是否包含用户名 + groups := strings.Split(channel.Group, ",") + for _, group := range groups { + if strings.Contains(group, username) { + filteredChannels = append(filteredChannels, channel) + break + } + } + } + channelData = append(channelData, filteredChannels...) + } else { + channelData = append(channelData, tagChannel...) + } } } } @@ -198,8 +259,26 @@ func SearchChannels(c *gin.Context) { }) return } - channelData = channels + + // 过滤非管理员可见的渠道 + if userRole < 100 { + filteredChannels := make([]*model.Channel, 0) + for _, channel := range channels { + // 检查渠道的分组是否包含用户名 + groups := strings.Split(channel.Group, ",") + for _, group := range groups { + if strings.Contains(group, username) { + filteredChannels = append(filteredChannels, channel) + break + } + } + } + channelData = filteredChannels + } else { + channelData = channels + } } + c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/controller/group.go b/controller/group.go index 2c725a4d703f..cd678029b6ce 100644 --- a/controller/group.go +++ b/controller/group.go @@ -1,15 +1,27 @@ package controller import ( - "github.com/gin-gonic/gin" "net/http" "one-api/model" "one-api/setting" + "strings" + + "github.com/gin-gonic/gin" ) func GetGroups(c *gin.Context) { + // 获取用户信息 + userRole := c.GetInt("role") + username := c.GetString("username") + groupNames := make([]string, 0) - for groupName, _ := range setting.GetGroupRatioCopy() { + for groupName := range setting.GetGroupRatioCopy() { + // 如果不是超级管理员(role < 100)，只能看到包含自己用户名的分组 + if userRole < 100 { + if !strings.Contains(groupName, username) { + continue + } + } groupNames = append(groupNames, groupName) } c.JSON(http.StatusOK, gin.H{ @@ -23,10 +35,22 @@ func GetUserGroups(c *gin.Context) { usableGroups := make(map[string]map[string]interface{}) userGroup := "" userId := c.GetInt("id") + userRole := c.GetInt("role") + username := c.GetString("username") userGroup, _ = model.GetUserGroup(userId, false) + + // 遍历所有分组及其比率 for groupName, ratio := range setting.GetGroupRatioCopy() { - // UserUsableGroups contains the groups that the user can use + // 获取用户可用的分组 userUsableGroups := setting.GetUserUsableGroups(userGroup) + + // 如果不是超级管理员(role < 100)，只能看到包含自己用户名的分组 + if userRole < 100 { + if !strings.Contains(groupName, username) { + continue + } + } + if desc, ok := userUsableGroups[groupName]; ok { usableGroups[groupName] = map[string]interface{}{ "ratio": ratio, @@ -34,6 +58,7 @@ func GetUserGroups(c *gin.Context) { } } } + c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/controller/misc.go b/controller/misc.go index a451b5e3faa5..e5588edbe4ef 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -270,3 +270,13 @@ func ResetPassword(c *gin.Context) { }) return } + +func Ping(c *gin.Context) { + c.Writer.Header().Set("Retry_request_id", "Retry_request_id") + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "pong", + "timestamp": common.GetTimestamp(), + }) + return +} diff --git a/controller/option.go b/controller/option.go index c82fbd7ead44..82ca43d3410a 100644 --- a/controller/option.go +++ b/controller/option.go @@ -4,6 +4,7 @@ import ( "encoding/json" "net/http" "one-api/common" + "one-api/middleware" "one-api/model" "one-api/setting" "strings" @@ -107,3 +108,26 @@ func UpdateOption(c *gin.Context) { }) return } + +// ToggleRequestLog 切换请求体日志的开关状态 +func ToggleRequestLog(c *gin.Context) { + var request struct { + Enable bool `json:"enable"` + } + if err := c.ShouldBindJSON(&request); err != nil { + c.JSON(200, gin.H{ + "success": false, + "message": "无效的请求参数", + }) + return + } + + middleware.EnableRequestBodyLogging = request.Enable + c.JSON(200, gin.H{ + "success": true, + "message": "请求体日志状态已更新", + "data": gin.H{ + "enable": middleware.EnableRequestBodyLogging, + }, + }) +} diff --git a/controller/playground.go b/controller/playground.go index a2b54790f752..ffcdaebf735d 100644 --- a/controller/playground.go +++ b/controller/playground.go @@ -59,7 +59,8 @@ func Playground(c *gin.Context) { c.Set("token_name", "playground-"+group) channel, err := model.CacheGetRandomSatisfiedChannel(group, playgroundRequest.Model, 0) if err != nil { - message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", group, playgroundRequest.Model) + groupId := setting.GetGroupId(group) + message := fmt.Sprintf("当前分组id %d 下对于模型 %s 无可用渠道", groupId, playgroundRequest.Model) openaiErr = service.OpenAIErrorWrapperLocal(errors.New(message), "get_playground_channel_failed", http.StatusInternalServerError) return } diff --git a/controller/relay.go b/controller/relay.go index 460599b54ee9..4a4f88eab75a 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -4,51 +4,117 @@ import ( "bytes" "errors" "fmt" - "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "io" "log" "net/http" "one-api/common" "one-api/dto" + "one-api/metrics" "one-api/middleware" "one-api/model" "one-api/relay" + relaycommon "one-api/relay/common" "one-api/relay/constant" relayconstant "one-api/relay/constant" "one-api/relay/helper" "one-api/service" + "strconv" "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" ) -func relayHandler(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode { +func relayInfoHandler(c *gin.Context, relayMode int) (*relaycommon.RelayInfo, interface{}, string, *dto.OpenAIErrorWithStatusCode) { + switch relayMode { + case relayconstant.RelayModeImagesGenerations: + relayInfo, request, err := relay.ImageInfo(c) + if err != nil { + return nil, nil, "", err + } + return relayInfo, request, request.Model, nil + case relayconstant.RelayModeAudioSpeech: + fallthrough + case relayconstant.RelayModeAudioTranslation: + fallthrough + case relayconstant.RelayModeAudioTranscription: + relayInfo, request, err := relay.AudioInfo(c) + if err != nil { + return nil, nil, "", err + } + return relayInfo, request, request.Model, nil + case relayconstant.RelayModeRerank: + relayInfo, request, err := relay.EmbeddingInfo(c) + if err != nil { + return nil, nil, "", err + } + return relayInfo, request, request.Model, nil + case relayconstant.RelayModeEmbeddings: + relayInfo, request, err := relay.EmbeddingInfo(c) + if err != nil { + return nil, nil, "", err + } + return relayInfo, request, request.Model, nil + default: + relayInfo, request, err := relay.TextInfo(c) + if err != nil { + return nil, nil, "", err + } + return relayInfo, request, request.Model, nil + } +} + +func relayExecuteHandler(c *gin.Context, relayMode int, relayInfo *relaycommon.RelayInfo, request interface{}) *dto.OpenAIErrorWithStatusCode { var err *dto.OpenAIErrorWithStatusCode switch relayMode { case relayconstant.RelayModeImagesGenerations: - err = relay.ImageHelper(c) + imageRequest, ok := request.(*dto.ImageRequest) + if !ok { + return service.OpenAIErrorWrapperLocal(fmt.Errorf("failed assert request: %d", relayMode), "invalid_request_type", http.StatusInternalServerError) + } + err = relay.ImageHelper(c, relayInfo, imageRequest) case relayconstant.RelayModeAudioSpeech: fallthrough case relayconstant.RelayModeAudioTranslation: fallthrough case relayconstant.RelayModeAudioTranscription: - err = relay.AudioHelper(c) + audioRequest, ok := request.(*dto.AudioRequest) + if !ok { + return service.OpenAIErrorWrapperLocal(fmt.Errorf("failed assert request: %d", relayMode), "invalid_request_type", http.StatusInternalServerError) + } + err = relay.AudioHelper(c, relayInfo, audioRequest) case relayconstant.RelayModeRerank: - err = relay.RerankHelper(c, relayMode) + rerankRequest, ok := request.(*dto.RerankRequest) + if !ok { + return service.OpenAIErrorWrapperLocal(fmt.Errorf("failed assert request: %d", relayMode), "invalid_request_type", http.StatusInternalServerError) + } + err = relay.RerankHelper(c, relayInfo, rerankRequest) case relayconstant.RelayModeEmbeddings: - err = relay.EmbeddingHelper(c) + embeddingRequest, ok := request.(*dto.EmbeddingRequest) + if !ok { + return service.OpenAIErrorWrapperLocal(fmt.Errorf("failed assert request: %d", relayMode), "invalid_request_type", http.StatusInternalServerError) + } + err = relay.EmbeddingHelper(c, relayInfo, embeddingRequest) default: - err = relay.TextHelper(c) + textRequest, ok := request.(*dto.GeneralOpenAIRequest) + if !ok { + return service.OpenAIErrorWrapperLocal(fmt.Errorf("failed assert request: %d", relayMode), "invalid_request_type", http.StatusInternalServerError) + } + err = relay.TextHelper(c, relayInfo, textRequest) } return err } func Relay(c *gin.Context) { + startTime := time.Now() relayMode := constant.Path2RelayMode(c.Request.URL.Path) requestId := c.GetString(common.RequestIdKey) group := c.GetString("group") originalModel := c.GetString("original_model") + tokenKey := c.GetString("token_key") + tokenName := c.GetString("token_name") var openaiErr *dto.OpenAIErrorWithStatusCode - for i := 0; i <= common.RetryTimes; i++ { channel, err := getChannel(c, group, originalModel, i) if err != nil { @@ -56,16 +122,41 @@ func Relay(c *gin.Context) { openaiErr = service.OpenAIErrorWrapperLocal(err, "get_channel_failed", http.StatusInternalServerError) break } - - openaiErr = relayRequest(c, relayMode, channel) - + // 设置 channel 信息到上下文 + c.Set("channel", strconv.Itoa(channel.Id)) + c.Set("channel_name", channel.Name) + fillRelayRequest(c, channel) + var ( + relayInfo *relaycommon.RelayInfo + request interface{} + requestModel string + ) + relayInfo, request, requestModel, openaiErr = relayInfoHandler(c, relayMode) + if i == 0 { + // e2e 用户请求计数 + metrics.IncrementRelayRequestE2ETotalCounter(strconv.Itoa(channel.Id), channel.Name, requestModel, group, tokenKey, tokenName, 1) + } else { + // 重试计数 + channelTag := "" + if channel.Tag != nil { + channelTag = *channel.Tag + } + metrics.IncrementRelayRetryCounter(strconv.Itoa(channel.Id), channel.Name, channelTag, channel.GetBaseURL(), requestModel, group, 1) + } if openaiErr == nil { - return // 成功处理请求，直接返回 + openaiErr = executeRelayRequest(c, relayMode, relayInfo, request) + if openaiErr == nil { + metrics.IncrementRelayRequestE2ESuccessCounter(strconv.Itoa(channel.Id), channel.Name, requestModel, group, tokenKey, tokenName, 1) + metrics.ObserveRelayRequestE2EDuration(strconv.Itoa(channel.Id), channel.Name, requestModel, group, tokenKey, tokenName, time.Since(startTime).Seconds()) + return + } } go processChannelError(c, channel.Id, channel.Type, channel.Name, channel.GetAutoBan(), openaiErr) if !shouldRetry(c, openaiErr, common.RetryTimes-i) { + // e2e 失败计数 + metrics.IncrementRelayRequestE2EFailedCounter(strconv.Itoa(channel.Id), channel.Name, requestModel, group, strconv.Itoa(openaiErr.StatusCode), tokenKey, tokenName, 1) break } } @@ -80,6 +171,33 @@ func Relay(c *gin.Context) { common.LogError(c, fmt.Sprintf("origin 429 error: %s", openaiErr.Error.Message)) openaiErr.Error.Message = "当前分组上游负载已饱和，请稍后再试" } + + // 处理自定义的 NewAPI batch 错误码 + if openaiErr.StatusCode == dto.StatusNewAPIBatchRateLimitExceeded { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "当前服务端限速已满，请稍后再试" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchTimeout { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "未等待到结果，请稍后使用Retry_request_id再次查询" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchInternal { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "服务内部错误，请稍后再试" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchSubmitted { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "批量请求已提交，但是结果还未出来，请使用Retry_request_id查询结果" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchAccepted { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "批量请求已接受，正在处理中，请稍后使用Retry_request_id查询结果" + } + if openaiErr.StatusCode == dto.StatusRequestConflict { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "请求冲突，有其他请求使用了这个Retry_request_id，请稍后再试" + } + openaiErr.Error.Message = common.MessageWithRequestId(openaiErr.Error.Message, requestId) c.JSON(openaiErr.StatusCode, gin.H{ "error": openaiErr.Error, @@ -143,16 +261,44 @@ func WssRelay(c *gin.Context) { if openaiErr.StatusCode == http.StatusTooManyRequests { openaiErr.Error.Message = "当前分组上游负载已饱和，请稍后再试" } + // 处理自定义的 NewAPI batch 错误码 + if openaiErr.StatusCode == dto.StatusNewAPIBatchRateLimitExceeded { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "当前服务端限速已满，请稍后再试" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchTimeout { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "未等待到结果，请稍后使用Retry_request_id再次查询" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchInternal { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "服务内部错误，请稍后再试" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchSubmitted { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "批量请求已提交，但是结果还未出来，请使用Retry_request_id查询结果" + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchAccepted { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "批量请求已接受，正在处理中，请稍后使用Retry_request_id查询结果" + } + if openaiErr.StatusCode == dto.StatusRequestConflict { + common.LogError(c, fmt.Sprintf("origin %d error: %s", openaiErr.StatusCode, openaiErr.Error.Message)) + openaiErr.Error.Message = "请求冲突，有其他请求使用了这个Retry_request_id，请稍后再试" + } openaiErr.Error.Message = common.MessageWithRequestId(openaiErr.Error.Message, requestId) helper.WssError(c, ws, openaiErr.Error) } } -func relayRequest(c *gin.Context, relayMode int, channel *model.Channel) *dto.OpenAIErrorWithStatusCode { +func fillRelayRequest(c *gin.Context, channel *model.Channel) { addUsedChannel(c, channel.Id) requestBody, _ := common.GetRequestBody(c) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) - return relayHandler(c, relayMode) +} + +func executeRelayRequest(c *gin.Context, relayMode int, relayInfo *relaycommon.RelayInfo, request interface{}) *dto.OpenAIErrorWithStatusCode { + return relayExecuteHandler(c, relayMode, relayInfo, request) } func wssRequest(c *gin.Context, ws *websocket.Conn, relayMode int, channel *model.Channel) *dto.OpenAIErrorWithStatusCode { @@ -175,10 +321,12 @@ func getChannel(c *gin.Context, group, originalModel string, retryCount int) (*m if !autoBan { autoBanInt = 0 } + channelTag := c.GetString("channel_tag") return &model.Channel{ Id: c.GetInt("channel_id"), Type: c.GetInt("channel_type"), Name: c.GetString("channel_name"), + Tag: &channelTag, AutoBan: &autoBanInt, }, nil } @@ -206,9 +354,34 @@ func shouldRetry(c *gin.Context, openaiErr *dto.OpenAIErrorWithStatusCode, retry if openaiErr.StatusCode == http.StatusTooManyRequests { return true } + // 处理自定义的 NewAPI batch 错误码 + if openaiErr.StatusCode == dto.StatusNewAPIBatchRateLimitExceeded { + return false + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchTimeout { + return false + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchInternal { + return false + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchSubmitted { + return false + } + if openaiErr.StatusCode == dto.StatusNewAPIBatchAccepted { + return false + } + if openaiErr.StatusCode == dto.StatusRequestConflict { + return false + } + if openaiErr.StatusCode == 307 { return true } + if strings.Contains(openaiErr.Error.Message, "deadline exceeded") || strings.Contains(openaiErr.Error.Message, "request canceled") || strings.Contains(openaiErr.Error.Message, "copy_response_body_failed") { + common.LogInfo(c, fmt.Sprintf("客户端请求下游超时，不再重试 : %s", openaiErr.Error.Message)) + return false + } + if openaiErr.StatusCode/100 == 5 { // 超时不重试 if openaiErr.StatusCode == 504 || openaiErr.StatusCode == 524 { diff --git a/controller/token.go b/controller/token.go index a88032797384..195d734b4aac 100644 --- a/controller/token.go +++ b/controller/token.go @@ -1,15 +1,19 @@ package controller import ( - "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/model" "strconv" + "strings" + + "github.com/gin-gonic/gin" ) func GetAllTokens(c *gin.Context) { userId := c.GetInt("id") + userRole := c.GetInt("role") + username := c.GetString("username") p, _ := strconv.Atoi(c.Query("p")) size, _ := strconv.Atoi(c.Query("size")) if p < 0 { @@ -28,6 +32,45 @@ func GetAllTokens(c *gin.Context) { }) return } + + // 如果不是超级管理员，过滤掉不属于自己分组的 token + if userRole < 100 { + filteredTokens := make([]*model.Token, 0) + for _, token := range tokens { + if strings.Contains(token.Group, username) { + filteredTokens = append(filteredTokens, token) + } + } + tokens = filteredTokens + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": tokens, + }) + return +} + +func RootGetAllTokens(c *gin.Context) { + p, _ := strconv.Atoi(c.Query("p")) + size, _ := strconv.Atoi(c.Query("size")) + if p < 0 { + p = 0 + } + if size <= 0 { + size = common.ItemsPerPage + } else if size > 100 { + size = 100 + } + tokens, err := model.GetAllTokens(p*size, size) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/controller/usedata.go b/controller/usedata.go index 270eadf32324..df1753a2e85e 100644 --- a/controller/usedata.go +++ b/controller/usedata.go @@ -1,17 +1,40 @@ package controller import ( + "fmt" "github.com/gin-gonic/gin" "net/http" "one-api/model" "strconv" + "time" ) func GetAllQuotaDates(c *gin.Context) { startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) username := c.Query("username") - dates, err := model.GetAllQuotaDates(startTimestamp, endTimestamp, username) + token_name := c.Query("token_name") + dates, err := model.GetAllQuotaDates(startTimestamp, endTimestamp, username, token_name) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": dates, + }) + return +} + +func GetBilling(c *gin.Context) { + startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) + endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) + username := c.Query("username") + dates, err := model.GetAllQuotaDates(startTimestamp, endTimestamp, username, "") if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -54,3 +77,82 @@ func GetUserQuotaDates(c *gin.Context) { }) return } + +func ExportBillingExcel(c *gin.Context) { + // 从查询参数获取时间范围 + startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) + endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) + username := c.Query("user_name") + tokenname := c.Query("token_name") + // 判断时间跨度是否超过 1 个月 + if endTimestamp-startTimestamp > 2592000 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "时间跨度不能超过 1 个月", + }) + return + } + if tokenname != "" && username == "" { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "令牌名称和用户名称需要同时填写", + }) + } + // 转换时间戳为时间格式 + startTime := time.Unix(startTimestamp, 0) + if startTime.IsZero() { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "message": "无效的开始时间格式", + }) + return + } + + endTime := time.Unix(endTimestamp, 0) + if endTime.IsZero() { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "message": "无效的结束时间格式", + }) + return + } + + // 获取Excel数据 + excelBytes, err := model.GetBillingAndExportExcel(startTime.Unix(), endTime.Unix(), username, tokenname) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + // 设置文件名 + filename := fmt.Sprintf("billing_%s_%s.xlsx", + startTime.Format("20060102"), + endTime.Format("20060102")) + if username != "" { + filename = fmt.Sprintf("%s_billing_%s_%s.xlsx", + username, + startTime.Format("20060102"), + endTime.Format("20060102")) + if tokenname != "" { + filename = fmt.Sprintf("%s_%s_billing_%s_%s.xlsx", + username, + tokenname, + startTime.Format("20060102"), + endTime.Format("20060102")) + } + } + + // 设置响应头 + c.Header("Content-Description", "File Transfer") + c.Header("Content-Type", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filename)) + c.Header("Content-Transfer-Encoding", "binary") + c.Header("Expires", "0") + c.Header("Cache-Control", "must-revalidate") + c.Header("Pragma", "public") + + // 写入响应 + c.Data(http.StatusOK, "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", excelBytes) +} diff --git a/controller/user_rate_limit.go b/controller/user_rate_limit.go new file mode 100644 index 000000000000..52e874758715 --- /dev/null +++ b/controller/user_rate_limit.go @@ -0,0 +1,45 @@ +package controller + +import ( + "net/http" + "one-api/model" + + "github.com/gin-gonic/gin" +) + +// GetSpecificUserRateLimitConfig 获取特定用户的限速配置 +func GetSpecificUserRateLimitConfig(c *gin.Context) { + username := c.Query("username") + groupName := c.Query("group_name") + modelName := c.Query("model_name") + + if username == "" || groupName == "" || modelName == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "message": "用户名、分组名、模型名不能为空", + }) + return + } + + config, err := model.GetUserRateLimitConfig(username, groupName, modelName) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{ + "success": false, + "message": "配置不存在", + }) + return + } + + // 只返回需要的字段 + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": gin.H{ + "group_name": config.GroupName, + "username": config.Username, + "model_name": config.ModelName, + "current_rate_limit": config.CurrentRateLimit, + "is_rate_limit_enabled": config.IsRateLimitEnabled, + }, + }) +} diff --git a/dto/claude.go b/dto/claude.go new file mode 100644 index 000000000000..8068feb8fd64 --- /dev/null +++ b/dto/claude.go @@ -0,0 +1,218 @@ +package dto + +import "encoding/json" + +type ClaudeMetadata struct { + UserId string `json:"user_id"` +} + +type ClaudeMediaMessage struct { + Type string `json:"type,omitempty"` + Text *string `json:"text,omitempty"` + Model string `json:"model,omitempty"` + Source *ClaudeMessageSource `json:"source,omitempty"` + Usage *ClaudeUsage `json:"usage,omitempty"` + StopReason *string `json:"stop_reason,omitempty"` + PartialJson *string `json:"partial_json,omitempty"` + Role string `json:"role,omitempty"` + Thinking string `json:"thinking,omitempty"` + Signature string `json:"signature,omitempty"` + Delta string `json:"delta,omitempty"` + // tool_calls + Id string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Input any `json:"input,omitempty"` + Content json.RawMessage `json:"content,omitempty"` + ToolUseId string `json:"tool_use_id,omitempty"` +} + +func (c *ClaudeMediaMessage) SetText(s string) { + c.Text = &s +} + +func (c *ClaudeMediaMessage) GetText() string { + if c.Text == nil { + return "" + } + return *c.Text +} + +func (c *ClaudeMediaMessage) IsStringContent() bool { + var content string + return json.Unmarshal(c.Content, &content) == nil +} + +func (c *ClaudeMediaMessage) GetStringContent() string { + var content string + if err := json.Unmarshal(c.Content, &content); err == nil { + return content + } + return "" +} + +func (c *ClaudeMediaMessage) GetJsonRowString() string { + jsonContent, _ := json.Marshal(c) + return string(jsonContent) +} + +func (c *ClaudeMediaMessage) SetContent(content any) { + jsonContent, _ := json.Marshal(content) + c.Content = jsonContent +} + +func (c *ClaudeMediaMessage) ParseMediaContent() []ClaudeMediaMessage { + var mediaContent []ClaudeMediaMessage + if err := json.Unmarshal(c.Content, &mediaContent); err == nil { + return mediaContent + } + return make([]ClaudeMediaMessage, 0) +} + +type ClaudeMessageSource struct { + Type string `json:"type"` + MediaType string `json:"media_type,omitempty"` + Data any `json:"data,omitempty"` + Url string `json:"url,omitempty"` +} + +type ClaudeMessage struct { + Role string `json:"role"` + Content any `json:"content"` +} + +func (c *ClaudeMessage) IsStringContent() bool { + _, ok := c.Content.(string) + return ok +} + +func (c *ClaudeMessage) GetStringContent() string { + if c.IsStringContent() { + return c.Content.(string) + } + return "" +} + +func (c *ClaudeMessage) SetStringContent(content string) { + c.Content = content +} + +func (c *ClaudeMessage) ParseContent() ([]ClaudeMediaMessage, error) { + // map content to []ClaudeMediaMessage + // parse to json + jsonContent, _ := json.Marshal(c.Content) + var contentList []ClaudeMediaMessage + err := json.Unmarshal(jsonContent, &contentList) + if err != nil { + return make([]ClaudeMediaMessage, 0), err + } + return contentList, nil +} + +type Tool struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema map[string]interface{} `json:"input_schema"` +} + +type InputSchema struct { + Type string `json:"type"` + Properties any `json:"properties,omitempty"` + Required any `json:"required,omitempty"` +} + +type ClaudeRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt,omitempty"` + System any `json:"system,omitempty"` + Messages []ClaudeMessage `json:"messages,omitempty"` + MaxTokens uint `json:"max_tokens,omitempty"` + MaxTokensToSample uint `json:"max_tokens_to_sample,omitempty"` + StopSequences []string `json:"stop_sequences,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP float64 `json:"top_p,omitempty"` + TopK int `json:"top_k,omitempty"` + //ClaudeMetadata `json:"metadata,omitempty"` + Stream bool `json:"stream,omitempty"` + Tools any `json:"tools,omitempty"` + ToolChoice any `json:"tool_choice,omitempty"` + Thinking *Thinking `json:"thinking,omitempty"` +} + +type Thinking struct { + Type string `json:"type"` + BudgetTokens int `json:"budget_tokens"` +} + +func (c *ClaudeRequest) IsStringSystem() bool { + _, ok := c.System.(string) + return ok +} + +func (c *ClaudeRequest) GetStringSystem() string { + if c.IsStringSystem() { + return c.System.(string) + } + return "" +} + +func (c *ClaudeRequest) SetStringSystem(system string) { + c.System = system +} + +func (c *ClaudeRequest) ParseSystem() []ClaudeMediaMessage { + // map content to []ClaudeMediaMessage + // parse to json + jsonContent, _ := json.Marshal(c.System) + var contentList []ClaudeMediaMessage + if err := json.Unmarshal(jsonContent, &contentList); err == nil { + return contentList + } + return make([]ClaudeMediaMessage, 0) +} + +type ClaudeError struct { + Type string `json:"type,omitempty"` + Message string `json:"message,omitempty"` +} + +type ClaudeErrorWithStatusCode struct { + Error ClaudeError `json:"error"` + StatusCode int `json:"status_code"` + LocalError bool +} + +type ClaudeResponse struct { + Id string `json:"id,omitempty"` + Type string `json:"type"` + Role string `json:"role,omitempty"` + Content []ClaudeMediaMessage `json:"content,omitempty"` + Completion string `json:"completion,omitempty"` + StopReason string `json:"stop_reason,omitempty"` + Model string `json:"model,omitempty"` + Error *ClaudeError `json:"error,omitempty"` + Usage *ClaudeUsage `json:"usage,omitempty"` + Index *int `json:"index,omitempty"` + ContentBlock *ClaudeMediaMessage `json:"content_block,omitempty"` + Delta *ClaudeMediaMessage `json:"delta,omitempty"` + Message *ClaudeMediaMessage `json:"message,omitempty"` +} + +// set index +func (c *ClaudeResponse) SetIndex(i int) { + c.Index = &i +} + +// get index +func (c *ClaudeResponse) GetIndex() int { + if c.Index == nil { + return 0 + } + return *c.Index +} + +type ClaudeUsage struct { + InputTokens int `json:"input_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + OutputTokens int `json:"output_tokens"` +} diff --git a/dto/error.go b/dto/error.go index b347f6a159ef..eaee8ac52097 100644 --- a/dto/error.go +++ b/dto/error.go @@ -53,3 +53,13 @@ func (e GeneralErrorResponse) ToMessage() string { } return "" } + +// 自定义HTTP状态码 (使用非标准状态码范围) +const ( + StatusNewAPIBatchRateLimitExceeded = 499 // 自定义限流状态码 + StatusNewAPIBatchTimeout = 598 // 自定义超时状态码 + StatusNewAPIBatchInternal = 599 // 自定义内部错误状态码 + StatusNewAPIBatchSubmitted = 203 // 批量请求已提交，需要重试获取结果 + StatusNewAPIBatchAccepted = 202 // 批量请求已接受，正在处理中 + StatusRequestConflict = 409 // 请求冲突，如分布式锁获取失败 +) diff --git a/dto/openai_request.go b/dto/openai_request.go index 812e14a5ec73..b4a327d69918 100644 --- a/dto/openai_request.go +++ b/dto/openai_request.go @@ -2,7 +2,10 @@ package dto import ( "encoding/json" + "fmt" "strings" + + "one-api/common" ) type ResponseFormat struct { @@ -42,6 +45,7 @@ type GeneralOpenAIRequest struct { ResponseFormat *ResponseFormat `json:"response_format,omitempty"` EncodingFormat any `json:"encoding_format,omitempty"` Seed float64 `json:"seed,omitempty"` + LogitBias map[string]int `json:"logit_bias,omitempty"` Tools []ToolCallRequest `json:"tools,omitempty"` ToolChoice any `json:"tool_choice,omitempty"` User string `json:"user,omitempty"` @@ -51,6 +55,17 @@ type GeneralOpenAIRequest struct { Modalities any `json:"modalities,omitempty"` Audio any `json:"audio,omitempty"` ExtraBody any `json:"extra_body,omitempty"` + Thinking *ThinkingOptions `json:"thinking,omitempty"` + ThinkingConfig *ThinkingConfigs `json:"thinking_config,omitempty"` +} + +type ThinkingConfigs struct { + Enable bool `json:"enable,omitempty"` +} + +type ThinkingOptions struct { + Type string `json:"type"` + BudgetTokens int `json:"budget_tokens"` } type ToolCallRequest struct { @@ -116,17 +131,23 @@ type MediaContent struct { type MessageImageUrl struct { Url string `json:"url"` Detail string `json:"detail"` + Format string `json:"format,omitempty"` + Data string `json:"data,omitempty"` } type MessageInputAudio struct { - Data string `json:"data"` //base64 - Format string `json:"format"` + Data string `json:"data"` //base64 + Format string `json:"format"` + Fps float64 `json:"fps,omitempty"` + Url string `json:"url,omitempty"` } const ( ContentTypeText = "text" ContentTypeImageURL = "image_url" ContentTypeInputAudio = "input_audio" + ContentTypeVideoURL = "video_url" + ContentTypeYoutube = "youtube" ) func (m *Message) GetPrefix() bool { @@ -251,20 +272,52 @@ func (m *Message) ParseContent() []MediaContent { ImageUrl: MessageImageUrl{ Url: v, Detail: "high", + Format: "url", }, }) case map[string]interface{}: url, ok1 := v["url"].(string) detail, ok2 := v["detail"].(string) + format, ok3 := v["format"].(string) + data, ok4 := v["data"].(string) + if !ok2 { detail = "high" } - if ok1 { + if !ok3 { + format = "url" + } + + if format == "base64" { + var base64Data string + + if ok4 && data != "" { + base64Data = data + } else if ok1 && url != "" && strings.HasPrefix(url, "data:") { + parts := strings.Split(url, ",") + if len(parts) == 2 { + base64Data = parts[1] + } + } + + if base64Data != "" { + contentList = append(contentList, MediaContent{ + Type: ContentTypeImageURL, + ImageUrl: MessageImageUrl{ + Data: base64Data, + Detail: detail, + Format: format, + Url: url, + }, + }) + } + } else if format == "url" && ok1 { contentList = append(contentList, MediaContent{ Type: ContentTypeImageURL, ImageUrl: MessageImageUrl{ Url: url, Detail: detail, + Format: format, }, }) } @@ -274,16 +327,93 @@ func (m *Message) ParseContent() []MediaContent { if audioData, ok := contentItem["input_audio"].(map[string]interface{}); ok { data, ok1 := audioData["data"].(string) format, ok2 := audioData["format"].(string) + url, ok3 := audioData["url"].(string) + fps, ok3 := audioData["fps"].(float64) + if !ok2 { + if mimeType, ok3 := audioData["mime_type"].(string); ok3 { + format = mimeType + ok2 = true + } + } + if !ok3 { + fps = 0 + } + common.SysLog(fmt.Sprintf("Parsing audio content: data_ok=%v, format_ok=%v, format=%s, data_length=%d", ok1, ok2, format, len(data))) if ok1 && ok2 { contentList = append(contentList, MediaContent{ Type: ContentTypeInputAudio, InputAudio: MessageInputAudio{ Data: data, Format: format, + Fps: fps, + Url: url, + }, + }) + } + } + case ContentTypeVideoURL: + if videoData, ok := contentItem["video_url"].(map[string]interface{}); ok { + url, ok1 := videoData["url"].(string) + format, ok2 := videoData["format"].(string) + fps, ok3 := videoData["fps"].(float64) + data, ok4 := videoData["data"].(string) + + if !ok2 { + format = "url" + } + if !ok3 { + fps = 0 + } + + common.SysLog(fmt.Sprintf("Parsing video content: url_ok=%v, format_ok=%v, format=%s, url=%s, data_ok=%v", ok1, ok2, format, replaceBase64InURL(url), ok4)) + + if format == "base64" { + var base64Data string + + if ok4 && data != "" { + base64Data = data + } else if ok1 && url != "" && strings.HasPrefix(url, "data:") { + parts := strings.Split(url, ",") + if len(parts) == 2 { + base64Data = parts[1] + } + } + + if base64Data != "" { + contentList = append(contentList, MediaContent{ + Type: ContentTypeVideoURL, + InputAudio: MessageInputAudio{ + Data: base64Data, + Format: format, + Fps: fps, + Url: url, + }, + }) + } + } else if format == "url" && ok1 && url != "" { + contentList = append(contentList, MediaContent{ + Type: ContentTypeVideoURL, + InputAudio: MessageInputAudio{ + Url: url, + Format: format, + Fps: fps, }, }) } } + case ContentTypeYoutube: + mimetype, ok1 := contentItem["mimetype"].(string) + url, ok2 := contentItem["url"].(string) + if ok1 && ok2 { + contentList = append(contentList, MediaContent{ + Type: ContentTypeYoutube, + Text: mimetype, + ImageUrl: MessageImageUrl{ + Url: url, + Detail: "high", + }, + }) + } } } } @@ -293,3 +423,13 @@ func (m *Message) ParseContent() []MediaContent { } return contentList } + +func replaceBase64InURL(url string) string { + if strings.HasPrefix(url, "data:") { + parts := strings.Split(url, ",") + if len(parts) == 2 { + return "data:..." + } + } + return url +} diff --git a/dto/realtime.go b/dto/realtime.go index e28d813e240d..367bbeb6f798 100644 --- a/dto/realtime.go +++ b/dto/realtime.go @@ -51,8 +51,9 @@ type InputTokenDetails struct { } type OutputTokenDetails struct { - TextTokens int `json:"text_tokens"` - AudioTokens int `json:"audio_tokens"` + TextTokens int `json:"text_tokens"` + AudioTokens int `json:"audio_tokens"` + ReasoningTokens int `json:"reasoning_tokens"` } type RealtimeSession struct { diff --git a/go.mod b/go.mod index c9da57c64f3b..8c09f403784a 100644 --- a/go.mod +++ b/go.mod @@ -22,15 +22,17 @@ require ( github.com/golang-jwt/jwt v3.2.2+incompatible github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.0 - github.com/jinzhu/copier v0.4.0 github.com/joho/godotenv v1.5.1 github.com/pkg/errors v0.9.1 github.com/pkoukk/tiktoken-go v0.1.7 + github.com/prometheus/client_golang v1.21.1 github.com/samber/lo v1.39.0 github.com/shirou/gopsutil v3.21.11+incompatible - golang.org/x/crypto v0.27.0 + github.com/volcengine/volcengine-go-sdk v1.1.17 + github.com/xuri/excelize/v2 v2.9.0 + golang.org/x/crypto v0.31.0 golang.org/x/image v0.23.0 - golang.org/x/net v0.28.0 + golang.org/x/net v0.33.0 gorm.io/driver/mysql v1.4.3 gorm.io/driver/postgres v1.5.2 gorm.io/gorm v1.25.2 @@ -42,6 +44,7 @@ require ( github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.5 // indirect github.com/aws/smithy-go v1.20.2 // indirect + github.com/beorn7/perks v1.0.1 // indirect github.com/bytedance/sonic v1.11.6 // indirect github.com/bytedance/sonic/loader v0.1.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect @@ -58,7 +61,6 @@ require ( github.com/go-playground/universal-translator v0.18.1 // indirect github.com/go-sql-driver/mysql v1.7.0 // indirect github.com/goccy/go-json v0.10.2 // indirect - github.com/google/go-cmp v0.6.0 // indirect github.com/gorilla/context v1.1.1 // indirect github.com/gorilla/securecookie v1.1.1 // indirect github.com/gorilla/sessions v1.2.1 // indirect @@ -68,26 +70,39 @@ require ( github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect + github.com/jmespath/go-jmespath v0.4.0 // indirect github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/compress v1.17.11 // indirect github.com/klauspost/cpuid/v2 v2.2.9 // indirect github.com/leodido/go-urn v1.4.0 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mitchellh/mapstructure v1.5.0 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/pelletier/go-toml/v2 v2.2.1 // indirect + github.com/prometheus/client_model v0.6.1 // indirect + github.com/prometheus/common v0.62.0 // indirect + github.com/prometheus/procfs v0.15.1 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/richardlehane/mscfb v1.0.4 // indirect + github.com/richardlehane/msoleps v1.0.4 // indirect github.com/tklauser/go-sysconf v0.3.12 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect + github.com/volcengine/volc-sdk-golang v1.0.23 // indirect + github.com/xuri/efp v0.0.0-20240408161823-9ad904a10d6d // indirect + github.com/xuri/nfp v0.0.0-20240318013403-ab9948c2c4a7 // indirect github.com/yusufpapurcu/wmi v1.2.3 // indirect golang.org/x/arch v0.12.0 // indirect golang.org/x/exp v0.0.0-20240404231335-c0f41cb1a7a0 // indirect golang.org/x/sync v0.10.0 // indirect - golang.org/x/sys v0.27.0 // indirect + golang.org/x/sys v0.28.0 // indirect golang.org/x/text v0.21.0 // indirect - google.golang.org/protobuf v1.34.2 // indirect + google.golang.org/protobuf v1.36.1 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect modernc.org/libc v1.22.5 // indirect modernc.org/mathutil v1.5.0 // indirect diff --git a/go.sum b/go.sum index 0194ca302909..f8016b2ddcf3 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= +github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/Calcium-Ion/go-epay v0.0.4 h1:C96M7WfRLadcIVscWzwLiYs8etI1wrDmtFMuK2zP22A= github.com/Calcium-Ion/go-epay v0.0.4/go.mod h1:cxo/ZOg8ClvE3VAnCmEzbuyAZINSq7kFEN9oHj5WQ2U= github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= @@ -6,6 +8,7 @@ github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0 h1:onfun1RA+Kc github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0/go.mod h1:4yg+jNTYlDEzBjhGS96v+zjyA3lfXlFd5CiTLIkPBLI= github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 h1:HblK3eJHq54yET63qPCTJnks3loDse5xRmmqHgHzwoI= github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6/go.mod h1:pbiaLIeYLUbgMY1kwEAdwO6UKD5ZNwdPGQlwokS9fe8= +github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY= github.com/aws/aws-sdk-go-v2 v1.26.1 h1:5554eUqIYVWpU0YmeeYZ0wU64H2VLBs8TlhRB2L+EkA= github.com/aws/aws-sdk-go-v2 v1.26.1/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2 h1:x6xsQXGSmW6frevwDA+vi/wqhp1ct18mVXYN08/93to= @@ -20,14 +23,18 @@ github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.7.4 h1:JgHnonzbnA3pbqj76w github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.7.4/go.mod h1:nZspkhg+9p8iApLFoyAqfyuMP0F38acy2Hm3r5r95Cg= github.com/aws/smithy-go v1.20.2 h1:tbp628ireGtzcHDDmLT/6ADHidqnwgF57XOXZe6tp4Q= github.com/aws/smithy-go v1.20.2/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/bytedance/gopkg v0.0.0-20220118071334-3db87571198b h1:LTGVFpNmNHhj0vhOlfgWueFJ32eK9blaIlHR2ciXOT0= github.com/bytedance/gopkg v0.0.0-20220118071334-3db87571198b/go.mod h1:2ZlV9BaUH4+NXIBF0aMdKKAnHTzqH+iMU4KUjAbL23Q= github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc0= github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4= github.com/bytedance/sonic/loader v0.1.1 h1:c+e5Pt1k/cy5wMveRDyk2X4B9hF4g7an8N3zCYjJFNM= github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= +github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y= github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg= @@ -42,6 +49,8 @@ github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxK github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= +github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= @@ -91,14 +100,31 @@ github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keLg81eXfW3O+oY= github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I= +github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= +github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8= +github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ= github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/context v1.1.1 h1:AWwleXJkX/nhcU9bZSnZoi3h/qGYqQAGhq6zZe/aQW8= @@ -117,30 +143,38 @@ github.com/jackc/pgx/v5 v5.7.1 h1:x7SYsPBYDkHDksogeSmZZ5xzThcTgRz++I5E+ePFUcs= github.com/jackc/pgx/v5 v5.7.1/go.mod h1:e7O26IywZZ+naJtWWos6i6fvWK+29etgITqrqHLfoZA= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= -github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= +github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= +github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= +github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= +github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc= +github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/klauspost/cpuid/v2 v2.2.9 h1:66ze0taIn2H33fBvCkXuv9BmCwDfafmiIVpKV9kKGuY= github.com/klauspost/cpuid/v2 v2.2.9/go.mod h1:rqkxqrZ1EhYM9G+hXH7YdowN5R5RGN6NK4QwQ3WMXF8= github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= -github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII= github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= @@ -157,6 +191,10 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0= github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU= github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= @@ -173,12 +211,27 @@ github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQ github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_golang v1.21.1 h1:DOvXXTqVzvkIewV/CDPFdejpMCGeMcbGCQ8YOmu+Ibk= +github.com/prometheus/client_golang v1.21.1/go.mod h1:U9NM32ykUErtVBxdvD3zfi+EuFkkaBvMb09mIfe0Zgg= +github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= +github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E= +github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY= +github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ2Io= +github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I= +github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc= +github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk= github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/richardlehane/mscfb v1.0.4 h1:WULscsljNPConisD5hR0+OyZjwK46Pfyr6mPu5ZawpM= +github.com/richardlehane/mscfb v1.0.4/go.mod h1:YzVpcZg9czvAuhk9T+a3avCpcFPMUWm7gK3DypaEsUk= +github.com/richardlehane/msoleps v1.0.1/go.mod h1:BWev5JBpU9Ko2WAgmZEuiz4/u3ZYTKbjLycmwiWUfWg= +github.com/richardlehane/msoleps v1.0.4 h1:WuESlvhX3gH2IHcd8UqyCuFY5yiq/GR/yqaSM/9/g00= +github.com/richardlehane/msoleps v1.0.4/go.mod h1:BWev5JBpU9Ko2WAgmZEuiz4/u3ZYTKbjLycmwiWUfWg= github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= -github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8= github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE= +github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= +github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= github.com/samber/lo v1.39.0 h1:4gTz1wUhNYLhFSKl6O+8peW0v2F4BCY034GRpU9WnuA= github.com/samber/lo v1.39.0/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA= github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI= @@ -189,14 +242,16 @@ github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpE github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= +github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= -github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= @@ -209,6 +264,16 @@ github.com/ugorji/go/codec v1.1.7/go.mod h1:Ax+UKWsSmolVDwsd+7N3ZtXu+yMGCf907BLY github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY= github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +github.com/volcengine/volc-sdk-golang v1.0.23 h1:anOslb2Qp6ywnsbyq9jqR0ljuO63kg9PY+4OehIk5R8= +github.com/volcengine/volc-sdk-golang v1.0.23/go.mod h1:AfG/PZRUkHJ9inETvbjNifTDgut25Wbkm2QoYBTbvyU= +github.com/volcengine/volcengine-go-sdk v1.1.17 h1:Izrcx/FERzGvpY3ufPjt4GR7Ak6y94aMVXbnLmeuw2g= +github.com/volcengine/volcengine-go-sdk v1.1.17/go.mod h1:EyKoi6t6eZxoPNGr2GdFCZti2Skd7MO3eUzx7TtSvNo= +github.com/xuri/efp v0.0.0-20240408161823-9ad904a10d6d h1:llb0neMWDQe87IzJLS4Ci7psK/lVsjIS2otl+1WyRyY= +github.com/xuri/efp v0.0.0-20240408161823-9ad904a10d6d/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI= +github.com/xuri/excelize/v2 v2.9.0 h1:1tgOaEq92IOEumR1/JfYS/eR0KHOCsRv/rYXXh6YJQE= +github.com/xuri/excelize/v2 v2.9.0/go.mod h1:uqey4QBZ9gdMeWApPLdhm9x+9o2lq4iVmjiLfBS5hdE= +github.com/xuri/nfp v0.0.0-20240318013403-ab9948c2c4a7 h1:hPVCafDV85blFTabnqKgNhDCkJX25eik94Si9cTER4A= +github.com/xuri/nfp v0.0.0-20240318013403-ab9948c2c4a7/go.mod h1:WwHg+CVyzlv/TX9xqBFXEZAuxOPxn2k1GNHwG41IIUQ= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= @@ -216,19 +281,34 @@ github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQ golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= golang.org/x/arch v0.12.0 h1:UsYJhbzPYGsT0HbEdmYcqtCv8UNGvnaL561NnIUvaKg= golang.org/x/arch v0.12.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= -golang.org/x/crypto v0.27.0 h1:GXm2NjJrPaiv/h1tb2UH8QfgC/hOf/+z0p6PT8o1w7A= -golang.org/x/crypto v0.27.0/go.mod h1:1Xngt8kV6Dvbssa53Ziq6Eqn0HqbZi5Z6R0ZpwQzt70= +golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= +golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= +golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20240404231335-c0f41cb1a7a0 h1:985EYyeCOxTpcgOTJpflJUwOeEz0CQOdPt73OzpE9F8= golang.org/x/exp v0.0.0-20240404231335-c0f41cb1a7a0/go.mod h1:/lliqkxwWAhPjf5oSOIJup2XcqJaw8RGS6k3TGEc7GI= golang.org/x/image v0.23.0 h1:HseQ7c2OpPKTPVzNjG5fwJsOTCiiwS4QdsYi5XU6H68= golang.org/x/image v0.23.0/go.mod h1:wJJBTdLfCCf3tiHa1fNxpZmUI4mmoZvwMCPP0ddoNKY= +golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= +golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= +golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= +golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= -golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE= -golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= +golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I= +golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4= +golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ= golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -239,22 +319,45 @@ golang.org/x/sys v0.0.0-20220110181412-a018aaa089fe/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s= -golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= +golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= +golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= +google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= +google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= +google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= +google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo= +google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= +google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= +google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= -google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= -google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= +google.golang.org/protobuf v1.36.1 h1:yBPeRvTftaleIgM3PZ/WBIZ7XM/eEYAaEyCwvyjq/gk= +google.golang.org/protobuf v1.36.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= @@ -275,6 +378,8 @@ gorm.io/driver/postgres v1.5.2/go.mod h1:fmpX0m2I1PKuR7mKZiEluwrP3hbs+ps7JIGMUBp gorm.io/gorm v1.23.8/go.mod h1:l2lP/RyAtc1ynaTjFksBde/O8v9oOGIApu2/xRitmZk= gorm.io/gorm v1.25.2 h1:gs1o6Vsa+oVKG/a9ElL3XgyGfghFfkKA2SInQaCyMho= gorm.io/gorm v1.25.2/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k= +honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= +honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= modernc.org/libc v1.22.5 h1:91BNch/e5B0uPbJFgqbxXuOnxBQjlS//icfQEGmvyjE= modernc.org/libc v1.22.5/go.mod h1:jj+Z7dTNX8fBScMVNRAYZ/jF91K8fdT2hYMThc3YjBY= modernc.org/mathutil v1.5.0 h1:rV0Ko/6SfM+8G+yKiyI830l3Wuz1zRutdslNoQ0kfiQ= diff --git a/main.go b/main.go index 495057cf1efc..c85be780eb5e 100644 --- a/main.go +++ b/main.go @@ -2,18 +2,24 @@ package main import ( "embed" + "flag" "fmt" "log" "net/http" "one-api/common" "one-api/constant" "one-api/controller" + "one-api/metrics" "one-api/middleware" "one-api/model" "one-api/router" "one-api/service" "os" "strconv" + "time" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/bytedance/gopkg/util/gopool" "github.com/gin-contrib/sessions" @@ -22,6 +28,8 @@ import ( "github.com/joho/godotenv" _ "net/http/pprof" + + "one-api/relay/channel/volcengine" ) //go:embed web/dist @@ -31,13 +39,47 @@ var buildFS embed.FS var indexPage []byte func main() { - err := godotenv.Load(".env") - if err != nil { - common.SysLog("Support for .env file is disabled") + // 添加命令行参数支持 + configFile := flag.String("config", "", "path to config file") + flag.Parse() + + // 打印时区和时间信息 + common.PrintTimeInfo() + + // 根据是否指定配置文件决定加载哪个文件 + if *configFile != "" { + err := godotenv.Load(*configFile) + if err != nil { + common.SysLog(fmt.Sprintf("Failed to load config file %s: %v", *configFile, err)) + } + } else { + err := godotenv.Load(".env") + if err != nil { + common.SysLog("Support for .env file is disabled") + } } common.LoadEnv() + // 读取透传日志配置 + if os.Getenv("LOG_PASSTHROUGH_ENABLED") == "true" { + common.LogPassthroughEnabled = true + common.SysLog("log passthrough enabled") + } + + // 读取日志采样比例配置 + if os.Getenv("LOG_SAMPLE_RATIO") != "" { + ratio, err := strconv.Atoi(os.Getenv("LOG_SAMPLE_RATIO")) + if err != nil { + common.FatalLog("failed to parse LOG_SAMPLE_RATIO: " + err.Error()) + } + if ratio < 0 || ratio > 100 { + common.FatalLog("LOG_SAMPLE_RATIO must be between 0 and 100") + } + common.LogSampleRatio = ratio + common.SysLog(fmt.Sprintf("log sample ratio set to %d%%", ratio)) + } + common.SetupLogger() common.SysLog("New API " + common.Version + " started") if os.Getenv("GIN_MODE") != "debug" { @@ -47,7 +89,7 @@ func main() { common.SysLog("running in debug mode") } // Initialize SQL Database - err = model.InitDB() + err := model.InitDB() if err != nil { common.FatalLog("failed to initialize database: " + err.Error()) } @@ -56,6 +98,36 @@ func main() { if err != nil { common.FatalLog("failed to initialize database: " + err.Error()) } + // Initialize Central Control Database + err = model.InitCentralDB() + if err != nil { + common.FatalLog("failed to initialize central control database: " + err.Error()) + } + err = model.InitLogTable() + if err != nil { + common.FatalLog("failed to initialize database: " + err.Error()) + } + model.GetLogTableName(time.Now().Unix()) + // 每5分钟执行一次GetLogTableName + go func() { + ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + for range ticker.C { + model.GetLogTableName(time.Now().Unix()) + } + }() + + // 初始化请求持久化存储 + if os.Getenv("REQUEST_PERSISTENCE_ENABLED") == "true" { + model.RequestPersistenceEnabled = true + common.SysLog("request persistence enabled") + err = model.InitRequestPersistence() + if err != nil { + common.FatalLog("failed to initialize request persistence: " + err.Error()) + } + model.StartTableCheckRoutine() + } + defer func() { err := model.CloseDB() if err != nil { @@ -69,10 +141,26 @@ func main() { common.FatalLog("failed to initialize Redis: " + err.Error()) } + // Initialize Keep-Alive Manager for Redis keys + if err := volcengine.InitKeepAliveManager(); err != nil { + common.FatalLog("failed to initialize keep-alive manager: " + err.Error()) + } + // 在应用关闭时清理保活管理器 + defer func() { + if err := volcengine.ShutdownKeepAliveManager(); err != nil { + common.SysError("failed to shutdown keep-alive manager: " + err.Error()) + } + }() + // Initialize constants constant.InitEnv() // Initialize options model.InitOptionMap() + model.InitGroups() + + // 初始化batch请求平均耗时 + volcengine.InitBatchRequestAverageDuration() + if common.RedisEnabled { // for compatibility with old versions common.MemoryCacheEnabled = true @@ -89,7 +177,6 @@ func main() { // 数据看板 go model.UpdateQuotaData() - if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" { frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY")) if err != nil { @@ -112,6 +199,21 @@ func main() { controller.UpdateTaskBulk() }) } + if os.Getenv("ENABLE_METRICS") != "" { + register := prometheus.NewRegistry() + metrics.RegisterMetrics(register) + gatherersRegistry := prometheus.Gatherers{register} + go func() { + http.Handle("/metrics", promhttp.HandlerFor(gatherersRegistry, promhttp.HandlerOpts{})) + metricsPort := "9090" + if os.Getenv("METRICS_PORT") != "" { + metricsPort = os.Getenv("METRICS_PORT") + } + log.Println(http.ListenAndServe(fmt.Sprintf("0.0.0.0:%s", metricsPort), + nil)) + }() + + } if os.Getenv("BATCH_UPDATE_ENABLED") == "true" { common.BatchUpdateEnabled = true common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s") @@ -141,6 +243,7 @@ func main() { })) // This will cause SSE not to work!!! //server.Use(gzip.Gzip(gzip.DefaultCompression)) + server.Use(middleware.RequestLogger()) server.Use(middleware.RequestId()) middleware.SetUpLogger(server) // Initialize session store diff --git a/metrics/metrics.go b/metrics/metrics.go new file mode 100644 index 000000000000..a3a1afc31cba --- /dev/null +++ b/metrics/metrics.go @@ -0,0 +1,215 @@ +package metrics + +import ( + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" +) + +const ( + Namespace = "new_api" +) + +func RegisterMetrics(registry prometheus.Registerer) { + // channel + registry.MustRegister(relayRequestTotalCounter) + registry.MustRegister(relayRequestSuccessCounter) + registry.MustRegister(relayRequestFailedCounter) + registry.MustRegister(relayRequestRetryCounter) + registry.MustRegister(relayRequestDurationObsever) + // e2e + registry.MustRegister(relayRequestE2ETotalCounter) + registry.MustRegister(relayRequestE2ESuccessCounter) + registry.MustRegister(relayRequestE2EFailedCounter) + registry.MustRegister(relayRequestE2EDurationObsever) + // batch + registry.MustRegister(batchRequestCounter) + registry.MustRegister(batchRequestDurationObsever) + // token metrics + registry.MustRegister(inputTokensCounter) + registry.MustRegister(outputTokensCounter) + registry.MustRegister(cacheHitTokensCounter) + registry.MustRegister(inferenceTokensCounter) + // error log metrics + registry.MustRegister(errorLogCounter) +} + +var ( + relayRequestTotalCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_total", + Help: "Total number of relay request total", + }, []string{"channel", "channel_name", "tag", "base_url", "model", "group"}) + relayRequestSuccessCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_success", + Help: "Total number of relay request success", + }, []string{"channel", "channel_name", "tag", "base_url", "model", "group", "code"}) + relayRequestFailedCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_failed", + Help: "Total number of relay request failed", + }, []string{"channel", "channel_name", "tag", "base_url", "model", "group", "code"}) + relayRequestRetryCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_retry", + Help: "Total number of relay request retry", + }, []string{"channel", "channel_name", "tag", "base_url", "model", "group"}) + relayRequestDurationObsever = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Subsystem: Namespace, + Name: "relay_request_duration", + Help: "Duration of relay request", + Buckets: prometheus.ExponentialBuckets(1, 2, 12), + }, + []string{"channel", "channel_name", "tag", "base_url", "model", "group"}, + ) + relayRequestE2ETotalCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_e2e_total", + Help: "Total number of relay request e2e total", + }, []string{"channel", "channel_name", "model", "group", "token_key", "token_name"}) + relayRequestE2ESuccessCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_e2e_success", + Help: "Total number of relay request e2e success", + }, []string{"channel", "channel_name", "model", "group", "token_key", "token_name"}) + relayRequestE2EFailedCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "relay_request_e2e_failed", + Help: "Total number of relay request e2e failed", + }, []string{"channel", "channel_name", "model", "group", "code", "token_key", "token_name"}) + relayRequestE2EDurationObsever = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Subsystem: Namespace, + Name: "relay_request_e2e_duration", + Help: "Duration of relay request e2e", + Buckets: prometheus.ExponentialBuckets(1, 2, 12), + }, + []string{"channel", "channel_name", "model", "group", "token_key", "token_name"}, + ) + // Batch request metrics + batchRequestCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "batch_request_total", + Help: "Total number of batch requests by status code", + }, []string{"channel", "channel_name", "tag", "base_url", "model", "group", "code", "retry_header"}) + batchRequestDurationObsever = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Subsystem: Namespace, + Name: "batch_request_duration", + Help: "Duration of batch request", + Buckets: prometheus.ExponentialBuckets(1, 2, 12), + }, + []string{"channel", "channel_name", "tag", "base_url", "model", "group", "code", "retry_header"}, + ) + // Token metrics + inputTokensCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "input_tokens_total", + Help: "Total number of input tokens processed", + }, []string{"channel", "channel_name", "model", "group", "user_id", "user_name", "token_name"}) + + outputTokensCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "output_tokens_total", + Help: "Total number of output tokens generated", + }, []string{"channel", "channel_name", "model", "group", "user_id", "user_name", "token_name"}) + + cacheHitTokensCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "cache_hit_tokens_total", + Help: "Total number of tokens served from cache", + }, []string{"channel", "channel_name", "model", "group", "user_id", "user_name", "token_name"}) + + inferenceTokensCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "inference_tokens_total", + Help: "Total number of tokens processed during inference", + }, []string{"channel", "channel_name", "model", "group", "user_id", "user_name", "token_name"}) + + errorLogCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Subsystem: Namespace, + Name: "error_log_total", + Help: "Total number of error logs", + }, []string{"channel", "channel_name", "error_code", "error_type", "model", "group", "token_name"}) +) + +func IncrementRelayRequestTotalCounter(channel, channelName, tag, baseURL, model, group string, add float64) { + relayRequestTotalCounter.WithLabelValues(channel, channelName, tag, baseURL, model, group).Add(add) +} + +func IncrementRelayRequestSuccessCounter(channel, channelName, tag, baseURL, model, group, statusCode string, add float64) { + relayRequestSuccessCounter.WithLabelValues(channel, channelName, tag, baseURL, model, group, statusCode).Add(add) +} + +func IncrementRelayRequestFailedCounter(channel, channelName, tag, baseURL, model, group, code string, add float64) { + relayRequestFailedCounter.WithLabelValues(channel, channelName, tag, baseURL, model, group, code).Add(add) +} + +func IncrementRelayRetryCounter(channel, channelName, tag, baseURL, model, group string, add float64) { + relayRequestRetryCounter.WithLabelValues(channel, channelName, tag, baseURL, model, group).Add(add) +} + +func ObserveRelayRequestDuration(channel, channelName, tag, baseURL, model, group string, duration float64) { + relayRequestDurationObsever.WithLabelValues(channel, channelName, tag, baseURL, model, group).Observe(duration) +} + +func IncrementRelayRequestE2ETotalCounter(channel, channelName, model, group, tokenKey, tokenName string, add float64) { + relayRequestE2ETotalCounter.WithLabelValues(channel, channelName, model, group, tokenKey, tokenName).Add(add) +} + +func IncrementRelayRequestE2ESuccessCounter(channel, channelName, model, group, tokenKey, tokenName string, add float64) { + relayRequestE2ESuccessCounter.WithLabelValues(channel, channelName, model, group, tokenKey, tokenName).Add(add) +} + +func IncrementRelayRequestE2EFailedCounter(channel, channelName, model, group, code, tokenKey, tokenName string, add float64) { + relayRequestE2EFailedCounter.WithLabelValues(channel, channelName, model, group, code, tokenKey, tokenName).Add(add) +} + +func ObserveRelayRequestE2EDuration(channel, channelName, model, group, tokenKey, tokenName string, duration float64) { + relayRequestE2EDurationObsever.WithLabelValues(channel, channelName, model, group, tokenKey, tokenName).Observe(duration) +} + +// Batch request metrics functions +func IncrementBatchRequestCounter(channel, channelName, tag, baseURL, model, group, code, retryHeader string, add float64) { + batchRequestCounter.WithLabelValues(channel, channelName, tag, baseURL, model, group, code, retryHeader).Add(add) +} + +func ObserveBatchRequestDuration(channel, channelName, tag, baseURL, model, group, code, retryHeader string, duration float64) { + batchRequestDurationObsever.WithLabelValues(channel, channelName, tag, baseURL, model, group, code, retryHeader).Observe(duration) +} + +// Token metrics functions +func IncrementInputTokens(channel, channelName, model, group, userId, userName, tokenName string, add float64) { + inputTokensCounter.WithLabelValues(channel, channelName, model, group, userId, userName, tokenName).Add(add) +} + +func IncrementOutputTokens(channel, channelName, model, group, userId, userName, tokenName string, add float64) { + outputTokensCounter.WithLabelValues(channel, channelName, model, group, userId, userName, tokenName).Add(add) +} + +func IncrementCacheHitTokens(channel, channelName, model, group, userId, userName, tokenName string, add float64) { + cacheHitTokensCounter.WithLabelValues(channel, channelName, model, group, userId, userName, tokenName).Add(add) +} + +func IncrementInferenceTokens(channel, channelName, model, group, userId, userName, tokenName string, add float64) { + inferenceTokensCounter.WithLabelValues(channel, channelName, model, group, userId, userName, tokenName).Add(add) +} + +// Error log metrics function +func IncrementErrorLog(channel, channelName, errorCode, errorType, model, group, tokenName string, add float64) { + errorLogCounter.WithLabelValues(channel, channelName, errorCode, errorType, model, group, tokenName).Add(add) +} diff --git a/middleware/auth.go b/middleware/auth.go index a589f52ccf04..4fab42d9d062 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -1,13 +1,15 @@ package middleware import ( - "github.com/gin-contrib/sessions" - "github.com/gin-gonic/gin" + "fmt" "net/http" "one-api/common" "one-api/model" "strconv" "strings" + + "github.com/gin-contrib/sessions" + "github.com/gin-gonic/gin" ) func validUserInfo(username string, role int) bool { @@ -116,6 +118,10 @@ func authHelper(c *gin.Context, minRole int) { c.Abort() return } + + // 添加日志打印 + common.SysLog(fmt.Sprintf("[Auth Info] UserID: %v, Role: %v, Username: %v", id, role, username)) + c.Set("username", username) c.Set("role", role) c.Set("id", id) diff --git a/middleware/distributor.go b/middleware/distributor.go index 49fcf59b8a13..eeacaea46510 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -1,8 +1,10 @@ package middleware import ( - "errors" + "bytes" + "encoding/json" "fmt" + "io" "net/http" "one-api/common" "one-api/constant" @@ -44,12 +46,14 @@ func Distribute() func(c *gin.Context) { if tokenGroup != "" { // check common.UserUsableGroups[userGroup] if _, ok := setting.GetUserUsableGroups(userGroup)[tokenGroup]; !ok { - abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("令牌分组 %s 已被禁用", tokenGroup)) + tokenGroupId := setting.GetGroupId(tokenGroup) + abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("令牌分组id %d 已被禁用", tokenGroupId)) return } // check group in common.GroupRatio if !setting.ContainsGroupRatio(tokenGroup) { - abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup)) + tokenGroupId := setting.GetGroupId(tokenGroup) + abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组id %d 已被弃用", tokenGroupId)) return } userGroup = tokenGroup @@ -97,7 +101,8 @@ func Distribute() func(c *gin.Context) { if shouldSelectChannel { channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model, 0) if err != nil { - message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model) + userGroupId := setting.GetGroupId(userGroup) + message := fmt.Sprintf("当前分组id %d 下对于模型 %s 无可用渠道", userGroupId, modelRequest.Model) // 如果错误，但是渠道不为空，说明是数据库一致性问题 if channel != nil { common.SysError(fmt.Sprintf("渠道不存在：%d", channel.Id)) @@ -108,7 +113,8 @@ func Distribute() func(c *gin.Context) { return } if channel == nil { - abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道（数据库一致性已被破坏）", userGroup, modelRequest.Model)) + userGroupId := setting.GetGroupId(userGroup) + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("当前分组id %d 下对于模型 %s 无可用渠道（数据库一致性已被破坏）", userGroupId, modelRequest.Model)) return } } @@ -134,7 +140,7 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { midjourneyRequest := dto.MidjourneyRequest{} err = common.UnmarshalBodyReusable(c, &midjourneyRequest) if err != nil { - return nil, false, err + return nil, false, fmt.Errorf("无效的请求, %s", err.Error()) } midjourneyModel, mjErr, success := service.GetMjRequestModel(relayMode, &midjourneyRequest) if mjErr != nil { @@ -163,11 +169,40 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { c.Set("platform", string(constant.TaskPlatformSuno)) c.Set("relay_mode", relayMode) } else if !strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") { - err = common.UnmarshalBodyReusable(c, &modelRequest) - } - if err != nil { - return nil, false, errors.New("无效的请求, " + err.Error()) + // 检查请求体是否为空 + body, err := io.ReadAll(c.Request.Body) + if err != nil { + return nil, false, fmt.Errorf("无效的请求, 读取请求体失败: %s", err.Error()) + } + // 重置请求体 + c.Request.Body = io.NopCloser(bytes.NewBuffer(body)) + + // 如果请求体为空，根据路径设置默认模型 + if len(body) == 0 { + if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") { + modelRequest.Model = "text-moderation-stable" + } else if strings.HasSuffix(c.Request.URL.Path, "embeddings") { + modelRequest.Model = c.Param("model") + } else if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") { + modelRequest.Model = "dall-e" + } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") { + modelRequest.Model = "tts-1" + } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") { + modelRequest.Model = "whisper-1" + } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") { + modelRequest.Model = "whisper-1" + } else if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") { + modelRequest.Model = c.Query("model") + } + } else { + // 请求体不为空，尝试解析 JSON + err = json.Unmarshal(body, &modelRequest) + if err != nil { + return nil, false, fmt.Errorf("无效的请求, JSON 解析失败: %s", err.Error()) + } + } } + if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") { //wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01 modelRequest.Model = c.Query("model") @@ -212,6 +247,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode c.Set("channel_name", channel.Name) c.Set("channel_type", channel.Type) c.Set("channel_setting", channel.GetSetting()) + c.Set("channel_tag", channel.GetTag()) if nil != channel.OpenAIOrganization && "" != *channel.OpenAIOrganization { c.Set("channel_organization", *channel.OpenAIOrganization) } @@ -220,6 +256,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode c.Set("status_code_mapping", channel.GetStatusCodeMapping()) c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key)) c.Set("base_url", channel.GetBaseURL()) + c.Set("endpoint", channel.GetEndpoint()) // TODO: api_version统一 switch channel.Type { case common.ChannelTypeAzure: diff --git a/middleware/mock.go b/middleware/mock.go new file mode 100644 index 000000000000..df3edfee682f --- /dev/null +++ b/middleware/mock.go @@ -0,0 +1,59 @@ +package middleware + +import ( + "net/http" + + "github.com/gin-gonic/gin" +) + +const ( + TestTrafficHeader = "X-Test-Traffic" +) + +func MockResponse() gin.HandlerFunc { + return func(c *gin.Context) { + // 检查是否是测试流量 + if c.GetHeader(TestTrafficHeader) == "true" { + // 构造mock响应数据 + mockResponse := map[string]interface{}{ + "id": "chatcmpl-7f757876e4a24f75a4b0025b4d8a0e62", + "model": "gemini-2.5-pro-preview-05-06", + "object": "chat.completion", + "created": 1748702089, + "choices": []map[string]interface{}{ + { + "index": 0, + "message": map[string]interface{}{ + "role": "assistant", + "content": "测试结果是1 + 1 = 2", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]interface{}{ + "prompt_tokens": 6, + "completion_tokens": 7, + "total_tokens": 154, + "prompt_tokens_details": map[string]interface{}{ + "cached_tokens": 0, + "text_tokens": 0, + "audio_tokens": 0, + "image_tokens": 0, + }, + "completion_tokens_details": map[string]interface{}{ + "text_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 141, + }, + }, + } + + // 设置响应头 + c.Header("Content-Type", "application/json") + c.JSON(http.StatusOK, mockResponse) + c.Abort() + return + } + c.Next() + } +} diff --git a/middleware/request-id.go b/middleware/request-id.go index e623be7a269d..20fd086fe3de 100644 --- a/middleware/request-id.go +++ b/middleware/request-id.go @@ -1,18 +1,42 @@ package middleware import ( - "context" - "github.com/gin-gonic/gin" + "crypto/md5" + "crypto/rand" + "encoding/hex" "one-api/common" + "strconv" + "time" + + "github.com/gin-gonic/gin" ) func RequestId() func(c *gin.Context) { return func(c *gin.Context) { - id := common.GetTimeString() + common.GetRandomString(8) + id := c.GetHeader(common.RequestIdKey) + if id == "" { + // 使用更安全的request ID生成方法 + id = GenerateUniqueRequestId() + c.Header(common.RequestIdKey, id) + } c.Set(common.RequestIdKey, id) - ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id) - c.Request = c.Request.WithContext(ctx) - c.Header(common.RequestIdKey, id) c.Next() } } + +// GenerateUniqueRequestId 生成唯一的request ID +func GenerateUniqueRequestId() string { + // 获取当前时间戳（纳秒精度） + timestamp := time.Now().UnixNano() + + // 生成16字节的随机数 + randomBytes := make([]byte, 16) + rand.Read(randomBytes) + + // 将时间戳和随机数组合 + combined := strconv.FormatInt(timestamp, 10) + hex.EncodeToString(randomBytes) + + // 使用MD5生成最终的request ID（32位十六进制字符串） + hash := md5.Sum([]byte(combined)) + return hex.EncodeToString(hash[:]) +} diff --git a/middleware/request-logger.go b/middleware/request-logger.go new file mode 100644 index 000000000000..542250f22b37 --- /dev/null +++ b/middleware/request-logger.go @@ -0,0 +1,139 @@ +package middleware + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "one-api/common" + "strings" + + "github.com/gin-gonic/gin" +) + +// EnableRequestBodyLogging 控制是否打印请求体 +var EnableRequestBodyLogging bool = false + +func RequestLogger() gin.HandlerFunc { + return func(c *gin.Context) { + // 获取请求头 + headers := make(map[string]string) + for k, v := range c.Request.Header { + // 跳过敏感信息 + if strings.EqualFold(k, "Authorization") || strings.EqualFold(k, "Cookie") { + headers[k] = "***" + continue + } + headers[k] = strings.Join(v, ", ") + } + + // 获取请求参数 因为param并且后面request会打印所以不在此处打印 + // var params interface{} + if c.Request.Method == "GET" { + // params = c.Request.URL.Query() + } else { + // 读取请求体 + body, err := io.ReadAll(c.Request.Body) + if err == nil { + // 尝试解析为JSON + var jsonBody interface{} + if err := json.Unmarshal(body, &jsonBody); err == nil { + // params = jsonBody + } else { + // params = string(body) + } + // 恢复请求体 + c.Request.Body = io.NopCloser(bytes.NewBuffer(body)) + } + } + + // 构建日志信息 + logInfo := fmt.Sprintf("Request: %s %s\tClient IP: %s\tHeaders: %s\t", + c.Request.Method, + c.Request.URL.Path, + c.ClientIP(), + formatMap(headers), + ) + + // 如果启用了请求体日志，则记录请求体 + if EnableRequestBodyLogging { + if c.Request.Method != "GET" { + body, err := io.ReadAll(c.Request.Body) + if err == nil { + // 尝试解析为JSON + var jsonBody interface{} + if err := json.Unmarshal(body, &jsonBody); err == nil { + logInfo += fmt.Sprintf("\tBody: %s", formatValue(jsonBody)) + } else { + logInfo += fmt.Sprintf("\tBody: %s", string(body)) + } + // 恢复请求体 + c.Request.Body = io.NopCloser(bytes.NewBuffer(body)) + } + } + } + + common.SysLog(logInfo) + c.Next() + } +} + +func formatMap(m map[string]string) string { + if len(m) == 0 { + return "{}" + } + var pairs []string + for k, v := range m { + pairs = append(pairs, fmt.Sprintf("%s: %s", k, v)) + } + return "{" + strings.Join(pairs, ", ") + "}" +} + +func formatValue(v interface{}) string { + if v == nil { + return "null" + } + switch val := v.(type) { + case string: + return val + case map[string]interface{}: + return formatMapInterface(val) + case []interface{}: + return formatArray(val) + default: + bytes, err := json.Marshal(v) + if err != nil { + return fmt.Sprintf("%v", v) + } + // 去掉换行符 + return strings.ReplaceAll(string(bytes), "\n", "") + } +} + +func formatMapInterface(m map[string]interface{}) string { + if len(m) == 0 { + return "{}" + } + var pairs []string + for k, v := range m { + // 处理值中的换行符 + valueStr := formatValue(v) + valueStr = strings.ReplaceAll(valueStr, "\n", "") + pairs = append(pairs, fmt.Sprintf("%s: %s", k, valueStr)) + } + return "{" + strings.Join(pairs, ", ") + "}" +} + +func formatArray(arr []interface{}) string { + if len(arr) == 0 { + return "[]" + } + var elements []string + for _, v := range arr { + // 处理值中的换行符 + valueStr := formatValue(v) + valueStr = strings.ReplaceAll(valueStr, "\n", "") + elements = append(elements, valueStr) + } + return "[" + strings.Join(elements, ", ") + "]" +} diff --git a/middleware/utils.go b/middleware/utils.go index 082f565718a5..f1f46346bc18 100644 --- a/middleware/utils.go +++ b/middleware/utils.go @@ -1,29 +1,58 @@ package middleware import ( + "bytes" + "encoding/json" "fmt" - "github.com/gin-gonic/gin" + "io" "one-api/common" + "strings" + + "github.com/gin-gonic/gin" ) func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string) { userId := c.GetInt("id") - c.JSON(statusCode, gin.H{ + + // 获取请求体内容 + var requestBody []byte + if c.Request.Body != nil { + requestBody, _ = io.ReadAll(c.Request.Body) + // 恢复请求体，以便后续处理 + c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) + } + + // 准备错误响应 + errorResponse := gin.H{ "error": gin.H{ "message": common.MessageWithRequestId(message, c.GetString(common.RequestIdKey)), "type": "new_api_error", }, - }) - c.Abort() - common.LogError(c.Request.Context(), fmt.Sprintf("user %d | %s", userId, message)) -} + } + + // 将错误响应转换为JSON字符串，确保中文正确显示 + var responseBuffer bytes.Buffer + encoder := json.NewEncoder(&responseBuffer) + encoder.SetEscapeHTML(false) + encoder.SetIndent("", "") + encoder.Encode(errorResponse) + responseStr := strings.TrimSpace(responseBuffer.String()) + + // 将请求体转换为紧凑的JSON格式 + if len(requestBody) > 0 { + var jsonObj interface{} + if err := json.Unmarshal(requestBody, &jsonObj); err == nil { + if prettyJSON, err := json.Marshal(jsonObj); err == nil { + requestStr := strings.ReplaceAll(string(prettyJSON), "\n", "") + common.LogError(c.Request.Context(), fmt.Sprintf("user %d | %s | request body: %s | response body: %s", + userId, + message, + requestStr, + responseStr)) + } + } + } -func abortWithMidjourneyMessage(c *gin.Context, statusCode int, code int, description string) { - c.JSON(statusCode, gin.H{ - "description": description, - "type": "new_api_error", - "code": code, - }) + c.JSON(statusCode, errorResponse) c.Abort() - common.LogError(c.Request.Context(), description) } diff --git a/model/central_control.go b/model/central_control.go new file mode 100644 index 000000000000..f3efef24b3f1 --- /dev/null +++ b/model/central_control.go @@ -0,0 +1,36 @@ +package model + +import ( + "time" +) + +// UserRateLimitConfig 用户限速配置表 +type UserRateLimitConfig struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement;comment:主键ID"` + SiteName string `json:"site_name" gorm:"type:varchar(100);not null;comment:站点名"` + Username string `json:"username" gorm:"type:varchar(100);not null;comment:用户名"` + UserId int64 `json:"user_id" gorm:"not null;comment:用户ID"` + GroupName string `json:"group_name" gorm:"type:varchar(100);not null;comment:分组名"` + GroupId int64 `json:"group_id" gorm:"not null;comment:分组ID"` + ModelName string `json:"model_name" gorm:"type:varchar(100);not null;comment:模型名"` + SuggestedRateLimit int `json:"suggested_rate_limit" gorm:"not null;default:60;comment:建议限速大小(rpm)"` + IsRateLimitEnabled bool `json:"is_rate_limit_enabled" gorm:"not null;default:false;comment:是否启用限速(1:启用, 0:禁用)"` + CurrentRateLimit int `json:"current_rate_limit" gorm:"not null;default:60;comment:当前限速(rpm)"` + CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime;comment:创建时间"` + UpdatedAt time.Time `json:"updated_at" gorm:"autoUpdateTime;comment:更新时间"` +} + +// TableName 指定表名 +func (UserRateLimitConfig) TableName() string { + return "user_rate_limit_config" +} + +// GetUserRateLimitConfig 获取用户限速配置 +func GetUserRateLimitConfig(username, groupName, modelName string) (*UserRateLimitConfig, error) { + var config UserRateLimitConfig + err := CENTRAL_DB.Where("site_name = ? AND username = ? AND group_name = ? AND model_name = ?", "newapi-prod-center", username, groupName, modelName).First(&config).Error + if err != nil { + return nil, err + } + return &config, nil +} diff --git a/model/channel.go b/model/channel.go index 6ff0901d9c30..7b1d7940501b 100644 --- a/model/channel.go +++ b/model/channel.go @@ -22,6 +22,7 @@ type Channel struct { TestTime int64 `json:"test_time" gorm:"bigint"` ResponseTime int `json:"response_time"` // in milliseconds BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"` + Endpoint *string `json:"endpoint" gorm:"column:endpoint;default:''"` Other string `json:"other"` Balance float64 `json:"balance"` // in USD BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"` @@ -162,7 +163,7 @@ func GetChannelById(id int, selectAll bool) (*Channel, error) { if selectAll { err = DB.First(&channel, "id = ?", id).Error } else { - err = DB.Omit("key").First(&channel, "id = ?", id).Error + err = DB.First(&channel, "id = ?", id).Error } return &channel, err } @@ -223,6 +224,13 @@ func (channel *Channel) GetBaseURL() string { return *channel.BaseURL } +func (channel *Channel) GetEndpoint() string { + if channel.Endpoint == nil { + return "" + } + return *channel.Endpoint +} + func (channel *Channel) GetModelMapping() string { if channel.ModelMapping == nil { return "" diff --git a/model/group.go b/model/group.go new file mode 100644 index 000000000000..e7589f7ec99e --- /dev/null +++ b/model/group.go @@ -0,0 +1,7 @@ +package model + +type Group struct { + Id int `json:"id" gorm:"primaryKey"` + Name string `json:"name" gorm:"type:varchar(255);uniqueIndex;not null;default:''"` + Ratio int `json:"ratio" gorm:"type:int;not null;default:0"` +} diff --git a/model/log.go b/model/log.go index 86850a55a44c..ad42b78de238 100644 --- a/model/log.go +++ b/model/log.go @@ -4,17 +4,21 @@ import ( "fmt" "one-api/common" "os" + "sort" "strings" "time" "github.com/gin-gonic/gin" + "sync" + "sync/atomic" + "github.com/bytedance/gopkg/util/gopool" - "gorm.io/gorm" ) type Log struct { Id int `json:"id" gorm:"index:idx_created_at_id,priority:1"` + RequestID string `json:"request_id" gorm:"request_id"` UserId int `json:"user_id" gorm:"index"` CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_id,priority:2;index:idx_created_at_type"` Type int `json:"type" gorm:"index:idx_created_at_type"` @@ -25,6 +29,7 @@ type Log struct { Quota int `json:"quota" gorm:"default:0"` PromptTokens int `json:"prompt_tokens" gorm:"default:0"` CompletionTokens int `json:"completion_tokens" gorm:"default:0"` + ThinkingTokens int `json:"thinking_tokens" gorm:"default:0"` UseTime int `json:"use_time" gorm:"default:0"` IsStream bool `json:"is_stream" gorm:"default:false"` ChannelId int `json:"channel" gorm:"index"` @@ -78,7 +83,7 @@ func RecordLog(userId int, logType int, content string) { log := &Log{ UserId: userId, Username: username, - CreatedAt: common.GetTimestamp(), + CreatedAt: common.GetBeijingTimestamp(), Type: logType, Content: content, } @@ -88,9 +93,60 @@ func RecordLog(userId int, logType int, content string) { } } -func RecordConsumeLog(c *gin.Context, userId int, channelId int, promptTokens int, completionTokens int, +// 添加新的全局变量 +var ( + currentLogTable atomic.Value + tableCreateLock sync.Mutex + nextDayTimestamp atomic.Int64 +) + +// 添加新的函数用于获取日志表名 +func GetLogTableName(timestamp int64) string { + // 获取下一天的时间戳 + next := nextDayTimestamp.Load() + if timestamp >= next { + tableCreateLock.Lock() + defer tableCreateLock.Unlock() + + // 双重检查 + if timestamp >= nextDayTimestamp.Load() { + // 计算新的表名 + t := common.GetBeijingTimeFromTimestamp(timestamp) + tableName := fmt.Sprintf("logs_%04d_%02d_%02d", t.Year(), t.Month(), t.Day()) + + // 创建新表 + newTableSQL := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s LIKE logs`, tableName) + if err := LOG_DB.Exec(newTableSQL).Error; err != nil { + common.SysError("failed to create new log table: " + err.Error()) + return "logs" + } + + // 更新下一天的时间戳 + nextDay := time.Date(t.Year(), t.Month(), t.Day()+1, 0, 0, 0, 0, common.BeijingLocation) + nextDayTimestamp.Store(nextDay.Unix()) + + // 存储当前表名 + currentLogTable.Store(tableName) + return tableName + } + } + + if current, ok := currentLogTable.Load().(string); ok && current != "" { + return current + } + return "logs" +} + +// 修改 RecordConsumeLog 函数中的相关部分 +func RecordConsumeLog(c *gin.Context, userId int, channelId int, promptTokens int, completionTokens int, thinkingTokens int, modelName string, tokenName string, quota int, content string, tokenId int, userQuota int, useTimeSeconds int, isStream bool, group string, other map[string]interface{}) { + // 如果是压测流量，不记录计费日志 + if c.GetHeader("X-Test-Traffic") == "true" { + common.LogInfo(c, "test traffic detected, skipping consume log") + return + } + common.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, 用户调用前余额=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, userQuota, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content)) if !common.LogConsumeEnabled { return @@ -98,72 +154,110 @@ func RecordConsumeLog(c *gin.Context, userId int, channelId int, promptTokens in username := c.GetString("username") otherStr := common.MapToJsonStr(other) log := &Log{ - UserId: userId, + UserId: common.GetOriginUserId(c, userId), + RequestID: c.GetString(common.RequestIdKey), Username: username, - CreatedAt: common.GetTimestamp(), + CreatedAt: common.GetBeijingTimestamp(), Type: LogTypeConsume, Content: content, PromptTokens: promptTokens, CompletionTokens: completionTokens, + ThinkingTokens: thinkingTokens, TokenName: tokenName, ModelName: modelName, Quota: quota, - ChannelId: channelId, + ChannelId: common.GetOriginChannelId(c, channelId), TokenId: tokenId, UseTime: useTimeSeconds, IsStream: isStream, Group: group, Other: otherStr, } - err := LOG_DB.Create(log).Error + tableName := GetLogTableName(log.CreatedAt) + if time.Now().In(common.BeijingLocation).Before(time.Date(2025, 3, 12, 23, 59, 59, 0, common.BeijingLocation)) { + tableName = "logs" + } + err := LOG_DB.Table(tableName).Create(log).Error if err != nil { common.LogError(c, "failed to record log: "+err.Error()) } if common.DataExportEnabled { gopool.Go(func() { - LogQuotaData(userId, username, modelName, quota, common.GetTimestamp(), promptTokens+completionTokens) + LogQuotaData(userId, tokenName, username, modelName, quota, common.GetBeijingTimestamp(), promptTokens+completionTokens) }) } } func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string) (logs []*Log, total int64, err error) { - var tx *gorm.DB - if logType == LogTypeUnknown { - tx = LOG_DB - } else { - tx = LOG_DB.Where("logs.type = ?", logType) + // 获取需要查询的所有表名 + tableNames := getTableNamesByTimeRange(startTimestamp, endTimestamp) + if len(tableNames) == 0 { + return nil, 0, nil } - if modelName != "" { - tx = tx.Where("logs.model_name like ?", modelName) - } - if username != "" { - tx = tx.Where("logs.username = ?", username) - } - if tokenName != "" { - tx = tx.Where("logs.token_name = ?", tokenName) - } - if startTimestamp != 0 { - tx = tx.Where("logs.created_at >= ?", startTimestamp) - } - if endTimestamp != 0 { - tx = tx.Where("logs.created_at <= ?", endTimestamp) - } - if channel != 0 { - tx = tx.Where("logs.channel_id = ?", channel) - } - if group != "" { - tx = tx.Where("logs."+groupCol+" = ?", group) + // 用于存储所有查询结果 + allLogs := make([]*Log, 0) + total = 0 + + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempTotal int64 + var tempLogs []*Log + var tx = LOG_DB.Table(tableName) + + if logType != LogTypeUnknown { + tx = tx.Where("type = ?", logType) + } + if modelName != "" { + tx = tx.Where("model_name like ?", modelName) + } + if username != "" { + tx = tx.Where("username = ?", username) + } + if tokenName != "" { + tx = tx.Where("token_name = ?", tokenName) + } + if startTimestamp != 0 { + tx = tx.Where("created_at >= ?", startTimestamp) + } + if endTimestamp != 0 { + tx = tx.Where("created_at <= ?", endTimestamp) + } + if channel != 0 { + tx = tx.Where("channel_id = ?", channel) + } + if group != "" { + tx = tx.Where(groupCol+" = ?", group) + } + + // 获取当前表的总数 + if err = tx.Count(&tempTotal).Error; err != nil { + return nil, 0, err + } + total += tempTotal + + // 获取当前表的数据 + if err = tx.Order("id desc").Find(&tempLogs).Error; err != nil { + return nil, 0, err + } + allLogs = append(allLogs, tempLogs...) } - err = tx.Model(&Log{}).Count(&total).Error - if err != nil { - return nil, 0, err + + // 对所有结果按时间倒序排序 + sort.Slice(allLogs, func(i, j int) bool { + return allLogs[i].CreatedAt > allLogs[j].CreatedAt + }) + + // 处理分页 + end := startIdx + num + if end > len(allLogs) { + end = len(allLogs) } - err = tx.Order("logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error - if err != nil { - return nil, 0, err + if startIdx < len(allLogs) { + logs = allLogs[startIdx:end] } + // 处理渠道信息 channelIds := make([]int, 0) channelMap := make(map[int]string) for _, log := range logs { @@ -187,54 +281,154 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName } } - return logs, total, err + return logs, total, nil } func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string) (logs []*Log, total int64, err error) { - var tx *gorm.DB - if logType == LogTypeUnknown { - tx = LOG_DB.Where("logs.user_id = ?", userId) - } else { - tx = LOG_DB.Where("logs.user_id = ? and logs.type = ?", userId, logType) + // 获取需要查询的所有表名 + tableNames := getTableNamesByTimeRange(startTimestamp, endTimestamp) + if len(tableNames) == 0 { + return nil, 0, nil } - if modelName != "" { - tx = tx.Where("logs.model_name like ?", modelName) - } - if tokenName != "" { - tx = tx.Where("logs.token_name = ?", tokenName) - } - if startTimestamp != 0 { - tx = tx.Where("logs.created_at >= ?", startTimestamp) - } - if endTimestamp != 0 { - tx = tx.Where("logs.created_at <= ?", endTimestamp) - } - if group != "" { - tx = tx.Where("logs."+groupCol+" = ?", group) + // 用于存储所有查询结果 + allLogs := make([]*Log, 0) + total = 0 + + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempTotal int64 + var tempLogs []*Log + var tx = LOG_DB.Table(tableName) + + if logType == LogTypeUnknown { + tx = tx.Where("user_id = ?", userId) + } else { + tx = tx.Where("user_id = ? and type = ?", userId, logType) + } + + if modelName != "" { + tx = tx.Where("model_name like ?", modelName) + } + if tokenName != "" { + tx = tx.Where("token_name = ?", tokenName) + } + if startTimestamp != 0 { + tx = tx.Where("created_at >= ?", startTimestamp) + } + if endTimestamp != 0 { + tx = tx.Where("created_at <= ?", endTimestamp) + } + if group != "" { + tx = tx.Where(groupCol+" = ?", group) + } + + // 获取当前表的总数 + if err = tx.Count(&tempTotal).Error; err != nil { + return nil, 0, err + } + total += tempTotal + + // 获取当前表的数据 + if err = tx.Order("id desc").Find(&tempLogs).Error; err != nil { + return nil, 0, err + } + allLogs = append(allLogs, tempLogs...) } - err = tx.Model(&Log{}).Count(&total).Error - if err != nil { - return nil, 0, err + + // 对所有结果按时间倒序排序 + sort.Slice(allLogs, func(i, j int) bool { + return allLogs[i].CreatedAt > allLogs[j].CreatedAt + }) + + // 处理分页 + end := startIdx + num + if end > len(allLogs) { + end = len(allLogs) } - err = tx.Order("logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error - if err != nil { - return nil, 0, err + if startIdx < len(allLogs) { + logs = allLogs[startIdx:end] } formatUserLogs(logs) - return logs, total, err + return logs, total, nil } func SearchAllLogs(keyword string) (logs []*Log, err error) { - err = LOG_DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error - return logs, err + // 获取当前时间 + now := time.Now() + // 获取一个月前的时间戳 + oneMonthAgo := now.AddDate(0, -1, 0) + + // 获取时间范围内的所有表名 + tableNames := getTableNamesByTimeRange(oneMonthAgo.Unix(), now.Unix()) + + // 用于存储所有查询结果 + allLogs := make([]*Log, 0) + + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempLogs []*Log + err = LOG_DB.Table(tableName). + Where("type = ? or content LIKE ?", keyword, keyword+"%"). + Order("id desc"). + Find(&tempLogs).Error + if err != nil { + return nil, err + } + allLogs = append(allLogs, tempLogs...) + } + + // 对所有结果按时间倒序排序 + sort.Slice(allLogs, func(i, j int) bool { + return allLogs[i].CreatedAt > allLogs[j].CreatedAt + }) + + // 只返回最近的 MaxRecentItems 条记录 + if len(allLogs) > common.MaxRecentItems { + allLogs = allLogs[:common.MaxRecentItems] + } + + return allLogs, nil } func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) { - err = LOG_DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error - formatUserLogs(logs) - return logs, err + // 获取当前时间 + now := time.Now() + // 获取一个月前的时间戳 + oneMonthAgo := now.AddDate(0, -1, 0) + + // 获取时间范围内的所有表名 + tableNames := getTableNamesByTimeRange(oneMonthAgo.Unix(), now.Unix()) + + // 用于存储所有查询结果 + allLogs := make([]*Log, 0) + + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempLogs []*Log + err = LOG_DB.Table(tableName). + Where("user_id = ? and type = ?", userId, keyword). + Order("id desc"). + Find(&tempLogs).Error + if err != nil { + return nil, err + } + allLogs = append(allLogs, tempLogs...) + } + + // 对所有结果按时间倒序排序 + sort.Slice(allLogs, func(i, j int) bool { + return allLogs[i].CreatedAt > allLogs[j].CreatedAt + }) + + // 只返回最近的 MaxRecentItems 条记录 + if len(allLogs) > common.MaxRecentItems { + allLogs = allLogs[:common.MaxRecentItems] + } + + formatUserLogs(allLogs) + return allLogs, nil } type Stat struct { @@ -243,74 +437,224 @@ type Stat struct { Tpm int `json:"tpm"` } +// 添加一个辅助函数用于获取时间范围内的所有表名 +func getTableNamesByTimeRange(startTimestamp, endTimestamp int64) []string { + if startTimestamp == 0 || endTimestamp == 0 { + return []string{"logs"} + } + + tables := make([]string, 0) + start := time.Unix(startTimestamp, 0) + end := time.Unix(endTimestamp, 0) + + // 如果在同一天，直接返回一个表名 + if start.Year() == end.Year() && start.Month() == end.Month() && start.Day() == end.Day() { + tableName := fmt.Sprintf("logs_%04d_%02d_%02d", start.Year(), start.Month(), start.Day()) + return []string{tableName} + } + + // 遍历日期范围内的每一天 + for d := start; !d.After(end); d = d.AddDate(0, 0, 1) { + tableName := fmt.Sprintf("logs_%04d_%02d_%02d", d.Year(), d.Month(), d.Day()) + tables = append(tables, tableName) + } + + return tables +} + +// 修改 SumUsedQuota 函数 func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat) { - tx := LOG_DB.Table("logs").Select("sum(quota) quota") + // 获取需要查询的所有表名 + tableNames := getTableNamesByTimeRange(startTimestamp, endTimestamp) + if len(tableNames) == 0 { + return stat + } + // 用于存储聚合结果 + var totalQuota int64 + // var totalRpm int64 + // var totalTpm int64 - // 为rpm和tpm创建单独的查询 - rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm") + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempStat Stat + + // 配额查询 + quotaQuery := LOG_DB.Table(tableName).Select("IFNULL(sum(quota), 0) as quota") + if username != "" { + quotaQuery = quotaQuery.Where("username = ?", username) + } + if tokenName != "" { + quotaQuery = quotaQuery.Where("token_name = ?", tokenName) + } + if startTimestamp != 0 { + quotaQuery = quotaQuery.Where("created_at >= ?", startTimestamp) + } + if endTimestamp != 0 { + quotaQuery = quotaQuery.Where("created_at <= ?", endTimestamp) + } + if modelName != "" { + quotaQuery = quotaQuery.Where("model_name like ?", modelName) + } + if channel != 0 { + quotaQuery = quotaQuery.Where("channel_id = ?", channel) + } + if group != "" { + quotaQuery = quotaQuery.Where(groupCol+" = ?", group) + } + quotaQuery = quotaQuery.Where("type = ?", LogTypeConsume) + quotaQuery.Scan(&tempStat) + + totalQuota += int64(tempStat.Quota) + } + + // RPM和TPM只需要查询最近的表 + rpmTpmQuery := LOG_DB.Table(tableNames[len(tableNames)-1]). + Select("count(*) rpm, IFNULL(sum(prompt_tokens), 0) + IFNULL(sum(completion_tokens), 0) tpm") if username != "" { - tx = tx.Where("username = ?", username) rpmTpmQuery = rpmTpmQuery.Where("username = ?", username) } if tokenName != "" { - tx = tx.Where("token_name = ?", tokenName) rpmTpmQuery = rpmTpmQuery.Where("token_name = ?", tokenName) } - if startTimestamp != 0 { - tx = tx.Where("created_at >= ?", startTimestamp) - } - if endTimestamp != 0 { - tx = tx.Where("created_at <= ?", endTimestamp) - } if modelName != "" { - tx = tx.Where("model_name like ?", modelName) rpmTpmQuery = rpmTpmQuery.Where("model_name like ?", modelName) } if channel != 0 { - tx = tx.Where("channel_id = ?", channel) rpmTpmQuery = rpmTpmQuery.Where("channel_id = ?", channel) } if group != "" { - tx = tx.Where(groupCol+" = ?", group) rpmTpmQuery = rpmTpmQuery.Where(groupCol+" = ?", group) } - tx = tx.Where("type = ?", LogTypeConsume) - rpmTpmQuery = rpmTpmQuery.Where("type = ?", LogTypeConsume) + rpmTpmQuery = rpmTpmQuery.Where("type = ?", LogTypeConsume). + Where("created_at >= ?", time.Now().Add(-60*time.Second).Unix()) - // 只统计最近60秒的rpm和tpm - rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", time.Now().Add(-60*time.Second).Unix()) + var tempStat Stat + rpmTpmQuery.Scan(&tempStat) - // 执行查询 - tx.Scan(&stat) - rpmTpmQuery.Scan(&stat) + // 合并结果 + stat.Quota = int(totalQuota) + stat.Rpm = tempStat.Rpm + stat.Tpm = tempStat.Tpm return stat } func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) { - tx := LOG_DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)") - if username != "" { - tx = tx.Where("username = ?", username) + // 获取需要查询的所有表名 + tableNames := getTableNamesByTimeRange(startTimestamp, endTimestamp) + if len(tableNames) == 0 { + return 0 } - if tokenName != "" { - tx = tx.Where("token_name = ?", tokenName) - } - if startTimestamp != 0 { - tx = tx.Where("created_at >= ?", startTimestamp) - } - if endTimestamp != 0 { - tx = tx.Where("created_at <= ?", endTimestamp) - } - if modelName != "" { - tx = tx.Where("model_name = ?", modelName) + + var totalTokens int64 + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempToken int64 + tx := LOG_DB.Table(tableName). + Select("IFNULL(sum(prompt_tokens), 0) + IFNULL(sum(completion_tokens), 0)") + + if username != "" { + tx = tx.Where("username = ?", username) + } + if tokenName != "" { + tx = tx.Where("token_name = ?", tokenName) + } + if startTimestamp != 0 { + tx = tx.Where("created_at >= ?", startTimestamp) + } + if endTimestamp != 0 { + tx = tx.Where("created_at <= ?", endTimestamp) + } + if modelName != "" { + tx = tx.Where("model_name = ?", modelName) + } + tx.Where("type = ?", LogTypeConsume).Scan(&tempToken) + + totalTokens += tempToken } - tx.Where("type = ?", LogTypeConsume).Scan(&token) - return token + + return int(totalTokens) } func DeleteOldLog(targetTimestamp int64) (int64, error) { result := LOG_DB.Where("created_at < ?", targetTimestamp).Delete(&Log{}) return result.RowsAffected, result.Error } + +// SELECT +// logs.channel_id, +// COALESCE(channels.name, logs.channel_name, '未知渠道') AS channel_name, +// logs.model_name, +// SUM(logs.prompt_tokens) AS total_prompt_tokens, +// SUM(logs.completion_tokens) AS total_completion_tokens +// FROM logs +// LEFT JOIN channels ON logs.channel_id = channels.id +// WHERE +// logs.created_at BETWEEN 1741338023 AND 1741341623 +// GROUP BY +// logs.channel_id, -- 渠道ID作为主分组键 +// channel_name, -- 直接使用SELECT中的别名（COALESCE表达式结果） +// logs.model_name -- 模型名称 +// ORDER BY +// logs.channel_id; + +func GetAllChannelBilling(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) { + // 获取需要查询的所有表名 + tableNames := getTableNamesByTimeRange(startTimestamp, endTimestamp) + if len(tableNames) == 0 { + return 0 + } + + var totalTokens int + // 遍历每个表进行查询 + for _, tableName := range tableNames { + var tempToken int + tx := LOG_DB.Table(tableName).Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)") + if username != "" { + tx = tx.Where("username = ?", username) + } + if tokenName != "" { + tx = tx.Where("token_name = ?", tokenName) + } + if startTimestamp != 0 { + tx = tx.Where("created_at >= ?", startTimestamp) + } + if endTimestamp != 0 { + tx = tx.Where("created_at <= ?", endTimestamp) + } + if modelName != "" { + tx = tx.Where("model_name = ?", modelName) + } + tx.Where("type = ?", LogTypeConsume).Scan(&tempToken) + totalTokens += tempToken + } + return totalTokens +} + +// 在 init 函数中初始化（添加新的 init 函数） +func init() { + // 设置初始的下一天时间戳 + now := common.GetBeijingTime() + nextDay := time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, common.BeijingLocation) + nextDayTimestamp.Store(nextDay.Unix()) + // 设置当前表名 + currentLogTable.Store(fmt.Sprintf("logs_%04d_%02d_%02d", now.Year(), now.Month(), now.Day())) +} + +func InitLogTable() error { + now := common.GetBeijingTime() + nextDay := time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, common.BeijingLocation) + nextDayTimestamp.Store(nextDay.Unix()) + + tableName := fmt.Sprintf("logs_%04d_%02d_%02d", now.Year(), now.Month(), now.Day()) + currentLogTable.Store(tableName) + + newTableSQL := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s LIKE logs`, tableName) + if err := LOG_DB.Exec(newTableSQL).Error; err != nil { + common.SysError("failed to create new log table: " + err.Error()) + return err + } + return nil +} diff --git a/model/main.go b/model/main.go index c0bf927c4dbc..cc4093fa771f 100644 --- a/model/main.go +++ b/model/main.go @@ -1,16 +1,17 @@ package model import ( - "github.com/glebarez/sqlite" - "gorm.io/driver/mysql" - "gorm.io/driver/postgres" - "gorm.io/gorm" "log" "one-api/common" "os" "strings" "sync" "time" + + "github.com/glebarez/sqlite" + "gorm.io/driver/mysql" + "gorm.io/driver/postgres" + "gorm.io/gorm" ) var groupCol string @@ -31,6 +32,9 @@ var DB *gorm.DB var LOG_DB *gorm.DB +// 中心控制库数据库连接 +var CENTRAL_DB *gorm.DB + func createRootAccountIfNeed() error { var user User //if user.Status != common.UserStatusEnabled { @@ -168,6 +172,36 @@ func InitLogDB() (err error) { return err } +// InitCentralDB 初始化中心控制库 +func InitCentralDB() (err error) { + if os.Getenv("CENTRAL_SQL_DSN") == "" { + // 如果没有配置中心控制库，使用主数据库 + CENTRAL_DB = DB + common.SysLog("CENTRAL_SQL_DSN not set, using main database for central control") + return nil + } + db, err := chooseDB("CENTRAL_SQL_DSN") + if err == nil { + if common.DebugEnabled { + db = db.Debug() + } + CENTRAL_DB = db + sqlDB, err := CENTRAL_DB.DB() + if err != nil { + return err + } + sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100)) + sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000)) + sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60))) + + common.SysLog("central control database connected") + return nil + } else { + common.FatalLog(err) + } + return err +} + func migrateDB() error { err := DB.AutoMigrate(&Channel{}) if err != nil { @@ -213,6 +247,10 @@ func migrateDB() error { if err != nil { return err } + err = DB.AutoMigrate(&Group{}) + if err != nil { + return err + } common.SysLog("database migrated") err = createRootAccountIfNeed() return err @@ -226,6 +264,17 @@ func migrateLOGDB() error { return nil } +func migrateCentralDB() error { + // 迁移用户限速配置表 + err := CENTRAL_DB.AutoMigrate(&UserRateLimitConfig{}) + if err != nil { + return err + } + + common.SysLog("central control database migrated") + return nil +} + func closeDB(db *gorm.DB) error { sqlDB, err := db.DB() if err != nil { @@ -242,6 +291,12 @@ func CloseDB() error { return err } } + if CENTRAL_DB != DB && CENTRAL_DB != LOG_DB { + err := closeDB(CENTRAL_DB) + if err != nil { + return err + } + } return closeDB(DB) } diff --git a/model/option.go b/model/option.go index fe12eab18519..3e2ce107224e 100644 --- a/model/option.go +++ b/model/option.go @@ -1,6 +1,7 @@ package model import ( + "encoding/json" "one-api/common" "one-api/setting" "one-api/setting/config" @@ -22,6 +23,19 @@ func AllOption() ([]*Option, error) { return options, err } +func InitGroups() { + common.GroupRWMutex.Lock() + defer common.GroupRWMutex.Unlock() + var groups []*Group + err := DB.Find(&groups).Error + if err != nil { + panic(err) + } + for _, group := range groups { + common.Groups[group.Name] = group.Id + } +} + func InitOptionMap() { common.OptionMapRWMutex.Lock() common.OptionMap = make(map[string]string) @@ -161,6 +175,31 @@ func UpdateOption(key string, value string) error { // If save value does not contain primary key, it will execute Create, // otherwise it will execute Update (with all fields). DB.Save(&option) + if key == "GroupRatio" { + groups := make(map[string]int) + err := json.Unmarshal([]byte(option.Value), &groups) + if err != nil { + return err + } + + // 将 GroupRatio 数据同步到 group 表中 + for groupName, ratio := range groups { + // 查询是否存在该记录 + var count int64 + DB.Table("groups").Where("name = ?", groupName).Count(&count) + + if count > 0 { + // 存在记录，更新 ratio + DB.Table("groups").Where("name = ?", groupName).Update("ratio", ratio) + } else { + // 不存在记录，创建新记录 + DB.Table("groups").Create(map[string]interface{}{ + "name": groupName, + "ratio": ratio, + }) + } + } + } // Update OptionMap return updateOptionMap(key, value) } diff --git a/model/text_request.go b/model/text_request.go new file mode 100644 index 000000000000..86ed59687962 --- /dev/null +++ b/model/text_request.go @@ -0,0 +1,120 @@ +package model + +import ( + "fmt" + "one-api/common" + "os" + "sync/atomic" + "time" + + "gorm.io/gorm" +) + +var ( + // RequestPersistenceEnabled 是否启用请求持久化存储 + RequestPersistenceEnabled = false + // RequestPersistenceDB 请求持久化存储的数据库连接 + RequestPersistenceDB *gorm.DB +) + +var ( + textRequestTableName atomic.Value +) + +// TextRequest 记录文本请求的输入输出 +type TextRequest struct { + Id int64 `json:"id" gorm:"primaryKey;type:bigint;autoIncrement"` + UserId int `json:"user_id" gorm:"index;type:int"` + RequestId string `json:"request_id" gorm:"index;type:varchar(100)"` + CreatedAt time.Time `json:"created_at" gorm:"index;type:datetime"` // 格式：2024-03-21 14:30:45 + Model string `json:"model" gorm:"type:varchar(100);index"` + RequestHeaders string `json:"request_headers" gorm:"type:text"` // 通常不会太大 + RequestBody string `json:"request_body" gorm:"type:longtext"` // 可能包含大量数据 + ResponseHeaders string `json:"response_headers" gorm:"type:text"` // 通常不会太大 + ResponseBody string `json:"response_body" gorm:"type:longtext"` // 可能包含大量数据 +} + +// InitRequestPersistence 初始化请求持久化存储 +func InitRequestPersistence() error { + if !RequestPersistenceEnabled { + return nil + } + + // 使用 REQUEST_PERSISTENCE_ENABLED_SQL_DSN 初始化数据库连接 + dsn := os.Getenv("REQUEST_PERSISTENCE_ENABLED_SQL_DSN") + if dsn == "" { + // 如果没有指定专门的数据库连接，使用主数据库 + RequestPersistenceDB = DB + } else { + // 使用指定的数据库连接 + db, err := chooseDB("REQUEST_PERSISTENCE_ENABLED_SQL_DSN") + if err != nil { + return fmt.Errorf("failed to initialize request persistence database: %v", err) + } + if common.DebugEnabled { + db = db.Debug() + } + RequestPersistenceDB = db + sqlDB, err := RequestPersistenceDB.DB() + if err != nil { + return err + } + sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100)) + sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000)) + sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60))) + } + + // 创建未来一周的表 + return createTablesForNextWeek() +} + +// createTablesForNextWeek 创建未来一周的表 +func createTablesForNextWeek() error { + now := time.Now() + // 设置当前表名，使用正确的日期格式 + textRequestTableName.Store(fmt.Sprintf("text_requests_%s", now.Format("20060102"))) + + // 创建未来一周的表 + for i := 0; i < 7; i++ { + date := now.AddDate(0, 0, i) + // 使用正确的日期格式 YYYYMMDD + tableName := fmt.Sprintf("text_requests_%s", date.Format("20060102")) + + // 检查表是否存在 + var count int64 + err := RequestPersistenceDB.Raw("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = ?", tableName).Count(&count).Error + if err != nil { + return fmt.Errorf("failed to check if table %s exists: %v", tableName, err) + } + + // 如果表不存在，则创建 + if count == 0 { + if err := RequestPersistenceDB.Table(tableName).Set("gorm:table_options", "ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci").AutoMigrate(&TextRequest{}); err != nil { + return fmt.Errorf("failed to create table %s: %v", tableName, err) + } + common.SysLog(fmt.Sprintf("created table %s", tableName)) + } + } + return nil +} + +// StartTableCheckRoutine 启动定时检查表的协程 +func StartTableCheckRoutine() { + go func() { + for { + // 计算下一个12点的时间 + now := time.Now() + next := time.Date(now.Year(), now.Month(), now.Day(), 12, 0, 0, 0, now.Location()) + if now.After(next) { + next = next.AddDate(0, 0, 1) + } + // 等待到下一个12点 + time.Sleep(next.Sub(now)) + + // 创建未来一周的表 + if err := createTablesForNextWeek(); err != nil { + common.SysError("failed to create tables for next week: " + err.Error()) + } + } + }() +} diff --git a/model/token.go b/model/token.go index 8587ea62a9f3..7c9f930394b3 100644 --- a/model/token.go +++ b/model/token.go @@ -16,6 +16,7 @@ type Token struct { Key string `json:"key" gorm:"type:char(48);uniqueIndex"` Status int `json:"status" gorm:"default:1"` Name string `json:"name" gorm:"index" ` + User string `json:"user"` CreatedTime int64 `json:"created_time" gorm:"bigint"` AccessedTime int64 `json:"accessed_time" gorm:"bigint"` ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired @@ -62,6 +63,25 @@ func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { return tokens, err } +func GetAllTokens(startIdx int, num int) ([]*Token, error) { + var tokens []*Token + var err error + var Users []*User + err = DB.Where("1 = 1").Order("id desc").Find(&Users).Error + if err != nil { + return nil, err + } + id2name := make(map[int]string) + for _, user := range Users { + id2name[user.Id] = user.Username + } + err = DB.Where("user_id is not null").Order("id desc").Limit(num).Offset(startIdx).Find(&tokens).Error + for _, token := range tokens { + token.User = id2name[token.UserId] + } + return tokens, err +} + func SearchUserTokens(userId int, keyword string, token string) (tokens []*Token, err error) { if token != "" { token = strings.Trim(token, "sk-") @@ -120,7 +140,7 @@ func GetTokenByIds(id int, userId int) (*Token, error) { } token := Token{Id: id, UserId: userId} var err error = nil - err = DB.First(&token, "id = ? and user_id = ?", id, userId).Error + err = DB.First(&token, "id = ? ", id).Error return &token, err } @@ -253,7 +273,7 @@ func DeleteTokenById(id int, userId int) (err error) { if id == 0 || userId == 0 { return errors.New("id 或 userId 为空！") } - token := Token{Id: id, UserId: userId} + token := Token{Id: id} err = DB.Where(token).First(&token).Error if err != nil { return err diff --git a/model/usedata.go b/model/usedata.go index 1255b0beddab..664efafa3117 100644 --- a/model/usedata.go +++ b/model/usedata.go @@ -2,8 +2,11 @@ package model import ( "fmt" + "github.com/xuri/excelize/v2" "gorm.io/gorm" "one-api/common" + "one-api/setting/operation_setting" + "sort" "sync" "time" ) @@ -13,6 +16,7 @@ type QuotaData struct { Id int `json:"id"` UserID int `json:"user_id" gorm:"index"` Username string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"` + TokenName string `json:"token_name" gorm:"size:256;default:''"` ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"` CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"` TokenUsed int `json:"token_used" gorm:"default:0"` @@ -20,6 +24,30 @@ type QuotaData struct { Quota int `json:"quota" gorm:"default:0"` } +type BillingData struct { + ChannelId int `json:"chanel_id"` + ChannelName string `json:"channel_name"` + ChannelTag string `json:"channel_tag"` + Count int `json:"count"` + ModelName string `json:"model_name"` + PromptTokens int `json:"prompt_tokens"` + CompletionsTokens int `json:"completions_tokens"` +} + +type BillingJsonData struct { + ChannelId int `json:"chanel_id"` + CurrentDate string `json:"current_date"` + ChannelName string `json:"channel_name"` + ChannelTag string `json:"channel_tag"` + Count int `json:"count"` + ModelName string `json:"model_name"` + PromptTokens float32 `json:"prompt_tokens"` + CompletionsTokens float32 `json:"completions_tokens"` + PromptPricing float32 `json:"prompt_pricing"` + CompletionsPricing float32 `json:"completions_pricing"` + /**/ Cost float32 `json:"cost"` +} + func UpdateQuotaData() { // recover defer func() { @@ -39,8 +67,8 @@ func UpdateQuotaData() { var CacheQuotaData = make(map[string]*QuotaData) var CacheQuotaDataLock = sync.Mutex{} -func logQuotaDataCache(userId int, username string, modelName string, quota int, createdAt int64, tokenUsed int) { - key := fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt) +func logQuotaDataCache(userId int, tokenName, username string, modelName string, quota int, createdAt int64, tokenUsed int) { + key := fmt.Sprintf("%d-%s-%s-%s-%d", userId, username, tokenName, modelName, createdAt) quotaData, ok := CacheQuotaData[key] if ok { quotaData.Count += 1 @@ -50,6 +78,7 @@ func logQuotaDataCache(userId int, username string, modelName string, quota int, quotaData = &QuotaData{ UserID: userId, Username: username, + TokenName: tokenName, ModelName: modelName, CreatedAt: createdAt, Count: 1, @@ -60,13 +89,13 @@ func logQuotaDataCache(userId int, username string, modelName string, quota int, CacheQuotaData[key] = quotaData } -func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64, tokenUsed int) { +func LogQuotaData(userId int, tokenName, username string, modelName string, quota int, createdAt int64, tokenUsed int) { // 只精确到小时 createdAt = createdAt - (createdAt % 3600) CacheQuotaDataLock.Lock() defer CacheQuotaDataLock.Unlock() - logQuotaDataCache(userId, username, modelName, quota, createdAt, tokenUsed) + logQuotaDataCache(userId, tokenName, username, modelName, quota, createdAt, tokenUsed) } func SaveQuotaDataCache() { @@ -79,13 +108,13 @@ func SaveQuotaDataCache() { // 3. 如果没有数据，就插入数据 for _, quotaData := range CacheQuotaData { quotaDataDB := &QuotaData{} - DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?", - quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt).First(quotaDataDB) + DB.Table("quota_data").Where("user_id = ? and token_name = ? and username = ? and model_name = ? and created_at = ?", + quotaData.UserID, quotaData.TokenName, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt).First(quotaDataDB) if quotaDataDB.Id > 0 { //quotaDataDB.Count += quotaData.Count //quotaDataDB.Quota += quotaData.Quota //DB.Table("quota_data").Save(quotaDataDB) - increaseQuotaData(quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.Count, quotaData.Quota, quotaData.CreatedAt, quotaData.TokenUsed) + increaseQuotaData(quotaData.UserID, quotaData.TokenName, quotaData.Username, quotaData.ModelName, quotaData.Count, quotaData.Quota, quotaData.CreatedAt, quotaData.TokenUsed) } else { DB.Table("quota_data").Create(quotaData) } @@ -94,9 +123,9 @@ func SaveQuotaDataCache() { common.SysLog(fmt.Sprintf("保存数据看板数据成功，共保存%d条数据", size)) } -func increaseQuotaData(userId int, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) { - err := DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?", - userId, username, modelName, createdAt).Updates(map[string]interface{}{ +func increaseQuotaData(userId int, tokenname, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) { + err := DB.Table("quota_data").Where("user_id = ? and token_name = ? and username = ? and model_name = ? and created_at = ?", + userId, tokenname, username, modelName, createdAt).Updates(map[string]interface{}{ "count": gorm.Expr("count + ?", count), "quota": gorm.Expr("quota + ?", quota), "token_used": gorm.Expr("token_used + ?", tokenUsed), @@ -106,11 +135,17 @@ func increaseQuotaData(userId int, username string, modelName string, count int, } } -func GetQuotaDataByUsername(username string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) { +func GetQuotaDataByUsername(username, tokenName string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) { var quotaDatas []*QuotaData // 从quota_data表中查询数据 - err = DB.Table("quota_data").Where("username = ? and created_at >= ? and created_at <= ?", username, startTime, endTime).Find("aDatas).Error - return quotaDatas, err + if tokenName != "" { + err = DB.Table("quota_data").Where("username = ? and token_name = ? and created_at >= ? and created_at <= ?", username, tokenName, startTime, endTime).Find("aDatas).Error + return quotaDatas, err + } else { + err = DB.Table("quota_data").Where("username = ? and created_at >= ? and created_at <= ?", username, startTime, endTime).Find("aDatas).Error + return quotaDatas, err + } + } func GetQuotaDataByUserId(userId int, startTime int64, endTime int64) (quotaData []*QuotaData, err error) { @@ -120,9 +155,9 @@ func GetQuotaDataByUserId(userId int, startTime int64, endTime int64) (quotaData return quotaDatas, err } -func GetAllQuotaDates(startTime int64, endTime int64, username string) (quotaData []*QuotaData, err error) { +func GetAllQuotaDates(startTime int64, endTime int64, username, tokenName string) (quotaData []*QuotaData, err error) { if username != "" { - return GetQuotaDataByUsername(username, startTime, endTime) + return GetQuotaDataByUsername(username, tokenName, startTime, endTime) } var quotaDatas []*QuotaData // 从quota_data表中查询数据 @@ -131,3 +166,263 @@ func GetAllQuotaDates(startTime int64, endTime int64, username string) (quotaDat err = DB.Table("quota_data").Select("model_name, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used, created_at").Where("created_at >= ? and created_at <= ?", startTime, endTime).Group("model_name, created_at").Find("aDatas).Error return quotaDatas, err } + +func GetBilling(startTime int64, endTime int64, userName, tokenname string) (billingJsonData []*BillingJsonData, err error) { + // 将时间戳转换为当天的开始时间（00:00:00） + if endTime > time.Now().Unix() { + endTime = time.Now().Unix() + } + currentTime := time.Unix(startTime, 0) + currentTime = time.Date(currentTime.Year(), currentTime.Month(), currentTime.Day(), 0, 0, 0, 0, currentTime.Location()) + endDateTime := time.Unix(endTime, 0) + + // 按天遍历时间范围 + for currentTime.Unix() <= endDateTime.Unix() { + dayStart := currentTime.Unix() + dayEnd := currentTime.Add(24 * time.Hour).Add(-time.Second).Unix() + tableName := fmt.Sprintf("logs_%04d_%02d_%02d", currentTime.Year(), currentTime.Month(), currentTime.Day()) + if dayEnd > endTime { + dayEnd = endTime + } + + var billingData []*BillingData + var tempBillingMap = make(map[string]*BillingData) // 用于临时存储聚合结果 + pageSize := 100000 + offset := 0 + + for { + var tempData []*struct { + ChannelId int + ChannelName string + ChannelTag string + ModelName string + PromptTokens int + CompletionTokens int + } + + if userName != "" { + // 分页查询原始日志数据 + if tokenname != "" { + err = DB.Table(tableName). + Select(fmt.Sprintf("%s.channel_id, channels.name as channel_name, channels.tag as channel_tag, "+ + "%s.model_name, %s.prompt_tokens, %s.completion_tokens", tableName, tableName, tableName, tableName)). + Joins(fmt.Sprintf("JOIN channels ON %s.channel_id = channels.id", tableName)). // 修复这里 + Where(fmt.Sprintf("%s.created_at BETWEEN ? AND ?", tableName), dayStart, dayEnd). + Where(fmt.Sprintf("%s.username = ?", tableName), userName). + Where(fmt.Sprintf("%s.token_name =?", tableName), tokenname). + Order(fmt.Sprintf("%s.id", tableName)). // 修复这里 + Limit(pageSize). + Offset(offset). + Find(&tempData).Error + } else { + err = DB.Table(tableName). + Select(fmt.Sprintf("%s.channel_id, channels.name as channel_name, channels.tag as channel_tag, "+ + "%s.model_name, %s.prompt_tokens, %s.completion_tokens", tableName, tableName, tableName, tableName)). + Joins(fmt.Sprintf("JOIN channels ON %s.channel_id = channels.id", tableName)). // 修复这里 + Where(fmt.Sprintf("%s.created_at BETWEEN ? AND ?", tableName), dayStart, dayEnd). + Where(fmt.Sprintf("%s.username = ?", tableName), userName). + Order(fmt.Sprintf("%s.id", tableName)). // 修复这里 + Limit(pageSize). + Offset(offset). + Find(&tempData).Error + } + } else { + // 分页查询原始日志数据 + err = DB.Table(tableName). + Select(fmt.Sprintf("%s.channel_id, channels.name as channel_name, channels.tag as channel_tag, "+ + "%s.model_name, %s.prompt_tokens, %s.completion_tokens", tableName, tableName, tableName, tableName)). + Joins(fmt.Sprintf("JOIN channels ON %s.channel_id = channels.id", tableName)). // 修复这里 + Where(fmt.Sprintf("%s.created_at BETWEEN ? AND ?", tableName), dayStart, dayEnd). + Order(fmt.Sprintf("%s.id", tableName)). // 修复这里 + Limit(pageSize). + Offset(offset). + Find(&tempData).Error + } + + if err != nil { + return nil, err + } + + // 如果没有更多数据，退出循环 + if len(tempData) == 0 { + break + } + + // 处理当前页的数据，进行内存聚合 + for _, item := range tempData { + key := fmt.Sprintf("%s_%s", item.ChannelTag, item.ModelName) + if _, ok := tempBillingMap[key]; !ok { + tempBillingMap[key] = &BillingData{ + ChannelId: item.ChannelId, + ChannelName: item.ChannelName, + ChannelTag: item.ChannelTag, + ModelName: item.ModelName, + Count: 0, + PromptTokens: 0, + CompletionsTokens: 0, + } + } + existing, _ := tempBillingMap[key] + // 已存在的记录，累加计数 + existing.Count++ + existing.PromptTokens += item.PromptTokens + existing.CompletionsTokens += item.CompletionTokens + } + + offset += pageSize + } + + // 将聚合结果转换为切片 + for _, data := range tempBillingMap { + billingData = append(billingData, data) + } + + sort.Slice(billingData, func(i, j int) bool { + if billingData[i].ChannelTag != billingData[j].ChannelTag { + return billingData[i].ChannelTag < billingData[j].ChannelTag + } else if billingData[i].ChannelId != billingData[j].ChannelId { + return billingData[i].ChannelId < billingData[j].ChannelId + } else { + return billingData[i].ModelName < billingData[j].ModelName + } + }) + + // 处理当天的数据 + for _, data := range billingData { + modelPrice1, ok1 := operation_setting.GetModelRatio(data.ModelName) + modelPrice := 1.0 + + if ok1 { + modelPrice = modelPrice1 + } + + billingJsonData = append(billingJsonData, &BillingJsonData{ + ChannelId: data.ChannelId, + ChannelName: data.ChannelName, + ChannelTag: data.ChannelTag, + CurrentDate: currentTime.Format("2006-01-02"), + Count: data.Count, + ModelName: data.ModelName, + PromptTokens: float32(data.PromptTokens), + CompletionsTokens: float32(data.CompletionsTokens), + PromptPricing: float32(modelPrice * 2), + CompletionsPricing: float32(modelPrice * 2 * operation_setting.GetCompletionRatio(data.ModelName)), + Cost: (float32(data.PromptTokens)*float32(modelPrice*2) + float32(data.CompletionsTokens)*float32(modelPrice*2*operation_setting.GetCompletionRatio(data.ModelName))) / 100_0000, + }) + } + + // 移动到下一天 + currentTime = currentTime.Add(24 * time.Hour) + } + + // 在返回之前对数据进行排序 + sort.Slice(billingJsonData, func(i, j int) bool { + // 首先按照 ChannelTag 排序 + if billingJsonData[i].ChannelTag != billingJsonData[j].ChannelTag { + return billingJsonData[i].ChannelTag < billingJsonData[j].ChannelTag + } + // ChannelTag 相同时，按照 CurrentDate 排序 + return billingJsonData[i].CurrentDate < billingJsonData[j].CurrentDate + }) + + return billingJsonData, nil +} + +func GetBillingAndExportExcel(startTime int64, endTime int64, userName string, tokenname string) ([]byte, error) { + billingData, err := GetBilling(startTime, endTime, userName, tokenname) + if err != nil { + return nil, err + } + + // 创建新的Excel文件 + f := excelize.NewFile() + defer f.Close() + + // 设置表头 + headers := []string{"渠道Tag（Tag相同则聚合）", "日期", "调用次数", "模型名字", + "提示Tokens", "补全Tokens", "提示价格", "补全价格", "金额"} + for i, header := range headers { + cell := fmt.Sprintf("%c1", 'A'+i) + f.SetCellValue("Sheet1", cell, header) + // 设置列宽为25 + f.SetColWidth("Sheet1", string('A'+i), string('A'+i), 25) + } + + row := 2 + currentChannelTag := "null" + var channelTotal float32 = 0 + + // 在 GetBillingAndExportExcel 函数开始处添加样式定义 + style, err := f.NewStyle(&excelize.Style{ + Fill: excelize.Fill{ + Type: "pattern", + Color: []string{"FFD699"}, // 橙色 + Pattern: 1, + }, + }) + if err != nil { + return nil, err + } + + // 写入数据 + for _, data := range billingData { + // 如果是新的渠道ID，且不是第一条数据 + if currentChannelTag != "null" && currentChannelTag != data.ChannelTag && channelTotal > 0 { + // 写入渠道总计行 + f.SetCellValue("Sheet1", fmt.Sprintf("A%d", row), currentChannelTag) + f.SetCellValue("Sheet1", fmt.Sprintf("B%d", row), "总计") + f.SetCellValue("Sheet1", fmt.Sprintf("C%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("D%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("E%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("F%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("G%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("H%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("I%d", row), channelTotal) + // 为整行设置样式 + for col := 'A'; col <= 'I'; col++ { + f.SetCellStyle("Sheet1", fmt.Sprintf("%c%d", col, row), fmt.Sprintf("%c%d", col, row), style) + } + row += 3 + channelTotal = 0 + } + + // 写入详细数据 + f.SetCellValue("Sheet1", fmt.Sprintf("A%d", row), data.ChannelTag) + f.SetCellValue("Sheet1", fmt.Sprintf("B%d", row), data.CurrentDate) + f.SetCellValue("Sheet1", fmt.Sprintf("C%d", row), data.Count) + f.SetCellValue("Sheet1", fmt.Sprintf("D%d", row), data.ModelName) + f.SetCellValue("Sheet1", fmt.Sprintf("E%d", row), data.PromptTokens) + f.SetCellValue("Sheet1", fmt.Sprintf("F%d", row), data.CompletionsTokens) + f.SetCellValue("Sheet1", fmt.Sprintf("G%d", row), data.PromptPricing) + f.SetCellValue("Sheet1", fmt.Sprintf("H%d", row), data.CompletionsPricing) + f.SetCellValue("Sheet1", fmt.Sprintf("I%d", row), data.Cost) + + channelTotal += data.Cost + currentChannelTag = data.ChannelTag + row++ + } + + // 写入最后一个渠道的总计行 + if channelTotal > 0 { + f.SetCellValue("Sheet1", fmt.Sprintf("A%d", row), currentChannelTag) + f.SetCellValue("Sheet1", fmt.Sprintf("B%d", row), "总计") + f.SetCellValue("Sheet1", fmt.Sprintf("C%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("D%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("E%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("F%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("G%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("H%d", row), "-") + f.SetCellValue("Sheet1", fmt.Sprintf("I%d", row), channelTotal) + for col := 'A'; col <= 'I'; col++ { + f.SetCellStyle("Sheet1", fmt.Sprintf("%c%d", col, row), fmt.Sprintf("%c%d", col, row), style) + } + } + + // 删除保存文件的代码，改为返回字节流 + buffer, err := f.WriteToBuffer() + if err != nil { + return nil, err + } + + return buffer.Bytes(), nil +} diff --git a/relay/channel/adapter.go b/relay/channel/adapter.go index c970fd4854ae..cf097c4508f3 100644 --- a/relay/channel/adapter.go +++ b/relay/channel/adapter.go @@ -1,11 +1,12 @@ package channel import ( - "github.com/gin-gonic/gin" "io" "net/http" "one-api/dto" relaycommon "one-api/relay/common" + + "github.com/gin-gonic/gin" ) type Adaptor interface { diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index a60bc6f1898b..fbaf536615de 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -1,15 +1,29 @@ package channel import ( + "bytes" + "context" "errors" "fmt" - "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "io" "net/http" + onecommon "one-api/common" "one-api/relay/common" "one-api/relay/constant" - "one-api/service" + "regexp" + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +// contextKey 是用于 context 值的自定义类型 +type contextKey string + +const ( + ginContextKey contextKey = "gin_context" ) func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) { @@ -24,6 +38,19 @@ func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Hea req.Set("Accept", "text/event-stream") } } + + // 添加自定义请求头 + for key, value := range info.Headers { + req.Set(key, value) + } + + // 添加指定的header - 从原始请求中获取 + if retryRequestId := c.GetHeader("retry_request_id"); retryRequestId != "" { + req.Set("retry_request_id", retryRequestId) + } + if retry := c.GetHeader("retry"); retry != "" { + req.Set("retry", retry) + } } func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) { @@ -91,23 +118,113 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody } func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { - var client *http.Client - var err error - if proxyURL, ok := info.ChannelSetting["proxy"]; ok { - client, err = service.NewProxyHttpClient(proxyURL.(string)) - if err != nil { - return nil, fmt.Errorf("new proxy http client failed: %w", err) + // Check if mock response is enabled and test traffic header is present + var response *http.Response + + if onecommon.MockResponseEnabled && c.GetHeader("X-Test-Traffic") == "true" { + + var responseBody string + if strings.Contains(strings.ToLower(info.UpstreamModelName), "gemini") { + responseBody = `{ + "candidates": [{ + "content": { + "parts": [{ + "text": "测试结果是1 + 1 = 2" + }], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 10, + "totalTokenCount": 20 + } + }` + } else { + responseBody = `{ + "id": "mock-response", + "model": "gpt-3.5-turbo", + "object": "chat.completion", + "choices": [{ + "message": { + "role": "assistant", + "content": "测试结果是1 + 1 = 2" + } + }] + }` } - } else { - client = service.GetHttpClient() + response = &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewBufferString(responseBody)), + } + // 设置正确的 Content-Type 头 + response.Header.Set("Content-Type", "application/json") + } - resp, err := client.Do(req) - if err != nil { - return nil, err + + // Create HTTP client + client := &http.Client{ + Timeout: time.Duration(onecommon.RelayTimeout) * time.Second, } - if resp == nil { - return nil, errors.New("resp is nil") + req.Header.Set(onecommon.RequestIdKey, c.GetString(onecommon.RequestIdKey)) + + // 添加来源标识和重试次数 + req.Header.Set("X-Origin-User-ID", strconv.Itoa(info.UserId)) + req.Header.Set("X-Origin-Channel-ID", strconv.Itoa(info.ChannelId)) + req.Header.Set("X-Retry-Count", strconv.Itoa(info.RetryCount)) + + // 添加指定的header - 从原始请求中获取 + if retryRequestId := c.GetHeader("retry_request_id"); retryRequestId != "" { + req.Header.Set("retry_request_id", retryRequestId) + } + if retry := c.GetHeader("retry"); retry != "" { + req.Header.Set("retry", retry) + } + + // 打印请求头 + requestId := c.GetString(onecommon.RequestIdKey) + ctx := context.WithValue(c.Request.Context(), onecommon.RequestIdKey, requestId) + ctx = context.WithValue(ctx, "gin_context", c) + onecommon.LogInfo(ctx, fmt.Sprintf("request headers: %v", req.Header)) + + // 读取并打印请求体 + if req.Body != nil { + bodyBytes, _ := io.ReadAll(req.Body) + if len(bodyBytes) > 0 { + // 只打印小于64KB的请求体 + if len(bodyBytes) < 64*1024 { + // 使用正则表达式替换base64数据 + bodyStr := string(bodyBytes) + replacedBody := replaceBase64InString(bodyStr) + onecommon.LogInfo(ctx, fmt.Sprintf("request body: %s", replacedBody)) + } else { + onecommon.LogInfo(ctx, fmt.Sprintf("request body too large (size: %d bytes), skipping print", len(bodyBytes))) + } + // 重新设置 body + req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + } + } + + var resp *http.Response + if response == nil { + var err error + resp, err = client.Do(req) + if err != nil { + return nil, err + } + if resp == nil { + return nil, errors.New("resp is nil") + } + } else { + resp = response } + + // 打印响应头 + onecommon.LogInfo(c, fmt.Sprintf("response headers: %v", resp.Header)) + _ = req.Body.Close() _ = c.Request.Body.Close() return resp, nil @@ -130,9 +247,41 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.TaskRelayInfo, if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) } + req.Header.Set(onecommon.RequestIdKey, c.GetString(onecommon.RequestIdKey)) resp, err := doRequest(c, req, info.RelayInfo) if err != nil { return nil, fmt.Errorf("do request failed: %w", err) } return resp, nil } + +func replaceBase64InString(s string) string { + // 替换data URL格式的base64数据 + // 匹配格式: "url": "data:video/mp4;base64,UklGRnoGAABXQVZFZm10IBAAAAABAAEAQB8AAEAfAAABAAgAZGF0YQoGAACBhYqFbF1fdJivrJBhNjVgodDbq2EcBj+a2/LDciUFLIHO8tiJNwgZaLvt559NEAxQp+PwtmMcBjiR1/LMeSwFJHfH8N2QQAoUXrTp66hVFApGn+DyvmwhBSuBzvLZiTYIG2m98OScTgwOUarm7blmGgU7k9n1unEiBC13yO/eizEIHWq+8+OWT" + // 替换为: "url": "data:video/mp4;base64,[BASE64_DATA_1896764_chars]" + + // 匹配data URL格式的正则表达式 + dataURLPattern := regexp.MustCompile(`"url":\s*"data:[^"]+;base64,[^"]+"`) + + // 替换函数 + replacer := func(match string) string { + // 提取MIME类型 + mimePattern := regexp.MustCompile(`data:([^;]+);base64,`) + mimeMatch := mimePattern.FindStringSubmatch(match) + if len(mimeMatch) > 1 { + mimeType := mimeMatch[1] + // 计算base64数据长度（大约） + base64Pattern := regexp.MustCompile(`base64,([^"]+)`) + base64Match := base64Pattern.FindStringSubmatch(match) + if len(base64Match) > 1 { + base64Data := base64Match[1] + // 计算字符数 + charCount := len(base64Data) + return fmt.Sprintf(`"url": "data:%s;base64,[BASE64_DATA_%d_chars]"`, mimeType, charCount) + } + } + return `"url": "data:image/png;base64,[BASE64_DATA_REPLACED]"` + } + + return dataURLPattern.ReplaceAllStringFunc(s, replacer) +} diff --git a/relay/channel/aws/adaptor.go b/relay/channel/aws/adaptor.go index 7f2a2841bc66..602464584bf8 100644 --- a/relay/channel/aws/adaptor.go +++ b/relay/channel/aws/adaptor.go @@ -2,13 +2,14 @@ package aws import ( "errors" - "github.com/gin-gonic/gin" "io" "net/http" "one-api/dto" "one-api/relay/channel/claude" relaycommon "one-api/relay/common" "one-api/setting/model_setting" + + "github.com/gin-gonic/gin" ) const ( @@ -40,6 +41,15 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error { model_setting.GetClaudeSettings().WriteHeaders(info.OriginModelName, req) + + // 添加指定的header - 从原始请求中获取 + if retryRequestId := c.GetHeader("retry_request_id"); retryRequestId != "" { + req.Set("retry_request_id", retryRequestId) + } + if retry := c.GetHeader("retry"); retry != "" { + req.Set("retry", retry) + } + return nil } @@ -50,7 +60,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re var claudeReq *claude.ClaudeRequest var err error - claudeReq, err = claude.RequestOpenAI2ClaudeMessage(*request) + claudeReq, err = claude.RequestOpenAI2ClaudeMessage(c, *request) if err != nil { return nil, err } diff --git a/relay/channel/claude/adaptor.go b/relay/channel/claude/adaptor.go index bf03e5f5fdaf..21a54476cdd0 100644 --- a/relay/channel/claude/adaptor.go +++ b/relay/channel/claude/adaptor.go @@ -67,7 +67,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re if a.RequestMode == RequestModeCompletion { return RequestOpenAI2ClaudeComplete(*request), nil } else { - return RequestOpenAI2ClaudeMessage(*request) + return RequestOpenAI2ClaudeMessage(c, *request) } } diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 09154bcbe51b..91f778be0c7f 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -11,6 +11,7 @@ import ( "one-api/relay/helper" "one-api/service" "one-api/setting/model_setting" + "strconv" "strings" "github.com/gin-gonic/gin" @@ -30,7 +31,6 @@ func stopReasonClaude2OpenAI(reason string) string { } func RequestOpenAI2ClaudeComplete(textRequest dto.GeneralOpenAIRequest) *ClaudeRequest { - claudeRequest := ClaudeRequest{ Model: textRequest.Model, Prompt: "", @@ -60,7 +60,7 @@ func RequestOpenAI2ClaudeComplete(textRequest dto.GeneralOpenAIRequest) *ClaudeR return &claudeRequest } -func RequestOpenAI2ClaudeMessage(textRequest dto.GeneralOpenAIRequest) (*ClaudeRequest, error) { +func RequestOpenAI2ClaudeMessage(c *gin.Context, textRequest dto.GeneralOpenAIRequest) (*ClaudeRequest, error) { claudeTools := make([]Tool, 0, len(textRequest.Tools)) for _, tool := range textRequest.Tools { @@ -105,12 +105,22 @@ func RequestOpenAI2ClaudeMessage(textRequest dto.GeneralOpenAIRequest) (*ClaudeR if claudeRequest.MaxTokens < 1280 { claudeRequest.MaxTokens = 1280 } - - // BudgetTokens 为 max_tokens 的 80% - claudeRequest.Thinking = &Thinking{ - Type: "enabled", - BudgetTokens: int(float64(claudeRequest.MaxTokens) * model_setting.GetClaudeSettings().ThinkingAdapterBudgetTokensPercentage), + claudeRequest.Thinking = &Thinking{Type: "enabled"} + // 支持用户覆盖 budget tokens + if textRequest.Thinking != nil { + // claude 对于 budget tokens 最小限制为 1024 + if textRequest.Thinking.BudgetTokens < 1024 { + textRequest.Thinking.BudgetTokens = 1024 + common.LogInfo(c, fmt.Sprintf("传入的 budget tokens %d 小于 1024，已设为 1024 ", textRequest.Thinking.BudgetTokens)) + } + claudeRequest.Thinking.BudgetTokens = textRequest.Thinking.BudgetTokens + common.LogInfo(c, fmt.Sprintf("用户自定义 budget tokens 长度: %d", claudeRequest.Thinking.BudgetTokens)) + } else { + // BudgetTokens 为 max_tokens 的 80% + claudeRequest.Thinking.BudgetTokens = int(float64(claudeRequest.MaxTokens) * model_setting.GetClaudeSettings().ThinkingAdapterBudgetTokensPercentage) + common.LogInfo(c, fmt.Sprintf("budget tokens 使用系统 max tokens 的 80%%: %d", claudeRequest.Thinking.BudgetTokens)) } + // TODO: 临时处理 // https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#important-considerations-when-using-extended-thinking claudeRequest.TopP = 0 @@ -560,3 +570,170 @@ func ClaudeHandler(c *gin.Context, resp *http.Response, requestMode int, info *r _, err = c.Writer.Write(jsonResponse) return nil, &usage } + +func ClaudeMessage2OpenAIRequest(claudeReq *ClaudeRequest) (*dto.GeneralOpenAIRequest, error) { + openaiReq := &dto.GeneralOpenAIRequest{ + Model: claudeReq.Model, + MaxTokens: claudeReq.MaxTokens, + Temperature: claudeReq.Temperature, + TopP: claudeReq.TopP, + Stream: claudeReq.Stream, + StreamOptions: &dto.StreamOptions{IncludeUsage: true}, + } + + if claudeReq.Thinking != nil { + openaiReq.Thinking = &dto.ThinkingOptions{ + Type: claudeReq.Thinking.Type, + BudgetTokens: claudeReq.Thinking.BudgetTokens, + } + } + + if claudeReq.Tools != nil { + openaiTools := make([]dto.ToolCallRequest, 0) + + switch tools := claudeReq.Tools.(type) { + case []Tool: + for _, claudeTool := range tools { + params := make(map[string]interface{}, 3) + for _, key := range []string{"type", "properties", "required"} { + if val, exist := claudeTool.InputSchema[key]; exist { + params[key] = val + } + } + + openaiTools = append(openaiTools, dto.ToolCallRequest{ + Type: "function", + Function: dto.FunctionRequest{ + Name: claudeTool.Name, + Description: claudeTool.Description, + Parameters: params, + }, + }) + } + + case []interface{}: // 处理通用类型 + for _, rawTool := range tools { + tool, ok := rawTool.(map[string]interface{}) + if !ok { + continue + } + + name, _ := tool["name"].(string) + desc, _ := tool["description"].(string) + schema, _ := tool["input_schema"].(map[string]interface{}) + + params := make(map[string]interface{}, 3) + if schema != nil { + for _, key := range []string{"type", "properties", "required"} { + if val, exist := schema[key]; exist { + params[key] = val + } + } + } + + openaiTools = append(openaiTools, dto.ToolCallRequest{ + Type: "function", + Function: dto.FunctionRequest{ + Name: name, + Description: desc, + Parameters: params, + }, + }) + } + } + + openaiReq.Tools = openaiTools + } + + if claudeReq.System != "" { + systemMsg := dto.Message{ + Role: "system", + Content: json.RawMessage([]byte(strconv.Quote(claudeReq.System))), + } + openaiReq.Messages = append(openaiReq.Messages, systemMsg) + } + + // 多模态 + for _, claudeMsg := range claudeReq.Messages { + openaiMsg := dto.Message{Role: claudeMsg.Role} + + switch content := claudeMsg.Content.(type) { + case string: // 纯文本 + openaiMsg.SetStringContent(content) + + case []ClaudeMediaMessage: // 复杂消息类型 + var mediaContents []dto.MediaContent + var toolCalls []dto.ToolCallRequest + + for _, media := range content { + switch media.Type { + case "text": + mediaContents = append(mediaContents, dto.MediaContent{ + Type: dto.ContentTypeText, + Text: media.Text, + }) + + case "image": + if media.Source != nil { + mediaContents = append(mediaContents, dto.MediaContent{ + Type: dto.ContentTypeImageURL, + ImageUrl: dto.MessageImageUrl{ + Url: fmt.Sprintf("data:%s;base64,%s", media.Source.MediaType, media.Source.Data), + Detail: "auto", + }, + }) + } + + case "tool_use": + args, _ := json.Marshal(media.Input) + toolCalls = append(toolCalls, dto.ToolCallRequest{ + ID: media.Id, + Function: dto.FunctionRequest{ + Name: media.Name, + Arguments: string(args), + }, + }) + + case "tool_result": + openaiMsg.ToolCallId = media.ToolUseId + openaiMsg.SetStringContent(media.Content) + } + } + + if len(toolCalls) > 0 { + toolCallData, _ := json.Marshal(toolCalls) + openaiMsg.ToolCalls = toolCallData + } + + if len(mediaContents) > 0 { + openaiMsg.SetMediaContent(mediaContents) + } + } + + openaiReq.Messages = append(openaiReq.Messages, openaiMsg) + } + + if len(claudeReq.StopSequences) > 0 { + if len(claudeReq.StopSequences) == 1 { + openaiReq.Stop = claudeReq.StopSequences[0] + } else { + stopSeq := make([]string, len(claudeReq.StopSequences)) + copy(stopSeq, claudeReq.StopSequences) + openaiReq.Stop = stopSeq + } + } + + if strings.Contains(claudeReq.Model, "claude-3") { + openaiReq.ResponseFormat = &dto.ResponseFormat{ + Type: "json_object", + JsonSchema: &dto.FormatJsonSchema{ + Schema: map[string]interface{}{ + "type": "object", + "properties": make(map[string]interface{}), + }, + }, + } + } + + return openaiReq, nil +} diff --git a/relay/channel/gemini/dto.go b/relay/channel/gemini/dto.go index bbcb1248d482..2a4004ea9150 100644 --- a/relay/channel/gemini/dto.go +++ b/relay/channel/gemini/dto.go @@ -43,6 +43,10 @@ type GeminiFileData struct { FileUri string `json:"fileUri,omitempty"` } +type GeminiVideoMetadata struct { + Fps float64 `json:"fps"` +} + type GeminiPart struct { Text string `json:"text,omitempty"` InlineData *GeminiInlineData `json:"inlineData,omitempty"` @@ -51,6 +55,7 @@ type GeminiPart struct { FileData *GeminiFileData `json:"fileData,omitempty"` ExecutableCode *GeminiPartExecutableCode `json:"executableCode,omitempty"` CodeExecutionResult *GeminiPartCodeExecutionResult `json:"codeExecutionResult,omitempty"` + VideoMetadata *GeminiVideoMetadata `json:"video_metadata,omitempty"` } type GeminiChatContent struct { @@ -71,15 +76,20 @@ type GeminiChatTool struct { } type GeminiChatGenerationConfig struct { - Temperature *float64 `json:"temperature,omitempty"` - TopP float64 `json:"topP,omitempty"` - TopK float64 `json:"topK,omitempty"` - MaxOutputTokens uint `json:"maxOutputTokens,omitempty"` - CandidateCount int `json:"candidateCount,omitempty"` - StopSequences []string `json:"stopSequences,omitempty"` - ResponseMimeType string `json:"responseMimeType,omitempty"` - ResponseSchema any `json:"responseSchema,omitempty"` - Seed int64 `json:"seed,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP float64 `json:"topP,omitempty"` + TopK float64 `json:"topK,omitempty"` + MaxOutputTokens uint `json:"maxOutputTokens,omitempty"` + CandidateCount int `json:"candidateCount,omitempty"` + StopSequences []string `json:"stopSequences,omitempty"` + ResponseMimeType string `json:"responseMimeType,omitempty"` + ResponseSchema any `json:"responseSchema,omitempty"` + Seed int64 `json:"seed,omitempty"` + ThinkingConfig *GeminiChatThinkingConfig `json:"thinkingConfig,omitempty"` +} + +type GeminiChatThinkingConfig struct { + ThinkingBudget int `json:"thinkingBudget"` } type GeminiChatCandidate struct { @@ -108,6 +118,7 @@ type GeminiUsageMetadata struct { PromptTokenCount int `json:"promptTokenCount"` CandidatesTokenCount int `json:"candidatesTokenCount"` TotalTokenCount int `json:"totalTokenCount"` + ThoughtsTokenCount int `json:"thoughtsTokenCount"` } // Imagen related structs diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index c1ce8219dccb..e2289f24460a 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -29,6 +29,14 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest) (*GeminiChatReque TopP: textRequest.TopP, MaxOutputTokens: textRequest.MaxTokens, Seed: int64(textRequest.Seed), + ThinkingConfig: func() *GeminiChatThinkingConfig { + if textRequest.Thinking != nil { + return &GeminiChatThinkingConfig{ + ThinkingBudget: textRequest.Thinking.BudgetTokens, + } + } + return nil + }(), }, } @@ -204,6 +212,28 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest) (*GeminiChatReque }, }) } + } else if part.Type == dto.ContentTypeInputAudio { + // 处理音频内容 + audioData := part.InputAudio.(dto.MessageInputAudio) + // 添加调试日志 + common.SysLog(fmt.Sprintf("Processing audio data: format=%s, data length=%d", audioData.Format, len(audioData.Data))) + // 将音频数据转换为Gemini的InlineData格式 + parts = append(parts, GeminiPart{ + InlineData: &GeminiInlineData{ + MimeType: audioData.Format, + Data: audioData.Data, + }, + VideoMetadata: &GeminiVideoMetadata{ + Fps: audioData.Fps, + }, + }) + } else if part.Type == dto.ContentTypeYoutube { + parts = append(parts, GeminiPart{ + FileData: &GeminiFileData{ + MimeType: part.Text, + FileUri: part.ImageUrl.(dto.MessageImageUrl).Url, + }, + }) } } @@ -553,6 +583,7 @@ func GeminiChatHandler(c *gin.Context, resp *http.Response, info *relaycommon.Re return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil } if len(geminiResponse.Candidates) == 0 { + common.SysError(fmt.Sprintf("no candidates returned: %s", string(responseBody))) return &dto.OpenAIErrorWithStatusCode{ Error: dto.OpenAIError{ Message: "No candidates returned", @@ -570,6 +601,7 @@ func GeminiChatHandler(c *gin.Context, resp *http.Response, info *relaycommon.Re CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount, TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount, } + usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount fullTextResponse.Usage = usage jsonResponse, err := json.Marshal(fullTextResponse) if err != nil { diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go index 6dbbb17e2165..c0ddab623f9c 100644 --- a/relay/channel/openai/adaptor.go +++ b/relay/channel/openai/adaptor.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "github.com/gin-gonic/gin" "io" "mime/multipart" "net/http" @@ -21,6 +20,8 @@ import ( relaycommon "one-api/relay/common" "one-api/relay/constant" "strings" + + "github.com/gin-gonic/gin" ) type Adaptor struct { @@ -75,6 +76,15 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info *relaycommon.RelayInfo) error { channel.SetupApiRequestHeader(info, c, header) + + // 添加用户ID和渠道ID到请求头 + if info.UserId != 0 { + header.Set("X-User-ID", fmt.Sprintf("%d", info.UserId)) + } + if info.ChannelId != 0 { + header.Set("X-Channel-ID", fmt.Sprintf("%d", info.ChannelId)) + } + if info.ChannelType == common.ChannelTypeAzure { header.Set("api-key", info.ApiKey) return nil @@ -136,6 +146,10 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re info.ReasoningEffort = request.ReasoningEffort info.UpstreamModelName = request.Model } + if strings.HasSuffix(request.Model, "-disable") { + request.Thinking = &dto.ThinkingOptions{Type: "disabled"} + request.Model = strings.TrimSuffix(request.Model, "-disable") + } if request.Model == "o1" || request.Model == "o1-2024-12-17" || strings.HasPrefix(request.Model, "o3") { //修改第一个Message的内容，将system改为developer if len(request.Messages) > 0 && request.Messages[0].Role == "system" { diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go new file mode 100644 index 000000000000..d53a38a6e458 --- /dev/null +++ b/relay/channel/openai/helper.go @@ -0,0 +1,128 @@ +package openai + +import ( + "encoding/json" + "one-api/common" + "one-api/dto" + relaycommon "one-api/relay/common" + relayconstant "one-api/relay/constant" + "one-api/service" + "strings" +) + +func ProcessStreamResponse(streamResponse dto.ChatCompletionsStreamResponse, responseTextBuilder *strings.Builder, toolCount *int) error { + for _, choice := range streamResponse.Choices { + responseTextBuilder.WriteString(choice.Delta.GetContentString()) + responseTextBuilder.WriteString(choice.Delta.GetReasoningContent()) + if choice.Delta.ToolCalls != nil { + if len(choice.Delta.ToolCalls) > *toolCount { + *toolCount = len(choice.Delta.ToolCalls) + } + for _, tool := range choice.Delta.ToolCalls { + responseTextBuilder.WriteString(tool.Function.Name) + responseTextBuilder.WriteString(tool.Function.Arguments) + } + } + } + return nil +} + +func processTokens(relayMode int, streamItems []string, responseTextBuilder *strings.Builder, toolCount *int) error { + streamResp := "[" + strings.Join(streamItems, ",") + "]" + + switch relayMode { + case relayconstant.RelayModeChatCompletions: + return processChatCompletions(streamResp, streamItems, responseTextBuilder, toolCount) + case relayconstant.RelayModeCompletions: + return processCompletions(streamResp, streamItems, responseTextBuilder) + } + return nil +} + +func processChatCompletions(streamResp string, streamItems []string, responseTextBuilder *strings.Builder, toolCount *int) error { + var streamResponses []dto.ChatCompletionsStreamResponse + if err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses); err != nil { + // 一次性解析失败，逐个解析 + common.SysError("error unmarshalling stream response: " + err.Error()) + for _, item := range streamItems { + var streamResponse dto.ChatCompletionsStreamResponse + if err := json.Unmarshal(common.StringToByteSlice(item), &streamResponse); err != nil { + return err + } + if err := ProcessStreamResponse(streamResponse, responseTextBuilder, toolCount); err != nil { + common.SysError("error processing stream response: " + err.Error()) + } + } + return nil + } + + // 批量处理所有响应 + for _, streamResponse := range streamResponses { + for _, choice := range streamResponse.Choices { + responseTextBuilder.WriteString(choice.Delta.GetContentString()) + responseTextBuilder.WriteString(choice.Delta.GetReasoningContent()) + if choice.Delta.ToolCalls != nil { + if len(choice.Delta.ToolCalls) > *toolCount { + *toolCount = len(choice.Delta.ToolCalls) + } + for _, tool := range choice.Delta.ToolCalls { + responseTextBuilder.WriteString(tool.Function.Name) + responseTextBuilder.WriteString(tool.Function.Arguments) + } + } + } + } + return nil +} + +func processCompletions(streamResp string, streamItems []string, responseTextBuilder *strings.Builder) error { + var streamResponses []dto.CompletionsStreamResponse + if err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses); err != nil { + // 一次性解析失败，逐个解析 + common.SysError("error unmarshalling stream response: " + err.Error()) + for _, item := range streamItems { + var streamResponse dto.CompletionsStreamResponse + if err := json.Unmarshal(common.StringToByteSlice(item), &streamResponse); err != nil { + continue + } + for _, choice := range streamResponse.Choices { + responseTextBuilder.WriteString(choice.Text) + } + } + return nil + } + + // 批量处理所有响应 + for _, streamResponse := range streamResponses { + for _, choice := range streamResponse.Choices { + responseTextBuilder.WriteString(choice.Text) + } + } + return nil +} + +func handleLastResponse(lastStreamData string, responseId *string, createAt *int64, + systemFingerprint *string, model *string, usage **dto.Usage, + containStreamUsage *bool, info *relaycommon.RelayInfo, + shouldSendLastResp *bool) error { + + var lastStreamResponse dto.ChatCompletionsStreamResponse + if err := json.Unmarshal(common.StringToByteSlice(lastStreamData), &lastStreamResponse); err != nil { + return err + } + + *responseId = lastStreamResponse.Id + *createAt = lastStreamResponse.Created + *systemFingerprint = lastStreamResponse.GetSystemFingerprint() + *model = lastStreamResponse.Model + + if service.ValidUsage(lastStreamResponse.Usage) { + *containStreamUsage = true + *usage = lastStreamResponse.Usage + if !info.ShouldIncludeUsage { + *shouldSendLastResp = false + } + } + + return nil +} diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 223ddd3d10bd..52635e721016 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -16,6 +16,7 @@ import ( "one-api/relay/helper" "one-api/service" "os" + "strconv" "strings" "github.com/bytedance/gopkg/util/gopool" @@ -282,23 +283,38 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model if err != nil { return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil } + common.LogInfo(c, fmt.Sprintf("response headers: %v", resp.Header)) + usageJson, _ := json.Marshal(simpleResponse.Usage) + common.LogInfo(c, fmt.Sprintf("raw response Usage: %s", string(usageJson))) if simpleResponse.Error.Type != "" { return &dto.OpenAIErrorWithStatusCode{ Error: simpleResponse.Error, StatusCode: resp.StatusCode, }, nil } + // Reset response body - resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) + // resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) // We shouldn't set the header before we parse the response body, because the parse part may fail. // And then we will have to send an error response, but in this case, the header has already been set. // So the httpClient will be confused by the response. // For example, Postman will report error, and we cannot check the response at all. + // for k, v := range resp.Header { + // c.Writer.Header().Set(k, v[0]) + // } for k, v := range resp.Header { + if k == "Content-Length" { + num, _ := strconv.Atoi(v[0]) + if num != len(responseBody) { + common.LogInfo(c, fmt.Sprintf("Content Length is %s but response body is %d", v[0], len(responseBody))) + c.Writer.Header().Set(k, strconv.Itoa(len(responseBody))) + continue + } + } c.Writer.Header().Set(k, v[0]) } c.Writer.WriteHeader(resp.StatusCode) - _, err = io.Copy(c.Writer, resp.Body) + _, err = io.Copy(c.Writer, bytes.NewReader(responseBody)) if err != nil { return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil } diff --git a/relay/channel/task/suno/adaptor.go b/relay/channel/task/suno/adaptor.go index 03d60516f06a..1eaa31fd2c82 100644 --- a/relay/channel/task/suno/adaptor.go +++ b/relay/channel/task/suno/adaptor.go @@ -5,7 +5,6 @@ import ( "context" "encoding/json" "fmt" - "github.com/gin-gonic/gin" "io" "net/http" "one-api/common" @@ -16,6 +15,8 @@ import ( "one-api/service" "strings" "time" + + "github.com/gin-gonic/gin" ) type TaskAdaptor struct { @@ -64,6 +65,15 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Header.Set("Accept", c.Request.Header.Get("Accept")) req.Header.Set("Authorization", "Bearer "+info.ApiKey) + + // 添加指定的header - 从原始请求中获取 + if retryRequestId := c.GetHeader("retry_request_id"); retryRequestId != "" { + req.Header.Set("retry_request_id", retryRequestId) + } + if retry := c.GetHeader("retry"); retry != "" { + req.Header.Set("retry", retry) + } + return nil } diff --git a/relay/channel/vertex/adaptor.go b/relay/channel/vertex/adaptor.go index 7ccd3f30dbdc..deff5acec603 100644 --- a/relay/channel/vertex/adaptor.go +++ b/relay/channel/vertex/adaptor.go @@ -124,7 +124,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re return nil, errors.New("request is nil") } if a.RequestMode == RequestModeClaude { - claudeReq, err := claude.RequestOpenAI2ClaudeMessage(*request) + claudeReq, err := claude.RequestOpenAI2ClaudeMessage(c, *request) if err != nil { return nil, err } diff --git a/relay/channel/volcengine/adaptor.go b/relay/channel/volcengine/adaptor.go index 3b57c67ca042..55098993a284 100644 --- a/relay/channel/volcengine/adaptor.go +++ b/relay/channel/volcengine/adaptor.go @@ -3,7 +3,6 @@ package volcengine import ( "errors" "fmt" - "github.com/gin-gonic/gin" "io" "net/http" "one-api/dto" @@ -12,6 +11,8 @@ import ( relaycommon "one-api/relay/common" "one-api/relay/constant" "strings" + + "github.com/gin-gonic/gin" ) type Adaptor struct { @@ -31,6 +32,10 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) { } func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { + if info.BaseUrl == "" { + info.BaseUrl = "https://ark.cn-beijing.volces.com" + fmt.Printf("no baseurl found, using %s\n", info.BaseUrl) + } switch info.RelayMode { case constant.RelayModeChatCompletions: if strings.HasPrefix(info.UpstreamModelName, "bot") { @@ -54,6 +59,11 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re if request == nil { return nil, errors.New("request is nil") } + + // Safely handle the Thinking field + if request.Thinking != nil { + request.Thinking = &dto.ThinkingOptions{Type: request.Thinking.Type} + } return request, nil } @@ -66,6 +76,10 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela } func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) { + // 检查模型名称是否包含 "batch"，如果是则使用批量接口 + if strings.Contains(strings.ToLower(info.OriginModelName), "batch") { + return DoBatchChatRequest(c, info, requestBody) + } return channel.DoApiRequest(a, c, info, requestBody) } diff --git a/relay/channel/volcengine/batchchat.go b/relay/channel/volcengine/batchchat.go new file mode 100644 index 000000000000..85ded735e908 --- /dev/null +++ b/relay/channel/volcengine/batchchat.go @@ -0,0 +1,1241 @@ +package volcengine + +import ( + "context" + "encoding/json" + "fmt" + "io" + "math/rand" + "net/http" + "one-api/common" + "one-api/dto" + "one-api/metrics" + relaycommon "one-api/relay/common" + "os" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/gin-gonic/gin" + "github.com/volcengine/volcengine-go-sdk/service/arkruntime" + "github.com/volcengine/volcengine-go-sdk/service/arkruntime/model" +) + +// 全局配置变量 +const ( + MaxParallelRequests = 2000 + + // 限速器默认大小为 MaxParallelRequests 的 3 倍 + RateLimiterSize = MaxParallelRequests * 3 + // 异步调用超时时间配置 + MinAsyncTimeout = 1 * time.Second + MaxAsyncTimeout = 3 * time.Second + // 限流等待时间配置 - 用于等待可用请求槽位的超时时间 + MinRateLimitWaitTime = 100 * time.Millisecond + MaxRateLimitWaitTime = 1000 * time.Millisecond + + // 重试响应延迟时间配置 - 用于控制重试请求的返回时间 + MinRetryResponseDelay = 0 * time.Millisecond + MaxRetryResponseDelay = 0 * time.Millisecond + + // CreateBatchChatCompletion 调用的超时时间 + BatchCompletionTimeout = 24 * time.Hour + + // 子协程最大存活时间 + SubGoroutineMaxLifetime = 24 * time.Hour + + // 分布式锁过期时间 + DistributedLockExpiration = 24 * time.Hour +) + +// 异步调用的超时时间 - 可通过环境变量VOLCENGINE_ASYNC_CALL_TIMEOUT配置，默认30秒 +var AsyncCallTimeout = time.Duration(common.GetEnvOrDefault("VOLCENGINE_ASYNC_CALL_TIMEOUT", 30)) * time.Second + +// 客户端缓存 +var ( + clientCache = make(map[string]*arkruntime.Client) + clientMutex sync.RWMutex +) + +// 请求计数器 +var ( + requestCounter int64 = 0 +) + +// 限速器 +var ( + rateLimiter = make(chan struct{}, RateLimiterSize) +) + +// 建议重试时间相关变量 +var ( + batchRequestAvgDuration float64 = 60.0 // 默认30秒 + batchRequestAvgDurationMutex sync.RWMutex +) + +// NewBatchClient 创建一个新的批量请求客户端实例 +func NewBatchClient(apiKey string) *arkruntime.Client { + return arkruntime.NewClientWithApiKey( + apiKey, + arkruntime.WithBatchMaxParallel(MaxParallelRequests), // 使用全局变量设置发起请求的最大并发数量 + ) +} + +// GetBatchClient 根据 channel ID 获取或创建客户端实例 +func GetBatchClient(channelId string, apiKey string) *arkruntime.Client { + clientMutex.RLock() + if client, exists := clientCache[channelId]; exists { + clientMutex.RUnlock() + return client + } + clientMutex.RUnlock() + + // 如果缓存中没有，创建新的客户端 + clientMutex.Lock() + defer clientMutex.Unlock() + + // 双重检查，防止并发创建 + if client, exists := clientCache[channelId]; exists { + return client + } + + client := NewBatchClient(apiKey) + clientCache[channelId] = client + return client +} + +// acquireRequestSlot 获取请求槽位，如果达到上限则返回错误 +func acquireRequestSlot() error { + // 获取当前计数器值 + currentCount := atomic.LoadInt64(&requestCounter) + + // 如果已达到上限，返回错误 + if currentCount >= int64(RateLimiterSize) { + return fmt.Errorf("request limit reached, please retry later") + } + + // 尝试获取限速器槽位 + select { + case rateLimiter <- struct{}{}: + // 成功获取槽位，增加计数器 + atomic.AddInt64(&requestCounter, 1) + return nil + default: + // 限速器已满，返回错误 + return fmt.Errorf("request limit reached, please retry later") + } +} + +// releaseRequestSlot 释放请求槽位 +func releaseRequestSlot() { + // 减少计数器 + atomic.AddInt64(&requestCounter, -1) + // 释放限速器槽位 + select { + case <-rateLimiter: + default: + // 如果限速器为空，忽略 + } +} + +// waitForAvailableSlot 等待可用的请求槽位 +func waitForAvailableSlot(ctx context.Context) error { + ticker := time.NewTicker(100 * time.Millisecond) // 每100ms检查一次 + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + if err := acquireRequestSlot(); err == nil { + return nil + } + } + } +} + +// 从 context 获取异步调用超时时间，如果没有设置则使用随机计算的时间 +func getAsyncCallTimeout(ctx context.Context) time.Duration { + if timeout, ok := ctx.Value("async_call_timeout").(time.Duration); ok && timeout > 0 { + return timeout + } + // 使用随机等待时间作为异步调用超时时间 + return calculateRandomWaitTime() +} + +// 从 context 获取批量推理超时时间，如果没有设置则使用默认值 +func getBatchCompletionTimeout(ctx context.Context) time.Duration { + if timeout, ok := ctx.Value("batch_completion_timeout").(time.Duration); ok && timeout > 0 { + return timeout + } + return BatchCompletionTimeout +} + +// GetCurrentTimeouts 获取当前 context 中的超时配置 +func GetCurrentTimeouts(ctx context.Context) map[string]time.Duration { + return map[string]time.Duration{ + "async_call_timeout": getAsyncCallTimeout(ctx), + "batch_completion_timeout": getBatchCompletionTimeout(ctx), + "min_async_timeout": MinAsyncTimeout, + "max_async_timeout": MaxAsyncTimeout, + "min_rate_limit_wait": MinRateLimitWaitTime, + "max_rate_limit_wait": MaxRateLimitWaitTime, + "min_retry_response_delay": MinRetryResponseDelay, + "max_retry_response_delay": MaxRetryResponseDelay, + "default_async_timeout": AsyncCallTimeout, + "default_batch_timeout": BatchCompletionTimeout, + } +} + +// getRequestID 从gin context获取请求ID，使用middleware设置的ID +func getRequestID(c *gin.Context) string { + // 优先使用retry_request_id + requestID := c.GetHeader("retry_request_id") + if requestID == "" { + // 如果没有retry_request_id，则使用正常的requestID + requestID = c.GetHeader(common.RequestIdKey) + } + return requestID +} + +// isRetryRequest 检查是否为重试请求 +func isRetryRequest(c *gin.Context) bool { + retryHeader := c.GetHeader("retry") + return retryHeader == "true" +} + +func DoBatchChatRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) { + // 获取请求ID + requestID := getRequestID(c) + + // 获取 retry header 状态 + retryHeaderStatus := "false" + if isRetryRequest(c) { + retryHeaderStatus = "true" + } + + // 记录请求开始时间 + requestStartTime := time.Now() + + // 用于记录最终状态码的变量 + var finalStatusCode string + var finalError error + + // 使用 defer 确保在所有情况下都记录指标 + defer func() { + // 如果没有设置状态码，说明是正常流程 + if finalStatusCode == "" { + finalStatusCode = "success" + } + + metrics.IncrementBatchRequestCounter( + fmt.Sprintf("%d", info.ChannelId), + info.ChannelName, + info.ChannelTag, + info.BaseUrl, + info.UpstreamModelName, + info.Group, + finalStatusCode, + retryHeaderStatus, + 1, + ) + metrics.ObserveBatchRequestDuration( + fmt.Sprintf("%d", info.ChannelId), + info.ChannelName, + info.ChannelTag, + info.BaseUrl, + info.UpstreamModelName, + info.Group, + finalStatusCode, + retryHeaderStatus, + time.Since(requestStartTime).Seconds(), + ) + }() + + // 尝试获取分布式锁，避免重复执行 + lockKey := requestID + "_lock" + lockAcquired, err := TryAcquireLock(lockKey, DistributedLockExpiration) + if err != nil { + finalStatusCode = "lock_acquisition_error" + finalError = fmt.Errorf("failed to acquire lock: %w", err) + return nil, finalError + } + + if !lockAcquired { + finalStatusCode = "lock_already_acquired" + + // 返回内部错误响应 + errorResponse := gin.H{ + "error": gin.H{ + "message": fmt.Sprintf("request %s is already being processed, another request is in progress", requestID), + "type": "internal_error", + "code": "lock_acquisition_failed", + "request_id": requestID, + }, + } + errorJson, _ := json.Marshal(errorResponse) + response := &http.Response{ + StatusCode: dto.StatusRequestConflict, + Body: io.NopCloser(strings.NewReader(string(errorJson))), + Header: make(http.Header), + } + response.Header.Set("Content-Type", "application/json") + return response, nil + } + + // 确保在函数结束时释放锁 + defer func() { + if releaseErr := ReleaseLock(lockKey); releaseErr != nil { + common.LogError(c, fmt.Sprintf("Failed to release lock for request %s: %v", requestID, releaseErr)) + } + }() + + // 检查是否为重试请求 + if isRetryRequest(c) { + // 应用重试响应延迟 + retryDelay := calculateRetryResponseDelay() + if retryDelay > 0 { + common.LogInfo(c.Request.Context(), fmt.Sprintf("Applying retry response delay: %v for request %s", retryDelay, requestID)) + time.Sleep(retryDelay) + } + + // 从Redis获取结果，使用当前的requestID（可能是retry_request_id） + resultData, err := GetBatchResultFromRedis(requestID) + if err == nil { + // 检查Result是否为空且状态为pending，这种情况说明第一次请求可能超时了 + if resultData.Result == "" && resultData.Status == "pending" { + finalStatusCode = "retry_pending" + common.LogInfo(c.Request.Context(), fmt.Sprintf("Found pending request for %s, returning retry response", requestID)) + + // 返回重试提示 + errorResponse := gin.H{ + "error": gin.H{ + "message": "Request is still being processed, please retry later", + "type": "request_in_progress", + "code": "request_still_processing", + "request_id": requestID, + }, + } + errorJson, _ := json.Marshal(errorResponse) + + response := &http.Response{ + StatusCode: dto.StatusNewAPIBatchSubmitted, // 203 - 批量请求已提交，需要重试获取结果 + Body: io.NopCloser(strings.NewReader(string(errorJson))), + Header: make(http.Header), + } + response.Header.Set("Content-Type", "application/json") + response.Header.Set("Retry_request_id", requestID) + // 添加建议重试时间header + avgDuration := GetBatchRequestAverageDuration() + response.Header.Set("X-Suggested-Retry-After", fmt.Sprintf("%.0f", avgDuration)) + c.Writer.Header().Set("Retry_request_id", requestID) + c.Writer.Header().Set("X-Suggested-Retry-After", fmt.Sprintf("%.0f", avgDuration)) + return response, nil + } else if resultData.Result != "" { + finalStatusCode = "retry_cache_hit" + // 只有在Result不为空时才处理缓存结果 + // 删除Redis中的key + err = DeleteBatchResultFromRedis(requestID) + if err != nil { + common.LogError(c, err.Error()) + } + + // 先用火山引擎格式解析，再转换为SimpleResponse + openaiResponse, err := convertVolcEngineResponseToOpenAI([]byte(resultData.Result)) + if err != nil { + finalStatusCode = "retry_cache_convert_error" + finalError = fmt.Errorf("failed to convert cached response format: %w", err) + return nil, finalError + } + + // 将转换后的结果序列化为JSON + openaiResponseJson, err := json.Marshal(openaiResponse) + if err != nil { + finalStatusCode = "retry_cache_marshal_error" + finalError = fmt.Errorf("failed to marshal cached OpenAI response: %w", err) + return nil, finalError + } + + // 找到结果，返回并删除key + response := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(string(openaiResponseJson))), + Header: make(http.Header), + } + response.Header.Set("Content-Type", "application/json") + return response, nil + } else { + // Result为空但状态不是pending，说明是错误状态 + common.LogInfo(c.Request.Context(), fmt.Sprintf("Found error status for request %s, continuing with new request", requestID)) + // 删除Redis中的key，继续执行新建流程 + err = DeleteBatchResultFromRedis(requestID) + if err != nil { + common.LogError(c, err.Error()) + } + } + } + // 如果Redis中没有找到结果，继续执行新建流程 + } + + // 尝试获取请求槽位，如果无法立即获取则等待 + if err := acquireRequestSlot(); err != nil { + // 如果无法立即获取槽位，等待可用槽位 + // 随机计算等待时间：在100ms-1000ms之间随机选择 + waitTime := calculateRateLimitWaitTime() + + ctx, cancel := context.WithTimeout(c.Request.Context(), waitTime) + defer cancel() + + if waitErr := waitForAvailableSlot(ctx); waitErr != nil { + finalStatusCode = "rate_limit_exceeded" + // 等待超时，返回自定义限流错误 + errorResponse := gin.H{ + "error": gin.H{ + "message": "Request limit reached, please retry later", + "type": "new_api_batch_rate_limit_exceeded", + "code": "new_api_batch_rate_limit_exceeded", + }, + } + errorJson, _ := json.Marshal(errorResponse) + response := &http.Response{ + StatusCode: dto.StatusNewAPIBatchRateLimitExceeded, + Body: io.NopCloser(strings.NewReader(string(errorJson))), + Header: make(http.Header), + } + response.Header.Set("Content-Type", "application/json") + return response, nil + } + } + + // 确保在函数结束时释放槽位 + defer releaseRequestSlot() + + // 解析请求体 + var request dto.GeneralOpenAIRequest + if err := json.NewDecoder(requestBody).Decode(&request); err != nil { + finalStatusCode = "request_decode_error" + finalError = fmt.Errorf("failed to decode request body: %w", err) + return nil, finalError + } + // 转换为豆包批量请求格式 + batchRequest, err := convertToBatchRequest(c.Request.Context(), &request, info.Endpoint) + if err != nil { + finalStatusCode = "request_convert_error" + finalError = fmt.Errorf("failed to convert request: %w", err) + return nil, finalError + } + + // 检查是否有未支持的参数 + checkUnsupportedParameters(&request) + + // 使用 channel ID 获取或创建客户端实例 + client := GetBatchClient(fmt.Sprintf("%d", info.ChannelId), info.ApiKey) + + // 创建带超时的 context，用于异步调用的整体超时 + timeoutDuration := getAsyncCallTimeout(c.Request.Context()) + asyncCtx, asyncCancel := context.WithTimeout(c.Request.Context(), timeoutDuration) + defer asyncCancel() + + // 创建通道用于接收异步结果 + resultChan := make(chan interface{}, 1) + errChan := make(chan error, 1) + + // 异步发起批量推理请求 + go func() { + // 使用独立的context，不受外层asyncCtx影响 + independentCtx := context.Background() + + result, err := executeBatchRequestWithRedis(independentCtx, client, batchRequest, requestID) + + if err != nil { + common.LogError(c, fmt.Sprintf("Async batch request failed for requestID %s: %v", requestID, err)) + errChan <- err + return + } + resultChan <- result + }() + + // 等待结果或超时 + var result interface{} + select { + case result = <-resultChan: + // 成功获取结果 + common.LogInfo(c, fmt.Sprintf("Received result for requestID %s", requestID)) + case err := <-errChan: + // 发生错误 + finalStatusCode = "async_request_failed" + finalError = fmt.Errorf("batch request failed: %w", err) + common.LogError(c, fmt.Sprintf("batch request failed: %v", err)) + return nil, finalError + case <-asyncCtx.Done(): + // 超时 + if asyncCtx.Err() == context.DeadlineExceeded { + finalStatusCode = "async_timeout" + common.LogError(c, fmt.Sprintf("Async call timeout after %v for requestID %s", timeoutDuration, requestID)) + c.Writer.Header().Set("Retry_request_id", requestID) + + // 返回自定义状态码表示请求已提交 + errorResponse := gin.H{ + "error": gin.H{ + "message": fmt.Sprintf("Async call timeout after %v for requestID %s, please retry later to get the result", timeoutDuration, requestID), + "type": "request_submitted", + "code": "request_submitted", + "request_id": requestID, + }, + } + errorJson, _ := json.Marshal(errorResponse) + response := &http.Response{ + StatusCode: dto.StatusNewAPIBatchSubmitted, // 203 - 批量请求已提交，需要重试获取结果 + Body: io.NopCloser(strings.NewReader(string(errorJson))), + Header: make(http.Header), + } + response.Header.Set("Content-Type", "application/json") + response.Header.Set("Retry_request_id", requestID) + // 添加建议重试时间header + avgDuration := GetBatchRequestAverageDuration() + response.Header.Set("X-Suggested-Retry-After", fmt.Sprintf("%.0f", avgDuration)) + c.Writer.Header().Set("Retry_request_id", requestID) + c.Writer.Header().Set("X-Suggested-Retry-After", fmt.Sprintf("%.0f", avgDuration)) + return response, nil + } + finalStatusCode = "async_cancelled" + finalError = fmt.Errorf("async call cancelled: %w", asyncCtx.Err()) + common.LogError(c, fmt.Sprintf("Async call cancelled for requestID %s: %w", requestID, asyncCtx.Err())) + c.Writer.Header().Set("Retry_request_id", requestID) + return nil, finalError + } + + // 将结果转换为JSON + resultJson, err := json.Marshal(result) + if err != nil { + finalStatusCode = "result_marshal_error" + finalError = fmt.Errorf("failed to marshal result: %w", err) + return nil, finalError + } + + // 将火山引擎的响应转换为标准的OpenAI格式 + openaiResponse, err := convertVolcEngineResponseToOpenAI(resultJson) + if err != nil { + finalStatusCode = "response_convert_error" + finalError = fmt.Errorf("failed to convert response format: %w", err) + return nil, finalError + } + + // 将转换后的结果序列化为JSON + openaiResponseJson, err := json.Marshal(openaiResponse) + if err != nil { + finalStatusCode = "openai_response_marshal_error" + finalError = fmt.Errorf("failed to marshal OpenAI response: %w", err) + return nil, finalError + } + + // 创建HTTP响应 + response := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(string(openaiResponseJson))), + Header: make(http.Header), + } + response.Header.Set("Content-Type", "application/json") + response.Header.Set("Retry_request_id", requestID) + return response, nil +} + +// calculateRandomWaitTime 在 MinAsyncTimeout-MaxAsyncTimeout 之间随机计算异步调用超时时间 +func calculateRandomWaitTime() time.Duration { + // 计算时间范围（毫秒） + timeRange := int(MaxAsyncTimeout.Milliseconds() - MinAsyncTimeout.Milliseconds()) + + // 如果时间范围为0或负数，直接返回MinAsyncTimeout + if timeRange <= 0 { + return MinAsyncTimeout + } + + // 生成 MinAsyncTimeout-MaxAsyncTimeout 之间的随机等待时间（毫秒） + waitMilliseconds := rand.Intn(timeRange) + int(MinAsyncTimeout.Milliseconds()) + + return time.Duration(waitMilliseconds) * time.Millisecond +} + +// calculateRateLimitWaitTime 在 MinRateLimitWaitTime-MaxRateLimitWaitTime 之间随机计算限流等待时间 +func calculateRateLimitWaitTime() time.Duration { + // 计算时间范围（毫秒） + timeRange := int(MaxRateLimitWaitTime.Milliseconds() - MinRateLimitWaitTime.Milliseconds()) + + // 如果时间范围为0或负数，直接返回MinRateLimitWaitTime + if timeRange <= 0 { + return MinRateLimitWaitTime + } + + // 生成 MinRateLimitWaitTime-MaxRateLimitWaitTime 之间的随机等待时间（毫秒） + waitMilliseconds := rand.Intn(timeRange) + int(MinRateLimitWaitTime.Milliseconds()) + + return time.Duration(waitMilliseconds) * time.Millisecond +} + +// calculateRetryResponseDelay 在 MinRetryResponseDelay-MaxRetryResponseDelay 之间随机计算重试响应延迟时间 +func calculateRetryResponseDelay() time.Duration { + // 计算时间范围（毫秒） + timeRange := int(MaxRetryResponseDelay.Milliseconds() - MinRetryResponseDelay.Milliseconds()) + + // 如果时间范围为0或负数，直接返回MinRetryResponseDelay + if timeRange <= 0 { + return MinRetryResponseDelay + } + + // 生成 MinRetryResponseDelay-MaxRetryResponseDelay 之间的随机延迟时间（毫秒） + delayMilliseconds := rand.Intn(timeRange) + int(MinRetryResponseDelay.Milliseconds()) + + return time.Duration(delayMilliseconds) * time.Millisecond +} + +// GetCurrentRequestCount 获取当前请求计数（用于监控） +func GetCurrentRequestCount() int64 { + return atomic.LoadInt64(&requestCounter) +} + +// GetRateLimiterStatus 获取限速器状态（用于监控） +func GetRateLimiterStatus() (current, capacity int) { + return len(rateLimiter), cap(rateLimiter) +} + +func MustMarshalJson(v interface{}) string { + s, _ := json.Marshal(v) + return string(s) +} + +// checkUnsupportedParameters 检查请求中是否有未支持的参数 +func checkUnsupportedParameters(request *dto.GeneralOpenAIRequest) { + var unsupportedParams []string + + // 检查 tool_choice 参数（豆包批量推理可能不支持） + if request.ToolChoice != nil { + unsupportedParams = append(unsupportedParams, "tool_choice") + } + + // 检查 n 参数（批量推理不支持，批量推理本身就是多个请求） + if request.N > 0 { + unsupportedParams = append(unsupportedParams, "n") + } + + // 检查 stream 参数（批量推理不支持流式输出） + if request.Stream { + unsupportedParams = append(unsupportedParams, "stream") + } + + // 检查 user 参数（豆包SDK支持，但批量推理可能不支持） + if request.User != "" { + unsupportedParams = append(unsupportedParams, "user") + } + + // 检查 seed 参数（豆包可能不支持） + if request.Seed != 0 { + unsupportedParams = append(unsupportedParams, "seed") + } + + // 检查 response_format 参数（豆包SDK支持，但批量推理可能不支持） + if request.ResponseFormat != nil { + unsupportedParams = append(unsupportedParams, "response_format") + } + + // 检查 stream_options 参数（豆包SDK支持，但批量推理可能不支持） + if request.StreamOptions != nil { + unsupportedParams = append(unsupportedParams, "stream_options") + } + + // 检查 functions 参数（已废弃，使用tools替代） + if request.Functions != nil { + unsupportedParams = append(unsupportedParams, "functions") + } + + // 检查其他可能不支持的参数 + if request.Prompt != nil { + unsupportedParams = append(unsupportedParams, "prompt") + } + + if request.Prefix != nil { + unsupportedParams = append(unsupportedParams, "prefix") + } + + if request.Suffix != nil { + unsupportedParams = append(unsupportedParams, "suffix") + } + + if request.Input != nil { + unsupportedParams = append(unsupportedParams, "input") + } + + if request.Instruction != "" { + unsupportedParams = append(unsupportedParams, "instruction") + } + + if request.Size != "" { + unsupportedParams = append(unsupportedParams, "size") + } + + if request.EncodingFormat != nil { + unsupportedParams = append(unsupportedParams, "encoding_format") + } + + if request.Dimensions > 0 { + unsupportedParams = append(unsupportedParams, "dimensions") + } + + if request.Modalities != nil { + unsupportedParams = append(unsupportedParams, "modalities") + } + + if request.Audio != nil { + unsupportedParams = append(unsupportedParams, "audio") + } + + if request.ExtraBody != nil { + unsupportedParams = append(unsupportedParams, "extra_body") + } + + if request.Thinking != nil { + unsupportedParams = append(unsupportedParams, "thinking") + } + + if request.ThinkingConfig != nil { + unsupportedParams = append(unsupportedParams, "thinking_config") + } + + // 如果有未支持的参数，打印错误日志 + if len(unsupportedParams) > 0 { + fmt.Printf("Error: Unsupported parameters detected in batch request: %v\n", unsupportedParams) + fmt.Printf("These parameters are not supported by VolcEngine batch inference API and will be ignored.\n") + } +} + +// convertToBatchRequest 将 OpenAI 格式的请求转换为豆包批量请求格式 +// 支持多模态消息：文本、图片、视频形式 +func convertToBatchRequest(ctx context.Context, request *dto.GeneralOpenAIRequest, endpoint string) (*model.CreateChatCompletionRequest, error) { + // 获取消息内容 + if len(request.Messages) == 0 { + return nil, fmt.Errorf("no messages found in request") + } + + // 转换消息为豆包格式，支持多模态消息 + messages := make([]*model.ChatCompletionMessage, 0, len(request.Messages)) + for i, msg := range request.Messages { + // 解析消息内容 + contentParts := msg.ParseContent() + + // 记录多模态内容信息 + if len(contentParts) > 1 { + common.SysLog(fmt.Sprintf("Processing multimodal message %d with %d content parts", i, len(contentParts))) + for j, part := range contentParts { + common.SysLog(fmt.Sprintf(" Part %d: type=%s", j, part.Type)) + } + } else if len(contentParts) == 1 { + common.SysLog(fmt.Sprintf("Processing single content message %d: type=%s", i, contentParts[0].Type)) + } + + var messageContent *model.ChatCompletionMessageContent + + if len(contentParts) == 1 && contentParts[0].Type == "text" { + // 单文本消息 + text := contentParts[0].Text + messageContent = &model.ChatCompletionMessageContent{ + StringValue: &text, + } + } else { + // 多模态消息或复杂消息 + parts := make([]*model.ChatCompletionMessageContentPart, 0, len(contentParts)) + for _, part := range contentParts { + switch part.Type { + case "text": + // 文本内容 + if part.Text != "" { + parts = append(parts, &model.ChatCompletionMessageContentPart{ + Type: "text", + Text: part.Text, + }) + common.SysLog(fmt.Sprintf("Added text part: length=%d", len(part.Text))) + } + case "image_url": + // 图片内容 - 支持URL格式和base64格式 + if imageUrl, ok := part.ImageUrl.(dto.MessageImageUrl); ok { + detail := model.ImageURLDetail(imageUrl.Detail) + + // 检查Format是否为"url"或"base64" + switch imageUrl.Format { + case "url": + parts = append(parts, &model.ChatCompletionMessageContentPart{ + Type: "image_url", + ImageURL: &model.ChatMessageImageURL{ + URL: imageUrl.Url, + Detail: detail, + }, + }) + common.SysLog(fmt.Sprintf("Added image_url part: url=%s, detail=%s", replaceBase64InURL(imageUrl.Url), imageUrl.Detail)) + case "base64": + // 处理base64格式的图片 + // 从data字段中提取base64数据 + base64Data := imageUrl.Data + if base64Data != "" { + // 根据URL或格式判断图片类型，设置正确的MIME类型 + mimeType := "image/jpeg" // 默认MIME类型 + + // 如果URL包含文件扩展名，根据扩展名判断MIME类型 + if imageUrl.Url != "" { + url := strings.ToLower(imageUrl.Url) + if strings.Contains(url, ".jpg") || strings.Contains(url, ".jpeg") { + mimeType = "image/jpeg" + } else if strings.Contains(url, ".png") { + mimeType = "image/png" + } else if strings.Contains(url, ".gif") { + mimeType = "image/gif" + } else if strings.Contains(url, ".webp") { + mimeType = "image/webp" + } else if strings.Contains(url, ".bmp") { + mimeType = "image/bmp" + } else if strings.Contains(url, ".tiff") || strings.Contains(url, ".tif") { + mimeType = "image/tiff" + } else if strings.Contains(url, ".ico") { + mimeType = "image/x-icon" + } else if strings.Contains(url, ".dib") { + mimeType = "image/bmp" + } else if strings.Contains(url, ".icns") { + mimeType = "image/icns" + } else if strings.Contains(url, ".sgi") { + mimeType = "image/sgi" + } else if strings.Contains(url, ".j2c") || strings.Contains(url, ".j2k") || strings.Contains(url, ".jp2") || strings.Contains(url, ".jpc") || strings.Contains(url, ".jpf") || strings.Contains(url, ".jpx") { + mimeType = "image/jp2" + } else if strings.Contains(url, ".heic") { + mimeType = "image/heic" + } else if strings.Contains(url, ".heif") { + mimeType = "image/heif" + } + } + + // 构造data URL格式 + dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, base64Data) + parts = append(parts, &model.ChatCompletionMessageContentPart{ + Type: "image_url", + ImageURL: &model.ChatMessageImageURL{ + URL: dataURL, + Detail: detail, + }, + }) + common.SysLog(fmt.Sprintf("Added image_url part: base64 data length=%d, detail=%s, mime_type=%s", len(base64Data), imageUrl.Detail, mimeType)) + } else { + common.SysLog(fmt.Sprintf("Skipping empty base64 image_url part")) + } + default: + // 如果不是支持的格式，跳过该部分 + common.SysLog(fmt.Sprintf("Skipping unsupported image_url part: format=%s", imageUrl.Format)) + } + } + case "video_url": + // 视频内容 - 支持URL格式和base64格式 + common.SysLog(fmt.Sprintf("Processing video_url part: InputAudio type=%T", part.InputAudio)) + if videoUrl, ok := part.InputAudio.(dto.MessageInputAudio); ok { + common.SysLog(fmt.Sprintf("Video URL details: url=%s, format=%s, fps=%f", replaceBase64InURL(videoUrl.Url), videoUrl.Format, videoUrl.Fps)) + + // 检查Format是否为"url"或"base64" + switch videoUrl.Format { + case "url": + parts = append(parts, &model.ChatCompletionMessageContentPart{ + Type: "video_url", + VideoURL: &model.ChatMessageVideoURL{ + URL: videoUrl.Url, + FPS: &videoUrl.Fps, + }, + }) + common.SysLog(fmt.Sprintf("Added video_url part: url=%s, fps=%f", replaceBase64InURL(videoUrl.Url), videoUrl.Fps)) + case "base64": + // 处理base64格式的视频 + // 从data字段中提取base64数据 + base64Data := videoUrl.Data + if base64Data != "" { + // 根据URL或格式判断视频类型，设置正确的MIME类型 + mimeType := "video/mp4" // 默认MIME类型 + + // 如果URL包含文件扩展名，根据扩展名判断MIME类型 + if videoUrl.Url != "" { + url := strings.ToLower(videoUrl.Url) + if strings.Contains(url, ".mp4") { + mimeType = "video/mp4" + } else if strings.Contains(url, ".avi") { + mimeType = "video/avi" + } else if strings.Contains(url, ".mov") { + mimeType = "video/quicktime" + } + } + + // 构造data URL格式 + dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, base64Data) + parts = append(parts, &model.ChatCompletionMessageContentPart{ + Type: "video_url", + VideoURL: &model.ChatMessageVideoURL{ + URL: dataURL, + FPS: &videoUrl.Fps, + }, + }) + common.SysLog(fmt.Sprintf("Added video_url part: base64 data length=%d, fps=%f, mime_type=%s", len(base64Data), videoUrl.Fps, mimeType)) + } else { + common.SysLog("Skipping empty base64 video_url part") + } + default: + // 如果不是支持的格式，跳过该部分 + common.SysLog(fmt.Sprintf("Skipping unsupported video_url part: format=%s", videoUrl.Format)) + } + } else { + common.SysLog(fmt.Sprintf("Failed to cast InputAudio to MessageInputAudio: %v", part.InputAudio)) + } + default: + // 对于不支持的内容类型，直接返回错误 + return nil, fmt.Errorf("unsupported content type: %s", part.Type) + } + } + if len(parts) > 0 { + messageContent = &model.ChatCompletionMessageContent{ + ListValue: parts, + } + } + } + + if messageContent != nil { + messages = append(messages, &model.ChatCompletionMessage{ + Role: msg.Role, + Content: messageContent, + }) + } + } + + // 使用JSON序列化来更好地显示messages内容 + messagesJson, _ := json.Marshal(messages) + // 替换base64数据为占位符 + replacedMessagesJson := replaceBase64InString(string(messagesJson)) + common.LogInfo(ctx, fmt.Sprintf("Messages: %s", replacedMessagesJson)) + // 转换为豆包批量请求格式 + batchRequest := model.CreateChatCompletionRequest{ + Model: endpoint, + Messages: messages, + } + + // 设置 max_tokens，只有当请求中包含时才设置 + if request.MaxTokens > 0 { + maxTokens := int(request.MaxTokens) + batchRequest.MaxTokens = &maxTokens + } + + // 设置 stop 参数，只有当请求中包含时才设置 + if request.Stop != nil { + // 类型断言处理 stop 参数 + switch stop := request.Stop.(type) { + case string: + batchRequest.Stop = []string{stop} + case []string: + batchRequest.Stop = stop + case []interface{}: + stopStrings := make([]string, 0, len(stop)) + for _, s := range stop { + if str, ok := s.(string); ok { + stopStrings = append(stopStrings, str) + } + } + batchRequest.Stop = stopStrings + } + } + + // 设置 frequency_penalty，只有当请求中包含且不为0时才设置 + if request.FrequencyPenalty != 0 { + freqPenalty := float32(request.FrequencyPenalty) + batchRequest.FrequencyPenalty = &freqPenalty + } + + // 设置 presence_penalty，只有当请求中包含且不为0时才设置 + if request.PresencePenalty != 0 { + presPenalty := float32(request.PresencePenalty) + batchRequest.PresencePenalty = &presPenalty + } + + // 设置 temperature，只有当请求中包含时才设置 + if request.Temperature != nil { + temp := float32(*request.Temperature) + batchRequest.Temperature = &temp + } + + // 设置 top_p，只有当请求中包含且不为0时才设置 + if request.TopP != 0 { + topP := float32(request.TopP) + batchRequest.TopP = &topP + } + + // 设置 logprobs，只有当请求中包含且为true时才设置 + if request.LogProbs { + batchRequest.LogProbs = &request.LogProbs + } + + // 设置 top_logprobs，只有当请求中包含且大于0时才设置 + if request.TopLogProbs > 0 { + batchRequest.TopLogProbs = &request.TopLogProbs + } + + // 设置 logit_bias，只有当请求中包含时才设置 + if len(request.LogitBias) > 0 { + batchRequest.LogitBias = request.LogitBias + } + + // 设置 tools，只有当请求中包含时才设置 + if len(request.Tools) > 0 { + tools := make([]*model.Tool, 0, len(request.Tools)) + for _, tool := range request.Tools { + // 只支持function类型的工具 + if tool.Type == "function" { + chatTool := &model.Tool{ + Type: model.ToolTypeFunction, + Function: &model.FunctionDefinition{ + Name: tool.Function.Name, + Description: tool.Function.Description, + Parameters: tool.Function.Parameters, + }, + } + tools = append(tools, chatTool) + } + } + if len(tools) > 0 { + batchRequest.Tools = tools + } + } + + return &batchRequest, nil +} + +// executeBatchRequestWithRedis 执行批量请求并保存结果到Redis +func executeBatchRequestWithRedis(ctx context.Context, client *arkruntime.Client, batchRequest *model.CreateChatCompletionRequest, requestID string) (interface{}, error) { + // 在发起请求前先创建Redis key + if err := CreateBatchRequestKey(requestID); err != nil { + common.LogError(ctx, fmt.Sprintf("Failed to create initial Redis key for request %s: %v", requestID, err)) + // 即使创建Redis key失败，也继续执行请求，只是不保存结果 + } + + // 为 CreateBatchChatCompletion 创建独立的超时 context，不复用传入的ctx + apiCtx, apiCancel := context.WithTimeout(context.Background(), getBatchCompletionTimeout(ctx)) + defer apiCancel() + + // 使用JSON序列化来更好地显示batchRequest内容 + batchRequestJson, _ := json.Marshal(batchRequest) + // 替换base64数据为占位符 + replacedBatchRequestJson := replaceBase64InString(string(batchRequestJson)) + common.LogInfo(ctx, fmt.Sprintf("Batch chat completion request: %s", replacedBatchRequestJson)) + result, err := client.CreateBatchChatCompletion(apiCtx, batchRequest) + if err != nil { + common.LogError(ctx, err.Error()) + // 检查是否是超时错误 + if apiCtx.Err() == context.DeadlineExceeded { + timeoutMsg := fmt.Sprintf("batch completion timeout after %v", getBatchCompletionTimeout(ctx)) + if saveErr := SaveBatchErrorToRedis(requestID, timeoutMsg); saveErr != nil { + fmt.Printf("Failed to save timeout error to Redis for request %s: %v\n", requestID, saveErr) + } + } else { + if saveErr := SaveBatchErrorToRedis(requestID, err.Error()); saveErr != nil { + fmt.Printf("Failed to save error to Redis for request %s: %v\n", requestID, saveErr) + } + } + return nil, err + } + // 保存成功结果到Redis，子协程独立运行 + common.LogInfo(ctx, fmt.Sprintf("Batch chat completion result: %+v", result)) + if saveErr := SaveBatchResultToRedis(requestID, result, "completed"); saveErr != nil { + fmt.Printf("Failed to save result to Redis for request %s: %v\n", requestID, saveErr) + } + + return result, nil +} + +// convertVolcEngineResponseToOpenAI 将火山引擎的响应转换为标准的OpenAI格式 +func convertVolcEngineResponseToOpenAI(resultJson []byte) (*dto.SimpleResponse, error) { + // 添加调试日志 + fmt.Printf("Original volcengine response: %s\n", string(resultJson)) + + // 直接解析为火山引擎的原始格式 + var volcResponse map[string]interface{} + if err := json.Unmarshal(resultJson, &volcResponse); err != nil { + return nil, fmt.Errorf("failed to unmarshal volcengine response: %w", err) + } + + // 提取choices和usage + choices, ok := volcResponse["choices"].([]interface{}) + if !ok { + return nil, fmt.Errorf("invalid choices field in volcengine response") + } + + usageRaw, ok := volcResponse["usage"].(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("invalid usage field in volcengine response") + } + + // 转换usage + usage := dto.Usage{} + if promptTokens, ok := usageRaw["prompt_tokens"].(float64); ok { + usage.PromptTokens = int(promptTokens) + } + if completionTokens, ok := usageRaw["completion_tokens"].(float64); ok { + usage.CompletionTokens = int(completionTokens) + } + if totalTokens, ok := usageRaw["total_tokens"].(float64); ok { + usage.TotalTokens = int(totalTokens) + } + + // 转换为标准的OpenAI格式 + openaiResponse := &dto.SimpleResponse{ + Usage: usage, + Choices: []dto.OpenAITextResponseChoice{}, + } + + // 转换choices + for _, choice := range choices { + choiceMap, ok := choice.(map[string]interface{}) + if !ok { + continue + } + + message, ok := choiceMap["message"].(map[string]interface{}) + if !ok { + continue + } + + // 处理content字段 + content := "" + if contentRaw, exists := message["content"]; exists { + switch v := contentRaw.(type) { + case string: + content = v + case []interface{}: + // 如果是数组，提取文本内容 + for _, part := range v { + if partMap, ok := part.(map[string]interface{}); ok { + if partType, ok := partMap["type"].(string); ok && partType == "text" { + if text, ok := partMap["text"].(string); ok { + content += text + } + } + } + } + } + } + + // 处理reasoning_content字段 + if reasoningContent, exists := message["reasoning_content"]; exists { + if reasoningStr, ok := reasoningContent.(string); ok && reasoningStr != "" { + content = reasoningStr + "\n" + content + } + } + + // 构建转换后的choice + index := 0 + if indexRaw, ok := choiceMap["index"].(float64); ok { + index = int(indexRaw) + } + + finishReason := "" + if finishReasonRaw, ok := choiceMap["finish_reason"].(string); ok { + finishReason = finishReasonRaw + } + + role := "" + if roleRaw, ok := message["role"].(string); ok { + role = roleRaw + } + + convertedChoice := dto.OpenAITextResponseChoice{ + Index: index, + Message: dto.Message{ + Role: role, + }, + FinishReason: finishReason, + } + convertedChoice.Message.SetStringContent(content) + + openaiResponse.Choices = append(openaiResponse.Choices, convertedChoice) + } + + // 添加调试日志 + openaiResponseJson, _ := json.Marshal(openaiResponse) + fmt.Printf("Converted OpenAI response: %s\n", string(openaiResponseJson)) + + return openaiResponse, nil +} + +// GetBatchRequestAverageDuration 获取batch请求的平均耗时（秒） +// 这个函数返回一个估算的平均耗时，用于建议重试时间 +func GetBatchRequestAverageDuration() float64 { + batchRequestAvgDurationMutex.RLock() + defer batchRequestAvgDurationMutex.RUnlock() + return batchRequestAvgDuration +} + +// SetBatchRequestAverageDuration 设置batch请求的平均耗时（秒） +func SetBatchRequestAverageDuration(duration float64) { + batchRequestAvgDurationMutex.Lock() + defer batchRequestAvgDurationMutex.Unlock() + if duration > 0 { + batchRequestAvgDuration = duration + } +} + +// InitBatchRequestAverageDuration 初始化batch请求的平均耗时 +// 从环境变量读取配置，如果没有配置则使用默认值 +func InitBatchRequestAverageDuration() { + if avgDurationStr := os.Getenv("BATCH_REQUEST_AVG_DURATION"); avgDurationStr != "" { + if avgDuration, err := strconv.ParseFloat(avgDurationStr, 64); err == nil && avgDuration > 0 { + SetBatchRequestAverageDuration(avgDuration) + } + } +} + +// replaceBase64InString 替换字符串中的base64数据为占位符 +func replaceBase64InString(input string) string { + // 替换data URL格式的base64数据 + // 匹配格式: "url": "data:video/mp4;base64,UklGRnoGAABXQVZFZm10IBAAAAABAAEAQB8AAEAfAAABAAgAZGF0YQoGAACBhYqFbF1fdJivrJBhNjVgodDbq2EcBj+a2/LDciUFLIHO8tiJNwgZaLvt559NEAxQp+PwtmMcBjiR1/LMeSwFJHfH8N2QQAoUXrTp66hVFApGn+DyvmwhBSuBzvLZiTYIG2m98OScTgwOUarm7blmGgU7k9n1unEiBC13yO/eizEIHWq+8+OWT" + // 替换为: "url": "data:video/mp4;base64,[BASE64_DATA_1896764_chars]" + + // 匹配data URL格式的正则表达式 + dataURLPattern := regexp.MustCompile(`"url":\s*"data:[^"]+;base64,[^"]+"`) + + // 替换函数 + replacer := func(match string) string { + // 提取MIME类型 + mimePattern := regexp.MustCompile(`data:([^;]+);base64,`) + mimeMatch := mimePattern.FindStringSubmatch(match) + if len(mimeMatch) > 1 { + mimeType := mimeMatch[1] + // 计算base64数据长度（大约） + base64Pattern := regexp.MustCompile(`base64,([^"]+)`) + base64Match := base64Pattern.FindStringSubmatch(match) + if len(base64Match) > 1 { + base64Data := base64Match[1] + // 计算字符数 + charCount := len(base64Data) + return fmt.Sprintf(`"url": "data:%s;base64,[BASE64_DATA_%d_chars]"`, mimeType, charCount) + } + } + return `"url": "data:image/png;base64,[BASE64_DATA_REPLACED]"` + } + + return dataURLPattern.ReplaceAllStringFunc(input, replacer) +} + +// replaceBase64InURL 替换URL中的base64数据为占位符 +func replaceBase64InURL(url string) string { + // 检查是否是data URL格式 + if strings.HasPrefix(url, "data:") && strings.Contains(url, ";base64,") { + // 提取MIME类型 + parts := strings.Split(url, ";base64,") + if len(parts) == 2 { + mimeType := strings.TrimPrefix(parts[0], "data:") + base64Data := parts[1] + // 计算字符数 + charCount := len(base64Data) + return fmt.Sprintf("data:%s;base64,[BASE64_DATA_%d_chars]", mimeType, charCount) + } + } + return url +} diff --git a/relay/channel/volcengine/keepalive.go b/relay/channel/volcengine/keepalive.go new file mode 100644 index 000000000000..79972600d998 --- /dev/null +++ b/relay/channel/volcengine/keepalive.go @@ -0,0 +1,546 @@ +package volcengine + +import ( + "context" + "fmt" + "math/rand" + "one-api/common" + "one-api/middleware" + "strings" + "sync" + "time" +) + +// KeepAliveManager 保活管理器 +type KeepAliveManager struct { + keys map[string]*KeepAliveKey + mutex sync.RWMutex + ctx context.Context + cancel context.CancelFunc + isRunning bool + interval time.Duration + expiration time.Duration +} + +// KeepAliveKey 保活key的信息 +type KeepAliveKey struct { + Key string `json:"key"` + CreatedAt time.Time `json:"created_at"` // 创建时间 + LastTouch time.Time `json:"last_touch"` // 最后触摸时间 + Expiration time.Duration `json:"expiration"` // 过期时间 + Status string `json:"status"` // active, inactive, error + ErrorCount int `json:"error_count"` // 错误计数 +} + +// 全局保活管理器实例 +var ( + keepAliveManager *KeepAliveManager + keepAliveOnce sync.Once +) + +// 默认配置 +const ( + DefaultKeepAliveInterval = 10 * time.Second // 默认保活间隔：10分钟 + DefaultKeepAliveExpiration = 10 * time.Minute // 默认过期时间：4小时 + MaxKeepAliveDuration = 10 * time.Minute // 最大保活时间：4小时 + MaxErrorCount = 5 // 最大错误次数 + KeepAliveTriggerTime = 5 * time.Minute // 保活触发时间：在key过期前5分钟开始保活 + MinKeepAliveInterval = 30 * time.Second // 随机保活时间范围 + MaxKeepAliveInterval = 2 * time.Minute // 随机保活时间范围 +) + +// GetKeepAliveManager 获取全局保活管理器实例（单例模式） +func GetKeepAliveManager() *KeepAliveManager { + keepAliveOnce.Do(func() { + keepAliveManager = NewKeepAliveManager(DefaultKeepAliveInterval, DefaultKeepAliveExpiration) + }) + return keepAliveManager +} + +// NewKeepAliveManager 创建新的保活管理器 +func NewKeepAliveManager(interval, expiration time.Duration) *KeepAliveManager { + ctx, cancel := context.WithCancel(context.Background()) + + manager := &KeepAliveManager{ + keys: make(map[string]*KeepAliveKey), + ctx: ctx, + cancel: cancel, + isRunning: false, + interval: interval, + expiration: expiration, + } + + return manager +} + +// Start 启动保活管理器 +func (kam *KeepAliveManager) Start() error { + kam.mutex.Lock() + defer kam.mutex.Unlock() + + if kam.isRunning { + return fmt.Errorf("keep-alive manager is already running") + } + + kam.isRunning = true + + // 启动保活协程 + go kam.keepAliveLoop() + + common.LogInfo(kam.ctx, fmt.Sprintf("Keep-alive manager started with interval: %v, expiration: %v", kam.interval, kam.expiration)) + return nil +} + +// Stop 停止保活管理器 +func (kam *KeepAliveManager) Stop() error { + kam.mutex.Lock() + defer kam.mutex.Unlock() + + if !kam.isRunning { + return fmt.Errorf("keep-alive manager is not running") + } + + kam.isRunning = false + kam.cancel() + + common.LogInfo(kam.ctx, "Keep-alive manager stopped") + return nil +} + +// AddKey 添加key到保活列表 +func (kam *KeepAliveManager) AddKey(key string, expiration time.Duration) error { + if key == "" { + return fmt.Errorf("key cannot be empty") + } + + if expiration <= 0 { + expiration = kam.expiration + } + + kam.mutex.Lock() + defer kam.mutex.Unlock() + + // 检查key是否已存在 + if _, exists := kam.keys[key]; exists { + return fmt.Errorf("key %s already exists in keep-alive list", key) + } + + // 创建新的保活key + keepAliveKey := &KeepAliveKey{ + Key: key, + CreatedAt: time.Now(), + LastTouch: time.Now(), + Expiration: expiration, + Status: "active", + ErrorCount: 0, + } + + kam.keys[key] = keepAliveKey + + common.LogInfo(kam.ctx, fmt.Sprintf("Added key %s to keep-alive list with expiration: %v", key, expiration)) + return nil +} + +// RemoveKey 从保活列表中移除key +func (kam *KeepAliveManager) RemoveKey(key string) error { + kam.mutex.Lock() + defer kam.mutex.Unlock() + + if _, exists := kam.keys[key]; !exists { + return fmt.Errorf("key %s not found in keep-alive list", key) + } + + delete(kam.keys, key) + + common.LogInfo(kam.ctx, fmt.Sprintf("Removed key %s from keep-alive list", key)) + return nil +} + +// GetKey 获取key的信息 +func (kam *KeepAliveManager) GetKey(key string) (*KeepAliveKey, error) { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + keepAliveKey, exists := kam.keys[key] + if !exists { + return nil, fmt.Errorf("key %s not found in keep-alive list", key) + } + + return keepAliveKey, nil +} + +// GetAllKeys 获取所有保活key的信息 +func (kam *KeepAliveManager) GetAllKeys() map[string]*KeepAliveKey { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + // 创建副本以避免并发访问问题 + result := make(map[string]*KeepAliveKey) + for key, value := range kam.keys { + result[key] = &KeepAliveKey{ + Key: value.Key, + CreatedAt: value.CreatedAt, + LastTouch: value.LastTouch, + Expiration: value.Expiration, + Status: value.Status, + ErrorCount: value.ErrorCount, + } + } + + return result +} + +// GetKeyCount 获取保活key的数量 +func (kam *KeepAliveManager) GetKeyCount() int { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + return len(kam.keys) +} + +// IsRunning 检查保活管理器是否正在运行 +func (kam *KeepAliveManager) IsRunning() bool { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + return kam.isRunning +} + +// keepAliveLoop 保活循环 +func (kam *KeepAliveManager) keepAliveLoop() { + for { + select { + case <-kam.ctx.Done(): + common.LogInfo(kam.ctx, "Keep-alive loop stopped") + return + default: + // 生成随机保活间隔时间 + randomInterval := kam.generateRandomInterval() + common.LogInfo(kam.ctx, fmt.Sprintf("Next keep-alive in %v", randomInterval)) + + // 等待随机时间 + select { + case <-kam.ctx.Done(): + common.LogInfo(kam.ctx, "Keep-alive loop stopped") + return + case <-time.After(randomInterval): + kam.performKeepAlive() + } + } + } +} + +// generateRandomInterval 生成随机保活间隔时间 +func (kam *KeepAliveManager) generateRandomInterval() time.Duration { + // 计算随机秒数 + minSeconds := int(MinKeepAliveInterval.Seconds()) + maxSeconds := int(MaxKeepAliveInterval.Seconds()) + randomSeconds := minSeconds + rand.Intn(maxSeconds-minSeconds+1) + + return time.Duration(randomSeconds) * time.Second +} + +// performKeepAlive 执行保活操作 +func (kam *KeepAliveManager) performKeepAlive() { + // 使用现有的RequestId生成逻辑创建ctx + requestID := middleware.GenerateUniqueRequestId() + ctx := context.WithValue(context.Background(), common.RequestIdKey, requestID) + + kam.mutex.RLock() + keys := make([]string, 0, len(kam.keys)) + for key := range kam.keys { + keys = append(keys, key) + } + kam.mutex.RUnlock() + + // 记录本轮保活开始 + common.LogInfo(ctx, fmt.Sprintf("Keep-alive round started, total keys: %d", len(keys))) + + // 统计变量 + var ( + successCount int + removedCount int + errorCount int + skippedCount int + ) + + // 逐个处理每个key + for i, key := range keys { + keyRequestID := fmt.Sprintf("%s-key-%d", requestID, i+1) + keyCtx := context.WithValue(ctx, "key_request_id", keyRequestID) + + result := kam.touchKey(keyCtx, key) + switch result { + case "success": + successCount++ + case "removed": + removedCount++ + case "error": + errorCount++ + case "skipped": + skippedCount++ + } + } + + // 记录本轮保活结束 + common.LogInfo(ctx, fmt.Sprintf("Keep-alive round completed, success: %d, removed: %d, errors: %d, skipped: %d", + successCount, removedCount, errorCount, skippedCount)) +} + +// touchKey 触摸单个key以保持活跃（内部方法，带详细日志） +func (kam *KeepAliveManager) touchKey(ctx context.Context, key string) string { + kam.mutex.Lock() + keepAliveKey, exists := kam.keys[key] + if !exists { + kam.mutex.Unlock() + common.LogInfo(ctx, fmt.Sprintf("Key %s not found in keep-alive list", key)) + return "removed" + } + kam.mutex.Unlock() + + // 记录开始处理key + common.LogInfo(ctx, fmt.Sprintf("Processing key %s, age: %v, error_count: %d", + key, time.Since(keepAliveKey.CreatedAt), keepAliveKey.ErrorCount)) + + // 检查是否超过最大保活时间 + if time.Since(keepAliveKey.CreatedAt) > MaxKeepAliveDuration { + common.LogInfo(ctx, fmt.Sprintf("Key %s has exceeded max keep-alive duration (%v), removing from keep-alive list", + key, MaxKeepAliveDuration)) + kam.RemoveKey(key) + return "removed" + } + + // 检查是否需要保活：只在key快到期的最后5分钟进行保活 + remainingTime := keepAliveKey.Expiration - time.Since(keepAliveKey.CreatedAt) + if remainingTime > KeepAliveTriggerTime { + common.LogInfo(ctx, fmt.Sprintf("Key %s has %v remaining, skipping keep-alive (trigger time: %v)", + key, remainingTime, KeepAliveTriggerTime)) + return "skipped" + } + + // 尝试触摸key + err := kam.touchRedisKey(key, keepAliveKey.Expiration) + + kam.mutex.Lock() + defer kam.mutex.Unlock() + + if err != nil { + // 检查是否是key不存在的错误 + if strings.Contains(err.Error(), "does not exist in Redis") { + common.LogInfo(ctx, fmt.Sprintf("Key %s no longer exists in Redis, removing from keep-alive list", + key)) + delete(kam.keys, key) + return "removed" + } + + // 其他错误，增加错误计数 + keepAliveKey.ErrorCount++ + keepAliveKey.Status = "error" + + common.LogError(ctx, fmt.Sprintf("Failed to touch key %s: %v (error count: %d)", + key, err, keepAliveKey.ErrorCount)) + + // 如果错误次数超过阈值，移除key + if keepAliveKey.ErrorCount >= MaxErrorCount { + common.LogError(ctx, fmt.Sprintf("Key %s exceeded max error count, removing from keep-alive list", + key)) + delete(kam.keys, key) + return "removed" + } + + return "error" + } else { + keepAliveKey.LastTouch = time.Now() + keepAliveKey.ErrorCount = 0 + keepAliveKey.Status = "active" + + common.LogInfo(ctx, fmt.Sprintf("Successfully touched key %s, new expiration: %v", + key, keepAliveKey.Expiration)) + + return "success" + } +} + +// touchRedisKey 触摸Redis中的key +func (kam *KeepAliveManager) touchRedisKey(key string, expiration time.Duration) error { + redisClient := getRedisClient() + ctx := context.Background() + + // 检查key是否存在 + exists, err := redisClient.Exists(ctx, key).Result() + if err != nil { + return fmt.Errorf("failed to check key existence: %w", err) + } + + if exists == 0 { + return fmt.Errorf("key %s does not exist in Redis, should be removed from keep-alive list", key) + } + + // 更新key的过期时间 + err = redisClient.Expire(ctx, key, expiration).Err() + if err != nil { + return fmt.Errorf("failed to update key expiration: %w", err) + } + + return nil +} + +// CleanupExpiredKeys 清理过期的key +func (kam *KeepAliveManager) CleanupExpiredKeys() int { + kam.mutex.Lock() + defer kam.mutex.Unlock() + + now := time.Now() + removedCount := 0 + + for key, keepAliveKey := range kam.keys { + if now.Sub(keepAliveKey.LastTouch) > keepAliveKey.Expiration { + delete(kam.keys, key) + removedCount++ + common.LogInfo(kam.ctx, fmt.Sprintf("Cleaned up expired key: %s", key)) + } + } + + if removedCount > 0 { + common.LogInfo(kam.ctx, fmt.Sprintf("Cleaned up %d expired keys", removedCount)) + } + + return removedCount +} + +// GetStats 获取保活管理器的统计信息 +func (kam *KeepAliveManager) GetStats() map[string]interface{} { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + stats := make(map[string]interface{}) + stats["is_running"] = kam.isRunning + stats["total_keys"] = len(kam.keys) + stats["interval"] = kam.interval.String() + stats["expiration"] = kam.expiration.String() + stats["max_keep_alive_duration"] = MaxKeepAliveDuration.String() + + // 统计不同状态的key数量 + statusCount := make(map[string]int) + // 统计保活时间分布 + ageDistribution := make(map[string]int) + now := time.Now() + + for _, key := range kam.keys { + statusCount[key.Status]++ + + // 计算key的年龄并分类 + age := now.Sub(key.CreatedAt) + switch { + case age < 1*time.Hour: + ageDistribution["<1h"]++ + case age < 2*time.Hour: + ageDistribution["1-2h"]++ + case age < 3*time.Hour: + ageDistribution["2-3h"]++ + case age < 4*time.Hour: + ageDistribution["3-4h"]++ + default: + ageDistribution[">4h"]++ + } + } + stats["status_count"] = statusCount + stats["age_distribution"] = ageDistribution + + return stats +} + +// 便捷函数，用于快速添加batch_result类型的key +func AddBatchResultKey(requestID string) error { + manager := GetKeepAliveManager() + key := "batch_result:" + requestID + return manager.AddKey(key, DefaultKeepAliveExpiration) +} + +// 便捷函数，用于快速移除batch_result类型的key +func RemoveBatchResultKey(requestID string) error { + manager := GetKeepAliveManager() + key := "batch_result:" + requestID + return manager.RemoveKey(key) +} + +// 便捷函数，用于获取batch_result类型key的剩余保活时间 +func GetBatchResultKeyRemainingTime(requestID string) (time.Duration, error) { + manager := GetKeepAliveManager() + key := "batch_result:" + requestID + return manager.GetKeyRemainingKeepAliveTime(key) +} + +// 便捷函数，用于获取batch_result类型key的年龄 +func GetBatchResultKeyAge(requestID string) (time.Duration, error) { + manager := GetKeepAliveManager() + key := "batch_result:" + requestID + return manager.GetKeyAge(key) +} + +// InitKeepAliveManager 初始化并启动保活管理器 +func InitKeepAliveManager() error { + manager := GetKeepAliveManager() + + // 如果已经运行，直接返回 + if manager.IsRunning() { + return nil + } + + // 启动保活管理器 + if err := manager.Start(); err != nil { + return fmt.Errorf("failed to start keep-alive manager: %w", err) + } + + common.LogInfo(context.Background(), "Keep-alive manager initialized and started successfully") + return nil +} + +// ShutdownKeepAliveManager 关闭保活管理器 +func ShutdownKeepAliveManager() error { + manager := GetKeepAliveManager() + + if !manager.IsRunning() { + return nil + } + + if err := manager.Stop(); err != nil { + return fmt.Errorf("failed to stop keep-alive manager: %w", err) + } + + common.LogInfo(context.Background(), "Keep-alive manager shutdown successfully") + return nil +} + +// GetKeyRemainingKeepAliveTime 获取key的剩余保活时间 +func (kam *KeepAliveManager) GetKeyRemainingKeepAliveTime(key string) (time.Duration, error) { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + keepAliveKey, exists := kam.keys[key] + if !exists { + return 0, fmt.Errorf("key %s not found in keep-alive list", key) + } + + elapsed := time.Since(keepAliveKey.CreatedAt) + remaining := MaxKeepAliveDuration - elapsed + + if remaining <= 0 { + return 0, nil + } + + return remaining, nil +} + +// GetKeyAge 获取key的年龄（从创建到现在的时间） +func (kam *KeepAliveManager) GetKeyAge(key string) (time.Duration, error) { + kam.mutex.RLock() + defer kam.mutex.RUnlock() + + keepAliveKey, exists := kam.keys[key] + if !exists { + return 0, fmt.Errorf("key %s not found in keep-alive list", key) + } + + return time.Since(keepAliveKey.CreatedAt), nil +} diff --git a/relay/channel/volcengine/redis_client.go b/relay/channel/volcengine/redis_client.go new file mode 100644 index 000000000000..2b6b0043c426 --- /dev/null +++ b/relay/channel/volcengine/redis_client.go @@ -0,0 +1,300 @@ +package volcengine + +import ( + "context" + "encoding/json" + "fmt" + "one-api/common" + "sync" + "time" + + "github.com/go-redis/redis/v8" +) + +// Redis 客户端 +var ( + redisClient *redis.Client + redisOnce sync.Once +) + +// BatchResultData 批量推理结果数据结构 +type BatchResultData struct { + Result string `json:"result"` + Timestamp int64 `json:"timestamp"` + Status string `json:"status"` + RequestID string `json:"request_id"` + Error string `json:"error,omitempty"` +} + +// getRedisClient 获取Redis客户端实例 +func getRedisClient() *redis.Client { + redisOnce.Do(func() { + // 使用项目统一的Redis客户端 + if common.RedisEnabled && common.RDB != nil { + redisClient = common.RDB + } else { + // 如果项目Redis未启用，创建一个默认的本地Redis客户端 + redisClient = redis.NewClient(&redis.Options{ + Addr: "localhost:6379", + Password: "", + DB: 0, + }) + } + }) + return redisClient +} + +// CreateBatchRequestKey 在发起请求前预先创建Redis key +func CreateBatchRequestKey(requestID string) error { + redisClient := getRedisClient() + ctx := context.Background() + + // 设置过期时间为10分钟 + expiration := 10 * time.Minute + + // 创建初始状态的数据结构 + resultData := BatchResultData{ + Result: "", + Timestamp: time.Now().Unix(), + Status: "pending", + RequestID: requestID, + } + + resultDataJson, err := json.Marshal(resultData) + if err != nil { + return fmt.Errorf("failed to marshal initial data: %w", err) + } + + // 写入Redis + key := "batch_result:" + requestID + err = redisClient.Set(ctx, key, string(resultDataJson), expiration).Err() + if err != nil { + return fmt.Errorf("failed to create initial key in Redis: %w", err) + } + + // 将key添加到保活管理器 + if err := AddBatchResultKey(requestID); err != nil { + // 即使添加到保活管理器失败，也不影响Redis key的创建 + fmt.Printf("Warning: Failed to add key %s to keep-alive manager: %v\n", key, err) + } + + fmt.Printf("Successfully created initial Redis key for request %s with status: pending\n", requestID) + return nil +} + +// SaveBatchResultToRedis 保存批量推理结果到Redis +func SaveBatchResultToRedis(requestID string, result interface{}, status string) error { + redisClient := getRedisClient() + ctx := context.Background() + + // 设置过期时间为24小时 + expiration := 24 * time.Hour + + // 将结果转换为JSON + var resultJson []byte + var err error + + if result == nil { + // 如果结果为nil，使用空字符串 + resultJson = []byte("") + } else { + resultJson, err = json.Marshal(result) + if err != nil { + return fmt.Errorf("failed to marshal result: %w", err) + } + } + + // 创建结果数据结构 + resultData := BatchResultData{ + Result: string(resultJson), + Timestamp: time.Now().Unix(), + Status: status, + RequestID: requestID, + } + + resultDataJson, err := json.Marshal(resultData) + if err != nil { + return fmt.Errorf("failed to marshal result data: %w", err) + } + + // 写入Redis + key := "batch_result:" + requestID + err = redisClient.Set(ctx, key, string(resultDataJson), expiration).Err() + if err != nil { + return fmt.Errorf("failed to write result to Redis: %w", err) + } + + fmt.Printf("Successfully wrote result to Redis for request %s with status: %s\n", requestID, status) + return nil +} + +// SaveBatchErrorToRedis 保存批量推理错误到Redis +func SaveBatchErrorToRedis(requestID string, errorMsg string) error { + redisClient := getRedisClient() + ctx := context.Background() + + // 设置过期时间为24小时 + expiration := 24 * time.Hour + + // 创建错误数据结构 + resultData := BatchResultData{ + Result: "", + Timestamp: time.Now().Unix(), + Status: "error", + RequestID: requestID, + Error: errorMsg, + } + + resultDataJson, err := json.Marshal(resultData) + if err != nil { + return fmt.Errorf("failed to marshal error data: %w", err) + } + + // 写入Redis + key := "batch_result:" + requestID + err = redisClient.Set(ctx, key, string(resultDataJson), expiration).Err() + if err != nil { + return fmt.Errorf("failed to write error to Redis: %w", err) + } + + fmt.Printf("Successfully wrote error to Redis for request %s\n", requestID) + return nil +} + +// GetBatchResultFromRedis 从Redis获取批量推理结果 +func GetBatchResultFromRedis(requestID string) (*BatchResultData, error) { + redisClient := getRedisClient() + ctx := context.Background() + + key := "batch_result:" + requestID + result, err := redisClient.Get(ctx, key).Result() + if err != nil { + if err == redis.Nil { + return nil, fmt.Errorf("result not found for request ID: %s", requestID) + } + return nil, fmt.Errorf("failed to get result from Redis: %w", err) + } + + var resultData BatchResultData + if err := json.Unmarshal([]byte(result), &resultData); err != nil { + return nil, fmt.Errorf("failed to unmarshal result data: %w", err) + } + + return &resultData, nil +} + +// DeleteBatchResultFromRedis 从Redis删除批量推理结果 +func DeleteBatchResultFromRedis(requestID string) error { + redisClient := getRedisClient() + ctx := context.Background() + + key := "batch_result:" + requestID + err := redisClient.Del(ctx, key).Err() + if err != nil { + return fmt.Errorf("failed to delete result from Redis: %w", err) + } + + fmt.Printf("Successfully deleted result from Redis for request %s\n", requestID) + return nil +} + +// ListBatchResultsFromRedis 列出所有批量推理结果 +func ListBatchResultsFromRedis() ([]string, error) { + redisClient := getRedisClient() + ctx := context.Background() + + pattern := "batch_result:*" + keys, err := redisClient.Keys(ctx, pattern).Result() + if err != nil { + return nil, fmt.Errorf("failed to list keys from Redis: %w", err) + } + + return keys, nil +} + +// GetBatchResultCount 获取批量推理结果数量 +func GetBatchResultCount() (int64, error) { + redisClient := getRedisClient() + ctx := context.Background() + + pattern := "batch_result:*" + count, err := redisClient.Keys(ctx, pattern).Result() + if err != nil { + return 0, fmt.Errorf("failed to count keys from Redis: %w", err) + } + + return int64(len(count)), nil +} + +// CleanExpiredBatchResults 清理过期的批量推理结果 +func CleanExpiredBatchResults() error { + redisClient := getRedisClient() + ctx := context.Background() + + // 获取所有批量推理结果的key + keys, err := ListBatchResultsFromRedis() + if err != nil { + return fmt.Errorf("failed to list keys: %w", err) + } + + // 检查每个key的TTL，如果小于等于0则删除 + for _, key := range keys { + ttl, err := redisClient.TTL(ctx, key).Result() + if err != nil { + fmt.Printf("Failed to get TTL for key %s: %v\n", key, err) + continue + } + + if ttl <= 0 { + err := redisClient.Del(ctx, key).Err() + if err != nil { + fmt.Printf("Failed to delete expired key %s: %v\n", key, err) + } else { + fmt.Printf("Successfully deleted expired key %s\n", key) + } + } + } + + return nil +} + +// PingRedis 测试Redis连接 +func PingRedis() error { + redisClient := getRedisClient() + ctx := context.Background() + + _, err := redisClient.Ping(ctx).Result() + if err != nil { + return fmt.Errorf("failed to ping Redis: %w", err) + } + + return nil +} + +// TryAcquireLock 尝试获取分布式锁 +func TryAcquireLock(lockKey string, expiration time.Duration) (bool, error) { + redisClient := getRedisClient() + ctx := context.Background() + + // 使用SET命令的NX和EX选项实现分布式锁 + result, err := redisClient.SetNX(ctx, "lock:"+lockKey, "locked", expiration).Result() + if err != nil { + return false, fmt.Errorf("failed to acquire lock: %w", err) + } + + return result, nil +} + +// ReleaseLock 释放分布式锁 +func ReleaseLock(lockKey string) error { + redisClient := getRedisClient() + ctx := context.Background() + + // 删除锁 + err := redisClient.Del(ctx, "lock:"+lockKey).Err() + if err != nil { + return fmt.Errorf("failed to release lock: %w", err) + } + + return nil +} diff --git a/relay/channel/xai/adaptor.go b/relay/channel/xai/adaptor.go new file mode 100644 index 000000000000..baecace19246 --- /dev/null +++ b/relay/channel/xai/adaptor.go @@ -0,0 +1,104 @@ +package xai + +import ( + "errors" + "fmt" + "github.com/gin-gonic/gin" + "io" + "net/http" + "one-api/dto" + "one-api/relay/channel" + relaycommon "one-api/relay/common" + "strings" +) + +type Adaptor struct { +} + +func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) { + //TODO implement me + //panic("implement me") + return nil, errors.New("not available") +} + +func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) { + //not available + return nil, errors.New("not available") +} + +func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) { + request.Size = "" + return request, nil +} + +func (a *Adaptor) Init(info *relaycommon.RelayInfo) { +} + +func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { + return fmt.Sprintf("%s/v1/chat/completions", info.BaseUrl), nil +} + +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error { + channel.SetupApiRequestHeader(info, c, req) + req.Set("Authorization", "Bearer "+info.ApiKey) + return nil +} + +func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + if strings.HasPrefix(request.Model, "grok-3-mini") { + if request.MaxCompletionTokens == 0 && request.MaxTokens != 0 { + request.MaxCompletionTokens = request.MaxTokens + request.MaxTokens = 0 + } + if strings.HasSuffix(request.Model, "-high") { + request.ReasoningEffort = "high" + request.Model = strings.TrimSuffix(request.Model, "-high") + } else if strings.HasSuffix(request.Model, "-low") { + request.ReasoningEffort = "low" + request.Model = strings.TrimSuffix(request.Model, "-low") + } else if strings.HasSuffix(request.Model, "-medium") { + request.ReasoningEffort = "medium" + request.Model = strings.TrimSuffix(request.Model, "-medium") + } + info.ReasoningEffort = request.ReasoningEffort + info.UpstreamModelName = request.Model + } + return request, nil +} + +func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) { + //not available + return nil, errors.New("not available") +} + +func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) { + return channel.DoApiRequest(a, c, info, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) { + if info.IsStream { + err, usage = xAIStreamHandler(c, resp, info) + } else { + err, usage = xAIHandler(c, resp, info) + } + //if _, ok := usage.(*dto.Usage); ok && usage != nil { + // usage.(*dto.Usage).CompletionTokens = usage.(*dto.Usage).TotalTokens - usage.(*dto.Usage).PromptTokens + //} + + return +} + +func (a *Adaptor) GetModelList() []string { + return ModelList +} + +func (a *Adaptor) GetChannelName() string { + return ChannelName +} diff --git a/relay/channel/xai/constants.go b/relay/channel/xai/constants.go new file mode 100644 index 000000000000..685fe3bba667 --- /dev/null +++ b/relay/channel/xai/constants.go @@ -0,0 +1,18 @@ +package xai + +var ModelList = []string{ + // grok-3 + "grok-3-beta", "grok-3-mini-beta", + // grok-3 mini + "grok-3-fast-beta", "grok-3-mini-fast-beta", + // extend grok-3-mini reasoning + "grok-3-mini-beta-high", "grok-3-mini-beta-low", "grok-3-mini-beta-medium", + "grok-3-mini-fast-beta-high", "grok-3-mini-fast-beta-low", "grok-3-mini-fast-beta-medium", + // image model + "grok-2-image", + // legacy models + "grok-2", "grok-2-vision", + "grok-beta", "grok-vision-beta", +} + +var ChannelName = "xai" diff --git a/relay/channel/xai/dto.go b/relay/channel/xai/dto.go new file mode 100644 index 000000000000..7036d5f1a0df --- /dev/null +++ b/relay/channel/xai/dto.go @@ -0,0 +1,14 @@ +package xai + +import "one-api/dto" + +// ChatCompletionResponse represents the response from XAI chat completion API +type ChatCompletionResponse struct { + Id string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []dto.ChatCompletionsStreamResponseChoice + Usage *dto.Usage `json:"usage"` + SystemFingerprint string `json:"system_fingerprint"` +} diff --git a/relay/channel/xai/text.go b/relay/channel/xai/text.go new file mode 100644 index 000000000000..9e4b22256c4e --- /dev/null +++ b/relay/channel/xai/text.go @@ -0,0 +1,119 @@ +package xai + +import ( + "bytes" + "encoding/json" + "github.com/gin-gonic/gin" + "io" + "net/http" + "one-api/common" + "one-api/dto" + "one-api/relay/channel/openai" + relaycommon "one-api/relay/common" + "one-api/relay/helper" + "one-api/service" + "strings" +) + +func streamResponseXAI2OpenAI(xAIResp *dto.ChatCompletionsStreamResponse, usage *dto.Usage) *dto.ChatCompletionsStreamResponse { + if xAIResp == nil { + return nil + } + if xAIResp.Usage != nil { + xAIResp.Usage.CompletionTokens = usage.CompletionTokens + } + openAIResp := &dto.ChatCompletionsStreamResponse{ + Id: xAIResp.Id, + Object: xAIResp.Object, + Created: xAIResp.Created, + Model: xAIResp.Model, + Choices: xAIResp.Choices, + Usage: xAIResp.Usage, + } + + return openAIResp +} + +func xAIStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) { + usage := &dto.Usage{} + var responseTextBuilder strings.Builder + var toolCount int + var containStreamUsage bool + + helper.SetEventStreamHeaders(c) + + helper.StreamScannerHandler(c, resp, info, func(data string) bool { + var xAIResp *dto.ChatCompletionsStreamResponse + err := json.Unmarshal([]byte(data), &xAIResp) + if err != nil { + common.SysError("error unmarshalling stream response: " + err.Error()) + return true + } + + // 把 xAI 的usage转换为 OpenAI 的usage + if xAIResp.Usage != nil { + containStreamUsage = true + usage.PromptTokens = xAIResp.Usage.PromptTokens + usage.TotalTokens = xAIResp.Usage.TotalTokens + usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens + } + + openaiResponse := streamResponseXAI2OpenAI(xAIResp, usage) + _ = openai.ProcessStreamResponse(*openaiResponse, &responseTextBuilder, &toolCount) + err = helper.ObjectData(c, openaiResponse) + if err != nil { + common.SysError(err.Error()) + } + return true + }) + + if !containStreamUsage { + usage, _ = service.ResponseText2Usage(responseTextBuilder.String(), info.UpstreamModelName, info.PromptTokens) + usage.CompletionTokens += toolCount * 7 + } + + helper.Done(c) + err := resp.Body.Close() + if err != nil { + //return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil + common.SysError("close_response_body_failed: " + err.Error()) + } + return nil, usage +} + +func xAIHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) { + responseBody, err := io.ReadAll(resp.Body) + var response *dto.TextResponse + err = json.NewDecoder(bytes.NewReader(responseBody)).Decode(&response) + if err != nil { + common.SysError("error unmarshalling stream response: " + err.Error()) + return nil, nil + } + response.Usage.CompletionTokens = response.Usage.TotalTokens - response.Usage.PromptTokens + response.Usage.CompletionTokenDetails.TextTokens = response.Usage.CompletionTokens - response.Usage.CompletionTokenDetails.ReasoningTokens + + // new body + encodeJson, err := json.Marshal(response) + if err != nil { + common.SysError("error marshalling stream response: " + err.Error()) + return nil, nil + } + + // set new body + resp.Body = io.NopCloser(bytes.NewBuffer(encodeJson)) + + for k, v := range resp.Header { + c.Writer.Header().Set(k, v[0]) + } + c.Writer.WriteHeader(resp.StatusCode) + _, err = io.Copy(c.Writer, resp.Body) + if err != nil { + return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil + } + err = resp.Body.Close() + if err != nil { + return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil + } + + return nil, &response.Usage +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index c1d3f4a4d4a4..8834c6eb7345 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -20,6 +20,8 @@ type ThinkingContentInfo struct { type RelayInfo struct { ChannelType int ChannelId int + ChannelTag string + ChannelName string TokenId int TokenKey string UserId int @@ -43,6 +45,7 @@ type RelayInfo struct { ApiKey string Organization string BaseUrl string + Endpoint string SupportStreamOptions bool ShouldIncludeUsage bool IsModelMapped bool @@ -58,6 +61,9 @@ type RelayInfo struct { UserSetting map[string]interface{} UserEmail string UserQuota int + Direct bool + RetryCount int + Headers map[string]string ThinkingContentInfo } @@ -86,7 +92,8 @@ func GenRelayInfo(c *gin.Context) *RelayInfo { channelType := c.GetInt("channel_type") channelId := c.GetInt("channel_id") channelSetting := c.GetStringMap("channel_setting") - + channelTag := c.GetString("channel_tag") + channelName := c.GetString("channel_name") tokenId := c.GetInt("token_id") tokenKey := c.GetString("token_key") userId := c.GetInt("id") @@ -104,9 +111,12 @@ func GenRelayInfo(c *gin.Context) *RelayInfo { isFirstResponse: true, RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path), BaseUrl: c.GetString("base_url"), + Endpoint: c.GetString("endpoint"), RequestURLPath: c.Request.URL.String(), ChannelType: channelType, ChannelId: channelId, + ChannelTag: channelTag, + ChannelName: channelName, TokenId: tokenId, TokenKey: tokenKey, UserId: userId, @@ -123,11 +133,17 @@ func GenRelayInfo(c *gin.Context) *RelayInfo { ApiKey: strings.TrimPrefix(c.Request.Header.Get("Authorization"), "Bearer "), Organization: c.GetString("channel_organization"), ChannelSetting: channelSetting, + Headers: make(map[string]string), ThinkingContentInfo: ThinkingContentInfo{ IsFirstThinkingContent: true, SendLastThinkingContent: false, }, } + // 使用直连模式 + if strings.HasPrefix(c.Request.URL.Path, "/v1/messages") { + info.Direct = true + } + if strings.HasPrefix(c.Request.URL.Path, "/pg") { info.IsPlayground = true info.RequestURLPath = strings.TrimPrefix(info.RequestURLPath, "/pg") diff --git a/relay/constant/api_type.go b/relay/constant/api_type.go index 8ccfee03c1ad..77812503defb 100644 --- a/relay/constant/api_type.go +++ b/relay/constant/api_type.go @@ -31,6 +31,8 @@ const ( APITypeVolcEngine APITypeBaiduV2 APITypeOpenRouter + APITypeXai + APITypeDoubaoOffline APITypeDummy // this one is only for count, do not add any channel after this ) @@ -89,6 +91,10 @@ func ChannelType2APIType(channelType int) (int, bool) { apiType = APITypeBaiduV2 case common.ChannelTypeOpenRouter: apiType = APITypeOpenRouter + case common.ChannelTypeXai: + apiType = APITypeXai + case common.ChannelTypeDoubaoOffline: + apiType = APITypeVolcEngine } if apiType == -1 { return APITypeOpenAI, false diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go index 845166c3183d..5e042fc70147 100644 --- a/relay/constant/relay_mode.go +++ b/relay/constant/relay_mode.go @@ -44,7 +44,7 @@ const ( func Path2RelayMode(path string) int { relayMode := RelayModeUnknown - if strings.HasPrefix(path, "/v1/chat/completions") || strings.HasPrefix(path, "/pg/chat/completions") { + if strings.HasPrefix(path, "/v1/chat/completions") || strings.HasPrefix(path, "/pg/chat/completions") || strings.HasPrefix(path, "/v1/messages") { relayMode = RelayModeChatCompletions } else if strings.HasPrefix(path, "/v1/completions") { relayMode = RelayModeCompletions diff --git a/relay/helper/common.go b/relay/helper/common.go index 2a72d30a9ce3..225da677d047 100644 --- a/relay/helper/common.go +++ b/relay/helper/common.go @@ -4,11 +4,12 @@ import ( "encoding/json" "errors" "fmt" - "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "net/http" "one-api/common" "one-api/dto" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" ) func SetEventStreamHeaders(c *gin.Context) { diff --git a/relay/relay-audio.go b/relay/relay-audio.go index b77ee80e058f..456ae80876e6 100644 --- a/relay/relay-audio.go +++ b/relay/relay-audio.go @@ -3,16 +3,20 @@ package relay import ( "errors" "fmt" - "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/dto" + "one-api/metrics" relaycommon "one-api/relay/common" relayconstant "one-api/relay/constant" "one-api/relay/helper" "one-api/service" "one-api/setting" + "strconv" "strings" + "time" + + "github.com/gin-gonic/gin" ) func getAndValidAudioRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.AudioRequest, error) { @@ -54,21 +58,40 @@ func getAndValidAudioRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto. return audioRequest, nil } -func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { +func AudioInfo(c *gin.Context) (*relaycommon.RelayInfo, *dto.AudioRequest, *dto.OpenAIErrorWithStatusCode) { relayInfo := relaycommon.GenRelayInfo(c) audioRequest, err := getAndValidAudioRequest(c, relayInfo) - if err != nil { common.LogError(c, fmt.Sprintf("getAndValidAudioRequest failed: %s", err.Error())) - return service.OpenAIErrorWrapper(err, "invalid_audio_request", http.StatusBadRequest) + return nil, nil, service.OpenAIErrorWrapper(err, "invalid_audio_request", http.StatusBadRequest) } - promptTokens := 0 + return relayInfo, audioRequest, nil +} + +func AudioHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, audioRequest *dto.AudioRequest) (openaiErr *dto.OpenAIErrorWithStatusCode) { + startTime := time.Now() + var funcErr *dto.OpenAIErrorWithStatusCode + var statusCode int = -1 + metrics.IncrementRelayRequestTotalCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, audioRequest.Model, relayInfo.Group, 1) + defer func() { + if funcErr != nil { + metrics.IncrementRelayRequestFailedCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, audioRequest.Model, relayInfo.Group, strconv.Itoa(funcErr.StatusCode), 1) + } else { + metrics.IncrementRelayRequestSuccessCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, audioRequest.Model, relayInfo.Group, strconv.Itoa(statusCode), 1) + metrics.ObserveRelayRequestDuration(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, audioRequest.Model, relayInfo.Group, time.Since(startTime).Seconds()) + } + }() + var ( + err error + promptTokens = 0 + ) preConsumedTokens := common.PreConsumedQuota if relayInfo.RelayMode == relayconstant.RelayModeAudioSpeech { promptTokens, err = service.CountTTSToken(audioRequest.Input, audioRequest.Model) if err != nil { - return service.OpenAIErrorWrapper(err, "count_audio_token_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "count_audio_token_failed", http.StatusInternalServerError) + return funcErr } preConsumedTokens = promptTokens relayInfo.PromptTokens = promptTokens @@ -76,11 +99,13 @@ func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { priceData, err := helper.ModelPriceHelper(c, relayInfo, preConsumedTokens, 0) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + return funcErr } preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo) if openaiErr != nil { + funcErr = openaiErr return openaiErr } defer func() { @@ -91,25 +116,29 @@ func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { err = helper.ModelMappedHelper(c, relayInfo) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + return funcErr } audioRequest.Model = relayInfo.UpstreamModelName adaptor := GetAdaptor(relayInfo.ApiType) if adaptor == nil { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + return funcErr } adaptor.Init(relayInfo) ioReader, err := adaptor.ConvertAudioRequest(c, relayInfo, *audioRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + return funcErr } resp, err := adaptor.DoRequest(c, relayInfo, ioReader) if err != nil { - return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + return funcErr } statusCodeMappingStr := c.GetString("status_code_mapping") @@ -118,6 +147,7 @@ func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { httpResp = resp.(*http.Response) if httpResp.StatusCode != http.StatusOK { openaiErr = service.RelayErrorHandler(httpResp) + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr @@ -126,11 +156,15 @@ func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo) if openaiErr != nil { + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr } + // 设置状态码用于指标记录 + statusCode = resp.(*http.Response).StatusCode + postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") return nil diff --git a/relay/relay-image.go b/relay/relay-image.go index 90b423f97074..a638c59dfce2 100644 --- a/relay/relay-image.go +++ b/relay/relay-image.go @@ -5,17 +5,21 @@ import ( "encoding/json" "errors" "fmt" - "github.com/gin-gonic/gin" "io" "net/http" "one-api/common" "one-api/dto" + "one-api/metrics" "one-api/model" relaycommon "one-api/relay/common" "one-api/relay/helper" "one-api/service" "one-api/setting" + "strconv" "strings" + "time" + + "github.com/gin-gonic/gin" ) func getAndValidImageRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.ImageRequest, error) { @@ -70,25 +74,42 @@ func getAndValidImageRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto. return imageRequest, nil } -func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { +func ImageInfo(c *gin.Context) (*relaycommon.RelayInfo, *dto.ImageRequest, *dto.OpenAIErrorWithStatusCode) { relayInfo := relaycommon.GenRelayInfo(c) imageRequest, err := getAndValidImageRequest(c, relayInfo) if err != nil { common.LogError(c, fmt.Sprintf("getAndValidImageRequest failed: %s", err.Error())) - return service.OpenAIErrorWrapper(err, "invalid_image_request", http.StatusBadRequest) + return nil, nil, service.OpenAIErrorWrapper(err, "invalid_image_request", http.StatusBadRequest) } + return relayInfo, imageRequest, nil +} - err = helper.ModelMappedHelper(c, relayInfo) +func ImageHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, imageRequest *dto.ImageRequest) (openaiErr *dto.OpenAIErrorWithStatusCode) { + startTime := time.Now() + var funcErr *dto.OpenAIErrorWithStatusCode + var statusCode int = -1 + metrics.IncrementRelayRequestTotalCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, imageRequest.Model, relayInfo.Group, 1) + defer func() { + if funcErr != nil { + metrics.IncrementRelayRequestFailedCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, imageRequest.Model, relayInfo.Group, strconv.Itoa(funcErr.StatusCode), 1) + } else { + metrics.IncrementRelayRequestSuccessCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, imageRequest.Model, relayInfo.Group, strconv.Itoa(statusCode), 1) + metrics.ObserveRelayRequestDuration(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, imageRequest.Model, relayInfo.Group, time.Since(startTime).Seconds()) + } + }() + err := helper.ModelMappedHelper(c, relayInfo) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + return funcErr } imageRequest.Model = relayInfo.UpstreamModelName priceData, err := helper.ModelPriceHelper(c, relayInfo, 0, 0) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + return funcErr } if !priceData.UsePrice { // modelRatio 16 = modelPrice $0.04 @@ -97,6 +118,9 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { } userQuota, err := model.GetUserQuota(relayInfo.UserId, false) + if err != nil { + common.LogError(c, fmt.Sprintf("get_user_quota_failed: %s", err.Error())) + } sizeRatio := 1.0 // Size @@ -122,12 +146,14 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { quota := int(priceData.ModelPrice * priceData.GroupRatio * common.QuotaPerUnit) if userQuota-quota < 0 { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("image pre-consumed quota failed, user quota: %s, need quota: %s", common.FormatQuota(userQuota), common.FormatQuota(quota)), "insufficient_user_quota", http.StatusForbidden) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("image pre-consumed quota failed, user quota: %s, need quota: %s", common.FormatQuota(userQuota), common.FormatQuota(quota)), "insufficient_user_quota", http.StatusForbidden) + return funcErr } adaptor := GetAdaptor(relayInfo.ApiType) if adaptor == nil { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + return funcErr } adaptor.Init(relayInfo) @@ -135,12 +161,14 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { convertedRequest, err := adaptor.ConvertImageRequest(c, relayInfo, *imageRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + return funcErr } jsonData, err := json.Marshal(convertedRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + return funcErr } requestBody = bytes.NewBuffer(jsonData) @@ -148,7 +176,8 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { resp, err := adaptor.DoRequest(c, relayInfo, requestBody) if err != nil { - return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + return funcErr } var httpResp *http.Response if resp != nil { @@ -156,19 +185,24 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream") if httpResp.StatusCode != http.StatusOK { openaiErr := service.RelayErrorHandler(httpResp) + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr } } - _, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo) + _, openaiErr = adaptor.DoResponse(c, httpResp, relayInfo) if openaiErr != nil { + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr } + // 设置状态码用于指标记录 + statusCode = resp.(*http.Response).StatusCode + usage := &dto.Usage{ PromptTokens: imageRequest.N, TotalTokens: imageRequest.N, diff --git a/relay/relay-mj.go b/relay/relay-mj.go index a7018456316d..a97c8f29e18c 100644 --- a/relay/relay-mj.go +++ b/relay/relay-mj.go @@ -208,7 +208,7 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse { other := make(map[string]interface{}) other["model_price"] = modelPrice other["group_ratio"] = groupRatio - model.RecordConsumeLog(c, userId, channelId, 0, 0, modelName, tokenName, + model.RecordConsumeLog(c, userId, channelId, 0, 0, 0, modelName, tokenName, quota, logContent, tokenId, userQuota, 0, false, group, other) model.UpdateUserUsedQuotaAndRequestCount(userId, quota) channelId := c.GetInt("channel_id") @@ -510,7 +510,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons other := make(map[string]interface{}) other["model_price"] = modelPrice other["group_ratio"] = groupRatio - model.RecordConsumeLog(c, userId, channelId, 0, 0, modelName, tokenName, + model.RecordConsumeLog(c, userId, channelId, 0, 0, 0, modelName, tokenName, quota, logContent, tokenId, userQuota, 0, false, group, other) model.UpdateUserUsedQuotaAndRequestCount(userId, quota) channelId := c.GetInt("channel_id") diff --git a/relay/relay-text.go b/relay/relay-text.go index ddf6767d1756..f8a2637678f6 100644 --- a/relay/relay-text.go +++ b/relay/relay-text.go @@ -5,22 +5,25 @@ import ( "encoding/json" "errors" "fmt" - "github.com/bytedance/gopkg/util/gopool" "io" "math" "net/http" "one-api/common" "one-api/constant" "one-api/dto" + "one-api/metrics" "one-api/model" relaycommon "one-api/relay/common" relayconstant "one-api/relay/constant" "one-api/relay/helper" "one-api/service" "one-api/setting" + "strconv" "strings" "time" + "github.com/bytedance/gopkg/util/gopool" + "github.com/gin-gonic/gin" ) @@ -66,28 +69,57 @@ func getAndValidateTextRequest(c *gin.Context, relayInfo *relaycommon.RelayInfo) return textRequest, nil } -func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { - +func TextInfo(c *gin.Context) (*relaycommon.RelayInfo, *dto.GeneralOpenAIRequest, *dto.OpenAIErrorWithStatusCode) { relayInfo := relaycommon.GenRelayInfo(c) + if relayInfo.Direct { + // support claude direct + if strings.HasPrefix(relayInfo.OriginModelName, "claude") { + textRequest, err := getAndValidateDirectRequest(c, relayInfo) + if err != nil { + common.LogError(c, fmt.Sprintf("getAndValidateDirectRequest failed: %s", err.Error())) + return nil, nil, service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) + } + return relayInfo, textRequest, nil + } + } + // get & validate textRequest 获取并验证文本请求 textRequest, err := getAndValidateTextRequest(c, relayInfo) if err != nil { common.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error())) - return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) + return nil, nil, service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) } + return relayInfo, textRequest, nil +} + +func TextHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest *dto.GeneralOpenAIRequest) (openaiErr *dto.OpenAIErrorWithStatusCode) { + startTime := common.GetBeijingTime() + var funcErr *dto.OpenAIErrorWithStatusCode + var statusCode int = -1 + metrics.IncrementRelayRequestTotalCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, textRequest.Model, relayInfo.Group, 1) + defer func() { + if funcErr != nil { + metrics.IncrementRelayRequestFailedCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, textRequest.Model, relayInfo.Group, strconv.Itoa(openaiErr.StatusCode), 1) + } else { + metrics.IncrementRelayRequestSuccessCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, textRequest.Model, relayInfo.Group, strconv.Itoa(statusCode), 1) + metrics.ObserveRelayRequestDuration(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, textRequest.Model, relayInfo.Group, time.Since(startTime).Seconds()) + } + }() if setting.ShouldCheckPromptSensitive() { words, err := checkRequestSensitive(textRequest, relayInfo) if err != nil { + funcErr = service.OpenAIErrorWrapperLocal(err, "sensitive_words_detected", http.StatusBadRequest) common.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", "))) - return service.OpenAIErrorWrapperLocal(err, "sensitive_words_detected", http.StatusBadRequest) + return funcErr } } - err = helper.ModelMappedHelper(c, relayInfo) + err := helper.ModelMappedHelper(c, relayInfo) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + return funcErr } textRequest.Model = relayInfo.UpstreamModelName @@ -98,22 +130,30 @@ func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { promptTokens = value.(int) relayInfo.PromptTokens = promptTokens } else { - promptTokens, err = getPromptTokens(textRequest, relayInfo) + promptTokens, err = getPromptTokens(c, textRequest, relayInfo) // count messages token error 计算promptTokens错误 if err != nil { - return service.OpenAIErrorWrapper(err, "count_token_messages_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "count_token_messages_failed", http.StatusInternalServerError) + return funcErr } c.Set("prompt_tokens", promptTokens) } + // Record input tokens metric + tokenName := c.GetString("token_name") + userName := c.GetString("username") + metrics.IncrementInputTokens(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, textRequest.Model, relayInfo.Group, strconv.Itoa(relayInfo.UserId), userName, tokenName, float64(promptTokens)) + priceData, err := helper.ModelPriceHelper(c, relayInfo, promptTokens, int(textRequest.MaxTokens)) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + return funcErr } // pre-consume quota 预消耗配额 preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo) if openaiErr != nil { + funcErr = openaiErr return openaiErr } defer func() { @@ -145,7 +185,8 @@ func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { adaptor := GetAdaptor(relayInfo.ApiType) if adaptor == nil { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + return funcErr } adaptor.Init(relayInfo) var requestBody io.Reader @@ -162,53 +203,139 @@ func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { convertedRequest, err := adaptor.ConvertRequest(c, relayInfo, textRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + return funcErr } jsonData, err := json.Marshal(convertedRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + return funcErr } + // if len(jsonData) <= 2048 { + // common.LogInfo(c, fmt.Sprintf("========>>> request data: %s", string(jsonData))) + // } requestBody = bytes.NewBuffer(jsonData) + // 如果请求中包含 X-Test-Traffic 头，则添加到 relayInfo 中 + if c.GetHeader("X-Test-Traffic") == "true" { + relayInfo.Headers = make(map[string]string) + relayInfo.Headers["X-Test-Traffic"] = "true" + } + + // // 如果请求中包含 retry_request_id 头，则添加到 relayInfo 中 + // if retryRequestId := c.GetHeader("retry_request_id"); retryRequestId != "" { + // if relayInfo.Headers == nil { + // relayInfo.Headers = make(map[string]string) + // } + // relayInfo.Headers["retry_request_id"] = retryRequestId + // } + + // // 如果请求中包含 retry 头，则添加到 relayInfo 中 + // if retry := c.GetHeader("retry"); retry != "" { + // if relayInfo.Headers == nil { + // relayInfo.Headers = make(map[string]string) + // } + // relayInfo.Headers["retry"] = retry + // } + statusCodeMappingStr := c.GetString("status_code_mapping") var httpResp *http.Response resp, err := adaptor.DoRequest(c, relayInfo, requestBody) + if err != nil { - return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + return funcErr } if resp != nil { httpResp = resp.(*http.Response) + // // 直接设置到 gin 的响应头 + // c.Writer.Header().Set("X-Origin-User-ID", strconv.Itoa(relayInfo.UserId)) + // c.Writer.Header().Set("X-Origin-Channel-ID", strconv.Itoa(relayInfo.ChannelId)) + // c.Writer.Header().Set("X-Retry-Count", strconv.Itoa(relayInfo.RetryCount)) + // if c.GetHeader("Retry_request_id") != "" { + // c.Writer.Header().Set("Retry_request_id", c.GetHeader("Retry_request_id")) + // } relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream") if httpResp.StatusCode != http.StatusOK { + for k, v := range httpResp.Header { + if k == "Content-Length" { + continue + } + c.Writer.Header().Set(k, v[0]) + // common.LogInfo(c, fmt.Sprintf("set header %s = %s", k, v[0])) + } openaiErr = service.RelayErrorHandler(httpResp) + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr } } + // 读取响应体并创建副本 + responseBodyBytes, err := io.ReadAll(httpResp.Body) + if err != nil { + funcErr = service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) + return funcErr + } + // 为adaptor创建一个新的响应体 + httpResp.Body = io.NopCloser(bytes.NewBuffer(responseBodyBytes)) + common.LogInfo(c, fmt.Sprintf("response body: %s", string(responseBodyBytes))) usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo) if openaiErr != nil { + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) + common.LogError(c, fmt.Sprintf("doResponse failed: %+v", openaiErr)) return openaiErr } + common.LogInfo(c, fmt.Sprintf("response status code: %d, Usage: %+v", httpResp.StatusCode, usage)) + // 设置状态码用于指标记录 + statusCode = resp.(*http.Response).StatusCode + + // Store request and response data together if persistence is enabled and status code is 200 + if model.RequestPersistenceEnabled && httpResp.StatusCode == http.StatusOK && !(c.GetHeader("X-Test-Traffic") == "true") { + // 读取请求数据 + requestHeaders, _ := json.Marshal(c.Request.Header) + requestBodyBytes, _ := io.ReadAll(c.Request.Body) + c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBodyBytes)) + + // 读取响应数据 + responseHeaders, _ := json.Marshal(httpResp.Header) + + // 创建并保存记录，使用请求开始时间 + textRequest := &model.TextRequest{ + UserId: common.GetOriginUserId(c, c.GetInt("id")), + CreatedAt: common.GetBeijingTime(), + RequestId: c.GetString(common.RequestIdKey), + Model: textRequest.Model, + RequestHeaders: string(requestHeaders), + RequestBody: string(requestBodyBytes), + ResponseHeaders: string(responseHeaders), + ResponseBody: string(responseBodyBytes), + } + tableName := fmt.Sprintf("text_requests_%s", startTime.Format("20060102")) + if err := model.RequestPersistenceDB.Table(tableName).Save(textRequest).Error; err != nil { + common.SysError("failed to save text request: " + err.Error()) + } + } if strings.HasPrefix(relayInfo.OriginModelName, "gpt-4o-audio") { service.PostAudioConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") } else { postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") } + return nil } -func getPromptTokens(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) (int, error) { +func getPromptTokens(ctx *gin.Context, textRequest *dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) (int, error) { var promptTokens int var err error switch info.RelayMode { case relayconstant.RelayModeChatCompletions: - promptTokens, err = service.CountTokenChatRequest(info, *textRequest) + promptTokens, err = service.CountTokenChatRequest(ctx, info, *textRequest) case relayconstant.RelayModeCompletions: promptTokens, err = service.CountTokenInput(textRequest.Prompt, textRequest.Model) case relayconstant.RelayModeModerations: @@ -298,6 +425,11 @@ func returnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, us func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, preConsumedQuota int, userQuota int, priceData helper.PriceData, extraContent string) { + // 如果是压测流量，不记录计费日志 + if ctx.GetHeader("X-Test-Traffic") == "true" { + common.LogInfo(ctx, "test traffic detected, skipping consume log") + return + } if usage == nil { usage = &dto.Usage{ PromptTokens: relayInfo.PromptTokens, @@ -309,10 +441,13 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix() promptTokens := usage.PromptTokens cacheTokens := usage.PromptTokensDetails.CachedTokens + completionTokens := usage.CompletionTokens + thinkingTokens := usage.CompletionTokenDetails.ReasoningTokens modelName := relayInfo.OriginModelName tokenName := ctx.GetString("token_name") + userName := ctx.GetString("username") completionRatio := priceData.CompletionRatio cacheRatio := priceData.CacheRatio ratio := priceData.ModelRatio * priceData.GroupRatio @@ -320,6 +455,10 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, groupRatio := priceData.GroupRatio modelPrice := priceData.ModelPrice + if usage.CompletionTokens+usage.PromptTokens < usage.TotalTokens { + completionTokens = completionTokens + thinkingTokens + } + quota := 0 if !priceData.UsePrice { quota = (promptTokens - cacheTokens) + int(math.Round(float64(cacheTokens)*cacheRatio)) @@ -331,8 +470,8 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, } else { quota = int(modelPrice * common.QuotaPerUnit * groupRatio) } - totalTokens := promptTokens + completionTokens - + totalTokens := promptTokens + completionTokens + thinkingTokens + var logContent string if !priceData.UsePrice { logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，分组倍率 %.2f", modelRatio, completionRatio, groupRatio) @@ -375,11 +514,16 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, if extraContent != "" { logContent += ", " + extraContent } + + // Record token metrics + metrics.IncrementOutputTokens(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, modelName, relayInfo.Group, strconv.Itoa(relayInfo.UserId), userName, tokenName, float64(completionTokens)) + + if cacheTokens > 0 { + metrics.IncrementCacheHitTokens(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, modelName, relayInfo.Group, strconv.Itoa(relayInfo.UserId), userName, tokenName, float64(cacheTokens)) + } + + metrics.IncrementInferenceTokens(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, modelName, relayInfo.Group, strconv.Itoa(relayInfo.UserId), userName, tokenName, float64(thinkingTokens)) other := service.GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, cacheTokens, cacheRatio, modelPrice) - model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, promptTokens, completionTokens, logModel, + model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, promptTokens, completionTokens, thinkingTokens, logModel, tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other) - - //if quota != 0 { - // - //} } diff --git a/relay/relay_adaptor.go b/relay/relay_adaptor.go index 00cff3168b35..6245d749cb12 100644 --- a/relay/relay_adaptor.go +++ b/relay/relay_adaptor.go @@ -26,6 +26,7 @@ import ( "one-api/relay/channel/tencent" "one-api/relay/channel/vertex" "one-api/relay/channel/volcengine" + "one-api/relay/channel/xai" "one-api/relay/channel/xunfei" "one-api/relay/channel/zhipu" "one-api/relay/channel/zhipu_4v" @@ -86,6 +87,10 @@ func GetAdaptor(apiType int) channel.Adaptor { return &baidu_v2.Adaptor{} case constant.APITypeOpenRouter: return &openrouter.Adaptor{} + case constant.APITypeXai: + return &xai.Adaptor{} + case constant.APITypeDoubaoOffline: + return &volcengine.Adaptor{} } return nil } diff --git a/relay/relay_direct.go b/relay/relay_direct.go new file mode 100644 index 000000000000..7b28daf5266d --- /dev/null +++ b/relay/relay_direct.go @@ -0,0 +1,31 @@ +package relay + +import ( + "errors" + "github.com/gin-gonic/gin" + "math" + "one-api/common" + "one-api/dto" + "one-api/relay/channel/claude" + relaycommon "one-api/relay/common" + "strings" +) + +func getAndValidateDirectRequest(c *gin.Context, relayInfo *relaycommon.RelayInfo) (*dto.GeneralOpenAIRequest, error) { + if strings.HasPrefix(relayInfo.OriginModelName, "claude") { + directRequest := &claude.ClaudeRequest{} + err := common.UnmarshalBodyReusable(c, directRequest) + if err != nil { + return nil, err + } + if directRequest.MaxTokens > math.MaxInt32/2 { + return nil, errors.New("max_tokens is invalid") + } + if directRequest.Model == "" { + return nil, errors.New("model is required") + } + + return claude.ClaudeMessage2OpenAIRequest(directRequest) + } + return nil, errors.New("direct model not support") +} diff --git a/relay/relay_embedding.go b/relay/relay_embedding.go index e5bfa8636deb..8dd76439162e 100644 --- a/relay/relay_embedding.go +++ b/relay/relay_embedding.go @@ -4,14 +4,18 @@ import ( "bytes" "encoding/json" "fmt" - "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/dto" + "one-api/metrics" relaycommon "one-api/relay/common" relayconstant "one-api/relay/constant" "one-api/relay/helper" "one-api/service" + "strconv" + "time" + + "github.com/gin-gonic/gin" ) func getEmbeddingPromptToken(embeddingRequest dto.EmbeddingRequest) int { @@ -32,24 +36,42 @@ func validateEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, embed return nil } -func EmbeddingHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) { +func EmbeddingInfo(c *gin.Context) (*relaycommon.RelayInfo, *dto.EmbeddingRequest, *dto.OpenAIErrorWithStatusCode) { relayInfo := relaycommon.GenRelayInfo(c) var embeddingRequest *dto.EmbeddingRequest err := common.UnmarshalBodyReusable(c, &embeddingRequest) if err != nil { common.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error())) - return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) + return nil, nil, service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) } + return relayInfo, embeddingRequest, nil +} + +func EmbeddingHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, embeddingRequest *dto.EmbeddingRequest) (openaiErr *dto.OpenAIErrorWithStatusCode) { + startTime := time.Now() + var funcErr *dto.OpenAIErrorWithStatusCode + var statusCode int = -1 + metrics.IncrementRelayRequestTotalCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, embeddingRequest.Model, relayInfo.Group, 1) + defer func() { + if funcErr != nil { + metrics.IncrementRelayRequestFailedCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, embeddingRequest.Model, relayInfo.Group, strconv.Itoa(funcErr.StatusCode), 1) + } else { + metrics.IncrementRelayRequestSuccessCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, embeddingRequest.Model, relayInfo.Group, strconv.Itoa(statusCode), 1) + metrics.ObserveRelayRequestDuration(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, embeddingRequest.Model, relayInfo.Group, time.Since(startTime).Seconds()) + } + }() - err = validateEmbeddingRequest(c, relayInfo, *embeddingRequest) + err := validateEmbeddingRequest(c, relayInfo, *embeddingRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "invalid_embedding_request", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(err, "invalid_embedding_request", http.StatusBadRequest) + return funcErr } err = helper.ModelMappedHelper(c, relayInfo) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + return funcErr } embeddingRequest.Model = relayInfo.UpstreamModelName @@ -59,11 +81,13 @@ func EmbeddingHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) priceData, err := helper.ModelPriceHelper(c, relayInfo, promptToken, 0) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + return funcErr } // pre-consume quota 预消耗配额 preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo) if openaiErr != nil { + funcErr = openaiErr return openaiErr } defer func() { @@ -74,24 +98,28 @@ func EmbeddingHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) adaptor := GetAdaptor(relayInfo.ApiType) if adaptor == nil { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + return funcErr } adaptor.Init(relayInfo) convertedRequest, err := adaptor.ConvertEmbeddingRequest(c, relayInfo, *embeddingRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + return funcErr } jsonData, err := json.Marshal(convertedRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + return funcErr } requestBody := bytes.NewBuffer(jsonData) statusCodeMappingStr := c.GetString("status_code_mapping") resp, err := adaptor.DoRequest(c, relayInfo, requestBody) if err != nil { - return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + return funcErr } var httpResp *http.Response @@ -99,6 +127,7 @@ func EmbeddingHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) httpResp = resp.(*http.Response) if httpResp.StatusCode != http.StatusOK { openaiErr = service.RelayErrorHandler(httpResp) + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr @@ -107,10 +136,15 @@ func EmbeddingHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo) if openaiErr != nil { + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr } + + // 设置状态码用于指标记录 + statusCode = resp.(*http.Response).StatusCode + postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") return nil } diff --git a/relay/relay_rerank.go b/relay/relay_rerank.go index a37613871142..e3c3f0832953 100644 --- a/relay/relay_rerank.go +++ b/relay/relay_rerank.go @@ -4,13 +4,17 @@ import ( "bytes" "encoding/json" "fmt" - "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/dto" + "one-api/metrics" relaycommon "one-api/relay/common" "one-api/relay/helper" "one-api/service" + "strconv" + "time" + + "github.com/gin-gonic/gin" ) func getRerankPromptToken(rerankRequest dto.RerankRequest) int { @@ -24,25 +28,44 @@ func getRerankPromptToken(rerankRequest dto.RerankRequest) int { return token } -func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWithStatusCode) { +func RerankInfo(c *gin.Context) (*relaycommon.RelayInfo, *dto.RerankRequest, *dto.OpenAIErrorWithStatusCode) { relayInfo := relaycommon.GenRelayInfo(c) var rerankRequest *dto.RerankRequest err := common.UnmarshalBodyReusable(c, &rerankRequest) if err != nil { common.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error())) - return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) + return nil, nil, service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) } + return relayInfo, rerankRequest, nil +} + +func RerankHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, rerankRequest *dto.RerankRequest) (openaiErr *dto.OpenAIErrorWithStatusCode) { + startTime := time.Now() + var funcErr *dto.OpenAIErrorWithStatusCode + var statusCode int = -1 + metrics.IncrementRelayRequestTotalCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, rerankRequest.Model, relayInfo.Group, 1) + defer func() { + if funcErr != nil { + metrics.IncrementRelayRequestFailedCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, rerankRequest.Model, relayInfo.Group, strconv.Itoa(funcErr.StatusCode), 1) + } else { + metrics.IncrementRelayRequestSuccessCounter(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, rerankRequest.Model, relayInfo.Group, strconv.Itoa(statusCode), 1) + metrics.ObserveRelayRequestDuration(strconv.Itoa(relayInfo.ChannelId), relayInfo.ChannelName, relayInfo.ChannelTag, relayInfo.BaseUrl, rerankRequest.Model, relayInfo.Group, time.Since(startTime).Seconds()) + } + }() if rerankRequest.Query == "" { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("query is empty"), "invalid_query", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("query is empty"), "invalid_query", http.StatusBadRequest) + return funcErr } if len(rerankRequest.Documents) == 0 { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("documents is empty"), "invalid_documents", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("documents is empty"), "invalid_documents", http.StatusBadRequest) + return funcErr } - err = helper.ModelMappedHelper(c, relayInfo) + err := helper.ModelMappedHelper(c, relayInfo) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_mapped_error", http.StatusInternalServerError) + return funcErr } rerankRequest.Model = relayInfo.UpstreamModelName @@ -52,11 +75,13 @@ func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWith priceData, err := helper.ModelPriceHelper(c, relayInfo, promptToken, 0) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError) + return funcErr } // pre-consume quota 预消耗配额 preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo) if openaiErr != nil { + funcErr = openaiErr return openaiErr } defer func() { @@ -67,23 +92,27 @@ func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWith adaptor := GetAdaptor(relayInfo.ApiType) if adaptor == nil { - return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest) + return funcErr } adaptor.Init(relayInfo) convertedRequest, err := adaptor.ConvertRerankRequest(c, relayInfo.RelayMode, *rerankRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError) + return funcErr } jsonData, err := json.Marshal(convertedRequest) if err != nil { - return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError) + return funcErr } requestBody := bytes.NewBuffer(jsonData) statusCodeMappingStr := c.GetString("status_code_mapping") resp, err := adaptor.DoRequest(c, relayInfo, requestBody) if err != nil { - return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + funcErr = service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) + return funcErr } var httpResp *http.Response @@ -91,6 +120,7 @@ func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWith httpResp = resp.(*http.Response) if httpResp.StatusCode != http.StatusOK { openaiErr = service.RelayErrorHandler(httpResp) + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr @@ -99,10 +129,15 @@ func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWith usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo) if openaiErr != nil { + funcErr = openaiErr // reset status code 重置状态码 service.ResetStatusCode(openaiErr, statusCodeMappingStr) return openaiErr } + + // 设置状态码用于指标记录 + statusCode = resp.(*http.Response).StatusCode + postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") return nil } diff --git a/relay/relay_task.go b/relay/relay_task.go index 26874ba6ef17..ee66647eafd4 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "github.com/gin-gonic/gin" "io" "net/http" "one-api/common" @@ -17,6 +16,8 @@ import ( "one-api/service" "one-api/setting" "one-api/setting/operation_setting" + + "github.com/gin-gonic/gin" ) /* @@ -123,7 +124,7 @@ func RelayTaskSubmit(c *gin.Context, relayMode int) (taskErr *dto.TaskError) { other := make(map[string]interface{}) other["model_price"] = modelPrice other["group_ratio"] = groupRatio - model.RecordConsumeLog(c, relayInfo.UserId, relayInfo.ChannelId, 0, 0, + model.RecordConsumeLog(c, relayInfo.UserId, relayInfo.ChannelId, 0, 0, 0, modelName, tokenName, quota, logContent, relayInfo.TokenId, userQuota, 0, false, relayInfo.Group, other) model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) diff --git a/router/api-router.go b/router/api-router.go index bc3f5d9fe3e2..cdcc4d2e14e1 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -13,6 +13,7 @@ func SetApiRouter(router *gin.Engine) { apiRouter.Use(gzip.Gzip(gzip.DefaultCompression)) apiRouter.Use(middleware.GlobalAPIRateLimit()) { + apiRouter.GET("/ping", controller.Ping) apiRouter.GET("/status", controller.GetStatus) apiRouter.GET("/models", middleware.UserAuth(), controller.DashboardListModels) apiRouter.GET("/status/test", middleware.AdminAuth(), controller.TestStatus) @@ -77,6 +78,7 @@ func SetApiRouter(router *gin.Engine) { optionRoute.GET("/", controller.GetOptions) optionRoute.PUT("/", controller.UpdateOption) optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio) + optionRoute.POST("/request_log", controller.ToggleRequestLog) } channelRoute := apiRouter.Group("/channel") channelRoute.Use(middleware.AdminAuth()) @@ -113,6 +115,11 @@ func SetApiRouter(router *gin.Engine) { tokenRoute.PUT("/", controller.UpdateToken) tokenRoute.DELETE("/:id", controller.DeleteToken) } + allTokenRoute := apiRouter.Group("/alltoken") + allTokenRoute.Use(middleware.RootAuth()) + { + allTokenRoute.GET("/", controller.RootGetAllTokens) + } redemptionRoute := apiRouter.Group("/redemption") redemptionRoute.Use(middleware.AdminAuth()) { @@ -132,8 +139,16 @@ func SetApiRouter(router *gin.Engine) { logRoute.GET("/self", middleware.UserAuth(), controller.GetUserLogs) logRoute.GET("/self/search", middleware.UserAuth(), controller.SearchUserLogs) + // 用户限速配置接口 + userRateLimitRoute := apiRouter.Group("/user_rate_limit") + userRateLimitRoute.Use(middleware.TokenAuth()) + { + userRateLimitRoute.GET("/config", controller.GetSpecificUserRateLimitConfig) + } + dataRoute := apiRouter.Group("/data") dataRoute.GET("/", middleware.AdminAuth(), controller.GetAllQuotaDates) + dataRoute.GET("/billing", controller.ExportBillingExcel) dataRoute.GET("/self", middleware.UserAuth(), controller.GetUserQuotaDates) logRoute.Use(middleware.CORS()) diff --git a/router/relay-router.go b/router/relay-router.go index 32e0c682686b..54799467e1e8 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -1,15 +1,20 @@ package router import ( - "github.com/gin-gonic/gin" "one-api/controller" "one-api/middleware" "one-api/relay" + + "github.com/gin-gonic/gin" ) func SetRelayRouter(router *gin.Engine) { router.Use(middleware.CORS()) router.Use(middleware.DecompressRequestMiddleware()) + + // 添加ping路由用于测试 + router.GET("/ping", controller.Ping) + // https://platform.openai.com/docs/api-reference/introduction modelsRouter := router.Group("/v1/models") modelsRouter.Use(middleware.TokenAuth()) @@ -59,6 +64,7 @@ func SetRelayRouter(router *gin.Engine) { httpRouter.DELETE("/models/:model", controller.RelayNotImplemented) httpRouter.POST("/moderations", controller.Relay) httpRouter.POST("/rerank", controller.Relay) + httpRouter.POST("/messages", controller.Relay) } relayMjRouter := router.Group("/mj") diff --git a/service/error.go b/service/error.go index c76013496d71..564e6255a31c 100644 --- a/service/error.go +++ b/service/error.go @@ -30,8 +30,7 @@ func OpenAIErrorWrapper(err error, code string, statusCode int) *dto.OpenAIError text := err.Error() lowerText := strings.ToLower(text) if strings.Contains(lowerText, "post") || strings.Contains(lowerText, "dial") || strings.Contains(lowerText, "http") { - common.SysLog(fmt.Sprintf("error: %s", text)) - text = "请求上游地址失败" + text = fmt.Sprintf("请求上游地址失败，错误信息：%s, code is %s", text, code) } openAIError := dto.OpenAIError{ Message: text, diff --git a/service/quota.go b/service/quota.go index e4499ff9487a..41a2ae0a0f0a 100644 --- a/service/quota.go +++ b/service/quota.go @@ -3,7 +3,6 @@ package service import ( "errors" "fmt" - "github.com/bytedance/gopkg/util/gopool" "math" "one-api/common" constant2 "one-api/constant" @@ -16,6 +15,8 @@ import ( "strings" "time" + "github.com/bytedance/gopkg/util/gopool" + "github.com/gin-gonic/gin" ) @@ -170,12 +171,17 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod logContent += ", " + extraContent } other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice) - model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.InputTokens, usage.OutputTokens, logModel, + model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.InputTokens, usage.OutputTokens, usage.OutputTokenDetails.ReasoningTokens, logModel, tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other) } func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, preConsumedQuota int, userQuota int, priceData helper.PriceData, extraContent string) { + // 如果是压测流量，不记录计费日志 + if ctx.GetHeader("X-Test-Traffic") == "true" { + common.LogInfo(ctx, "test traffic detected, skipping consume log") + return + } useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix() textInputTokens := usage.PromptTokensDetails.TextTokens @@ -244,7 +250,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, logContent += ", " + extraContent } other := GenerateAudioOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice) - model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.PromptTokens, usage.CompletionTokens, logModel, + model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.PromptTokens, usage.CompletionTokens, usage.CompletionTokenDetails.ReasoningTokens, logModel, tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other) } diff --git a/service/token_counter.go b/service/token_counter.go index a6b8e86a177d..be4b356eaf72 100644 --- a/service/token_counter.go +++ b/service/token_counter.go @@ -14,6 +14,7 @@ import ( "strings" "unicode/utf8" + "github.com/gin-gonic/gin" "github.com/pkoukk/tiktoken-go" ) @@ -33,7 +34,7 @@ func InitTokenEncoders() { if err != nil { common.FatalLog(fmt.Sprintf("failed to get gpt-4o token encoder: %s", err.Error())) } - for model, _ := range operation_setting.GetDefaultModelRatioMap() { + for model := range operation_setting.GetDefaultModelRatioMap() { if strings.HasPrefix(model, "gpt-3.5") { tokenEncoderMap[model] = cl100TokenEncoder } else if strings.HasPrefix(model, "gpt-4") { @@ -162,13 +163,17 @@ func getImageToken(info *relaycommon.RelayInfo, imageUrl *dto.MessageImageUrl, m return tiles*tileTokens + baseTokens, nil } -func CountTokenChatRequest(info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) (int, error) { +func CountTokenChatRequest(ctx *gin.Context, info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) (int, error) { + if request.Messages == nil { + return 0, errors.New("messages is required") + } tkm := 0 - msgTokens, err := CountTokenMessages(info, request.Messages, request.Model, request.Stream) + msgTokens, err := CountTokenMessages(ctx, info, request.Messages, request.Model, request.Stream) if err != nil { return 0, err } tkm += msgTokens + if request.Tools != nil { openaiTools := request.Tools countStr := "" @@ -258,7 +263,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent, return textToken, audioToken, nil } -func CountTokenMessages(info *relaycommon.RelayInfo, messages []dto.Message, model string, stream bool) (int, error) { +func CountTokenMessages(ctx *gin.Context, info *relaycommon.RelayInfo, messages []dto.Message, model string, stream bool) (int, error) { //recover when panic tokenEncoder := getTokenEncoder(model) // Reference: @@ -293,7 +298,7 @@ func CountTokenMessages(info *relaycommon.RelayInfo, messages []dto.Message, mod return 0, err } tokenNum += imageTokenNum - log.Printf("image token num: %d", imageTokenNum) + common.LogInfo(ctx, fmt.Sprintf("image token num: %d", imageTokenNum)) } else if m.Type == dto.ContentTypeInputAudio { // TODO: 音频token数量计算 tokenNum += 100 diff --git a/setting/operation_setting/model-ratio.go b/setting/operation_setting/model-ratio.go index d9312e6c8818..f0a920a6f1f9 100644 --- a/setting/operation_setting/model-ratio.go +++ b/setting/operation_setting/model-ratio.go @@ -2,6 +2,7 @@ package operation_setting import ( "encoding/json" + "fmt" "one-api/common" "strings" "sync" @@ -206,6 +207,121 @@ var defaultModelRatio = map[string]float64{ "llama-3-sonar-large-32k-online": 1 / 1000 * USD, } +var newModelRation = map[string]float64{ + "gemini-2.0-flash-lite": 4, + "deepseek-v3": 4, + "deepseek-reasoner": 4, + "gemini-2.0-flash-thinking-exp-01-21": 4, + "claude-3-5-sonnet-latest": 5, + "grok-2-latest": 5, + "grok-2": 5, + "grok-2-1212": 5, + "grok-2-vision-1212": 5, + "grok-vision-beta": 3, + "grok-beta": 3, + "gpt-4o-2024-11-20": 4, + "deepseek-chat": 4, + "gemini-1.5-pro": 4, + "gemini-exp-1114": 4, + "gemini-exp-1121": 4, + "gemini-exp-1206": 4, + "gemini-1.5-flash-latest": 4, + "gemini-1.5-pro-latest": 4, + "gemini-1.5-pro-001": 4, + "gpt-4-all": 4, + "gpt-4o": 4, + "gpt-4o-2024-08-06": 4, + "gpt-4o-2024-08-06-1": 4, + "gpt-4o-mini": 4, + "gpt-4o-mini-2024-07-18": 4, + "gpt-4o-all": 3, + "net-gpt-4": 1, + "mixtral-8x7b-instruct": 1, + "Meta-Llama-3.1-405B-Instruct": 1, + "Llama-3.1-405B": 1, + "gemini-1.5-pro-exp-0801": 4, + "o1-mini": 4, + "o1-mini-2024-09-12": 4, + "o1-preview": 4, + "o1-preview-2024-09-12": 4, + "gemini-1.5-pro-exp-0827": 4, + "gemini-1.5-flash-exp-0827": 4, + "GPT-4o-TL": 3, + "gemma2-27b-it": 1, + "gemma2-9b-it": 1, + "gemini-1.5-flash-002": 4, + "gemini-1.5-flash-8b": 4, + "gemini-1.5-flash": 4, + "gemini-1.5-pro-002": 4, + "gemini-2.0-flash-thinking-exp-1219": 4, + "gemini-2.0-flash-thinking-exp": 4, + "gemini-2.0-flash-exp": 4, + "gpt-4o-realtime-preview": 4, + "deepseek-r1": 4, + "gpt-4o-realtime-preview-2024-10-01": 4, + "cursor-3-5-sonnet-20240620": 5, + "claude-3-7-sonnet-20250219": 5, + "claude-3-7-sonnet-20250219-thinking": 5, +} + +var completionRation = map[string]float64{ + + "gemini-2.0-flash-lite": 4, + "deepseek-v3": 4, + "deepseek-reasoner": 4, + "gemini-2.0-flash-thinking-exp-01-21": 4, + "claude-3-5-sonnet-latest": 5, + "grok-2-latest": 5, + "grok-2": 5, + "grok-2-1212": 5, + "grok-2-vision-1212": 5, + "grok-vision-beta": 3, + "grok-beta": 3, + "gpt-4o-2024-11-20": 4, + "deepseek-chat": 4, + "gemini-1.5-pro": 4, + "gemini-exp-1114": 4, + "gemini-exp-1121": 4, + "gemini-exp-1206": 4, + "gemini-1.5-flash-latest": 4, + "gemini-1.5-pro-latest": 4, + "gemini-1.5-pro-001": 4, + "gpt-4-all": 4, + "gpt-4o": 4, + "gpt-4o-2024-08-06": 4, + "gpt-4o-2024-08-06-1": 4, + "gpt-4o-mini": 4, + "gpt-4o-mini-2024-07-18": 4, + "gpt-4o-all": 3, + "net-gpt-4": 1, + "mixtral-8x7b-instruct": 1, + "Meta-Llama-3.1-405B-Instruct": 1, + "Llama-3.1-405B": 1, + "gemini-1.5-pro-exp-0801": 4, + "o1-mini": 4, + "o1-mini-2024-09-12": 4, + "o1-preview": 4, + "o1-preview-2024-09-12": 4, + "gemini-1.5-pro-exp-0827": 4, + "gemini-1.5-flash-exp-0827": 4, + "GPT-4o-TL": 3, + "gemma2-27b-it": 1, + "gemma2-9b-it": 1, + "gemini-1.5-flash-002": 4, + "gemini-1.5-flash-8b": 4, + "gemini-1.5-flash": 4, + "gemini-1.5-pro-002": 4, + "gemini-2.0-flash-thinking-exp-1219": 4, + "gemini-2.0-flash-thinking-exp": 4, + "gemini-2.0-flash-exp": 4, + "gpt-4o-realtime-preview": 4, + "deepseek-r1": 4, + "gpt-4o-realtime-preview-2024-10-01": 4, + "cursor-3-5-sonnet-20240620": 5, + "claude-3-7-sonnet-20250219": 5, + "claude-3-7-sonnet-20250219-thinking": 5, +} + var defaultModelPrice = map[string]float64{ "suno_music": 0.1, "suno_lyrics": 0.01, @@ -342,6 +458,10 @@ func GetDefaultModelRatioMap() map[string]float64 { return defaultModelRatio } +func GetNewModelRationMap() map[string]float64 { + return newModelRation +} + func GetCompletionRatioMap() map[string]float64 { CompletionRatioMutex.Lock() defer CompletionRatioMutex.Unlock() @@ -360,21 +480,26 @@ func CompletionRatio2JSONString() string { return string(jsonBytes) } -func UpdateCompletionRatioByJSONString(jsonStr string) error { +func UpdateCompletionRatioByJSONString(jsonStr string) (err error) { CompletionRatioMutex.Lock() defer CompletionRatioMutex.Unlock() CompletionRatio = make(map[string]float64) - return json.Unmarshal([]byte(jsonStr), &CompletionRatio) + // common.SysLog("Updating completion ratio, " + jsonStr) + err = json.Unmarshal([]byte(jsonStr), &CompletionRatio) + common.SysLog("Updated completion ratio success, " + fmt.Sprintf("%v", CompletionRatio["gemini-2.5-pro-preview-03-25"])) + return err } func GetCompletionRatio(name string) float64 { GetCompletionRatioMap() - - if strings.Contains(name, "/") { - if ratio, ok := CompletionRatio[name]; ok { - return ratio - } + if ratio, ok := CompletionRatio[name]; ok { + return ratio } + //if strings.Contains(name, "/") { + // if ratio, ok := CompletionRatio[name]; ok { + // return ratio + // } + //} lowercaseName := strings.ToLower(name) if strings.HasPrefix(name, "gpt-4-gizmo") { name = "gpt-4-gizmo-*" @@ -465,6 +590,9 @@ func GetCompletionRatio(name string) float64 { if ratio, ok := CompletionRatio[name]; ok { return ratio } + if ratio, ok := completionRation[name]; ok { + return ratio + } return 1 } diff --git a/setting/user_usable_group.go b/setting/user_usable_group.go index 7082b6836688..327a240c1221 100644 --- a/setting/user_usable_group.go +++ b/setting/user_usable_group.go @@ -50,3 +50,11 @@ func GroupInUserUsableGroups(groupName string) bool { _, ok := userUsableGroups[groupName] return ok } + +func GetGroupId(groupName string) int { + id, ok := common.Groups[groupName] + if !ok { + return -1 + } + return id +} diff --git a/web/src/components/LogsTable.js b/web/src/components/LogsTable.js index abf28297e742..8e0b4bcc8d15 100644 --- a/web/src/components/LogsTable.js +++ b/web/src/components/LogsTable.js @@ -2,7 +2,7 @@ import React, { useContext, useEffect, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { API, - copy, + copy, getTodayEndTimestamp, getTodayStartTimestamp, isAdmin, showError, @@ -499,7 +499,7 @@ const LogsTable = () => { token_name: '', model_name: '', start_timestamp: timestamp2string(getTodayStartTimestamp()), - end_timestamp: timestamp2string(now.getTime() / 1000 + 3600), + end_timestamp: timestamp2string(getTodayEndTimestamp()), channel: '', group: '', }); @@ -720,7 +720,12 @@ const LogsTable = () => { url = encodeURI(url); const res = await API.get(url); const { success, message, data } = res.data; + if (success) { + if (!data.items || data.items.length === 0) { + setLoading(false); + return; + } const newPageData = data.items; setActivePage(data.page); setPageSize(data.page_size); diff --git a/web/src/components/TokensTable.js b/web/src/components/TokensTable.js index 2d82df35cdf5..78bd297a42f1 100644 --- a/web/src/components/TokensTable.js +++ b/web/src/components/TokensTable.js @@ -5,6 +5,7 @@ import { showError, showSuccess, timestamp2string, + isRoot } from '../helpers'; import { ITEMS_PER_PAGE } from '../constants'; @@ -81,6 +82,10 @@ const TokensTable = () => { title: t('名称'), dataIndex: 'name', }, + { + title: t('创建者'), + dataIndex: 'user', + }, { title: t('状态'), dataIndex: 'status', @@ -348,7 +353,10 @@ const TokensTable = () => { ); const loadTokens = async (startIdx) => { setLoading(true); - const res = await API.get(`/api/token/?p=${startIdx}&size=${pageSize}`); + const apiUrl = isRoot() + ? `/api/alltoken/?p=${startIdx}&size=${pageSize}` + : `/api/token/?p=${startIdx}&size=${pageSize}`; + const res = await API.get(apiUrl); const { success, message, data } = res.data; if (success) { if (startIdx === 0) { diff --git a/web/src/components/UsersTable.js b/web/src/components/UsersTable.js index a1e43b455265..e436861349cc 100644 --- a/web/src/components/UsersTable.js +++ b/web/src/components/UsersTable.js @@ -264,7 +264,13 @@ const UsersTable = () => { } const loadUsers = async (startIdx, pageSize) => { - const res = await API.get(`/api/user/?p=${startIdx}&page_size=${pageSize}`); + const currentUser = JSON.parse(localStorage.getItem('user')); + let url = `/api/user/?p=${startIdx}&page_size=${pageSize}`; + if (currentUser?.role < 100) { + // 非超级管理员只能看到同组用户 + url += `&group=${currentUser.group}`; + } + const res = await API.get(url); const { success, message, data } = res.data; if (success) { const newPageData = data.items; @@ -329,20 +335,30 @@ const UsersTable = () => { const searchUsers = async (startIdx, pageSize, searchKeyword, searchGroup) => { if (searchKeyword === '' && searchGroup === '') { - // if keyword is blank, load files instead. - await loadUsers(startIdx, pageSize); - return; + // if keyword is blank, load files instead. + await loadUsers(startIdx, pageSize); + return; } setSearching(true); - const res = await API.get(`/api/user/search?keyword=${searchKeyword}&group=${searchGroup}&p=${startIdx}&page_size=${pageSize}`); + const currentUser = JSON.parse(localStorage.getItem('user')); + let url = `/api/user/search?keyword=${searchKeyword}&p=${startIdx}&page_size=${pageSize}`; + + // 如果选择了特定分组就用选择的分组，否则非超级管理员用自己的分组 + if (searchGroup) { + url += `&group=${searchGroup}`; + } else if (currentUser?.role < 100) { + url += `&group=${currentUser.group}`; + } + + const res = await API.get(url); const { success, message, data } = res.data; if (success) { - const newPageData = data.items; - setActivePage(data.page); - setUserCount(data.total); - setUserFormat(newPageData); + const newPageData = data.items; + setActivePage(data.page); + setUserCount(data.total); + setUserFormat(newPageData); } else { - showError(message); + showError(message); } setSearching(false); }; @@ -383,16 +399,21 @@ const UsersTable = () => { const fetchGroups = async () => { try { let res = await API.get(`/api/group/`); - // add 'all' option - // res.data.data.unshift('all'); if (res === undefined) { return; } + let groups = res.data.data; + const currentUser = JSON.parse(localStorage.getItem('user')); + // 如果不是超级管理员，只显示用户名完全匹配的分组 + if (currentUser?.role < 100) { + const usernamePattern = new RegExp(`^${currentUser.username}$|^${currentUser.username}_|_${currentUser.username}$|_${currentUser.username}_`, 'i'); + groups = groups.filter(group => usernamePattern.test(group)); + } setGroupOptions( - res.data.data.map((group) => ({ + groups.map((group) => ({ label: group, value: group, - })), + })) ); } catch (error) { showError(error.message); diff --git a/web/src/constants/channel.constants.js b/web/src/constants/channel.constants.js index 5738d656bb1f..60641cca44d6 100644 --- a/web/src/constants/channel.constants.js +++ b/web/src/constants/channel.constants.js @@ -108,5 +108,10 @@ export const CHANNEL_OPTIONS = [ value: 44, color: 'purple', label: '嵌入模型：MokaAI M3E' + }, + { + value: 100, + color: 'blue', + label: '豆包离线' } ]; diff --git a/web/src/helpers/render.js b/web/src/helpers/render.js index d4add44f3487..f132ddd08367 100644 --- a/web/src/helpers/render.js +++ b/web/src/helpers/render.js @@ -206,7 +206,7 @@ export function renderNumber(num) { return (num / 1000).toFixed(1) + 'k'; } else { return num; - } + } } export function renderQuotaNumberWithDigit(num, digits = 2) { diff --git a/web/src/helpers/utils.js b/web/src/helpers/utils.js index a40b2079a208..aaaab092c5c7 100644 --- a/web/src/helpers/utils.js +++ b/web/src/helpers/utils.js @@ -156,6 +156,12 @@ export function getTodayStartTimestamp() { return Math.floor(now.getTime() / 1000); } +export function getTodayEndTimestamp() { + var now = new Date(); + now.setHours(23, 59, 59, 59); + return Math.floor(now.getTime() / 1000); +} + export function timestamp2string(timestamp) { let date = new Date(timestamp * 1000); let year = date.getFullYear().toString(); diff --git a/web/src/pages/Channel/EditChannel.js b/web/src/pages/Channel/EditChannel.js index bfc611fe7fd5..824dc4898b46 100644 --- a/web/src/pages/Channel/EditChannel.js +++ b/web/src/pages/Channel/EditChannel.js @@ -77,6 +77,7 @@ const EditChannel = (props) => { openai_organization: '', max_input_tokens: 0, base_url: '', + endpoint: '', other: '', model_mapping: '', status_code_mapping: '', @@ -537,7 +538,7 @@ const EditChannel = (props) => { value={inputs.name} autoComplete="new-password" /> - {inputs.type !== 3 && inputs.type !== 8 && inputs.type !== 22 && inputs.type !== 36 && inputs.type !== 45 && ( + {inputs.type !== 3 && inputs.type !== 8 && inputs.type !== 22 && inputs.type !== 36 && inputs.type !== 45 && inputs.type !== 100 && ( <>
{t('代理站地址')}： @@ -556,6 +557,25 @@ const EditChannel = (props) => { )} + {inputs.type === 100 && ( + <> +
+ {t('接入点地址')}： +
+ + { + handleInputChange('endpoint', value); + }} + value={inputs.endpoint} + autoComplete="new-password" + /> + + + )}
{t('密钥')}：
diff --git a/web/src/pages/Detail/index.js b/web/src/pages/Detail/index.js index ce82a22efb8f..ccc579a314bf 100644 --- a/web/src/pages/Detail/index.js +++ b/web/src/pages/Detail/index.js @@ -6,6 +6,7 @@ import { VChart } from "@visactor/react-vchart"; import { API, isAdmin, + isRoot, showError, timestamp2string, timestamp2string1, @@ -31,19 +32,19 @@ const Detail = (props) => { const [styleState, styleDispatch] = useContext(StyleContext); const [inputs, setInputs] = useState({ username: '', + tokenname : '', token_name: '', model_name: '', - start_timestamp: - localStorage.getItem('data_export_default_time') === 'hour' - ? timestamp2string(now.getTime() / 1000 - 86400) - : localStorage.getItem('data_export_default_time') === 'week' - ? timestamp2string(now.getTime() / 1000 - 86400 * 30) - : timestamp2string(now.getTime() / 1000 - 86400 * 7), - end_timestamp: timestamp2string(now.getTime() / 1000 + 3600), + start_timestamp: timestamp2string( + new Date(now.getFullYear(), now.getMonth(), now.getDate(), 0, 0, 0).getTime() / 1000 + ), + end_timestamp: timestamp2string( + new Date(now.getFullYear(), now.getMonth(), now.getDate(), 23, 59, 59).getTime() / 1000 + ), channel: '', data_export_default_time: '', }); - const { username, model_name, start_timestamp, end_timestamp, channel } = + const { username, tokenname,model_name, start_timestamp, end_timestamp, channel } = inputs; const isAdminUser = isAdmin(); const initialized = useRef(false); @@ -194,16 +195,28 @@ const Detail = (props) => { setInputs((inputs) => ({ ...inputs, [name]: value })); }; + // 添加校验函数 + const validateInputs = () => { + if (!username && tokenname) { + showError('选择令牌时必须指定用户名'); + return false; + } + return true; + }; + const loadQuotaData = async () => { + if (!validateInputs()) { + return; + } setLoading(true); try { let url = ''; let localStartTimestamp = Date.parse(start_timestamp) / 1000; let localEndTimestamp = Date.parse(end_timestamp) / 1000; if (isAdminUser) { - url = `/api/data/?username=${username}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&default_time=${dataExportDefaultTime}`; + url = `/api/data/?username=${username}&token_name=${tokenname}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&default_time=${dataExportDefaultTime}`; } else { - url = `/api/data/self/?start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&default_time=${dataExportDefaultTime}`; + url = `/api/data/self/?token_name=${tokenname}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}&default_time=${dataExportDefaultTime}`; } const res = await API.get(url); const { success, message, data } = res.data; @@ -228,8 +241,53 @@ const Detail = (props) => { } }; - const refresh = async () => { - await loadQuotaData(); + const exportBillingData = async () => { + if (!validateInputs()) { + return; + } + setLoading(true); + try { + let url = ''; + let localStartTimestamp = Date.parse(start_timestamp) / 1000; + let localEndTimestamp = Date.parse(end_timestamp) / 1000; + if (isAdminUser) { + url = `/api/data/billing/?user_name=${username}&token_name=${tokenname}&start_timestamp=${localStartTimestamp}&end_timestamp=${localEndTimestamp}`; + + const response = await fetch(url); + if (!response.ok) throw new Error('导出失败'); + + const blob = await response.blob(); + const downloadUrl = window.URL.createObjectURL(blob); + const link = document.createElement('a'); + link.href = downloadUrl; + + // 从Content-Disposition获取文件名，如果没有则使用默认名称 + const filename = response.headers.get('Content-Disposition')?.split('filename=')[1] || + `billing_${timestamp2string1(localStartTimestamp, 'day')}_${timestamp2string1(localEndTimestamp, 'day')}.xlsx`; + + link.setAttribute('download', filename); + document.body.appendChild(link); + link.click(); + link.remove(); + window.URL.revokeObjectURL(downloadUrl); + } + } catch (error) { + showError(error.message); + } finally { + setLoading(false); + } + }; + + const refresh = async () => { + await loadQuotaData(); + }; + + const exportBilling = async () => { + if (isRoot()) { + await exportBillingData(); + }else{ + + } }; const initChart = async () => { @@ -455,6 +513,15 @@ const Detail = (props) => { name='username' onChange={(value) => handleInputChange(value, 'username')} /> + handleInputChange(value, 'tokenname')} + /> )} + diff --git a/web/src/pages/Token/EditToken.js b/web/src/pages/Token/EditToken.js index 8fcde7106afd..74de4c0f64e5 100644 --- a/web/src/pages/Token/EditToken.js +++ b/web/src/pages/Token/EditToken.js @@ -1,4 +1,4 @@ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useState, useContext } from 'react'; import { useNavigate } from 'react-router-dom'; import { API, @@ -54,6 +54,7 @@ const EditToken = (props) => { const [groups, setGroups] = useState([]); const navigate = useNavigate(); const { t } = useTranslation(); + const handleInputChange = (name, value) => { setInputs((inputs) => ({ ...inputs, [name]: value })); }; @@ -93,18 +94,25 @@ const EditToken = (props) => { } }; - const loadGroups = async () => { - let res = await API.get(`/api/user/self/groups`); - const { success, message, data } = res.data; - if (success) { - let localGroupOptions = Object.entries(data).map(([group, info]) => ({ - label: info.desc, - value: group, - ratio: info.ratio - })); - setGroups(localGroupOptions); - } else { - showError(t(message)); + const fetchGroups = async () => { + try { + let res = await API.get(`/api/user/self/groups`); + const { success, message, data } = res.data; + if (success) { + let groupsData = []; + for (let key in data) { + groupsData.push({ + label: `${key} (${data[key].desc}) - 倍率: ${data[key].ratio}`, + value: key, + ratio: data[key].ratio + }); + } + setGroups(groupsData); + } else { + showError(message); + } + } catch (error) { + showError(error.message); } }; @@ -140,7 +148,7 @@ const EditToken = (props) => { }); } loadModels(); - loadGroups(); + fetchGroups(); }, [isEdit]); // 新增 state 变量 tokenCount 来记录用户想要创建的令牌数量，默认为 1 @@ -197,6 +205,13 @@ const EditToken = (props) => { showError(t(message)); } } else { + // 验证分组是否已选择 + if (!inputs.group) { + showError(t('请选择令牌分组！')); + setLoading(false); + return; + } + // 处理新增多个令牌的情况 let successCount = 0; // 记录成功创建的令牌数量 for (let i = 0; i < tokenCount; i++) { @@ -438,12 +453,12 @@ const EditToken = (props) => { disabled={!model_limits_enabled} />
- {t('令牌分组，默认为用户的分组')} + {t('令牌分组（必选）')}
{groups.length > 0 ?