| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509 |
- package controller
- import (
- "encoding/json"
- "fmt"
- "net/http"
- "one-api/common"
- "one-api/constant"
- "one-api/dto"
- "one-api/model"
- "strconv"
- "strings"
- "github.com/gin-gonic/gin"
- )
- type OpenAIModel struct {
- ID string `json:"id"`
- Object string `json:"object"`
- Created int64 `json:"created"`
- OwnedBy string `json:"owned_by"`
- Permission []struct {
- ID string `json:"id"`
- Object string `json:"object"`
- Created int64 `json:"created"`
- AllowCreateEngine bool `json:"allow_create_engine"`
- AllowSampling bool `json:"allow_sampling"`
- AllowLogprobs bool `json:"allow_logprobs"`
- AllowSearchIndices bool `json:"allow_search_indices"`
- AllowView bool `json:"allow_view"`
- AllowFineTuning bool `json:"allow_fine_tuning"`
- Organization string `json:"organization"`
- Group string `json:"group"`
- IsBlocking bool `json:"is_blocking"`
- } `json:"permission"`
- Root string `json:"root"`
- Parent string `json:"parent"`
- }
- type OpenAIModelsResponse struct {
- Data []OpenAIModel `json:"data"`
- Success bool `json:"success"`
- }
- func parseStatusFilter(statusParam string) int {
- switch strings.ToLower(statusParam) {
- case "enabled", "1":
- return common.ChannelStatusEnabled
- case "disabled", "0":
- return 0
- default:
- return -1
- }
- }
- func clearChannelInfo(channel *model.Channel) {
- if channel.ChannelInfo.IsMultiKey {
- channel.ChannelInfo.MultiKeyDisabledReason = nil
- channel.ChannelInfo.MultiKeyDisabledTime = nil
- }
- }
- func GetAllChannels(c *gin.Context) {
- pageInfo := common.GetPageQuery(c)
- channelData := make([]*model.Channel, 0)
- idSort, _ := strconv.ParseBool(c.Query("id_sort"))
- enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode"))
- statusParam := c.Query("status")
- // statusFilter: -1 all, 1 enabled, 0 disabled (include auto & manual)
- statusFilter := parseStatusFilter(statusParam)
- // type filter
- typeStr := c.Query("type")
- typeFilter := -1
- if typeStr != "" {
- if t, err := strconv.Atoi(typeStr); err == nil {
- typeFilter = t
- }
- }
- var total int64
- if enableTagMode {
- tags, err := model.GetPaginatedTags(pageInfo.GetStartIdx(), pageInfo.GetPageSize())
- if err != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
- return
- }
- for _, tag := range tags {
- if tag == nil || *tag == "" {
- continue
- }
- tagChannels, err := model.GetChannelsByTag(*tag, idSort)
- if err != nil {
- continue
- }
- filtered := make([]*model.Channel, 0)
- for _, ch := range tagChannels {
- if statusFilter == common.ChannelStatusEnabled && ch.Status != common.ChannelStatusEnabled {
- continue
- }
- if statusFilter == 0 && ch.Status == common.ChannelStatusEnabled {
- continue
- }
- if typeFilter >= 0 && ch.Type != typeFilter {
- continue
- }
- filtered = append(filtered, ch)
- }
- channelData = append(channelData, filtered...)
- }
- total, _ = model.CountAllTags()
- } else {
- baseQuery := model.DB.Model(&model.Channel{})
- if typeFilter >= 0 {
- baseQuery = baseQuery.Where("type = ?", typeFilter)
- }
- if statusFilter == common.ChannelStatusEnabled {
- baseQuery = baseQuery.Where("status = ?", common.ChannelStatusEnabled)
- } else if statusFilter == 0 {
- baseQuery = baseQuery.Where("status != ?", common.ChannelStatusEnabled)
- }
- baseQuery.Count(&total)
- order := "priority desc"
- if idSort {
- order = "id desc"
- }
- err := baseQuery.Order(order).Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Omit("key").Find(&channelData).Error
- if err != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
- return
- }
- }
- for _, datum := range channelData {
- clearChannelInfo(datum)
- }
- countQuery := model.DB.Model(&model.Channel{})
- if statusFilter == common.ChannelStatusEnabled {
- countQuery = countQuery.Where("status = ?", common.ChannelStatusEnabled)
- } else if statusFilter == 0 {
- countQuery = countQuery.Where("status != ?", common.ChannelStatusEnabled)
- }
- var results []struct {
- Type int64
- Count int64
- }
- _ = countQuery.Select("type, count(*) as count").Group("type").Find(&results).Error
- typeCounts := make(map[int64]int64)
- for _, r := range results {
- typeCounts[r.Type] = r.Count
- }
- common.ApiSuccess(c, gin.H{
- "items": channelData,
- "total": total,
- "page": pageInfo.GetPage(),
- "page_size": pageInfo.GetPageSize(),
- "type_counts": typeCounts,
- })
- return
- }
- func FetchUpstreamModels(c *gin.Context) {
- id, err := strconv.Atoi(c.Param("id"))
- if err != nil {
- common.ApiError(c, err)
- return
- }
- channel, err := model.GetChannelById(id, true)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- baseURL := constant.ChannelBaseURLs[channel.Type]
- if channel.GetBaseURL() != "" {
- baseURL = channel.GetBaseURL()
- }
- var url string
- switch channel.Type {
- case constant.ChannelTypeGemini:
- // curl https://example.com/v1beta/models?key=$GEMINI_API_KEY
- url = fmt.Sprintf("%s/v1beta/openai/models", baseURL) // Remove key in url since we need to use AuthHeader
- case constant.ChannelTypeAli:
- url = fmt.Sprintf("%s/compatible-mode/v1/models", baseURL)
- default:
- url = fmt.Sprintf("%s/v1/models", baseURL)
- }
- // 获取响应体 - 根据渠道类型决定是否添加 AuthHeader
- var body []byte
- key := strings.Split(channel.Key, "\n")[0]
- if channel.Type == constant.ChannelTypeGemini {
- body, err = GetResponseBody("GET", url, channel, GetAuthHeader(key)) // Use AuthHeader since Gemini now forces it
- } else {
- body, err = GetResponseBody("GET", url, channel, GetAuthHeader(key))
- }
- if err != nil {
- common.ApiError(c, err)
- return
- }
- var result OpenAIModelsResponse
- if err = json.Unmarshal(body, &result); err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": fmt.Sprintf("解析响应失败: %s", err.Error()),
- })
- return
- }
- var ids []string
- for _, model := range result.Data {
- id := model.ID
- if channel.Type == constant.ChannelTypeGemini {
- id = strings.TrimPrefix(id, "models/")
- }
- ids = append(ids, id)
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": ids,
- })
- }
- func FixChannelsAbilities(c *gin.Context) {
- success, fails, err := model.FixAbility()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": gin.H{
- "success": success,
- "fails": fails,
- },
- })
- }
- func SearchChannels(c *gin.Context) {
- keyword := c.Query("keyword")
- group := c.Query("group")
- modelKeyword := c.Query("model")
- statusParam := c.Query("status")
- statusFilter := parseStatusFilter(statusParam)
- idSort, _ := strconv.ParseBool(c.Query("id_sort"))
- enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode"))
- channelData := make([]*model.Channel, 0)
- if enableTagMode {
- tags, err := model.SearchTags(keyword, group, modelKeyword, idSort)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- for _, tag := range tags {
- if tag != nil && *tag != "" {
- tagChannel, err := model.GetChannelsByTag(*tag, idSort)
- if err == nil {
- channelData = append(channelData, tagChannel...)
- }
- }
- }
- } else {
- channels, err := model.SearchChannels(keyword, group, modelKeyword, idSort)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- channelData = channels
- }
- if statusFilter == common.ChannelStatusEnabled || statusFilter == 0 {
- filtered := make([]*model.Channel, 0, len(channelData))
- for _, ch := range channelData {
- if statusFilter == common.ChannelStatusEnabled && ch.Status != common.ChannelStatusEnabled {
- continue
- }
- if statusFilter == 0 && ch.Status == common.ChannelStatusEnabled {
- continue
- }
- filtered = append(filtered, ch)
- }
- channelData = filtered
- }
- // calculate type counts for search results
- typeCounts := make(map[int64]int64)
- for _, channel := range channelData {
- typeCounts[int64(channel.Type)]++
- }
- typeParam := c.Query("type")
- typeFilter := -1
- if typeParam != "" {
- if tp, err := strconv.Atoi(typeParam); err == nil {
- typeFilter = tp
- }
- }
- if typeFilter >= 0 {
- filtered := make([]*model.Channel, 0, len(channelData))
- for _, ch := range channelData {
- if ch.Type == typeFilter {
- filtered = append(filtered, ch)
- }
- }
- channelData = filtered
- }
- page, _ := strconv.Atoi(c.DefaultQuery("p", "1"))
- pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
- if page < 1 {
- page = 1
- }
- if pageSize <= 0 {
- pageSize = 20
- }
- total := len(channelData)
- startIdx := (page - 1) * pageSize
- if startIdx > total {
- startIdx = total
- }
- endIdx := startIdx + pageSize
- if endIdx > total {
- endIdx = total
- }
- pagedData := channelData[startIdx:endIdx]
- for _, datum := range pagedData {
- clearChannelInfo(datum)
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": gin.H{
- "items": pagedData,
- "total": total,
- "type_counts": typeCounts,
- },
- })
- return
- }
- func GetChannel(c *gin.Context) {
- id, err := strconv.Atoi(c.Param("id"))
- if err != nil {
- common.ApiError(c, err)
- return
- }
- channel, err := model.GetChannelById(id, false)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- if channel != nil {
- clearChannelInfo(channel)
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": channel,
- })
- return
- }
- // GetChannelKey 验证2FA后获取渠道密钥
- func GetChannelKey(c *gin.Context) {
- type GetChannelKeyRequest struct {
- Code string `json:"code" binding:"required"`
- }
- var req GetChannelKeyRequest
- if err := c.ShouldBindJSON(&req); err != nil {
- common.ApiError(c, fmt.Errorf("参数错误: %v", err))
- return
- }
- userId := c.GetInt("id")
- channelId, err := strconv.Atoi(c.Param("id"))
- if err != nil {
- common.ApiError(c, fmt.Errorf("渠道ID格式错误: %v", err))
- return
- }
- // 获取2FA记录并验证
- twoFA, err := model.GetTwoFAByUserId(userId)
- if err != nil {
- common.ApiError(c, fmt.Errorf("获取2FA信息失败: %v", err))
- return
- }
- if twoFA == nil || !twoFA.IsEnabled {
- common.ApiError(c, fmt.Errorf("用户未启用2FA,无法查看密钥"))
- return
- }
- // 统一的2FA验证逻辑
- if !validateTwoFactorAuth(twoFA, req.Code) {
- common.ApiError(c, fmt.Errorf("验证码或备用码错误,请重试"))
- return
- }
- // 获取渠道信息(包含密钥)
- channel, err := model.GetChannelById(channelId, true)
- if err != nil {
- common.ApiError(c, fmt.Errorf("获取渠道信息失败: %v", err))
- return
- }
- if channel == nil {
- common.ApiError(c, fmt.Errorf("渠道不存在"))
- return
- }
- // 记录操作日志
- model.RecordLog(userId, model.LogTypeSystem, fmt.Sprintf("查看渠道密钥信息 (渠道ID: %d)", channelId))
- // 统一的成功响应格式
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "验证成功",
- "data": map[string]interface{}{
- "key": channel.Key,
- },
- })
- }
- // validateTwoFactorAuth 统一的2FA验证函数
- func validateTwoFactorAuth(twoFA *model.TwoFA, code string) bool {
- // 尝试验证TOTP
- if cleanCode, err := common.ValidateNumericCode(code); err == nil {
- if isValid, _ := twoFA.ValidateTOTPAndUpdateUsage(cleanCode); isValid {
- return true
- }
- }
- // 尝试验证备用码
- if isValid, err := twoFA.ValidateBackupCodeAndUpdateUsage(code); err == nil && isValid {
- return true
- }
- return false
- }
- // validateChannel 通用的渠道校验函数
- func validateChannel(channel *model.Channel, isAdd bool) error {
- // 校验 channel settings
- if err := channel.ValidateSettings(); err != nil {
- return fmt.Errorf("渠道额外设置[channel setting] 格式错误:%s", err.Error())
- }
- // 如果是添加操作,检查 channel 和 key 是否为空
- if isAdd {
- if channel == nil || channel.Key == "" {
- return fmt.Errorf("channel cannot be empty")
- }
- // 检查模型名称长度是否超过 255
- for _, m := range channel.GetModels() {
- if len(m) > 255 {
- return fmt.Errorf("模型名称过长: %s", m)
- }
- }
- }
- // VertexAI 特殊校验
- if channel.Type == constant.ChannelTypeVertexAi {
- if channel.Other == "" {
- return fmt.Errorf("部署地区不能为空")
- }
- regionMap, err := common.StrToMap(channel.Other)
- if err != nil {
- return fmt.Errorf("部署地区必须是标准的Json格式,例如{\"default\": \"us-central1\", \"region2\": \"us-east1\"}")
- }
- if regionMap["default"] == nil {
- return fmt.Errorf("部署地区必须包含default字段")
- }
- }
- return nil
- }
- type AddChannelRequest struct {
- Mode string `json:"mode"`
- MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
- BatchAddSetKeyPrefix2Name bool `json:"batch_add_set_key_prefix_2_name"`
- Channel *model.Channel `json:"channel"`
- }
- func getVertexArrayKeys(keys string) ([]string, error) {
- if keys == "" {
- return nil, nil
- }
- var keyArray []interface{}
- err := common.Unmarshal([]byte(keys), &keyArray)
- if err != nil {
- return nil, fmt.Errorf("批量添加 Vertex AI 必须使用标准的JsonArray格式,例如[{key1}, {key2}...],请检查输入: %w", err)
- }
- cleanKeys := make([]string, 0, len(keyArray))
- for _, key := range keyArray {
- var keyStr string
- switch v := key.(type) {
- case string:
- keyStr = strings.TrimSpace(v)
- default:
- bytes, err := json.Marshal(v)
- if err != nil {
- return nil, fmt.Errorf("Vertex AI key JSON 编码失败: %w", err)
- }
- keyStr = string(bytes)
- }
- if keyStr != "" {
- cleanKeys = append(cleanKeys, keyStr)
- }
- }
- if len(cleanKeys) == 0 {
- return nil, fmt.Errorf("批量添加 Vertex AI 的 keys 不能为空")
- }
- return cleanKeys, nil
- }
- func AddChannel(c *gin.Context) {
- addChannelRequest := AddChannelRequest{}
- err := c.ShouldBindJSON(&addChannelRequest)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- // 使用统一的校验函数
- if err := validateChannel(addChannelRequest.Channel, true); err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- addChannelRequest.Channel.CreatedTime = common.GetTimestamp()
- keys := make([]string, 0)
- switch addChannelRequest.Mode {
- case "multi_to_single":
- addChannelRequest.Channel.ChannelInfo.IsMultiKey = true
- addChannelRequest.Channel.ChannelInfo.MultiKeyMode = addChannelRequest.MultiKeyMode
- if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi && addChannelRequest.Channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey {
- array, err := getVertexArrayKeys(addChannelRequest.Channel.Key)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(array)
- addChannelRequest.Channel.Key = strings.Join(array, "\n")
- } else {
- cleanKeys := make([]string, 0)
- for _, key := range strings.Split(addChannelRequest.Channel.Key, "\n") {
- if key == "" {
- continue
- }
- key = strings.TrimSpace(key)
- cleanKeys = append(cleanKeys, key)
- }
- addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(cleanKeys)
- addChannelRequest.Channel.Key = strings.Join(cleanKeys, "\n")
- }
- keys = []string{addChannelRequest.Channel.Key}
- case "batch":
- if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi && addChannelRequest.Channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey {
- // multi json
- keys, err = getVertexArrayKeys(addChannelRequest.Channel.Key)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- } else {
- keys = strings.Split(addChannelRequest.Channel.Key, "\n")
- }
- case "single":
- keys = []string{addChannelRequest.Channel.Key}
- default:
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "不支持的添加模式",
- })
- return
- }
- channels := make([]model.Channel, 0, len(keys))
- for _, key := range keys {
- if key == "" {
- continue
- }
- localChannel := addChannelRequest.Channel
- localChannel.Key = key
- if addChannelRequest.BatchAddSetKeyPrefix2Name && len(keys) > 1 {
- keyPrefix := localChannel.Key
- if len(localChannel.Key) > 8 {
- keyPrefix = localChannel.Key[:8]
- }
- localChannel.Name = fmt.Sprintf("%s %s", localChannel.Name, keyPrefix)
- }
- channels = append(channels, *localChannel)
- }
- err = model.BatchInsertChannels(channels)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- })
- return
- }
- func DeleteChannel(c *gin.Context) {
- id, _ := strconv.Atoi(c.Param("id"))
- channel := model.Channel{Id: id}
- err := channel.Delete()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- })
- return
- }
- func DeleteDisabledChannel(c *gin.Context) {
- rows, err := model.DeleteDisabledChannel()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": rows,
- })
- return
- }
- type ChannelTag struct {
- Tag string `json:"tag"`
- NewTag *string `json:"new_tag"`
- Priority *int64 `json:"priority"`
- Weight *uint `json:"weight"`
- ModelMapping *string `json:"model_mapping"`
- Models *string `json:"models"`
- Groups *string `json:"groups"`
- }
- func DisableTagChannels(c *gin.Context) {
- channelTag := ChannelTag{}
- err := c.ShouldBindJSON(&channelTag)
- if err != nil || channelTag.Tag == "" {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "参数错误",
- })
- return
- }
- err = model.DisableChannelByTag(channelTag.Tag)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- })
- return
- }
- func EnableTagChannels(c *gin.Context) {
- channelTag := ChannelTag{}
- err := c.ShouldBindJSON(&channelTag)
- if err != nil || channelTag.Tag == "" {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "参数错误",
- })
- return
- }
- err = model.EnableChannelByTag(channelTag.Tag)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- })
- return
- }
- func EditTagChannels(c *gin.Context) {
- channelTag := ChannelTag{}
- err := c.ShouldBindJSON(&channelTag)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "参数错误",
- })
- return
- }
- if channelTag.Tag == "" {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "tag不能为空",
- })
- return
- }
- err = model.EditChannelByTag(channelTag.Tag, channelTag.NewTag, channelTag.ModelMapping, channelTag.Models, channelTag.Groups, channelTag.Priority, channelTag.Weight)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- })
- return
- }
- type ChannelBatch struct {
- Ids []int `json:"ids"`
- Tag *string `json:"tag"`
- }
- func DeleteChannelBatch(c *gin.Context) {
- channelBatch := ChannelBatch{}
- err := c.ShouldBindJSON(&channelBatch)
- if err != nil || len(channelBatch.Ids) == 0 {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "参数错误",
- })
- return
- }
- err = model.BatchDeleteChannels(channelBatch.Ids)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": len(channelBatch.Ids),
- })
- return
- }
- type PatchChannel struct {
- model.Channel
- MultiKeyMode *string `json:"multi_key_mode"`
- KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加
- }
- func UpdateChannel(c *gin.Context) {
- channel := PatchChannel{}
- err := c.ShouldBindJSON(&channel)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- // 使用统一的校验函数
- if err := validateChannel(&channel.Channel, false); err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- // Preserve existing ChannelInfo to ensure multi-key channels keep correct state even if the client does not send ChannelInfo in the request.
- originChannel, err := model.GetChannelById(channel.Id, true)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- // Always copy the original ChannelInfo so that fields like IsMultiKey and MultiKeySize are retained.
- channel.ChannelInfo = originChannel.ChannelInfo
- // If the request explicitly specifies a new MultiKeyMode, apply it on top of the original info.
- if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" {
- channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode)
- }
- // 处理多key模式下的密钥追加/覆盖逻辑
- if channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey {
- switch *channel.KeyMode {
- case "append":
- // 追加模式:将新密钥添加到现有密钥列表
- if originChannel.Key != "" {
- var newKeys []string
- var existingKeys []string
- // 解析现有密钥
- if strings.HasPrefix(strings.TrimSpace(originChannel.Key), "[") {
- // JSON数组格式
- var arr []json.RawMessage
- if err := json.Unmarshal([]byte(strings.TrimSpace(originChannel.Key)), &arr); err == nil {
- existingKeys = make([]string, len(arr))
- for i, v := range arr {
- existingKeys[i] = string(v)
- }
- }
- } else {
- // 换行分隔格式
- existingKeys = strings.Split(strings.Trim(originChannel.Key, "\n"), "\n")
- }
- // 处理 Vertex AI 的特殊情况
- if channel.Type == constant.ChannelTypeVertexAi && channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey {
- // 尝试解析新密钥为JSON数组
- if strings.HasPrefix(strings.TrimSpace(channel.Key), "[") {
- array, err := getVertexArrayKeys(channel.Key)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "追加密钥解析失败: " + err.Error(),
- })
- return
- }
- newKeys = array
- } else {
- // 单个JSON密钥
- newKeys = []string{channel.Key}
- }
- // 合并密钥
- allKeys := append(existingKeys, newKeys...)
- channel.Key = strings.Join(allKeys, "\n")
- } else {
- // 普通渠道的处理
- inputKeys := strings.Split(channel.Key, "\n")
- for _, key := range inputKeys {
- key = strings.TrimSpace(key)
- if key != "" {
- newKeys = append(newKeys, key)
- }
- }
- // 合并密钥
- allKeys := append(existingKeys, newKeys...)
- channel.Key = strings.Join(allKeys, "\n")
- }
- }
- case "replace":
- // 覆盖模式:直接使用新密钥(默认行为,不需要特殊处理)
- }
- }
- err = channel.Update()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- channel.Key = ""
- clearChannelInfo(&channel.Channel)
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": channel,
- })
- return
- }
- func FetchModels(c *gin.Context) {
- var req struct {
- BaseURL string `json:"base_url"`
- Type int `json:"type"`
- Key string `json:"key"`
- }
- if err := c.ShouldBindJSON(&req); err != nil {
- c.JSON(http.StatusBadRequest, gin.H{
- "success": false,
- "message": "Invalid request",
- })
- return
- }
- baseURL := req.BaseURL
- if baseURL == "" {
- baseURL = constant.ChannelBaseURLs[req.Type]
- }
- client := &http.Client{}
- url := fmt.Sprintf("%s/v1/models", baseURL)
- request, err := http.NewRequest("GET", url, nil)
- if err != nil {
- c.JSON(http.StatusInternalServerError, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- // remove line breaks and extra spaces.
- key := strings.TrimSpace(req.Key)
- // If the key contains a line break, only take the first part.
- key = strings.Split(key, "\n")[0]
- request.Header.Set("Authorization", "Bearer "+key)
- response, err := client.Do(request)
- if err != nil {
- c.JSON(http.StatusInternalServerError, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- //check status code
- if response.StatusCode != http.StatusOK {
- c.JSON(http.StatusInternalServerError, gin.H{
- "success": false,
- "message": "Failed to fetch models",
- })
- return
- }
- defer response.Body.Close()
- var result struct {
- Data []struct {
- ID string `json:"id"`
- } `json:"data"`
- }
- if err := json.NewDecoder(response.Body).Decode(&result); err != nil {
- c.JSON(http.StatusInternalServerError, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- var models []string
- for _, model := range result.Data {
- models = append(models, model.ID)
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "data": models,
- })
- }
- func BatchSetChannelTag(c *gin.Context) {
- channelBatch := ChannelBatch{}
- err := c.ShouldBindJSON(&channelBatch)
- if err != nil || len(channelBatch.Ids) == 0 {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "参数错误",
- })
- return
- }
- err = model.BatchSetChannelTag(channelBatch.Ids, channelBatch.Tag)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": len(channelBatch.Ids),
- })
- return
- }
- func GetTagModels(c *gin.Context) {
- tag := c.Query("tag")
- if tag == "" {
- c.JSON(http.StatusBadRequest, gin.H{
- "success": false,
- "message": "tag不能为空",
- })
- return
- }
- channels, err := model.GetChannelsByTag(tag, false) // Assuming false for idSort is fine here
- if err != nil {
- c.JSON(http.StatusInternalServerError, gin.H{
- "success": false,
- "message": err.Error(),
- })
- return
- }
- var longestModels string
- maxLength := 0
- // Find the longest models string among all channels with the given tag
- for _, channel := range channels {
- if channel.Models != "" {
- currentModels := strings.Split(channel.Models, ",")
- if len(currentModels) > maxLength {
- maxLength = len(currentModels)
- longestModels = channel.Models
- }
- }
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": longestModels,
- })
- return
- }
- // CopyChannel handles cloning an existing channel with its key.
- // POST /api/channel/copy/:id
- // Optional query params:
- //
- // suffix - string appended to the original name (default "_复制")
- // reset_balance - bool, when true will reset balance & used_quota to 0 (default true)
- func CopyChannel(c *gin.Context) {
- id, err := strconv.Atoi(c.Param("id"))
- if err != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid id"})
- return
- }
- suffix := c.DefaultQuery("suffix", "_复制")
- resetBalance := true
- if rbStr := c.DefaultQuery("reset_balance", "true"); rbStr != "" {
- if v, err := strconv.ParseBool(rbStr); err == nil {
- resetBalance = v
- }
- }
- // fetch original channel with key
- origin, err := model.GetChannelById(id, true)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
- return
- }
- // clone channel
- clone := *origin // shallow copy is sufficient as we will overwrite primitives
- clone.Id = 0 // let DB auto-generate
- clone.CreatedTime = common.GetTimestamp()
- clone.Name = origin.Name + suffix
- clone.TestTime = 0
- clone.ResponseTime = 0
- if resetBalance {
- clone.Balance = 0
- clone.UsedQuota = 0
- }
- // insert
- if err := model.BatchInsertChannels([]model.Channel{clone}); err != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
- return
- }
- model.InitChannelCache()
- // success
- c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"id": clone.Id}})
- }
- // MultiKeyManageRequest represents the request for multi-key management operations
- type MultiKeyManageRequest struct {
- ChannelId int `json:"channel_id"`
- Action string `json:"action"` // "disable_key", "enable_key", "delete_disabled_keys", "get_key_status"
- KeyIndex *int `json:"key_index,omitempty"` // for disable_key and enable_key actions
- Page int `json:"page,omitempty"` // for get_key_status pagination
- PageSize int `json:"page_size,omitempty"` // for get_key_status pagination
- Status *int `json:"status,omitempty"` // for get_key_status filtering: 1=enabled, 2=manual_disabled, 3=auto_disabled, nil=all
- }
- // MultiKeyStatusResponse represents the response for key status query
- type MultiKeyStatusResponse struct {
- Keys []KeyStatus `json:"keys"`
- Total int `json:"total"`
- Page int `json:"page"`
- PageSize int `json:"page_size"`
- TotalPages int `json:"total_pages"`
- // Statistics
- EnabledCount int `json:"enabled_count"`
- ManualDisabledCount int `json:"manual_disabled_count"`
- AutoDisabledCount int `json:"auto_disabled_count"`
- }
- type KeyStatus struct {
- Index int `json:"index"`
- Status int `json:"status"` // 1: enabled, 2: disabled
- DisabledTime int64 `json:"disabled_time,omitempty"`
- Reason string `json:"reason,omitempty"`
- KeyPreview string `json:"key_preview"` // first 10 chars of key for identification
- }
- // ManageMultiKeys handles multi-key management operations
- func ManageMultiKeys(c *gin.Context) {
- request := MultiKeyManageRequest{}
- err := c.ShouldBindJSON(&request)
- if err != nil {
- common.ApiError(c, err)
- return
- }
- channel, err := model.GetChannelById(request.ChannelId, true)
- if err != nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "渠道不存在",
- })
- return
- }
- if !channel.ChannelInfo.IsMultiKey {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "该渠道不是多密钥模式",
- })
- return
- }
- lock := model.GetChannelPollingLock(channel.Id)
- lock.Lock()
- defer lock.Unlock()
- switch request.Action {
- case "get_key_status":
- keys := channel.GetKeys()
- // Default pagination parameters
- page := request.Page
- pageSize := request.PageSize
- if page <= 0 {
- page = 1
- }
- if pageSize <= 0 {
- pageSize = 50 // Default page size
- }
- // Statistics for all keys (unchanged by filtering)
- var enabledCount, manualDisabledCount, autoDisabledCount int
- // Build all key status data first
- var allKeyStatusList []KeyStatus
- for i, key := range keys {
- status := 1 // default enabled
- var disabledTime int64
- var reason string
- if channel.ChannelInfo.MultiKeyStatusList != nil {
- if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
- status = s
- }
- }
- // Count for statistics (all keys)
- switch status {
- case 1:
- enabledCount++
- case 2:
- manualDisabledCount++
- case 3:
- autoDisabledCount++
- }
- if status != 1 {
- if channel.ChannelInfo.MultiKeyDisabledTime != nil {
- disabledTime = channel.ChannelInfo.MultiKeyDisabledTime[i]
- }
- if channel.ChannelInfo.MultiKeyDisabledReason != nil {
- reason = channel.ChannelInfo.MultiKeyDisabledReason[i]
- }
- }
- // Create key preview (first 10 chars)
- keyPreview := key
- if len(key) > 10 {
- keyPreview = key[:10] + "..."
- }
- allKeyStatusList = append(allKeyStatusList, KeyStatus{
- Index: i,
- Status: status,
- DisabledTime: disabledTime,
- Reason: reason,
- KeyPreview: keyPreview,
- })
- }
- // Apply status filter if specified
- var filteredKeyStatusList []KeyStatus
- if request.Status != nil {
- for _, keyStatus := range allKeyStatusList {
- if keyStatus.Status == *request.Status {
- filteredKeyStatusList = append(filteredKeyStatusList, keyStatus)
- }
- }
- } else {
- filteredKeyStatusList = allKeyStatusList
- }
- // Calculate pagination based on filtered results
- filteredTotal := len(filteredKeyStatusList)
- totalPages := (filteredTotal + pageSize - 1) / pageSize
- if totalPages == 0 {
- totalPages = 1
- }
- if page > totalPages {
- page = totalPages
- }
- // Calculate range for current page
- start := (page - 1) * pageSize
- end := start + pageSize
- if end > filteredTotal {
- end = filteredTotal
- }
- // Get the page data
- var pageKeyStatusList []KeyStatus
- if start < filteredTotal {
- pageKeyStatusList = filteredKeyStatusList[start:end]
- }
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "",
- "data": MultiKeyStatusResponse{
- Keys: pageKeyStatusList,
- Total: filteredTotal, // Total of filtered results
- Page: page,
- PageSize: pageSize,
- TotalPages: totalPages,
- EnabledCount: enabledCount, // Overall statistics
- ManualDisabledCount: manualDisabledCount, // Overall statistics
- AutoDisabledCount: autoDisabledCount, // Overall statistics
- },
- })
- return
- case "disable_key":
- if request.KeyIndex == nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "未指定要禁用的密钥索引",
- })
- return
- }
- keyIndex := *request.KeyIndex
- if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "密钥索引超出范围",
- })
- return
- }
- if channel.ChannelInfo.MultiKeyStatusList == nil {
- channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
- }
- if channel.ChannelInfo.MultiKeyDisabledTime == nil {
- channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
- }
- if channel.ChannelInfo.MultiKeyDisabledReason == nil {
- channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
- }
- channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled
- err = channel.Update()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "密钥已禁用",
- })
- return
- case "enable_key":
- if request.KeyIndex == nil {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "未指定要启用的密钥索引",
- })
- return
- }
- keyIndex := *request.KeyIndex
- if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "密钥索引超出范围",
- })
- return
- }
- // 从状态列表中删除该密钥的记录,使其回到默认启用状态
- if channel.ChannelInfo.MultiKeyStatusList != nil {
- delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
- }
- if channel.ChannelInfo.MultiKeyDisabledTime != nil {
- delete(channel.ChannelInfo.MultiKeyDisabledTime, keyIndex)
- }
- if channel.ChannelInfo.MultiKeyDisabledReason != nil {
- delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex)
- }
- err = channel.Update()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": "密钥已启用",
- })
- return
- case "enable_all_keys":
- // 清空所有禁用状态,使所有密钥回到默认启用状态
- var enabledCount int
- if channel.ChannelInfo.MultiKeyStatusList != nil {
- enabledCount = len(channel.ChannelInfo.MultiKeyStatusList)
- }
- channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
- channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
- channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
- err = channel.Update()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": fmt.Sprintf("已启用 %d 个密钥", enabledCount),
- })
- return
- case "disable_all_keys":
- // 禁用所有启用的密钥
- if channel.ChannelInfo.MultiKeyStatusList == nil {
- channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
- }
- if channel.ChannelInfo.MultiKeyDisabledTime == nil {
- channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
- }
- if channel.ChannelInfo.MultiKeyDisabledReason == nil {
- channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
- }
- var disabledCount int
- for i := 0; i < channel.ChannelInfo.MultiKeySize; i++ {
- status := 1 // default enabled
- if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
- status = s
- }
- // 只禁用当前启用的密钥
- if status == 1 {
- channel.ChannelInfo.MultiKeyStatusList[i] = 2 // disabled
- disabledCount++
- }
- }
- if disabledCount == 0 {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "没有可禁用的密钥",
- })
- return
- }
- err = channel.Update()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount),
- })
- return
- case "delete_disabled_keys":
- keys := channel.GetKeys()
- var remainingKeys []string
- var deletedCount int
- var newStatusList = make(map[int]int)
- var newDisabledTime = make(map[int]int64)
- var newDisabledReason = make(map[int]string)
- newIndex := 0
- for i, key := range keys {
- status := 1 // default enabled
- if channel.ChannelInfo.MultiKeyStatusList != nil {
- if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
- status = s
- }
- }
- // 只删除自动禁用(status == 3)的密钥,保留启用(status == 1)和手动禁用(status == 2)的密钥
- if status == 3 {
- deletedCount++
- } else {
- remainingKeys = append(remainingKeys, key)
- // 保留非自动禁用密钥的状态信息,重新索引
- if status != 1 {
- newStatusList[newIndex] = status
- if channel.ChannelInfo.MultiKeyDisabledTime != nil {
- if t, exists := channel.ChannelInfo.MultiKeyDisabledTime[i]; exists {
- newDisabledTime[newIndex] = t
- }
- }
- if channel.ChannelInfo.MultiKeyDisabledReason != nil {
- if r, exists := channel.ChannelInfo.MultiKeyDisabledReason[i]; exists {
- newDisabledReason[newIndex] = r
- }
- }
- }
- newIndex++
- }
- }
- if deletedCount == 0 {
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "没有需要删除的自动禁用密钥",
- })
- return
- }
- // Update channel with remaining keys
- channel.Key = strings.Join(remainingKeys, "\n")
- channel.ChannelInfo.MultiKeySize = len(remainingKeys)
- channel.ChannelInfo.MultiKeyStatusList = newStatusList
- channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
- channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
- err = channel.Update()
- if err != nil {
- common.ApiError(c, err)
- return
- }
- model.InitChannelCache()
- c.JSON(http.StatusOK, gin.H{
- "success": true,
- "message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount),
- "data": deletedCount,
- })
- return
- default:
- c.JSON(http.StatusOK, gin.H{
- "success": false,
- "message": "不支持的操作",
- })
- return
- }
- }
|