Skip to content
Merged
10 changes: 10 additions & 0 deletions common/gin.go
Original file line number Diff line number Diff line change
Expand Up @@ -76,3 +76,13 @@ func GetContextKeyStringMap(c *gin.Context, key constant.ContextKey) map[string]
func GetContextKeyTime(c *gin.Context, key constant.ContextKey) time.Time {
return c.GetTime(string(key))
}

func GetContextKeyType[T any](c *gin.Context, key constant.ContextKey) (T, bool) {
if value, ok := c.Get(string(key)); ok {
if v, ok := value.(T); ok {
return v, true
}
}
var t T
return t, false
}
13 changes: 13 additions & 0 deletions common/str.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package common

import (
"encoding/base64"
"encoding/json"
"math/rand"
"strconv"
Expand Down Expand Up @@ -68,3 +69,15 @@ func StringToByteSlice(s string) []byte {
tmp2 := [3]uintptr{tmp1[0], tmp1[1], tmp1[1]}
return *(*[]byte)(unsafe.Pointer(&tmp2))
}

func EncodeBase64(str string) string {
return base64.StdEncoding.EncodeToString([]byte(str))
}

func GetJsonString(data any) string {
if data == nil {
return ""
}
b, _ := json.Marshal(data)
return string(b)
}
7 changes: 0 additions & 7 deletions constant/channel_setting.go

This file was deleted.

16 changes: 0 additions & 16 deletions constant/user_setting.go

This file was deleted.

15 changes: 13 additions & 2 deletions controller/channel-test.go
Original file line number Diff line number Diff line change
Expand Up @@ -173,8 +173,19 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr
consumedTime := float64(milliseconds) / 1000.0
other := service.GenerateTextOtherInfo(c, info, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.CompletionRatio,
usage.PromptTokensDetails.CachedTokens, priceData.CacheRatio, priceData.ModelPrice, priceData.GroupRatioInfo.GroupSpecialRatio)
model.RecordConsumeLog(c, 1, channel.Id, usage.PromptTokens, usage.CompletionTokens, info.OriginModelName, "模型测试",
quota, "模型测试", 0, quota, int(consumedTime), false, info.UsingGroup, other)
model.RecordConsumeLog(c, 1, model.RecordConsumeLogParams{
ChannelId: channel.Id,
PromptTokens: usage.PromptTokens,
CompletionTokens: usage.CompletionTokens,
ModelName: info.OriginModelName,
TokenName: "模型测试",
Quota: quota,
Content: "模型测试",
UseTimeSeconds: int(consumedTime),
IsStream: false,
Group: info.UsingGroup,
Other: other,
})
common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody)))
return nil, nil
}
Expand Down
23 changes: 21 additions & 2 deletions controller/channel.go
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ func FetchUpstreamModels(c *gin.Context) {
}

func FixChannelsAbilities(c *gin.Context) {
count, err := model.FixAbility()
success, fails, err := model.FixAbility()
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
Expand All @@ -239,7 +239,10 @@ func FixChannelsAbilities(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": count,
"data": gin.H{
"success": success,
"fails": fails,
},
})
}

