Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'songquanpeng/main'

Browse Source 
# Conflicts:
#	controller/relay-text.go
#	middleware/distributor.go
This commit is contained in:
quzard 2023-06-20 20:15:57 +08:00
commit bd7a9b6576
22 changed files with 245 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
common/constants.go
View File

@ -15,6 +15,8 @@ var Footer = ""
var Logo = ""
var TopUpLink = ""
var ChatLink = ""
var QuotaPerUnit = 500 * 1000.0 // $0.002 / 1K tokens
var DisplayInCurrencyEnabled = false
var UsingSQLite = false

8
common/logger.go
View File

@ -42,3 +42,11 @@ func FatalLog(v ...any) {
_, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v)
os.Exit(1)
}
func LogQuota(quota int) string {
if DisplayInCurrencyEnabled {
return fmt.Sprintf("＄%.6f 额度", float64(quota)/QuotaPerUnit)
} else {
return fmt.Sprintf("%d 点额度", quota)
}
}

15
common/redis.go
View File

@ -37,3 +37,18 @@ func ParseRedisOption() *redis.Options {
}
return opt
}
func RedisSet(key string, value string, expiration time.Duration) error {
ctx := context.Background()
return RDB.Set(ctx, key, value, expiration).Err()
}
func RedisGet(key string) (string, error) {
ctx := context.Background()
return RDB.Get(ctx, key).Result()
}
func RedisDel(key string) error {
ctx := context.Background()
return RDB.Del(ctx, key).Err()
}

28
controller/billing.go
View File

@ -2,6 +2,7 @@ package controller
import (
"github.com/gin-gonic/gin"
"one-api/common"
"one-api/model"
)
@ -18,23 +19,38 @@ func GetSubscription(c *gin.Context) {
})
return
}
amount := float64(quota)
if common.DisplayInCurrencyEnabled {
amount /= common.QuotaPerUnit
}
subscription := OpenAISubscriptionResponse{
Object: "billing_subscription",
HasPaymentMethod: true,
SoftLimitUSD: float64(quota),
HardLimitUSD: float64(quota),
SystemHardLimitUSD: float64(quota),
SoftLimitUSD: amount,
HardLimitUSD: amount,
SystemHardLimitUSD: amount,
}
c.JSON(200, subscription)
return
}
func GetUsage(c *gin.Context) {
//userId := c.GetInt("id")
// TODO: get usage from database
userId := c.GetInt("id")
quota, err := model.GetUserUsedQuota(userId)
if err != nil {
openAIError := OpenAIError{
Message: err.Error(),
Type: "one_api_error",
}
c.JSON(200, gin.H{
"error": openAIError,
})
return
}
amount := float64(quota)
usage := OpenAIUsageResponse{
Object: "list",
TotalUsage: 0,
TotalUsage: amount,
}
c.JSON(200, usage)
return

32
controller/misc.go
View File

@ -14,21 +14,23 @@ func GetStatus(c *gin.Context) {
"success": true,
"message": "",
"data": gin.H{
"version": common.Version,
"start_time": common.StartTime,
"email_verification": common.EmailVerificationEnabled,
"github_oauth": common.GitHubOAuthEnabled,
"github_client_id": common.GitHubClientId,
"system_name": common.SystemName,
"logo": common.Logo,
"footer_html": common.Footer,
"wechat_qrcode": common.WeChatAccountQRCodeImageURL,
"wechat_login": common.WeChatAuthEnabled,
"server_address": common.ServerAddress,
"turnstile_check": common.TurnstileCheckEnabled,
"turnstile_site_key": common.TurnstileSiteKey,
"top_up_link": common.TopUpLink,
"chat_link": common.ChatLink,
"version": common.Version,
"start_time": common.StartTime,
"email_verification": common.EmailVerificationEnabled,
"github_oauth": common.GitHubOAuthEnabled,
"github_client_id": common.GitHubClientId,
"system_name": common.SystemName,
"logo": common.Logo,
"footer_html": common.Footer,
"wechat_qrcode": common.WeChatAccountQRCodeImageURL,
"wechat_login": common.WeChatAuthEnabled,
"server_address": common.ServerAddress,
"turnstile_check": common.TurnstileCheckEnabled,
"turnstile_site_key": common.TurnstileSiteKey,
"top_up_link": common.TopUpLink,
"chat_link": common.ChatLink,
"quota_per_unit": common.QuotaPerUnit,
"display_in_currency": common.DisplayInCurrencyEnabled,
},
})
return

