feat: Support Cloudflare AI

This commit is contained in:
Martial BE
2024-04-16 18:08:56 +08:00
parent 5606a104f6
commit 344555418e
14 changed files with 606 additions and 50 deletions

View File

@@ -104,6 +104,10 @@ func InitDB() (err error) {
return nil
}
common.SysLog("database migration started")
// err = MigrateDB(DB)
// if err != nil {
// return err
// }
err = db.AutoMigrate(&Channel{})
if err != nil {
return err
@@ -157,6 +161,24 @@ func InitDB() (err error) {
return err
}
// func MigrateDB(db *gorm.DB) error {
// if DB.Migrator().HasConstraint(&Price{}, "model") {
// fmt.Println("----Price model has constraint----")
// // 如果是主键,移除主键约束
// err := db.Migrator().DropConstraint(&Price{}, "model")
// if err != nil {
// return err
// }
// // 修改字段长度
// err = db.Migrator().AlterColumn(&Price{}, "model")
// if err != nil {
// return err
// }
// }
// return nil
// }
func CloseDB() error {
sqlDB, err := DB.DB()
if err != nil {

View File

@@ -15,7 +15,7 @@ const (
)
type Price struct {
Model string `json:"model" gorm:"type:varchar(30);primaryKey" binding:"required"`
Model string `json:"model" gorm:"type:varchar(100)" binding:"required"`
Type string `json:"type" gorm:"default:'tokens'" binding:"required"`
ChannelType int `json:"channel_type" gorm:"default:0" binding:"gte=0"`
Input float64 `json:"input" gorm:"default:0" binding:"gte=0"`
@@ -98,11 +98,7 @@ func DeleteAllPrices(tx *gorm.DB) error {
}
func (price *Price) Delete() error {
err := DB.Delete(price).Error
if err != nil {
return err
}
return err
return DB.Where("model = ?", price.Model).Delete(&Price{}).Error
}
type ModelType struct {
@@ -289,6 +285,16 @@ func GetDefaultPrice() []*Price {
"yi-34b-chat-200k": {[]float64{0.8571, 0.8571}, common.ChannelTypeLingyi},
// 6 元 / 1M tokens 0.006 / 1k tokens
"yi-vl-plus": {[]float64{0.4286, 0.4286}, common.ChannelTypeLingyi},
"@cf/stabilityai/stable-diffusion-xl-base-1.0": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@cf/lykon/dreamshaper-8-lcm": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@cf/bytedance/stable-diffusion-xl-lightning": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@cf/qwen/qwen1.5-7b-chat-awq": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@cf/qwen/qwen1.5-14b-chat-awq": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@hf/thebloke/deepseek-coder-6.7b-base-awq": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@hf/google/gemma-7b-it": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@hf/thebloke/llama-2-13b-chat-awq": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
"@cf/openai/whisper": {[]float64{0, 0}, common.ChannelTypeCloudflareAI},
}
var prices []*Price