GithubProxy/coai
1
0
Fork 0
mirror of https://github.com/coaidev/coai.git synced 2025-05-19 04:50:14 +09:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: add input token and quota calculation for API request pre-checks

Browse Source
This commit is contained in:
Sh1n3zZ 2025-02-24 20:24:18 +08:00
parent 4d841c4a33
commit 9f6cec4298
No known key found for this signature in database
GPG Key ID: 696702CF723B0452
6 changed files with 24 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
addition/generation/api.go
View File

@ -5,8 +5,9 @@ import (
"chat/globals"
"chat/utils"
"fmt"
"github.com/gin-gonic/gin"
"strings"
"github.com/gin-gonic/gin"
)
type WebsocketGenerationForm struct {
@ -53,7 +54,7 @@ func GenerateAPI(c *gin.Context) {
return
}
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, form.Model)
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, form.Model, []globals.Message{})
if check != nil {
conn.Send(globals.GenerationSegmentResponse{
Message: check.Error(),

21
auth/rule.go
View File

@ -4,6 +4,10 @@ import (
"chat/channel"
"database/sql"
"fmt"
"chat/globals"
"chat/utils"
"github.com/go-redis/redis/v8"
)
@ -11,10 +15,11 @@ const (
ErrNotAuthenticated = "not authenticated error (model: %s)"
ErrNotSetPrice = "the price of the model is not set (model: %s)"
ErrNotEnoughQuota = "user quota is not enough error (model: %s, minimum quota: %0.2f, your quota: %0.2f)"
ErrEstimatedCost = "estimated cost exceeds user quota (model: %s, estimated cost: %0.2f, your quota: %0.2f)"
)
// CanEnableModel returns whether the model can be enabled (without subscription)
func CanEnableModel(db *sql.DB, user *User, model string) error {
func CanEnableModel(db *sql.DB, user *User, model string, messages []globals.Message) error {
isAuth := user != nil
isAdmin := isAuth && user.IsAdmin(db)
@ -37,21 +42,23 @@ func CanEnableModel(db *sql.DB, user *User, model string) error {
return fmt.Errorf(ErrNotAuthenticated, model)
}
// return if the user is authenticated and has enough quota
limit := charge.GetLimit()
// Calculate estimated input cost
inputTokens := utils.NumTokensFromMessages(messages, model, false)
estimatedInputCost := float32(inputTokens) / 1000 * charge.GetInput()
// Get user's current quota
quota := user.GetQuota(db)
if quota < limit {
return fmt.Errorf(ErrNotEnoughQuota, model, limit, quota)
if quota < estimatedInputCost {
return fmt.Errorf(ErrEstimatedCost, model, estimatedInputCost, quota)
}
return nil
}
func CanEnableModelWithSubscription(db *sql.DB, cache *redis.Client, user *User, model string) (canEnable error, usePlan bool) {
func CanEnableModelWithSubscription(db *sql.DB, cache *redis.Client, user *User, model string, messages []globals.Message) (canEnable error, usePlan bool) {
// use subscription quota first
if user != nil && HandleSubscriptionUsage(db, cache, user, model) {
return nil, true
}
return CanEnableModel(db, user, model), false
return CanEnableModel(db, user, model, messages), false
}

2
manager/chat.go
View File

@ -202,7 +202,7 @@ func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conve
model := instance.GetModel()
segment := adapter.ClearMessages(model, web.ToChatSearched(instance, restart))
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, model)
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, model, segment)
conn.Send(globals.ChatSegmentResponse{
Conversation: instance.GetId(),
})

8
manager/chat_completions.go
View File

@ -28,12 +28,12 @@ func supportRelayPlan() bool {
return channel.SystemInstance.SupportRelayPlan()
}
func checkEnableState(db *sql.DB, cache *redis.Client, user *auth.User, model string) (state error, plan bool) {
func checkEnableState(db *sql.DB, cache *redis.Client, user *auth.User, model string, messages []globals.Message) (state error, plan bool) {
if supportRelayPlan() {
return auth.CanEnableModelWithSubscription(db, cache, user, model)
return auth.CanEnableModelWithSubscription(db, cache, user, model, messages)
}
return auth.CanEnableModel(db, user, model), false
return auth.CanEnableModel(db, user, model, messages), false
}
func ChatRelayAPI(c *gin.Context) {
@ -80,7 +80,7 @@ func ChatRelayAPI(c *gin.Context) {
form.Official = true
}
check, plan := checkEnableState(db, cache, user, form.Model)
check, plan := checkEnableState(db, cache, user, form.Model, messages)
if check != nil {
sendErrorResponse(c, check, "quota_exceeded_error")
return

2
manager/completions.go
View File

@ -28,7 +28,7 @@ func NativeChatHandler(c *gin.Context, user *auth.User, model string, message []
db := utils.GetDBFromContext(c)
cache := utils.GetCacheFromContext(c)
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, model)
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, model, segment)
if check != nil {
return check.Error(), 0

2
manager/images.go
View File

@ -54,7 +54,7 @@ func ImagesRelayAPI(c *gin.Context) {
form.Model = strings.TrimSuffix(form.Model, "-official")
}
check := auth.CanEnableModel(db, user, form.Model)
check := auth.CanEnableModel(db, user, form.Model, []globals.Message{})
if check != nil {
sendErrorResponse(c, check, "quota_exceeded_error")
return