中文 | English
您好，你正在进行%s邮箱验证。
"+ "您的验证码为: %s
"+ "验证码 %d 分钟内有效，如果不是本人操作，请忽略。
", common.SystemName, code, common.VerificationValidMinutes) err := common.SendEmail(subject, email, content) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } func SendPasswordResetEmail(c *gin.Context) { email := c.Query("email") if err := common.Validate.Var(email, "required,email"); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无效的参数", }) return } if !model.IsEmailAlreadyTaken(email) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该邮箱地址未注册", }) return } code := common.GenerateVerificationCode(0) common.RegisterVerificationCodeWithKey(email, code, common.PasswordResetPurpose) link := fmt.Sprintf("%s/user/reset?email=%s&token=%s", setting.ServerAddress, email, code) subject := fmt.Sprintf("%s密码重置", common.SystemName) content := fmt.Sprintf("您好，你正在进行%s密码重置。
"+ "点击 此处 进行密码重置。
"+ "如果链接无法点击，请尝试点击下面的链接或将其复制到浏览器中打开：
%s
重置链接 %d 分钟内有效，如果不是本人操作，请忽略。
", common.SystemName, link, link, common.VerificationValidMinutes) err := common.SendEmail(subject, email, content) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } type PasswordResetRequest struct { Email string `json:"email"` Token string `json:"token"` } func ResetPassword(c *gin.Context) { var req PasswordResetRequest err := json.NewDecoder(c.Request.Body).Decode(&req) if req.Email == "" || req.Token == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无效的参数", }) return } if !common.VerifyCodeWithKey(req.Email, req.Token, common.PasswordResetPurpose) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "重置链接非法或已过期", }) return } password := common.GenerateVerificationCode(12) err = model.ResetUserPasswordByEmail(req.Email, password) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } common.DeleteKey(req.Email, common.PasswordResetPurpose) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", "data": password, }) return } ``` ## /controller/model.go ```go path="/controller/model.go" package controller import ( "fmt" "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/constant" "one-api/dto" "one-api/model" "one-api/relay" "one-api/relay/channel/ai360" "one-api/relay/channel/lingyiwanwu" "one-api/relay/channel/minimax" "one-api/relay/channel/moonshot" relaycommon "one-api/relay/common" relayconstant "one-api/relay/constant" ) // https://platform.openai.com/docs/api-reference/models/list var openAIModels []dto.OpenAIModels var openAIModelsMap map[string]dto.OpenAIModels var channelId2Models map[int][]string func getPermission() []dto.OpenAIModelPermission { var permission []dto.OpenAIModelPermission permission = append(permission, dto.OpenAIModelPermission{ Id: "modelperm-LwHkVFn8AcMItP432fKKDIKJ", Object: "model_permission", Created: 1626777600, AllowCreateEngine: true, AllowSampling: true, AllowLogprobs: true, AllowSearchIndices: false, AllowView: true, AllowFineTuning: false, Organization: "*", Group: nil, IsBlocking: false, }) return permission } func init() { // https://platform.openai.com/docs/models/model-endpoint-compatibility permission := getPermission() for i := 0; i < relayconstant.APITypeDummy; i++ { if i == relayconstant.APITypeAIProxyLibrary { continue } adaptor := relay.GetAdaptor(i) channelName := adaptor.GetChannelName() modelNames := adaptor.GetModelList() for _, modelName := range modelNames { openAIModels = append(openAIModels, dto.OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, OwnedBy: channelName, Permission: permission, Root: modelName, Parent: nil, }) } } for _, modelName := range ai360.ModelList { openAIModels = append(openAIModels, dto.OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, OwnedBy: ai360.ChannelName, Permission: permission, Root: modelName, Parent: nil, }) } for _, modelName := range moonshot.ModelList { openAIModels = append(openAIModels, dto.OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, OwnedBy: moonshot.ChannelName, Permission: permission, Root: modelName, Parent: nil, }) } for _, modelName := range lingyiwanwu.ModelList { openAIModels = append(openAIModels, dto.OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, OwnedBy: lingyiwanwu.ChannelName, Permission: permission, Root: modelName, Parent: nil, }) } for _, modelName := range minimax.ModelList { openAIModels = append(openAIModels, dto.OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, OwnedBy: minimax.ChannelName, Permission: permission, Root: modelName, Parent: nil, }) } for modelName, _ := range constant.MidjourneyModel2Action { openAIModels = append(openAIModels, dto.OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, OwnedBy: "midjourney", Permission: permission, Root: modelName, Parent: nil, }) } openAIModelsMap = make(map[string]dto.OpenAIModels) for _, aiModel := range openAIModels { openAIModelsMap[aiModel.Id] = aiModel } channelId2Models = make(map[int][]string) for i := 1; i <= common.ChannelTypeDummy; i++ { apiType, success := relayconstant.ChannelType2APIType(i) if !success || apiType == relayconstant.APITypeAIProxyLibrary { continue } meta := &relaycommon.RelayInfo{ChannelType: i} adaptor := relay.GetAdaptor(apiType) adaptor.Init(meta) channelId2Models[i] = adaptor.GetModelList() } } func ListModels(c *gin.Context) { userOpenAiModels := make([]dto.OpenAIModels, 0) permission := getPermission() modelLimitEnable := c.GetBool("token_model_limit_enabled") if modelLimitEnable { s, ok := c.Get("token_model_limit") var tokenModelLimit map[string]bool if ok { tokenModelLimit = s.(map[string]bool) } else { tokenModelLimit = map[string]bool{} } for allowModel, _ := range tokenModelLimit { if _, ok := openAIModelsMap[allowModel]; ok { userOpenAiModels = append(userOpenAiModels, openAIModelsMap[allowModel]) } else { userOpenAiModels = append(userOpenAiModels, dto.OpenAIModels{ Id: allowModel, Object: "model", Created: 1626777600, OwnedBy: "custom", Permission: permission, Root: allowModel, Parent: nil, }) } } } else { userId := c.GetInt("id") userGroup, err := model.GetUserGroup(userId, true) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "get user group failed", }) return } group := userGroup tokenGroup := c.GetString("token_group") if tokenGroup != "" { group = tokenGroup } models := model.GetGroupModels(group) for _, s := range models { if _, ok := openAIModelsMap[s]; ok { userOpenAiModels = append(userOpenAiModels, openAIModelsMap[s]) } else { userOpenAiModels = append(userOpenAiModels, dto.OpenAIModels{ Id: s, Object: "model", Created: 1626777600, OwnedBy: "custom", Permission: permission, Root: s, Parent: nil, }) } } } c.JSON(200, gin.H{ "success": true, "data": userOpenAiModels, }) } func ChannelListModels(c *gin.Context) { c.JSON(200, gin.H{ "success": true, "data": openAIModels, }) } func DashboardListModels(c *gin.Context) { c.JSON(200, gin.H{ "success": true, "data": channelId2Models, }) } func EnabledListModels(c *gin.Context) { c.JSON(200, gin.H{ "success": true, "data": model.GetEnabledModels(), }) } func RetrieveModel(c *gin.Context) { modelId := c.Param("model") if aiModel, ok := openAIModelsMap[modelId]; ok { c.JSON(200, aiModel) } else { openAIError := dto.OpenAIError{ Message: fmt.Sprintf("The model '%s' does not exist", modelId), Type: "invalid_request_error", Param: "model", Code: "model_not_found", } c.JSON(200, gin.H{ "error": openAIError, }) } } ``` ## /controller/oidc.go ```go path="/controller/oidc.go" package controller import ( "encoding/json" "errors" "fmt" "net/http" "net/url" "one-api/common" "one-api/model" "one-api/setting" "one-api/setting/system_setting" "strconv" "strings" "time" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" ) type OidcResponse struct { AccessToken string `json:"access_token"` IDToken string `json:"id_token"` RefreshToken string `json:"refresh_token"` TokenType string `json:"token_type"` ExpiresIn int `json:"expires_in"` Scope string `json:"scope"` } type OidcUser struct { OpenID string `json:"sub"` Email string `json:"email"` Name string `json:"name"` PreferredUsername string `json:"preferred_username"` Picture string `json:"picture"` } func getOidcUserInfoByCode(code string) (*OidcUser, error) { if code == "" { return nil, errors.New("无效的参数") } values := url.Values{} values.Set("client_id", system_setting.GetOIDCSettings().ClientId) values.Set("client_secret", system_setting.GetOIDCSettings().ClientSecret) values.Set("code", code) values.Set("grant_type", "authorization_code") values.Set("redirect_uri", fmt.Sprintf("%s/oauth/oidc", setting.ServerAddress)) formData := values.Encode() req, err := http.NewRequest("POST", system_setting.GetOIDCSettings().TokenEndpoint, strings.NewReader(formData)) if err != nil { return nil, err } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") client := http.Client{ Timeout: 5 * time.Second, } res, err := client.Do(req) if err != nil { common.SysLog(err.Error()) return nil, errors.New("无法连接至 OIDC 服务器，请稍后重试！") } defer res.Body.Close() var oidcResponse OidcResponse err = json.NewDecoder(res.Body).Decode(&oidcResponse) if err != nil { return nil, err } if oidcResponse.AccessToken == "" { common.SysError("OIDC 获取 Token 失败，请检查设置！") return nil, errors.New("OIDC 获取 Token 失败，请检查设置！") } req, err = http.NewRequest("GET", system_setting.GetOIDCSettings().UserInfoEndpoint, nil) if err != nil { return nil, err } req.Header.Set("Authorization", "Bearer "+oidcResponse.AccessToken) res2, err := client.Do(req) if err != nil { common.SysLog(err.Error()) return nil, errors.New("无法连接至 OIDC 服务器，请稍后重试！") } defer res2.Body.Close() if res2.StatusCode != http.StatusOK { common.SysError("OIDC 获取用户信息失败！请检查设置！") return nil, errors.New("OIDC 获取用户信息失败！请检查设置！") } var oidcUser OidcUser err = json.NewDecoder(res2.Body).Decode(&oidcUser) if err != nil { return nil, err } if oidcUser.OpenID == "" || oidcUser.Email == "" { common.SysError("OIDC 获取用户信息为空！请检查设置！") return nil, errors.New("OIDC 获取用户信息为空！请检查设置！") } return &oidcUser, nil } func OidcAuth(c *gin.Context) { session := sessions.Default(c) state := c.Query("state") if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "state is empty or not same", }) return } username := session.Get("username") if username != nil { OidcBind(c) return } if !system_setting.GetOIDCSettings().Enabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "管理员未开启通过 OIDC 登录以及注册", }) return } code := c.Query("code") oidcUser, err := getOidcUserInfoByCode(code) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } user := model.User{ OidcId: oidcUser.OpenID, } if model.IsOidcIdAlreadyTaken(user.OidcId) { err := user.FillUserByOidcId() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } else { if common.RegisterEnabled { user.Email = oidcUser.Email if oidcUser.PreferredUsername != "" { user.Username = oidcUser.PreferredUsername } else { user.Username = "oidc_" + strconv.Itoa(model.GetMaxUserId()+1) } if oidcUser.Name != "" { user.DisplayName = oidcUser.Name } else { user.DisplayName = "OIDC User" } err := user.Insert(0) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } else { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "管理员关闭了新用户注册", }) return } } if user.Status != common.UserStatusEnabled { c.JSON(http.StatusOK, gin.H{ "message": "用户已被封禁", "success": false, }) return } setupLogin(&user, c) } func OidcBind(c *gin.Context) { if !system_setting.GetOIDCSettings().Enabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "管理员未开启通过 OIDC 登录以及注册", }) return } code := c.Query("code") oidcUser, err := getOidcUserInfoByCode(code) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } user := model.User{ OidcId: oidcUser.OpenID, } if model.IsOidcIdAlreadyTaken(user.OidcId) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该 OIDC 账户已被绑定", }) return } session := sessions.Default(c) id := session.Get("id") // id := c.GetInt("id") // critical bug! user.Id = id.(int) err = user.FillUserById() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } user.OidcId = oidcUser.OpenID err = user.Update(false) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "bind", }) return } ``` ## /controller/option.go ```go path="/controller/option.go" package controller import ( "encoding/json" "net/http" "one-api/common" "one-api/model" "one-api/setting" "one-api/setting/system_setting" "strings" "github.com/gin-gonic/gin" ) func GetOptions(c *gin.Context) { var options []*model.Option common.OptionMapRWMutex.Lock() for k, v := range common.OptionMap { if strings.HasSuffix(k, "Token") || strings.HasSuffix(k, "Secret") || strings.HasSuffix(k, "Key") { continue } options = append(options, &model.Option{ Key: k, Value: common.Interface2String(v), }) } common.OptionMapRWMutex.Unlock() c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", "data": options, }) return } func UpdateOption(c *gin.Context) { var option model.Option err := json.NewDecoder(c.Request.Body).Decode(&option) if err != nil { c.JSON(http.StatusBadRequest, gin.H{ "success": false, "message": "无效的参数", }) return } switch option.Key { case "GitHubOAuthEnabled": if option.Value == "true" && common.GitHubClientId == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 GitHub OAuth，请先填入 GitHub Client Id 以及 GitHub Client Secret！", }) return } case "oidc.enabled": if option.Value == "true" && system_setting.GetOIDCSettings().ClientId == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 OIDC 登录，请先填入 OIDC Client Id 以及 OIDC Client Secret！", }) return } case "LinuxDOOAuthEnabled": if option.Value == "true" && common.LinuxDOClientId == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 LinuxDO OAuth，请先填入 LinuxDO Client Id 以及 LinuxDO Client Secret！", }) return } case "EmailDomainRestrictionEnabled": if option.Value == "true" && len(common.EmailDomainWhitelist) == 0 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用邮箱域名限制，请先填入限制的邮箱域名！", }) return } case "WeChatAuthEnabled": if option.Value == "true" && common.WeChatServerAddress == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用微信登录，请先填入微信登录相关配置信息！", }) return } case "TurnstileCheckEnabled": if option.Value == "true" && common.TurnstileSiteKey == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 Turnstile 校验，请先填入 Turnstile 校验相关配置信息！", }) return } case "TelegramOAuthEnabled": if option.Value == "true" && common.TelegramBotToken == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 Telegram OAuth，请先填入 Telegram Bot Token！", }) return } case "GroupRatio": err = setting.CheckGroupRatio(option.Value) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } err = model.UpdateOption(option.Key, option.Value) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } ``` ## /controller/playground.go ```go path="/controller/playground.go" package controller import ( "errors" "fmt" "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/constant" "one-api/dto" "one-api/middleware" "one-api/model" "one-api/service" "one-api/setting" "time" ) func Playground(c *gin.Context) { var openaiErr *dto.OpenAIErrorWithStatusCode defer func() { if openaiErr != nil { c.JSON(openaiErr.StatusCode, gin.H{ "error": openaiErr.Error, }) } }() useAccessToken := c.GetBool("use_access_token") if useAccessToken { openaiErr = service.OpenAIErrorWrapperLocal(errors.New("暂不支持使用 access token"), "access_token_not_supported", http.StatusBadRequest) return } playgroundRequest := &dto.PlayGroundRequest{} err := common.UnmarshalBodyReusable(c, playgroundRequest) if err != nil { openaiErr = service.OpenAIErrorWrapperLocal(err, "unmarshal_request_failed", http.StatusBadRequest) return } if playgroundRequest.Model == "" { openaiErr = service.OpenAIErrorWrapperLocal(errors.New("请选择模型"), "model_required", http.StatusBadRequest) return } c.Set("original_model", playgroundRequest.Model) group := playgroundRequest.Group userGroup := c.GetString("group") if group == "" { group = userGroup } else { if !setting.GroupInUserUsableGroups(group) && group != userGroup { openaiErr = service.OpenAIErrorWrapperLocal(errors.New("无权访问该分组"), "group_not_allowed", http.StatusForbidden) return } c.Set("group", group) } 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) openaiErr = service.OpenAIErrorWrapperLocal(errors.New(message), "get_playground_channel_failed", http.StatusInternalServerError) return } middleware.SetupContextForSelectedChannel(c, channel, playgroundRequest.Model) c.Set(constant.ContextKeyRequestStartTime, time.Now()) Relay(c) } ``` ## /controller/pricing.go ```go path="/controller/pricing.go" package controller import ( "github.com/gin-gonic/gin" "one-api/model" "one-api/setting" "one-api/setting/operation_setting" ) func GetPricing(c *gin.Context) { pricing := model.GetPricing() userId, exists := c.Get("id") usableGroup := map[string]string{} groupRatio := map[string]float64{} for s, f := range setting.GetGroupRatioCopy() { groupRatio[s] = f } var group string if exists { user, err := model.GetUserCache(userId.(int)) if err == nil { group = user.Group } } usableGroup = setting.GetUserUsableGroups(group) // check groupRatio contains usableGroup for group := range setting.GetGroupRatioCopy() { if _, ok := usableGroup[group]; !ok { delete(groupRatio, group) } } c.JSON(200, gin.H{ "success": true, "data": pricing, "group_ratio": groupRatio, "usable_group": usableGroup, }) } func ResetModelRatio(c *gin.Context) { defaultStr := operation_setting.DefaultModelRatio2JSONString() err := model.UpdateOption("ModelRatio", defaultStr) if err != nil { c.JSON(200, gin.H{ "success": false, "message": err.Error(), }) return } err = operation_setting.UpdateModelRatioByJSONString(defaultStr) if err != nil { c.JSON(200, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(200, gin.H{ "success": true, "message": "重置模型倍率成功", }) } ``` ## /controller/redemption.go ```go path="/controller/redemption.go" package controller import ( "net/http" "one-api/common" "one-api/model" "strconv" "github.com/gin-gonic/gin" ) func GetAllRedemptions(c *gin.Context) { p, _ := strconv.Atoi(c.Query("p")) pageSize, _ := strconv.Atoi(c.Query("page_size")) if p < 0 { p = 0 } if pageSize < 1 { pageSize = common.ItemsPerPage } redemptions, total, err := model.GetAllRedemptions((p-1)*pageSize, pageSize) 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": gin.H{ "items": redemptions, "total": total, "page": p, "page_size": pageSize, }, }) return } func SearchRedemptions(c *gin.Context) { keyword := c.Query("keyword") p, _ := strconv.Atoi(c.Query("p")) pageSize, _ := strconv.Atoi(c.Query("page_size")) if p < 0 { p = 0 } if pageSize < 1 { pageSize = common.ItemsPerPage } redemptions, total, err := model.SearchRedemptions(keyword, (p-1)*pageSize, pageSize) 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": gin.H{ "items": redemptions, "total": total, "page": p, "page_size": pageSize, }, }) return } func GetRedemption(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } redemption, err := model.GetRedemptionById(id) 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": redemption, }) return } func AddRedemption(c *gin.Context) { redemption := model.Redemption{} err := c.ShouldBindJSON(&redemption) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if len(redemption.Name) == 0 || len(redemption.Name) > 20 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "兑换码名称长度必须在1-20之间", }) return } if redemption.Count <= 0 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "兑换码个数必须大于0", }) return } if redemption.Count > 100 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "一次兑换码批量生成的个数不能大于 100", }) return } var keys []string for i := 0; i < redemption.Count; i++ { key := common.GetUUID() cleanRedemption := model.Redemption{ UserId: c.GetInt("id"), Name: redemption.Name, Key: key, CreatedTime: common.GetTimestamp(), Quota: redemption.Quota, } err = cleanRedemption.Insert() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), "data": keys, }) return } keys = append(keys, key) } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", "data": keys, }) return } func DeleteRedemption(c *gin.Context) { id, _ := strconv.Atoi(c.Param("id")) err := model.DeleteRedemptionById(id) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } func UpdateRedemption(c *gin.Context) { statusOnly := c.Query("status_only") redemption := model.Redemption{} err := c.ShouldBindJSON(&redemption) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } cleanRedemption, err := model.GetRedemptionById(redemption.Id) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if statusOnly != "" { cleanRedemption.Status = redemption.Status } else { // If you add more fields, please also update redemption.Update() cleanRedemption.Name = redemption.Name cleanRedemption.Quota = redemption.Quota } err = cleanRedemption.Update() 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": cleanRedemption, }) return } ``` ## /controller/relay.go ```go path="/controller/relay.go" package controller 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/middleware" "one-api/model" "one-api/relay" "one-api/relay/constant" relayconstant "one-api/relay/constant" "one-api/relay/helper" "one-api/service" "strings" ) func relayHandler(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode { var err *dto.OpenAIErrorWithStatusCode switch relayMode { case relayconstant.RelayModeImagesGenerations, relayconstant.RelayModeImagesEdits: err = relay.ImageHelper(c) case relayconstant.RelayModeAudioSpeech: fallthrough case relayconstant.RelayModeAudioTranslation: fallthrough case relayconstant.RelayModeAudioTranscription: err = relay.AudioHelper(c) case relayconstant.RelayModeRerank: err = relay.RerankHelper(c, relayMode) case relayconstant.RelayModeEmbeddings: err = relay.EmbeddingHelper(c) default: err = relay.TextHelper(c) } if err != nil { // 保存错误日志到mysql中 userId := c.GetInt("id") tokenName := c.GetString("token_name") modelName := c.GetString("original_model") tokenId := c.GetInt("token_id") userGroup := c.GetString("group") channelId := c.GetInt("channel_id") other := make(map[string]interface{}) other["error_type"] = err.Error.Type other["error_code"] = err.Error.Code other["status_code"] = err.StatusCode other["channel_id"] = channelId other["channel_name"] = c.GetString("channel_name") other["channel_type"] = c.GetInt("channel_type") model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.Error.Message, tokenId, 0, false, userGroup, other) } return err } func Relay(c *gin.Context) { relayMode := constant.Path2RelayMode(c.Request.URL.Path) requestId := c.GetString(common.RequestIdKey) group := c.GetString("group") originalModel := c.GetString("original_model") var openaiErr *dto.OpenAIErrorWithStatusCode for i := 0; i <= common.RetryTimes; i++ { channel, err := getChannel(c, group, originalModel, i) if err != nil { common.LogError(c, err.Error()) openaiErr = service.OpenAIErrorWrapperLocal(err, "get_channel_failed", http.StatusInternalServerError) break } openaiErr = relayRequest(c, relayMode, channel) if openaiErr == nil { return // 成功处理请求，直接返回 } go processChannelError(c, channel.Id, channel.Type, channel.Name, channel.GetAutoBan(), openaiErr) if !shouldRetry(c, openaiErr, common.RetryTimes-i) { break } } useChannel := c.GetStringSlice("use_channel") if len(useChannel) > 1 { retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]")) common.LogInfo(c, retryLogStr) } if openaiErr != nil { if openaiErr.StatusCode == http.StatusTooManyRequests { common.LogError(c, fmt.Sprintf("origin 429 error: %s", openaiErr.Error.Message)) openaiErr.Error.Message = "当前分组上游负载已饱和，请稍后再试" } openaiErr.Error.Message = common.MessageWithRequestId(openaiErr.Error.Message, requestId) c.JSON(openaiErr.StatusCode, gin.H{ "error": openaiErr.Error, }) } } var upgrader = websocket.Upgrader{ Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol CheckOrigin: func(r *http.Request) bool { return true // 允许跨域 }, } func WssRelay(c *gin.Context) { // 将 HTTP 连接升级为 WebSocket 连接 ws, err := upgrader.Upgrade(c.Writer, c.Request, nil) defer ws.Close() if err != nil { openaiErr := service.OpenAIErrorWrapper(err, "get_channel_failed", http.StatusInternalServerError) helper.WssError(c, ws, openaiErr.Error) return } relayMode := constant.Path2RelayMode(c.Request.URL.Path) requestId := c.GetString(common.RequestIdKey) group := c.GetString("group") //wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01 originalModel := c.GetString("original_model") var openaiErr *dto.OpenAIErrorWithStatusCode for i := 0; i <= common.RetryTimes; i++ { channel, err := getChannel(c, group, originalModel, i) if err != nil { common.LogError(c, err.Error()) openaiErr = service.OpenAIErrorWrapperLocal(err, "get_channel_failed", http.StatusInternalServerError) break } openaiErr = wssRequest(c, ws, relayMode, channel) if openaiErr == nil { return // 成功处理请求，直接返回 } go processChannelError(c, channel.Id, channel.Type, channel.Name, channel.GetAutoBan(), openaiErr) if !shouldRetry(c, openaiErr, common.RetryTimes-i) { break } } useChannel := c.GetStringSlice("use_channel") if len(useChannel) > 1 { retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]")) common.LogInfo(c, retryLogStr) } if openaiErr != nil { if openaiErr.StatusCode == http.StatusTooManyRequests { openaiErr.Error.Message = "当前分组上游负载已饱和，请稍后再试" } openaiErr.Error.Message = common.MessageWithRequestId(openaiErr.Error.Message, requestId) helper.WssError(c, ws, openaiErr.Error) } } func RelayClaude(c *gin.Context) { //relayMode := constant.Path2RelayMode(c.Request.URL.Path) requestId := c.GetString(common.RequestIdKey) group := c.GetString("group") originalModel := c.GetString("original_model") var claudeErr *dto.ClaudeErrorWithStatusCode for i := 0; i <= common.RetryTimes; i++ { channel, err := getChannel(c, group, originalModel, i) if err != nil { common.LogError(c, err.Error()) claudeErr = service.ClaudeErrorWrapperLocal(err, "get_channel_failed", http.StatusInternalServerError) break } claudeErr = claudeRequest(c, channel) if claudeErr == nil { return // 成功处理请求，直接返回 } openaiErr := service.ClaudeErrorToOpenAIError(claudeErr) go processChannelError(c, channel.Id, channel.Type, channel.Name, channel.GetAutoBan(), openaiErr) if !shouldRetry(c, openaiErr, common.RetryTimes-i) { break } } useChannel := c.GetStringSlice("use_channel") if len(useChannel) > 1 { retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]")) common.LogInfo(c, retryLogStr) } if claudeErr != nil { claudeErr.Error.Message = common.MessageWithRequestId(claudeErr.Error.Message, requestId) c.JSON(claudeErr.StatusCode, gin.H{ "type": "error", "error": claudeErr.Error, }) } } func relayRequest(c *gin.Context, relayMode int, channel *model.Channel) *dto.OpenAIErrorWithStatusCode { addUsedChannel(c, channel.Id) requestBody, _ := common.GetRequestBody(c) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) return relayHandler(c, relayMode) } func wssRequest(c *gin.Context, ws *websocket.Conn, relayMode int, channel *model.Channel) *dto.OpenAIErrorWithStatusCode { addUsedChannel(c, channel.Id) requestBody, _ := common.GetRequestBody(c) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) return relay.WssHelper(c, ws) } func claudeRequest(c *gin.Context, channel *model.Channel) *dto.ClaudeErrorWithStatusCode { addUsedChannel(c, channel.Id) requestBody, _ := common.GetRequestBody(c) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) return relay.ClaudeHelper(c) } func addUsedChannel(c *gin.Context, channelId int) { useChannel := c.GetStringSlice("use_channel") useChannel = append(useChannel, fmt.Sprintf("%d", channelId)) c.Set("use_channel", useChannel) } func getChannel(c *gin.Context, group, originalModel string, retryCount int) (*model.Channel, error) { if retryCount == 0 { autoBan := c.GetBool("auto_ban") autoBanInt := 1 if !autoBan { autoBanInt = 0 } return &model.Channel{ Id: c.GetInt("channel_id"), Type: c.GetInt("channel_type"), Name: c.GetString("channel_name"), AutoBan: &autoBanInt, }, nil } channel, err := model.CacheGetRandomSatisfiedChannel(group, originalModel, retryCount) if err != nil { return nil, errors.New(fmt.Sprintf("获取重试渠道失败: %s", err.Error())) } middleware.SetupContextForSelectedChannel(c, channel, originalModel) return channel, nil } func shouldRetry(c *gin.Context, openaiErr *dto.OpenAIErrorWithStatusCode, retryTimes int) bool { if openaiErr == nil { return false } if openaiErr.LocalError { return false } if retryTimes <= 0 { return false } if _, ok := c.Get("specific_channel_id"); ok { return false } if openaiErr.StatusCode == http.StatusTooManyRequests { return true } if openaiErr.StatusCode == 307 { return true } if openaiErr.StatusCode/100 == 5 { // 超时不重试 if openaiErr.StatusCode == 504 || openaiErr.StatusCode == 524 { return false } return true } if openaiErr.StatusCode == http.StatusBadRequest { channelType := c.GetInt("channel_type") if channelType == common.ChannelTypeAnthropic { return true } return false } if openaiErr.StatusCode == 408 { // azure处理超时不重试 return false } if openaiErr.StatusCode/100 == 2 { return false } return true } func processChannelError(c *gin.Context, channelId int, channelType int, channelName string, autoBan bool, err *dto.OpenAIErrorWithStatusCode) { // 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况 // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously common.LogError(c, fmt.Sprintf("relay error (channel #%d, status code: %d): %s", channelId, err.StatusCode, err.Error.Message)) if service.ShouldDisableChannel(channelType, err) && autoBan { service.DisableChannel(channelId, channelName, err.Error.Message) } } func RelayMidjourney(c *gin.Context) { relayMode := c.GetInt("relay_mode") var err *dto.MidjourneyResponse switch relayMode { case relayconstant.RelayModeMidjourneyNotify: err = relay.RelayMidjourneyNotify(c) case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition: err = relay.RelayMidjourneyTask(c, relayMode) case relayconstant.RelayModeMidjourneyTaskImageSeed: err = relay.RelayMidjourneyTaskImageSeed(c) case relayconstant.RelayModeSwapFace: err = relay.RelaySwapFace(c) default: err = relay.RelayMidjourneySubmit(c, relayMode) } //err = relayMidjourneySubmit(c, relayMode) log.Println(err) if err != nil { statusCode := http.StatusBadRequest if err.Code == 30 { err.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。" statusCode = http.StatusTooManyRequests } c.JSON(statusCode, gin.H{ "description": fmt.Sprintf("%s %s", err.Description, err.Result), "type": "upstream_error", "code": err.Code, }) channelId := c.GetInt("channel_id") common.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", err.Description, err.Result))) } } func RelayNotImplemented(c *gin.Context) { err := dto.OpenAIError{ Message: "API not implemented", Type: "new_api_error", Param: "", Code: "api_not_implemented", } c.JSON(http.StatusNotImplemented, gin.H{ "error": err, }) } func RelayNotFound(c *gin.Context) { err := dto.OpenAIError{ Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path), Type: "invalid_request_error", Param: "", Code: "", } c.JSON(http.StatusNotFound, gin.H{ "error": err, }) } func RelayTask(c *gin.Context) { retryTimes := common.RetryTimes channelId := c.GetInt("channel_id") relayMode := c.GetInt("relay_mode") group := c.GetString("group") originalModel := c.GetString("original_model") c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)}) taskErr := taskRelayHandler(c, relayMode) if taskErr == nil { retryTimes = 0 } for i := 0; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && i < retryTimes; i++ { channel, err := model.CacheGetRandomSatisfiedChannel(group, originalModel, i) if err != nil { common.LogError(c, fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", err.Error())) break } channelId = channel.Id useChannel := c.GetStringSlice("use_channel") useChannel = append(useChannel, fmt.Sprintf("%d", channelId)) c.Set("use_channel", useChannel) common.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, i)) middleware.SetupContextForSelectedChannel(c, channel, originalModel) requestBody, err := common.GetRequestBody(c) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) taskErr = taskRelayHandler(c, relayMode) } useChannel := c.GetStringSlice("use_channel") if len(useChannel) > 1 { retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]")) common.LogInfo(c, retryLogStr) } if taskErr != nil { if taskErr.StatusCode == http.StatusTooManyRequests { taskErr.Message = "当前分组上游负载已饱和，请稍后再试" } c.JSON(taskErr.StatusCode, taskErr) } } func taskRelayHandler(c *gin.Context, relayMode int) *dto.TaskError { var err *dto.TaskError switch relayMode { case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID: err = relay.RelayTaskFetch(c, relayMode) default: err = relay.RelayTaskSubmit(c, relayMode) } return err } func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool { if taskErr == nil { return false } if retryTimes <= 0 { return false } if _, ok := c.Get("specific_channel_id"); ok { return false } if taskErr.StatusCode == http.StatusTooManyRequests { return true } if taskErr.StatusCode == 307 { return true } if taskErr.StatusCode/100 == 5 { // 超时不重试 if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 { return false } return true } if taskErr.StatusCode == http.StatusBadRequest { return false } if taskErr.StatusCode == 408 { // azure处理超时不重试 return false } if taskErr.LocalError { return false } if taskErr.StatusCode/100 == 2 { return false } return true } ``` ## /controller/setup.go ```go path="/controller/setup.go" package controller import ( "github.com/gin-gonic/gin" "one-api/common" "one-api/constant" "one-api/model" "one-api/setting/operation_setting" "time" ) type Setup struct { Status bool `json:"status"` RootInit bool `json:"root_init"` DatabaseType string `json:"database_type"` } type SetupRequest struct { Username string `json:"username"` Password string `json:"password"` ConfirmPassword string `json:"confirmPassword"` SelfUseModeEnabled bool `json:"SelfUseModeEnabled"` DemoSiteEnabled bool `json:"DemoSiteEnabled"` } func GetSetup(c *gin.Context) { setup := Setup{ Status: constant.Setup, } if constant.Setup { c.JSON(200, gin.H{ "success": true, "data": setup, }) return } setup.RootInit = model.RootUserExists() if common.UsingMySQL { setup.DatabaseType = "mysql" } if common.UsingPostgreSQL { setup.DatabaseType = "postgres" } if common.UsingSQLite { setup.DatabaseType = "sqlite" } c.JSON(200, gin.H{ "success": true, "data": setup, }) } func PostSetup(c *gin.Context) { // Check if setup is already completed if constant.Setup { c.JSON(400, gin.H{ "success": false, "message": "系统已经初始化完成", }) return } // Check if root user already exists rootExists := model.RootUserExists() var req SetupRequest err := c.ShouldBindJSON(&req) if err != nil { c.JSON(400, gin.H{ "success": false, "message": "请求参数有误", }) return } // If root doesn't exist, validate and create admin account if !rootExists { // Validate password if req.Password != req.ConfirmPassword { c.JSON(400, gin.H{ "success": false, "message": "两次输入的密码不一致", }) return } if len(req.Password) < 8 { c.JSON(400, gin.H{ "success": false, "message": "密码长度至少为8个字符", }) return } // Create root user hashedPassword, err := common.Password2Hash(req.Password) if err != nil { c.JSON(500, gin.H{ "success": false, "message": "系统错误: " + err.Error(), }) return } rootUser := model.User{ Username: req.Username, Password: hashedPassword, Role: common.RoleRootUser, Status: common.UserStatusEnabled, DisplayName: "Root User", AccessToken: nil, Quota: 100000000, } err = model.DB.Create(&rootUser).Error if err != nil { c.JSON(500, gin.H{ "success": false, "message": "创建管理员账号失败: " + err.Error(), }) return } } // Set operation modes operation_setting.SelfUseModeEnabled = req.SelfUseModeEnabled operation_setting.DemoSiteEnabled = req.DemoSiteEnabled // Save operation modes to database for persistence err = model.UpdateOption("SelfUseModeEnabled", boolToString(req.SelfUseModeEnabled)) if err != nil { c.JSON(500, gin.H{ "success": false, "message": "保存自用模式设置失败: " + err.Error(), }) return } err = model.UpdateOption("DemoSiteEnabled", boolToString(req.DemoSiteEnabled)) if err != nil { c.JSON(500, gin.H{ "success": false, "message": "保存演示站点模式设置失败: " + err.Error(), }) return } // Update setup status constant.Setup = true setup := model.Setup{ Version: common.Version, InitializedAt: time.Now().Unix(), } err = model.DB.Create(&setup).Error if err != nil { c.JSON(500, gin.H{ "success": false, "message": "系统初始化失败: " + err.Error(), }) return } c.JSON(200, gin.H{ "success": true, "message": "系统初始化成功", }) } func boolToString(b bool) string { if b { return "true" } return "false" } ``` ## /controller/task.go ```go path="/controller/task.go" package controller import ( "context" "encoding/json" "errors" "fmt" "github.com/gin-gonic/gin" "github.com/samber/lo" "io" "net/http" "one-api/common" "one-api/constant" "one-api/dto" "one-api/model" "one-api/relay" "sort" "strconv" "time" ) func UpdateTaskBulk() { //revocer //imageModel := "midjourney" for { time.Sleep(time.Duration(15) * time.Second) common.SysLog("任务进度轮询开始") ctx := context.TODO() allTasks := model.GetAllUnFinishSyncTasks(500) platformTask := make(map[constant.TaskPlatform][]*model.Task) for _, t := range allTasks { platformTask[t.Platform] = append(platformTask[t.Platform], t) } for platform, tasks := range platformTask { if len(tasks) == 0 { continue } taskChannelM := make(map[int][]string) taskM := make(map[string]*model.Task) nullTaskIds := make([]int64, 0) for _, task := range tasks { if task.TaskID == "" { // 统计失败的未完成任务 nullTaskIds = append(nullTaskIds, task.ID) continue } taskM[task.TaskID] = task taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.TaskID) } if len(nullTaskIds) > 0 { err := model.TaskBulkUpdateByID(nullTaskIds, map[string]any{ "status": "FAILURE", "progress": "100%", }) if err != nil { common.LogError(ctx, fmt.Sprintf("Fix null task_id task error: %v", err)) } else { common.LogInfo(ctx, fmt.Sprintf("Fix null task_id task success: %v", nullTaskIds)) } } if len(taskChannelM) == 0 { continue } UpdateTaskByPlatform(platform, taskChannelM, taskM) } common.SysLog("任务进度轮询完成") } } func UpdateTaskByPlatform(platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) { switch platform { case constant.TaskPlatformMidjourney: //_ = UpdateMidjourneyTaskAll(context.Background(), tasks) case constant.TaskPlatformSuno: _ = UpdateSunoTaskAll(context.Background(), taskChannelM, taskM) default: common.SysLog("未知平台") } } func UpdateSunoTaskAll(ctx context.Context, taskChannelM map[int][]string, taskM map[string]*model.Task) error { for channelId, taskIds := range taskChannelM { err := updateSunoTaskAll(ctx, channelId, taskIds, taskM) if err != nil { common.LogError(ctx, fmt.Sprintf("渠道 #%d 更新异步任务失败: %d", channelId, err.Error())) } } return nil } func updateSunoTaskAll(ctx context.Context, channelId int, taskIds []string, taskM map[string]*model.Task) error { common.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds))) if len(taskIds) == 0 { return nil } channel, err := model.CacheGetChannel(channelId) if err != nil { common.SysLog(fmt.Sprintf("CacheGetChannel: %v", err)) err = model.TaskBulkUpdate(taskIds, map[string]any{ "fail_reason": fmt.Sprintf("获取渠道信息失败，请联系管理员，渠道ID：%d", channelId), "status": "FAILURE", "progress": "100%", }) if err != nil { common.SysError(fmt.Sprintf("UpdateMidjourneyTask error2: %v", err)) } return err } adaptor := relay.GetTaskAdaptor(constant.TaskPlatformSuno) if adaptor == nil { return errors.New("adaptor not found") } resp, err := adaptor.FetchTask(*channel.BaseURL, channel.Key, map[string]any{ "ids": taskIds, }) if err != nil { common.SysError(fmt.Sprintf("Get Task Do req error: %v", err)) return err } if resp.StatusCode != http.StatusOK { common.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) return errors.New(fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) } defer resp.Body.Close() responseBody, err := io.ReadAll(resp.Body) if err != nil { common.SysError(fmt.Sprintf("Get Task parse body error: %v", err)) return err } var responseItems dto.TaskResponse[[]dto.SunoDataResponse] err = json.Unmarshal(responseBody, &responseItems) if err != nil { common.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody))) return err } if !responseItems.IsSuccess() { common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, 成功获取到任务数: %d", channelId, len(taskIds), string(responseBody))) return err } for _, responseItem := range responseItems.Data { task := taskM[responseItem.TaskID] if !checkTaskNeedUpdate(task, responseItem) { continue } task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status) task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason) task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime) task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime) task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime) if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure { common.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason) task.Progress = "100%" //err = model.CacheUpdateUserQuota(task.UserId) ? if err != nil { common.LogError(ctx, "error update user quota cache: "+err.Error()) } else { quota := task.Quota if quota != 0 { err = model.IncreaseUserQuota(task.UserId, quota, false) if err != nil { common.LogError(ctx, "fail to increase user quota: "+err.Error()) } logContent := fmt.Sprintf("异步任务执行失败 %s，补偿 %s", task.TaskID, common.LogQuota(quota)) model.RecordLog(task.UserId, model.LogTypeSystem, logContent) } } } if responseItem.Status == model.TaskStatusSuccess { task.Progress = "100%" } task.Data = responseItem.Data err = task.Update() if err != nil { common.SysError("UpdateMidjourneyTask task error: " + err.Error()) } } return nil } func checkTaskNeedUpdate(oldTask *model.Task, newTask dto.SunoDataResponse) bool { if oldTask.SubmitTime != newTask.SubmitTime { return true } if oldTask.StartTime != newTask.StartTime { return true } if oldTask.FinishTime != newTask.FinishTime { return true } if string(oldTask.Status) != newTask.Status { return true } if oldTask.FailReason != newTask.FailReason { return true } if oldTask.FinishTime != newTask.FinishTime { return true } if (oldTask.Status == model.TaskStatusFailure || oldTask.Status == model.TaskStatusSuccess) && oldTask.Progress != "100%" { return true } oldData, _ := json.Marshal(oldTask.Data) newData, _ := json.Marshal(newTask.Data) sort.Slice(oldData, func(i, j int) bool { return oldData[i] < oldData[j] }) sort.Slice(newData, func(i, j int) bool { return newData[i] < newData[j] }) if string(oldData) != string(newData) { return true } return false } func GetAllTask(c *gin.Context) { p, _ := strconv.Atoi(c.Query("p")) if p < 0 { p = 0 } startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) // 解析其他查询参数 queryParams := model.SyncTaskQueryParams{ Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp, } logs := model.TaskGetAllTasks(p*common.ItemsPerPage, common.ItemsPerPage, queryParams) if logs == nil { logs = make([]*model.Task, 0) } c.JSON(200, gin.H{ "success": true, "message": "", "data": logs, }) } func GetUserTask(c *gin.Context) { p, _ := strconv.Atoi(c.Query("p")) if p < 0 { p = 0 } userId := c.GetInt("id") startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) queryParams := model.SyncTaskQueryParams{ Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp, } logs := model.TaskGetAllUserTask(userId, p*common.ItemsPerPage, common.ItemsPerPage, queryParams) if logs == nil { logs = make([]*model.Task, 0) } c.JSON(200, gin.H{ "success": true, "message": "", "data": logs, }) } ``` ## /controller/telegram.go ```go path="/controller/telegram.go" package controller import ( "crypto/hmac" "crypto/sha256" "encoding/hex" "io" "net/http" "one-api/common" "one-api/model" "sort" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" ) func TelegramBind(c *gin.Context) { if !common.TelegramOAuthEnabled { c.JSON(200, gin.H{ "message": "管理员未开启通过 Telegram 登录以及注册", "success": false, }) return } params := c.Request.URL.Query() if !checkTelegramAuthorization(params, common.TelegramBotToken) { c.JSON(200, gin.H{ "message": "无效的请求", "success": false, }) return } telegramId := params["id"][0] if model.IsTelegramIdAlreadyTaken(telegramId) { c.JSON(200, gin.H{ "message": "该 Telegram 账户已被绑定", "success": false, }) return } session := sessions.Default(c) id := session.Get("id") user := model.User{Id: id.(int)} if err := user.FillUserById(); err != nil { c.JSON(200, gin.H{ "message": err.Error(), "success": false, }) return } if user.Id == 0 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户已注销", }) return } user.TelegramId = telegramId if err := user.Update(false); err != nil { c.JSON(200, gin.H{ "message": err.Error(), "success": false, }) return } c.Redirect(302, "/setting") } func TelegramLogin(c *gin.Context) { if !common.TelegramOAuthEnabled { c.JSON(200, gin.H{ "message": "管理员未开启通过 Telegram 登录以及注册", "success": false, }) return } params := c.Request.URL.Query() if !checkTelegramAuthorization(params, common.TelegramBotToken) { c.JSON(200, gin.H{ "message": "无效的请求", "success": false, }) return } telegramId := params["id"][0] user := model.User{TelegramId: telegramId} if err := user.FillUserByTelegramId(); err != nil { c.JSON(200, gin.H{ "message": err.Error(), "success": false, }) return } setupLogin(&user, c) } func checkTelegramAuthorization(params map[string][]string, token string) bool { strs := []string{} var hash = "" for k, v := range params { if k == "hash" { hash = v[0] continue } strs = append(strs, k+"="+v[0]) } sort.Strings(strs) var imploded = "" for _, s := range strs { if imploded != "" { imploded += "\n" } imploded += s } sha256hash := sha256.New() io.WriteString(sha256hash, token) hmachash := hmac.New(sha256.New, sha256hash.Sum(nil)) io.WriteString(hmachash, imploded) ss := hex.EncodeToString(hmachash.Sum(nil)) return hash == ss } ``` ## /controller/token.go ```go path="/controller/token.go" package controller import ( "github.com/gin-gonic/gin" "net/http" "one-api/common" "one-api/model" "strconv" ) func GetAllTokens(c *gin.Context) { userId := c.GetInt("id") 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.GetAllUserTokens(userId, 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": "", "data": tokens, }) return } func SearchTokens(c *gin.Context) { userId := c.GetInt("id") keyword := c.Query("keyword") token := c.Query("token") tokens, err := model.SearchUserTokens(userId, keyword, token) 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": tokens, }) return } func GetToken(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) userId := c.GetInt("id") if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } token, err := model.GetTokenByIds(id, userId) 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": token, }) return } func GetTokenStatus(c *gin.Context) { tokenId := c.GetInt("token_id") userId := c.GetInt("id") token, err := model.GetTokenByIds(tokenId, userId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } expiredAt := token.ExpiredTime if expiredAt == -1 { expiredAt = 0 } c.JSON(http.StatusOK, gin.H{ "object": "credit_summary", "total_granted": token.RemainQuota, "total_used": 0, // not supported currently "total_available": token.RemainQuota, "expires_at": expiredAt * 1000, }) } func AddToken(c *gin.Context) { token := model.Token{} err := c.ShouldBindJSON(&token) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if len(token.Name) > 30 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌名称过长", }) return } key, err := common.GenerateKey() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "生成令牌失败", }) common.SysError("failed to generate token key: " + err.Error()) return } cleanToken := model.Token{ UserId: c.GetInt("id"), Name: token.Name, Key: key, CreatedTime: common.GetTimestamp(), AccessedTime: common.GetTimestamp(), ExpiredTime: token.ExpiredTime, RemainQuota: token.RemainQuota, UnlimitedQuota: token.UnlimitedQuota, ModelLimitsEnabled: token.ModelLimitsEnabled, ModelLimits: token.ModelLimits, AllowIps: token.AllowIps, Group: token.Group, } err = cleanToken.Insert() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } func DeleteToken(c *gin.Context) { id, _ := strconv.Atoi(c.Param("id")) userId := c.GetInt("id") err := model.DeleteTokenById(id, userId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } func UpdateToken(c *gin.Context) { userId := c.GetInt("id") statusOnly := c.Query("status_only") token := model.Token{} err := c.ShouldBindJSON(&token) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if len(token.Name) > 30 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌名称过长", }) return } cleanToken, err := model.GetTokenByIds(token.Id, userId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if token.Status == common.TokenStatusEnabled { if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() && cleanToken.ExpiredTime != -1 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌已过期，无法启用，请先修改令牌过期时间，或者设置为永不过期", }) return } if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌可用额度已用尽，无法启用，请先修改令牌剩余额度，或者设置为无限额度", }) return } } if statusOnly != "" { cleanToken.Status = token.Status } else { // If you add more fields, please also update token.Update() cleanToken.Name = token.Name cleanToken.ExpiredTime = token.ExpiredTime cleanToken.RemainQuota = token.RemainQuota cleanToken.UnlimitedQuota = token.UnlimitedQuota cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled cleanToken.ModelLimits = token.ModelLimits cleanToken.AllowIps = token.AllowIps cleanToken.Group = token.Group } err = cleanToken.Update() 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": cleanToken, }) return } ``` ## /controller/topup.go ```go path="/controller/topup.go" package controller import ( "fmt" "log" "net/url" "one-api/common" "one-api/model" "one-api/service" "one-api/setting" "strconv" "sync" "time" "github.com/Calcium-Ion/go-epay/epay" "github.com/gin-gonic/gin" "github.com/samber/lo" "github.com/shopspring/decimal" ) type EpayRequest struct { Amount int64 `json:"amount"` PaymentMethod string `json:"payment_method"` TopUpCode string `json:"top_up_code"` } type AmountRequest struct { Amount int64 `json:"amount"` TopUpCode string `json:"top_up_code"` } func GetEpayClient() *epay.Client { if setting.PayAddress == "" || setting.EpayId == "" || setting.EpayKey == "" { return nil } withUrl, err := epay.NewClient(&epay.Config{ PartnerID: setting.EpayId, Key: setting.EpayKey, }, setting.PayAddress) if err != nil { return nil } return withUrl } func getPayMoney(amount int64, group string) float64 { dAmount := decimal.NewFromInt(amount) if !common.DisplayInCurrencyEnabled { dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) dAmount = dAmount.Div(dQuotaPerUnit) } topupGroupRatio := common.GetTopupGroupRatio(group) if topupGroupRatio == 0 { topupGroupRatio = 1 } dTopupGroupRatio := decimal.NewFromFloat(topupGroupRatio) dPrice := decimal.NewFromFloat(setting.Price) payMoney := dAmount.Mul(dPrice).Mul(dTopupGroupRatio) return payMoney.InexactFloat64() } func getMinTopup() int64 { minTopup := setting.MinTopUp if !common.DisplayInCurrencyEnabled { dMinTopup := decimal.NewFromInt(int64(minTopup)) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) minTopup = int(dMinTopup.Mul(dQuotaPerUnit).IntPart()) } return int64(minTopup) } func RequestEpay(c *gin.Context) { var req EpayRequest err := c.ShouldBindJSON(&req) if err != nil { c.JSON(200, gin.H{"message": "error", "data": "参数错误"}) return } if req.Amount < getMinTopup() { c.JSON(200, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", getMinTopup())}) return } id := c.GetInt("id") group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(200, gin.H{"message": "error", "data": "获取用户分组失败"}) return } payMoney := getPayMoney(req.Amount, group) if payMoney < 0.01 { c.JSON(200, gin.H{"message": "error", "data": "充值金额过低"}) return } payType := "wxpay" if req.PaymentMethod == "zfb" { payType = "alipay" } if req.PaymentMethod == "wx" { req.PaymentMethod = "wxpay" payType = "wxpay" } callBackAddress := service.GetCallbackAddress() returnUrl, _ := url.Parse(setting.ServerAddress + "/log") notifyUrl, _ := url.Parse(callBackAddress + "/api/user/epay/notify") tradeNo := fmt.Sprintf("%s%d", common.GetRandomString(6), time.Now().Unix()) tradeNo = fmt.Sprintf("USR%dNO%s", id, tradeNo) client := GetEpayClient() if client == nil { c.JSON(200, gin.H{"message": "error", "data": "当前管理员未配置支付信息"}) return } uri, params, err := client.Purchase(&epay.PurchaseArgs{ Type: payType, ServiceTradeNo: tradeNo, Name: fmt.Sprintf("TUC%d", req.Amount), Money: strconv.FormatFloat(payMoney, 'f', 2, 64), Device: epay.PC, NotifyUrl: notifyUrl, ReturnUrl: returnUrl, }) if err != nil { c.JSON(200, gin.H{"message": "error", "data": "拉起支付失败"}) return } amount := req.Amount if !common.DisplayInCurrencyEnabled { dAmount := decimal.NewFromInt(int64(amount)) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) amount = dAmount.Div(dQuotaPerUnit).IntPart() } topUp := &model.TopUp{ UserId: id, Amount: amount, Money: payMoney, TradeNo: tradeNo, CreateTime: time.Now().Unix(), Status: "pending", } err = topUp.Insert() if err != nil { c.JSON(200, gin.H{"message": "error", "data": "创建订单失败"}) return } c.JSON(200, gin.H{"message": "success", "data": params, "url": uri}) } // tradeNo lock var orderLocks sync.Map var createLock sync.Mutex // LockOrder 尝试对给定订单号加锁 func LockOrder(tradeNo string) { lock, ok := orderLocks.Load(tradeNo) if !ok { createLock.Lock() defer createLock.Unlock() lock, ok = orderLocks.Load(tradeNo) if !ok { lock = new(sync.Mutex) orderLocks.Store(tradeNo, lock) } } lock.(*sync.Mutex).Lock() } // UnlockOrder 释放给定订单号的锁 func UnlockOrder(tradeNo string) { lock, ok := orderLocks.Load(tradeNo) if ok { lock.(*sync.Mutex).Unlock() } } func EpayNotify(c *gin.Context) { params := lo.Reduce(lo.Keys(c.Request.URL.Query()), func(r map[string]string, t string, i int) map[string]string { r[t] = c.Request.URL.Query().Get(t) return r }, map[string]string{}) client := GetEpayClient() if client == nil { log.Println("易支付回调失败 未找到配置信息") _, err := c.Writer.Write([]byte("fail")) if err != nil { log.Println("易支付回调写入失败") return } } verifyInfo, err := client.Verify(params) if err == nil && verifyInfo.VerifyStatus { _, err := c.Writer.Write([]byte("success")) if err != nil { log.Println("易支付回调写入失败") } } else { _, err := c.Writer.Write([]byte("fail")) if err != nil { log.Println("易支付回调写入失败") } log.Println("易支付回调签名验证失败") return } if verifyInfo.TradeStatus == epay.StatusTradeSuccess { log.Println(verifyInfo) LockOrder(verifyInfo.ServiceTradeNo) defer UnlockOrder(verifyInfo.ServiceTradeNo) topUp := model.GetTopUpByTradeNo(verifyInfo.ServiceTradeNo) if topUp == nil { log.Printf("易支付回调未找到订单: %v", verifyInfo) return } if topUp.Status == "pending" { topUp.Status = "success" err := topUp.Update() if err != nil { log.Printf("易支付回调更新订单失败: %v", topUp) return } //user, _ := model.GetUserById(topUp.UserId, false) //user.Quota += topUp.Amount * 500000 dAmount := decimal.NewFromInt(int64(topUp.Amount)) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) quotaToAdd := int(dAmount.Mul(dQuotaPerUnit).IntPart()) err = model.IncreaseUserQuota(topUp.UserId, quotaToAdd, true) if err != nil { log.Printf("易支付回调更新用户失败: %v", topUp) return } log.Printf("易支付回调更新用户成功 %v", topUp) model.RecordLog(topUp.UserId, model.LogTypeTopup, fmt.Sprintf("使用在线充值成功，充值金额: %v，支付金额：%f", common.LogQuota(quotaToAdd), topUp.Money)) } } else { log.Printf("易支付异常回调: %v", verifyInfo) } } func RequestAmount(c *gin.Context) { var req AmountRequest err := c.ShouldBindJSON(&req) if err != nil { c.JSON(200, gin.H{"message": "error", "data": "参数错误"}) return } if req.Amount < getMinTopup() { c.JSON(200, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", getMinTopup())}) return } id := c.GetInt("id") group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(200, gin.H{"message": "error", "data": "获取用户分组失败"}) return } payMoney := getPayMoney(req.Amount, group) if payMoney <= 0.01 { c.JSON(200, gin.H{"message": "error", "data": "充值金额过低"}) return } c.JSON(200, gin.H{"message": "success", "data": strconv.FormatFloat(payMoney, 'f', 2, 64)}) } ``` ## /controller/usedata.go ```go path="/controller/usedata.go" package controller import ( "github.com/gin-gonic/gin" "net/http" "one-api/model" "strconv" ) 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) 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 GetUserQuotaDates(c *gin.Context) { userId := c.GetInt("id") startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) // 判断时间跨度是否超过 1 个月 if endTimestamp-startTimestamp > 2592000 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "时间跨度不能超过 1 个月", }) return } dates, err := model.GetQuotaDataByUserId(userId, startTimestamp, endTimestamp) 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 } ``` The content has been capped at 50000 tokens, and files over NaN bytes have been omitted. The user could consider applying other filters to refine the result. The better and more specific the context, the better the LLM can follow instructions. If the context seems verbose, the user can refine the filter using uithub. Thank you for using https://uithub.com - Perfect LLM context for any GitHub repo.