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

fix: fix relay plan

Browse Source 
fix: fix relay plan
Co-Authored-By: Minghan Zhang <112773885+zmh-program@users.noreply.github.com>
This commit is contained in:
Deng Junhai 2024-04-04 12:14:19 +08:00
parent a888fa62ef
commit 269c3a23fa
1 changed files with 20 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
manager/chat_completions.go
View File

@ -1,19 +1,22 @@
package manager package manager
import ( import (
"chat/adapter/common" adaptercommon "chat/adapter/common"
"chat/addition/web" "chat/addition/web"
"chat/admin" "chat/admin"
"chat/auth" "chat/auth"
"chat/channel" "chat/channel"
"chat/globals" "chat/globals"
"chat/utils" "chat/utils"
"database/sql"
"fmt" "fmt"
"github.com/gin-gonic/gin"
"io" "io"
"net/http" "net/http"
"strings" "strings"
"time" "time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
) )
const ( const (
@ -25,6 +28,14 @@ func supportRelayPlan() bool {
return channel.SystemInstance.SupportRelayPlan() return channel.SystemInstance.SupportRelayPlan()
} }
func checkEnableState(db *sql.DB, cache *redis.Client, user *auth.User, model string) (state error, plan bool) {
if supportRelayPlan() {
return auth.CanEnableModelWithSubscription(db, cache, user, model)
}
return auth.CanEnableModel(db, user, model), false
}
func ChatRelayAPI(c *gin.Context) { func ChatRelayAPI(c *gin.Context) {
if globals.CloseRelay { if globals.CloseRelay {
abortWithErrorResponse(c, fmt.Errorf("relay api is denied of access"), "access_denied_error") abortWithErrorResponse(c, fmt.Errorf("relay api is denied of access"), "access_denied_error")
@ -49,6 +60,7 @@ func ChatRelayAPI(c *gin.Context) {
} }
db := utils.GetDBFromContext(c) db := utils.GetDBFromContext(c)
cache := utils.GetCacheFromContext(c)
user := &auth.User{ user := &auth.User{
Username: username, Username: username,
} }
@ -68,20 +80,20 @@ func ChatRelayAPI(c *gin.Context) {
form.Official = true form.Official = true
} }
check := auth.CanEnableModel(db, user, form.Model) check, plan := checkEnableState(db, cache, user, form.Model)
if check != nil { if check != nil {
sendErrorResponse(c, check, "quota_exceeded_error") sendErrorResponse(c, check, "quota_exceeded_error")
return return
} }
if form.Stream { if form.Stream {
sendStreamTranshipmentResponse(c, form, messages, id, created, user, supportRelayPlan()) sendStreamTranshipmentResponse(c, form, messages, id, created, user, plan)
} else { } else {
sendTranshipmentResponse(c, form, messages, id, created, user, supportRelayPlan()) sendTranshipmentResponse(c, form, messages, id, created, user, plan)
} }
} }
func getChatProps(form RelayForm, messages []globals.Message, buffer *utils.Buffer, plan bool) *adaptercommon.ChatProps { func getChatProps(form RelayForm, messages []globals.Message, buffer *utils.Buffer) *adaptercommon.ChatProps {
return &adaptercommon.ChatProps{ return &adaptercommon.ChatProps{
Model: form.Model, Model: form.Model,
Message: messages, Message: messages,
@ -103,7 +115,7 @@ func sendTranshipmentResponse(c *gin.Context, form RelayForm, messages []globals
cache := utils.GetCacheFromContext(c) cache := utils.GetCacheFromContext(c)
buffer := utils.NewBuffer(form.Model, messages, channel.ChargeInstance.GetCharge(form.Model)) buffer := utils.NewBuffer(form.Model, messages, channel.ChargeInstance.GetCharge(form.Model))
hit, err := channel.NewChatRequestWithCache(cache, buffer, auth.GetGroup(db, user), getChatProps(form, messages, buffer, plan), func(data *globals.Chunk) error { hit, err := channel.NewChatRequestWithCache(cache, buffer, auth.GetGroup(db, user), getChatProps(form, messages, buffer), func(data *globals.Chunk) error {
buffer.WriteChunk(data) buffer.WriteChunk(data)
return nil return nil
}) })
@ -212,7 +224,7 @@ func sendStreamTranshipmentResponse(c *gin.Context, form RelayForm, messages []g
go func() { go func() {
buffer := utils.NewBuffer(form.Model, messages, charge) buffer := utils.NewBuffer(form.Model, messages, charge)
hit, err := channel.NewChatRequestWithCache( hit, err := channel.NewChatRequestWithCache(
cache, buffer, group, getChatProps(form, messages, buffer, plan), cache, buffer, group, getChatProps(form, messages, buffer),
func(data *globals.Chunk) error { func(data *globals.Chunk) error {
buffer.WriteChunk(data) buffer.WriteChunk(data)