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

feat: add official feature

Browse Source
This commit is contained in:
Zhang Minghan 2023-11-28 09:11:36 +08:00
parent 989e61cc62
commit e862bf4651
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
manager/transhipment.go
View File

@ -28,6 +28,7 @@ type TranshipmentForm struct {
TopK *int `json:"top_k"`
Tools *globals.FunctionTools
ToolChoice *interface{}
Official bool `json:"official"`
}
type Choice struct {
@ -49,7 +50,7 @@ type TranshipmentResponse struct {
Model string `json:"model"`
Choices []Choice `json:"choices"`
Usage Usage `json:"usage"`
Quota float32 `json:"quota"`
Quota *float32 `json:"quota,omitempty"`
}
type ChoiceDelta struct {
@ -65,7 +66,7 @@ type TranshipmentStreamResponse struct {
Model string `json:"model"`
Choices []ChoiceDelta `json:"choices"`
Usage Usage `json:"usage"`
Quota float32 `json:"quota"`
Quota *float32 `json:"quota,omitempty"`
}
func ModelAPI(c *gin.Context) {
@ -116,6 +117,11 @@ func TranshipmentAPI(c *gin.Context) {
}
}
if strings.HasSuffix(form.Model, "-official") {
form.Model = strings.TrimSuffix(form.Model, "-official")
form.Official = true
}
check, plan := auth.CanEnableModelWithSubscription(db, cache, user, form.Model)
if !check {
c.JSON(http.StatusForbidden, gin.H{
@ -185,7 +191,7 @@ func sendTranshipmentResponse(c *gin.Context, form TranshipmentForm, id string,
CompletionTokens: buffer.CountOutputToken(),
TotalTokens: buffer.CountToken(),
},
Quota: buffer.GetQuota(),
Quota: utils.Multi[*float32](form.Official, nil, utils.ToPtr(buffer.GetQuota())),
})
}
@ -210,7 +216,7 @@ func getStreamTranshipmentForm(id string, created int64, form TranshipmentForm,
CompletionTokens: utils.MultiF(end, func() int { return buffer.CountOutputToken() }, 0),
TotalTokens: utils.MultiF(end, func() int { return buffer.CountToken() }, 0),
},
Quota: buffer.GetQuota(),
Quota: utils.Multi[*float32](form.Official, nil, utils.ToPtr(buffer.GetQuota())),
}
}