2
controller/user.go
View File

@ -384,7 +384,7 @@ func UpdateUser(c *gin.Context) {
return
}
if originUser.Quota != updatedUser.Quota {
model.RecordLog(originUser.Id, model.LogTypeManage, fmt.Sprintf("管理员将用户额度从 %d 点修改为 %d 点", originUser.Quota, updatedUser.Quota))
model.RecordLog(originUser.Id, model.LogTypeManage, fmt.Sprintf("管理员将用户额度从 %s修改为 %s", common.LogQuota(originUser.Quota), common.LogQuota(updatedUser.Quota)))
}
c.JSON(http.StatusOK, gin.H{
"success": true,

6
main.go
View File

@ -47,12 +47,18 @@ func main() {
// Initialize options
model.InitOptionMap()
if common.RedisEnabled {
model.InitChannelCache()
}
if os.Getenv("SYNC_FREQUENCY") != "" {
frequency, err := strconv.Atoi(os.Getenv("SYNC_FREQUENCY"))
if err != nil {
common.FatalLog(err)
}
go model.SyncOptions(frequency)
if common.RedisEnabled {
go model.SyncChannelCache(frequency)
}
}
// Initialize HTTP server

4
middleware/distributor.go
View File

@ -17,7 +17,7 @@ type ModelRequest struct {
func Distribute() func(c *gin.Context) {
return func(c *gin.Context) {
userId := c.GetInt("id")
userGroup, _ := model.GetUserGroup(userId)
userGroup, _ := model.CacheGetUserGroup(userId)
c.Set("group", userGroup)
var channel *model.Channel
channelId, ok := c.Get("channelId")
@ -76,7 +76,7 @@ func Distribute() func(c *gin.Context) {
if strings.HasPrefix(modelRequest.Model, "gpt-35-turbo") {
modelRequest.Model = strings.Replace(modelRequest.Model, "gpt-35-turbo", "gpt-3.5-turbo", 1)
}
channel, err = model.GetRandomSatisfiedChannel(userGroup, modelRequest.Model)
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model)
if err != nil {
c.JSON(200, gin.H{
"error": gin.H{

99
model/cache.go Normal file
View File

@ -0,0 +1,99 @@
package model
import (
"encoding/json"
"fmt"
"one-api/common"
"sync"
"time"
)
const (
TokenCacheSeconds = 60 * 60
UserId2GroupCacheSeconds = 60 * 60
)
func CacheGetTokenByKey(key string) (*Token, error) {
var token Token
if !common.RedisEnabled {
err := DB.Where("`key` = ?", key).First(&token).Error
return &token, err
}
tokenObjectString, err := common.RedisGet(fmt.Sprintf("token:%s", key))
if err != nil {
err := DB.Where("`key` = ?", key).First(&token).Error
if err != nil {
return nil, err
}
jsonBytes, err := json.Marshal(token)
if err != nil {
return nil, err
}
err = common.RedisSet(fmt.Sprintf("token:%s", key), string(jsonBytes), TokenCacheSeconds*time.Second)
if err != nil {
common.SysError("Redis set token error: " + err.Error())
}
}
err = json.Unmarshal([]byte(tokenObjectString), &token)
return &token, err
}
func CacheGetUserGroup(id int) (group string, err error) {
if !common.RedisEnabled {
return GetUserGroup(id)
}
group, err = common.RedisGet(fmt.Sprintf("user_group:%d", id))
if err != nil {
group, err = GetUserGroup(id)
if err != nil {
return "", err
}
err = common.RedisSet(fmt.Sprintf("user_group:%d", id), group, UserId2GroupCacheSeconds*time.Second)
if err != nil {
common.SysError("Redis set user group error: " + err.Error())
}
}
return group, err
}
var channelId2channel map[int]*Channel
var channelSyncLock sync.RWMutex
var group2model2channels map[string]map[string][]*Channel
func InitChannelCache() {
channelSyncLock.Lock()
defer channelSyncLock.Unlock()
channelId2channel = make(map[int]*Channel)
var channels []*Channel
DB.Find(&channels)
for _, channel := range channels {
channelId2channel[channel.Id] = channel
}
var abilities []*Ability
DB.Find(&abilities)
groups := make(map[string]bool)
for _, ability := range abilities {
groups[ability.Group] = true
}
group2model2channels = make(map[string]map[string][]*Channel)
for group := range groups {
group2model2channels[group] = make(map[string][]*Channel)
// TODO: implement this
}
}
func SyncChannelCache(frequency int) {
for {
time.Sleep(time.Duration(frequency) * time.Second)
common.SysLog("Syncing channels from database")
InitChannelCache()
}
}
func CacheGetRandomSatisfiedChannel(group string, model string) (*Channel, error) {
if !common.RedisEnabled {
return GetRandomSatisfiedChannel(group, model)
}
// TODO: implement this
return nil, nil
}

6
model/option.go
View File

@ -35,6 +35,7 @@ func InitOptionMap() {
common.OptionMap["RegisterEnabled"] = strconv.FormatBool(common.RegisterEnabled)
common.OptionMap["AutomaticDisableChannelEnabled"] = strconv.FormatBool(common.AutomaticDisableChannelEnabled)
common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled)
common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled)
common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64)
common.OptionMap["SMTPServer"] = ""
common.OptionMap["SMTPFrom"] = ""
@ -64,6 +65,7 @@ func InitOptionMap() {
common.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
common.OptionMap["TopUpLink"] = common.TopUpLink
common.OptionMap["ChatLink"] = common.ChatLink
common.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(common.QuotaPerUnit, 'f', -1, 64)
common.OptionMapRWMutex.Unlock()
loadOptionsFromDatabase()
}
@ -140,6 +142,8 @@ func updateOptionMap(key string, value string) (err error) {
common.AutomaticDisableChannelEnabled = boolValue
case "LogConsumeEnabled":
common.LogConsumeEnabled = boolValue
case "DisplayInCurrencyEnabled":
common.DisplayInCurrencyEnabled = boolValue
}
}
switch key {
@ -196,6 +200,8 @@ func updateOptionMap(key string, value string) (err error) {
common.ChatLink = value
case "ChannelDisableThreshold":
common.ChannelDisableThreshold, _ = strconv.ParseFloat(value, 64)
case "QuotaPerUnit":
common.QuotaPerUnit, _ = strconv.ParseFloat(value, 64)
}
return err
}

2
model/redemption.go
View File

@ -66,7 +66,7 @@ func Redeem(key string, userId int) (quota int, err error) {
if err != nil {
common.SysError("更新兑换码状态失败：" + err.Error())
}
RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %d 点额度", redemption.Quota))
RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s", common.LogQuota(redemption.Quota)))
}()
return redemption.Quota, nil
}

3
model/token.go
View File

@ -36,8 +36,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
if key == "" {
return nil, errors.New("未提供 token")
}
token = &Token{}
err = DB.Where("`key` = ?", key).First(token).Error
token, err = CacheGetTokenByKey(key)
if err == nil {
if token.Status != common.TokenStatusEnabled {
return nil, errors.New("该 token 状态不可用")

11
model/user.go
View File

@ -93,16 +93,16 @@ func (user *User) Insert(inviterId int) error {
return result.Error
}
if common.QuotaForNewUser > 0 {
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %d 点额度", common.QuotaForNewUser))
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(common.QuotaForNewUser)))
}
if inviterId != 0 {
if common.QuotaForInvitee > 0 {
_ = IncreaseUserQuota(user.Id, common.QuotaForInvitee)
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %d 点额度", common.QuotaForInvitee))
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(common.QuotaForInvitee)))
}
if common.QuotaForInviter > 0 {
_ = IncreaseUserQuota(inviterId, common.QuotaForInviter)
RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %d 点额度", common.QuotaForInviter))
RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(common.QuotaForInviter)))
}
}
return nil
@ -256,6 +256,11 @@ func GetUserQuota(id int) (quota int, err error) {
return quota, err
}
func GetUserUsedQuota(id int) (quota int, err error) {
err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
return quota, err
}
func GetUserEmail(id int) (email string, err error) {
err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
return email, err

2
web/src/App.js
View File

@ -48,6 +48,8 @@ function App() {
localStorage.setItem('system_name', data.system_name);
localStorage.setItem('logo', data.logo);
localStorage.setItem('footer_html', data.footer_html);
localStorage.setItem('quota_per_unit', data.quota_per_unit);
localStorage.setItem('display_in_currency', data.display_in_currency);
if (data.chat_link) {
localStorage.setItem('chat_link', data.chat_link);
} else {

39
web/src/components/OperationSetting.js
View File

@ -13,9 +13,11 @@ const OperationSetting = () => {
GroupRatio: '',
TopUpLink: '',
ChatLink: '',
QuotaPerUnit: 0,
AutomaticDisableChannelEnabled: '',
ChannelDisableThreshold: 0,
LogConsumeEnabled: ''
LogConsumeEnabled: '',
DisplayInCurrencyEnabled: ''
});
const [originInputs, setOriginInputs] = useState({});
let [loading, setLoading] = useState(false);
@ -118,6 +120,9 @@ const OperationSetting = () => {
if (originInputs['ChatLink'] !== inputs.ChatLink) {
await updateOption('ChatLink', inputs.ChatLink);
}
if (originInputs['QuotaPerUnit'] !== inputs.QuotaPerUnit) {
await updateOption('QuotaPerUnit', inputs.QuotaPerUnit);
}
break;
}
};
@ -129,7 +134,7 @@ const OperationSetting = () => {
<Header as='h3'>
通用设置
</Header>
<Form.Group widths={2}>
<Form.Group widths={3}>
<Form.Input
label='充值链接'
name='TopUpLink'
@ -148,6 +153,30 @@ const OperationSetting = () => {
type='link'
placeholder='例如 ChatGPT Next Web 的部署地址'
/>
<Form.Input
label='额度汇率'
name='QuotaPerUnit'
onChange={handleInputChange}
autoComplete='new-password'
value={inputs.QuotaPerUnit}
type='number'
step='0.01'
placeholder='一单位货币能兑换的额度'
/>
</Form.Group>
<Form.Group inline>
<Form.Checkbox
checked={inputs.LogConsumeEnabled === 'true'}
label='启用额度消费日志记录'
name='LogConsumeEnabled'
onChange={handleInputChange}
/>
<Form.Checkbox
checked={inputs.DisplayInCurrencyEnabled === 'true'}
label='以货币形式显示额度'
name='DisplayInCurrencyEnabled'
onChange={handleInputChange}
/>
</Form.Group>
<Form.Button onClick={() => {
submitConfig('general').then();
@ -264,12 +293,6 @@ const OperationSetting = () => {
placeholder='为一个 JSON 文本，键为分组名称，值为倍率'
/>
</Form.Group>
<Form.Checkbox
checked={inputs.LogConsumeEnabled === 'true'}
label='启用额度消费日志记录'
name='LogConsumeEnabled'
onChange={handleInputChange}
/>
<Form.Button onClick={() => {
submitConfig('ratio').then();
}}>保存倍率设置</Form.Button>

3
web/src/components/RedemptionsTable.js
View File

@ -4,6 +4,7 @@ import { Link } from 'react-router-dom';
import { API, copy, showError, showInfo, showSuccess, showWarning, timestamp2string } from '../helpers';
import { ITEMS_PER_PAGE } from '../constants';
import { renderQuota } from '../helpers/render';
function renderTimestamp(timestamp) {
return (
@ -220,7 +221,7 @@ const RedemptionsTable = () => {
<Table.Cell>{redemption.id}</Table.Cell>
<Table.Cell>{redemption.name ? redemption.name : '无'}</Table.Cell>
<Table.Cell>{renderStatus(redemption.status)}</Table.Cell>
<Table.Cell>{redemption.quota}</Table.Cell>
<Table.Cell>{renderQuota(redemption.quota)}</Table.Cell>
<Table.Cell>{renderTimestamp(redemption.created_time)}</Table.Cell>
<Table.Cell>{redemption.redeemed_time ? renderTimestamp(redemption.redeemed_time) : "尚未兑换"} </Table.Cell>
<Table.Cell>

3
web/src/components/TokensTable.js
View File

@ -4,6 +4,7 @@ import { Link } from 'react-router-dom';
import { API, copy, showError, showSuccess, showWarning, timestamp2string } from '../helpers';
import { ITEMS_PER_PAGE } from '../constants';
import { renderQuota } from '../helpers/render';
function renderTimestamp(timestamp) {
return (
@ -220,7 +221,7 @@ const TokensTable = () => {
<Table.Row key={token.id}>
<Table.Cell>{token.name ? token.name : '无'}</Table.Cell>
<Table.Cell>{renderStatus(token.status)}</Table.Cell>
<Table.Cell>{token.unlimited_quota ? '无限制' : token.remain_quota}</Table.Cell>
<Table.Cell>{token.unlimited_quota ? '无限制' : renderQuota(token.remain_quota, 2)}</Table.Cell>
<Table.Cell>{renderTimestamp(token.created_time)}</Table.Cell>
<Table.Cell>{token.expired_time === -1 ? '永不过期' : renderTimestamp(token.expired_time)}</Table.Cell>
<Table.Cell>

6
web/src/components/UsersTable.js
View File

@ -4,7 +4,7 @@ import { Link } from 'react-router-dom';
import { API, showError, showSuccess } from '../helpers';
import { ITEMS_PER_PAGE } from '../constants';
import { renderGroup, renderNumber, renderText } from '../helpers/render';
import { renderGroup, renderNumber, renderQuota, renderText } from '../helpers/render';
function renderRole(role) {
switch (role) {
@ -244,8 +244,8 @@ const UsersTable = () => {
{user.email ? <Popup hoverable content={user.email} trigger={<span>{renderText(user.email, 24)}</span>} /> : '无'}
</Table.Cell>
<Table.Cell>
<Popup content='剩余额度' trigger={<Label>{renderNumber(user.quota)}</Label>} />
<Popup content='已用额度' trigger={<Label>{renderNumber(user.used_quota)}</Label>} />
<Popup content='剩余额度' trigger={<Label>{renderQuota(user.quota)}</Label>} />
<Popup content='已用额度' trigger={<Label>{renderQuota(user.used_quota)}</Label>} />
<Popup content='请求次数' trigger={<Label>{renderNumber(user.request_count)}</Label>} />
</Table.Cell>
<Table.Cell>{renderRole(user.role)}</Table.Cell>

11
web/src/helpers/render.js
View File

@ -35,4 +35,15 @@ export function renderNumber(num) {
} else {
return num;
}
}
export function renderQuota(quota, digits = 2) {
let quotaPerUnit = localStorage.getItem('quota_per_unit');
let displayInCurrency = localStorage.getItem('display_in_currency');
quotaPerUnit = parseFloat(quotaPerUnit);
displayInCurrency = displayInCurrency === 'true';
if (displayInCurrency) {
return '$' + (quota / quotaPerUnit).toFixed(digits);
}
return renderNumber(quota);
}

3
web/src/pages/Redemption/EditRedemption.js
View File

@ -2,6 +2,7 @@ import React, { useEffect, useState } from 'react';
import { Button, Form, Header, Segment } from 'semantic-ui-react';
import { useParams } from 'react-router-dom';
import { API, downloadTextAsFile, showError, showSuccess } from '../../helpers';
import { renderQuota } from '../../helpers/render';
const EditRedemption = () => {
const params = useParams();
@ -87,7 +88,7 @@ const EditRedemption = () => {
</Form.Field>
<Form.Field>
<Form.Input
label='额度'
label={`额度（等价金额 ${renderQuota(quota)}）`}
name='quota'
placeholder={'请输入单个兑换码中包含的额度'}
onChange={handleInputChange}

3
web/src/pages/Token/EditToken.js
View File

@ -2,6 +2,7 @@ import React, { useEffect, useState } from 'react';
import { Button, Form, Header, Message, Segment } from 'semantic-ui-react';
import { useParams } from 'react-router-dom';
import { API, showError, showSuccess, timestamp2string } from '../../helpers';
import { renderQuota } from '../../helpers/render';
const EditToken = () => {
const params = useParams();
@ -137,7 +138,7 @@ const EditToken = () => {
<Message>注意，令牌的额度仅用于限制令牌本身的最大额度使用量，实际的使用受到账户的剩余额度限制。</Message>
<Form.Field>
<Form.Input
label='额度'
label={`额度（等价金额 ${renderQuota(remain_quota)}）`}
name='remain_quota'
placeholder={'请输入额度'}
onChange={handleInputChange}

3
web/src/pages/TopUp/index.js
View File

@ -1,6 +1,7 @@
import React, { useEffect, useState } from 'react';
import { Button, Form, Grid, Header, Segment, Statistic } from 'semantic-ui-react';
import { API, showError, showInfo, showSuccess } from '../../helpers';
import { renderQuota } from '../../helpers/render';
const TopUp = () => {
const [redemptionCode, setRedemptionCode] = useState('');
@ -81,7 +82,7 @@ const TopUp = () => {
<Grid.Column>
<Statistic.Group widths='one'>
<Statistic>
<Statistic.Value>{userQuota.toLocaleString()}</Statistic.Value>
<Statistic.Value>{renderQuota(userQuota)}</Statistic.Value>
<Statistic.Label>剩余额度</Statistic.Label>
</Statistic>
</Statistic.Group>