Expand Down Expand Up @@ -387,6 +390,14 @@ func AddChannel(c *gin.Context) {
})
return
}
err = channel.ValidateSettings()
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "channel setting 格式错误:" + err.Error(),
})
return
}
channel.CreatedTime = common.GetTimestamp()
keys := strings.Split(channel.Key, "\n")
if channel.Type == constant.ChannelTypeVertexAi {
Expand Down Expand Up @@ -614,6 +625,14 @@ func UpdateChannel(c *gin.Context) {
})
return
}
err = channel.ValidateSettings()
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "channel setting 格式错误:" + err.Error(),
})
return
}
if channel.Type == constant.ChannelTypeVertexAi {
if channel.Other == "" {
c.JSON(http.StatusOK, gin.H{
Expand Down
2 changes: 1 addition & 1 deletion controller/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ func ListModels(c *gin.Context) {
}
} else {
userId := c.GetInt("id")
userGroup, err := model.GetUserGroup(userId, true)
userGroup, err := model.GetUserGroup(userId, false)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
Expand Down
27 changes: 14 additions & 13 deletions controller/user.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"net/http"
"net/url"
"one-api/common"
"one-api/dto"
"one-api/model"
"one-api/setting"
"strconv"
Expand Down Expand Up @@ -961,7 +962,7 @@ func UpdateUserSetting(c *gin.Context) {
}

// 验证预警类型
if req.QuotaWarningType != constant.NotifyTypeEmail && req.QuotaWarningType != constant.NotifyTypeWebhook {
if req.QuotaWarningType != dto.NotifyTypeEmail && req.QuotaWarningType != dto.NotifyTypeWebhook {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "无效的预警类型",
Expand All @@ -979,7 +980,7 @@ func UpdateUserSetting(c *gin.Context) {
}

// 如果是webhook类型,验证webhook地址
if req.QuotaWarningType == constant.NotifyTypeWebhook {
if req.QuotaWarningType == dto.NotifyTypeWebhook {
if req.WebhookUrl == "" {
c.JSON(http.StatusOK, gin.H{
"success": false,
Expand All @@ -998,7 +999,7 @@ func UpdateUserSetting(c *gin.Context) {
}

// 如果是邮件类型,验证邮箱地址
if req.QuotaWarningType == constant.NotifyTypeEmail && req.NotificationEmail != "" {
if req.QuotaWarningType == dto.NotifyTypeEmail && req.NotificationEmail != "" {
// 验证邮箱格式
if !strings.Contains(req.NotificationEmail, "@") {
c.JSON(http.StatusOK, gin.H{
Expand All @@ -1020,24 +1021,24 @@ func UpdateUserSetting(c *gin.Context) {
}

// 构建设置
settings := map[string]interface{}{
constant.UserSettingNotifyType: req.QuotaWarningType,
constant.UserSettingQuotaWarningThreshold: req.QuotaWarningThreshold,
"accept_unset_model_ratio_model": req.AcceptUnsetModelRatioModel,
constant.UserSettingRecordIpLog: req.RecordIpLog,
settings := dto.UserSetting{
NotifyType: req.QuotaWarningType,
QuotaWarningThreshold: req.QuotaWarningThreshold,
AcceptUnsetRatioModel: req.AcceptUnsetModelRatioModel,
RecordIpLog: req.RecordIpLog,
}

// 如果是webhook类型,添加webhook相关设置
if req.QuotaWarningType == constant.NotifyTypeWebhook {
settings[constant.UserSettingWebhookUrl] = req.WebhookUrl
if req.QuotaWarningType == dto.NotifyTypeWebhook {
settings.WebhookUrl = req.WebhookUrl
if req.WebhookSecret != "" {
settings[constant.UserSettingWebhookSecret] = req.WebhookSecret
settings.WebhookSecret = req.WebhookSecret
}
}

// 如果提供了通知邮箱,添加到设置中
if req.QuotaWarningType == constant.NotifyTypeEmail && req.NotificationEmail != "" {
settings[constant.UserSettingNotificationEmail] = req.NotificationEmail
if req.QuotaWarningType == dto.NotifyTypeEmail && req.NotificationEmail != "" {
settings.NotificationEmail = req.NotificationEmail
}

// 更新用户设置
Expand Down
7 changes: 7 additions & 0 deletions dto/channel_settings.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
package dto

type ChannelSettings struct {
ForceFormat bool `json:"force_format,omitempty"`
ThinkingToContent bool `json:"thinking_to_content,omitempty"`
Proxy string `json:"proxy"`
}
16 changes: 16 additions & 0 deletions dto/user_settings.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
package dto

type UserSetting struct {
NotifyType string `json:"notify_type,omitempty"` // QuotaWarningType 额度预警类型
QuotaWarningThreshold float64 `json:"quota_warning_threshold,omitempty"` // QuotaWarningThreshold 额度预警阈值
WebhookUrl string `json:"webhook_url,omitempty"` // WebhookUrl webhook地址
WebhookSecret string `json:"webhook_secret,omitempty"` // WebhookSecret webhook密钥
NotificationEmail string `json:"notification_email,omitempty"` // NotificationEmail 通知邮箱地址
AcceptUnsetRatioModel bool `json:"accept_unset_model_ratio_model,omitempty"` // AcceptUnsetRatioModel 是否接受未设置价格的模型
RecordIpLog bool `json:"record_ip_log,omitempty"` // 是否记录请求和错误日志IP
}

var (
NotifyTypeEmail = "email" // Email 邮件
NotifyTypeWebhook = "webhook" // Webhook
)
7 changes: 4 additions & 3 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@ func main() {
return
}

common.SetupLogger()
common.SysLog("New API " + common.Version + " started")
if os.Getenv("GIN_MODE") != "debug" {
gin.SetMode(gin.ReleaseMode)
Expand Down Expand Up @@ -69,9 +68,9 @@ func main() {
if r := recover(); r != nil {
common.SysError(fmt.Sprintf("InitChannelCache panic: %v, retrying once", r))
// Retry once
_, fixErr := model.FixAbility()
_, _, fixErr := model.FixAbility()
if fixErr != nil {
common.SysError(fmt.Sprintf("InitChannelCache failed: %s", fixErr.Error()))
common.FatalLog(fmt.Sprintf("InitChannelCache failed: %s", fixErr.Error()))
}
}
}()
Expand Down Expand Up @@ -172,6 +171,8 @@ func InitResources() error {
// 加载环境变量
common.InitEnv()

common.SetupLogger()

// Initialize model settings
ratio_setting.InitRatioSettings()

Expand Down
6 changes: 3 additions & 3 deletions middleware/distributor.go
Original file line number Diff line number Diff line change
Expand Up @@ -247,9 +247,9 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
}
c.Set("channel_id", channel.Id)
c.Set("channel_name", channel.Name)
c.Set("channel_type", channel.Type)
common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type)
c.Set("channel_create_time", channel.CreatedTime)
c.Set("channel_setting", channel.GetSetting())
common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting())
c.Set("param_override", channel.GetParamOverride())
if nil != channel.OpenAIOrganization && "" != *channel.OpenAIOrganization {
c.Set("channel_organization", *channel.OpenAIOrganization)
Expand All @@ -258,7 +258,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
c.Set("model_mapping", channel.GetModelMapping())
c.Set("status_code_mapping", channel.GetStatusCodeMapping())
c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
c.Set("base_url", channel.GetBaseURL())
common.SetContextKey(c, constant.ContextKeyBaseUrl, channel.GetBaseURL())
// TODO: api_version统一
switch channel.Type {
case constant.ChannelTypeAzure:
Expand Down
Loading
Loading