mirror of
https://github.com/coaidev/coai.git
synced 2025-05-20 05:20:15 +09:00
103 lines
2.5 KiB
Go
103 lines
2.5 KiB
Go
package generation
|
|
|
|
import (
|
|
"chat/auth"
|
|
"chat/globals"
|
|
"chat/utils"
|
|
"fmt"
|
|
"github.com/gin-gonic/gin"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
type WebsocketGenerationForm struct {
|
|
Token string `json:"token" binding:"required"`
|
|
Prompt string `json:"prompt" binding:"required"`
|
|
Model string `json:"model" binding:"required"`
|
|
}
|
|
|
|
func ProjectTarDownloadAPI(c *gin.Context) {
|
|
hash := strings.TrimSpace(c.Query("hash"))
|
|
c.Writer.Header().Add("Content-Disposition", "attachment; filename=code.tar.gz")
|
|
c.File(fmt.Sprintf("generation/data/out/%s.tar.gz", hash))
|
|
}
|
|
|
|
func ProjectZipDownloadAPI(c *gin.Context) {
|
|
hash := strings.TrimSpace(c.Query("hash"))
|
|
c.Writer.Header().Add("Content-Disposition", "attachment; filename=code.zip")
|
|
c.File(fmt.Sprintf("generation/data/out/%s.zip", hash))
|
|
}
|
|
|
|
func GenerateAPI(c *gin.Context) {
|
|
var conn *utils.WebSocket
|
|
if conn = utils.NewWebsocket(c); conn == nil {
|
|
return
|
|
}
|
|
defer conn.DeferClose()
|
|
|
|
var form *WebsocketGenerationForm
|
|
if form = utils.ReadForm[WebsocketGenerationForm](conn); form == nil {
|
|
return
|
|
}
|
|
|
|
user := auth.ParseToken(c, form.Token)
|
|
authenticated := user != nil
|
|
|
|
db := utils.GetDBFromContext(c)
|
|
cache := utils.GetCacheFromContext(c)
|
|
|
|
id := auth.GetId(db, user)
|
|
|
|
if !utils.IncrWithLimit(cache,
|
|
fmt.Sprintf(":generation:%s", utils.Multi[string](authenticated, strconv.FormatInt(id, 10), c.ClientIP())),
|
|
1,
|
|
30,
|
|
3600,
|
|
) {
|
|
conn.Send(globals.GenerationSegmentResponse{
|
|
End: true,
|
|
Error: "generation rate limit exceeded, the max generation rate is 30 per hour.",
|
|
})
|
|
return
|
|
}
|
|
|
|
reversible := auth.CanEnableSubscription(db, cache, user)
|
|
if !auth.CanEnableModelWithSubscription(db, user, form.Model, reversible) {
|
|
conn.Send(globals.GenerationSegmentResponse{
|
|
Message: "You don't have enough quota to use this model.",
|
|
Quota: 0,
|
|
End: true,
|
|
})
|
|
return
|
|
}
|
|
|
|
var instance *utils.Buffer
|
|
hash, err := CreateGenerationWithCache(form.Model, form.Prompt, reversible, func(buffer *utils.Buffer, data string) {
|
|
instance = buffer
|
|
conn.Send(globals.GenerationSegmentResponse{
|
|
End: false,
|
|
Message: data,
|
|
Quota: buffer.GetQuota(),
|
|
})
|
|
})
|
|
|
|
if instance != nil && !reversible && instance.GetQuota() > 0 && user != nil {
|
|
user.UseQuota(db, instance.GetQuota())
|
|
}
|
|
|
|
if err != nil {
|
|
conn.Send(globals.GenerationSegmentResponse{
|
|
End: true,
|
|
Error: err.Error(),
|
|
Quota: instance.GetQuota(),
|
|
})
|
|
return
|
|
}
|
|
|
|
conn.Send(globals.GenerationSegmentResponse{
|
|
End: true,
|
|
Hash: hash,
|
|
Quota: instance.GetQuota(),
|
|
})
|
|
}
|