xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-12-25 01:25:59 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add 'type' field for ChatModel, support Chat and Image model

Browse Source
This commit is contained in:
RockYang
2024-12-25 18:57:18 +08:00
parent cbf06eea24
commit acee2d9d81
35 changed files with 766 additions and 698 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
api/core/app_server.go
View File

@@ -15,12 +15,6 @@ import (
"geekai/store/model"
"geekai/utils"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/golang-jwt/jwt/v5"
"github.com/nfnt/resize"
"golang.org/x/image/webp"
"gorm.io/gorm"
"image"
"image/jpeg"
"io"
@@ -29,6 +23,13 @@ import (
"runtime/debug"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/golang-jwt/jwt/v5"
"github.com/nfnt/resize"
"golang.org/x/image/webp"
"gorm.io/gorm"
)
type AppServer struct {
@@ -228,6 +229,7 @@ func needLogin(c *gin.Context) bool {
c.Request.URL.Path == "/api/suno/detail" ||
c.Request.URL.Path == "/api/suno/play" ||
c.Request.URL.Path == "/api/download" ||
c.Request.URL.Path == "/api/dall/models" ||
strings.HasPrefix(c.Request.URL.Path, "/api/test") ||
strings.HasPrefix(c.Request.URL.Path, "/api/payment/notify/") ||
strings.HasPrefix(c.Request.URL.Path, "/api/user/clogin") ||

4
api/core/types/config.go
View File

@@ -133,7 +133,7 @@ type SystemConfig struct {
AdminTitle string `json:"admin_title,omitempty"` // 管理后台标题
Logo string `json:"logo,omitempty"` // 方形 Logo
InitPower int `json:"init_power,omitempty"` // 新用户注册赠送算力值
DailyPower int `json:"daily_power,omitempty"` // 每日赠送算力
DailyPower int `json:"daily_power,omitempty"` // 每日签到赠送算力
InvitePower int `json:"invite_power,omitempty"` // 邀请新用户赠送算力值
VipMonthPower int `json:"vip_month_power,omitempty"` // VIP 会员每月赠送的算力值
@@ -146,7 +146,7 @@ type SystemConfig struct {
MjPower int `json:"mj_power,omitempty"` // MJ 绘画消耗算力
MjActionPower int `json:"mj_action_power,omitempty"` // MJ 操作（放大，变换）消耗算力
SdPower int `json:"sd_power,omitempty"` // SD 绘画消耗算力
DallPower int `json:"dall_power,omitempty"` // DALL-E-3 绘图消耗算力
DallPower int `json:"dall_power,omitempty"` // DALL-E-3 绘图消耗算力
SunoPower int `json:"suno_power,omitempty"` // Suno 生成歌曲消耗算力
LumaPower int `json:"luma_power,omitempty"` // Luma 生成视频消耗算力
AdvanceVoicePower int `json:"advance_voice_power,omitempty"` // 高级语音对话消耗算力

5
api/core/types/task.go
View File

@@ -74,6 +74,8 @@ type SdTaskParams struct {
// DallTask DALL-E task
type DallTask struct {
ClientId string `json:"client_id"`
ModelId uint `json:"model_id"`
ModelName string `json:"model_name"`
Id uint `json:"id"`
UserId uint `json:"user_id"`
Prompt string `json:"prompt"`
@@ -81,8 +83,7 @@ type DallTask struct {
Quality string `json:"quality"`
Size string `json:"size"`
Style string `json:"style"`
Power int `json:"power"`
Power int `json:"power"`
TranslateModelId int `json:"translate_model_id"` // 提示词翻译模型ID
}

3
api/handler/admin/chat_model_handler.go
View File

@@ -43,6 +43,7 @@ func (h *ChatModelHandler) Save(c *gin.Context) {
Temperature float32 `json:"temperature"` // 模型温度
KeyId int `json:"key_id,omitempty"`
CreatedAt int64 `json:"created_at"`
Type string `json:"type"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
@@ -65,7 +66,7 @@ func (h *ChatModelHandler) Save(c *gin.Context) {
item.MaxContext = data.MaxContext
item.Temperature = data.Temperature
item.KeyId = data.KeyId
item.Type = data.Type
var res *gorm.DB
if data.Id > 0 {
res = h.DB.Save(&item)

2
api/handler/chat_model_handler.go
View File

@@ -30,7 +30,7 @@ func NewChatModelHandler(app *core.AppServer, db *gorm.DB) *ChatModelHandler {
func (h *ChatModelHandler) List(c *gin.Context) {
var items []model.ChatModel
var chatModels = make([]vo.ChatModel, 0)
session := h.DB.Session(&gorm.Session{}).Where("enabled", true)
session := h.DB.Session(&gorm.Session{}).Where("type", "chat").Where("enabled", true)
t := c.Query("type")
if t != "" {
session = session.Where("type", t)

21
api/handler/openai_handler.go → api/handler/chat_openai_handler.go
View File

@@ -17,10 +17,11 @@ import (
"geekai/store/model"
"geekai/store/vo"
"geekai/utils"
req2 "github.com/imroc/req/v3"
"io"
"strings"
"time"
req2 "github.com/imroc/req/v3"
)
type Usage struct {
@@ -172,16 +173,22 @@ func (h *ChatHandler) sendOpenAiMessage(
var apiRes types.BizVo
r, err := req2.C().R().SetHeader("Body-Type", "application/json").
SetHeader("Authorization", function.Token).
SetBody(params).
SetSuccessResult(&apiRes).Post(function.Action)
SetBody(params).Post(function.Action)
errMsg := ""
if err != nil {
errMsg = err.Error()
} else if r.IsErrorState() {
errMsg = r.Status
} else {
all, _ := io.ReadAll(r.Body)
err = json.Unmarshal(all, &apiRes)
if err != nil {
errMsg = err.Error()
} else if apiRes.Code != types.Success {
errMsg = apiRes.Message
}
}
if errMsg != "" || apiRes.Code != types.Success {
errMsg = "调用函数工具出错：" + apiRes.Message + errMsg
if errMsg != "" {
errMsg = "调用函数工具出错：" + errMsg
contents = append(contents, errMsg)
} else {
errMsg = utils.InterfaceToString(apiRes.Data)

79
api/handler/dalle_handler.go
View File

@@ -8,6 +8,7 @@ package handler
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
@@ -17,14 +18,13 @@ import (
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"gorm.io/gorm"
)
type DallJobHandler struct {
BaseHandler
redis *redis.Client
dallService *dalle.Service
uploader *oss.UploaderManager
userService *service.UserService
@@ -42,49 +42,49 @@ func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service,
}
}
func (h *DallJobHandler) preCheck(c *gin.Context) bool {
user, err := h.GetLoginUser(c)
if err != nil {
resp.NotAuth(c)
return false
}
if user.Power < h.App.SysConfig.DallPower {
resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
return false
}
return true
}
// Image 创建一个绘画任务
func (h *DallJobHandler) Image(c *gin.Context) {
if !h.preCheck(c) {
return
}
var data types.DallTask
if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
resp.ERROR(c, types.InvalidArgs)
return
}
var chatModel model.ChatModel
if res := h.DB.Where("id = ?", data.ModelId).First(&chatModel); res.Error != nil {
resp.ERROR(c, "模型不存在")
return
}
// 检查用户剩余算力
user, err := h.GetLoginUser(c)
if err != nil {
resp.NotAuth(c)
return
}
if user.Power < chatModel.Power {
resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
return
}
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
task := types.DallTask{
ClientId: data.ClientId,
UserId: uint(userId),
ModelId: chatModel.Id,
ModelName: chatModel.Value,
Prompt: data.Prompt,
Quality: data.Quality,
Size: data.Size,
Style: data.Style,
Power: h.App.SysConfig.DallPower,
TranslateModelId: h.App.SysConfig.TranslateModelId,
Power: chatModel.Power,
}
job := model.DallJob{
UserId: uint(userId),
Prompt: data.Prompt,
Power: task.Power,
Power: chatModel.Power,
TaskInfo: utils.JsonEncode(task),
}
res := h.DB.Create(&job)
@@ -95,6 +95,17 @@ func (h *DallJobHandler) Image(c *gin.Context) {
task.Id = job.Id
h.dallService.PushTask(task)
// 扣减算力
err = h.userService.DecreasePower(int(user.Id), chatModel.Power, model.PowerLog{
Type: types.PowerConsume,
Model: chatModel.Value,
Remark: fmt.Sprintf("绘画提示词：%s", utils.CutWords(task.Prompt, 10)),
})
if err != nil {
resp.ERROR(c, "error with decrease power: "+err.Error())
return
}
resp.SUCCESS(c)
}
@@ -210,3 +221,25 @@ func (h *DallJobHandler) Publish(c *gin.Context) {
resp.SUCCESS(c)
}
func (h *DallJobHandler) GetModels(c *gin.Context) {
var models []model.ChatModel
err := h.DB.Where("type", "img").Where("enabled", true).Find(&models).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
var modelVos []vo.ChatModel
for _, v := range models {
var modelVo vo.ChatModel
err := utils.CopyObject(v, &modelVo)
if err != nil {
continue
}
modelVo.Id = v.Id
modelVos = append(modelVos, modelVo)
}
resp.SUCCESS(c, modelVos)
}

67
api/handler/function_handler.go
View File

@@ -12,6 +12,7 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/service/dalle"
"geekai/service/oss"
"geekai/store/model"
@@ -32,6 +33,7 @@ type FunctionHandler struct {
config types.ApiConfig
uploadManager *oss.UploaderManager
dallService *dalle.Service
userService *service.UserService
}
func NewFunctionHandler(
@@ -39,7 +41,8 @@ func NewFunctionHandler(
db *gorm.DB,
config *types.AppConfig,
manager *oss.UploaderManager,
dallService *dalle.Service) *FunctionHandler {
dallService *dalle.Service,
userService *service.UserService) *FunctionHandler {
return &FunctionHandler{
BaseHandler: BaseHandler{
App: server,
@@ -48,6 +51,7 @@ func NewFunctionHandler(
config: config.ApiConfig,
uploadManager: manager,
dallService: dallService,
userService: userService,
}
}
@@ -152,8 +156,12 @@ func (h *FunctionHandler) ZaoBao(c *gin.Context) {
SetHeader("AppId", h.config.AppId).
SetHeader("Authorization", fmt.Sprintf("Bearer %s", h.config.Token)).
SetSuccessResult(&res).Get(url)
if err != nil || r.IsErrorState() {
resp.ERROR(c, fmt.Sprintf("%v%v", err, r.Err))
if err != nil {
resp.ERROR(c, fmt.Sprintf("%v", err))
return
}
if r.IsErrorState() {
resp.ERROR(c, fmt.Sprintf("%v", r.Err))
return
}
@@ -167,7 +175,7 @@ func (h *FunctionHandler) ZaoBao(c *gin.Context) {
for _, v := range res.Data.Items {
builder = append(builder, v.Title)
}
builder = append(builder, fmt.Sprintf("%s", res.Data.Title))
builder = append(builder, res.Data.Title)
resp.SUCCESS(c, strings.Join(builder, "\n\n"))
}
@@ -199,33 +207,48 @@ func (h *FunctionHandler) Dall3(c *gin.Context) {
// create dall task
prompt := utils.InterfaceToString(params["prompt"])
job := model.DallJob{
UserId: user.Id,
Prompt: prompt,
Power: h.App.SysConfig.DallPower,
task := types.DallTask{
UserId: user.Id,
Prompt: prompt,
ModelId: 0,
ModelName: "dall-e-3",
TranslateModelId: h.App.SysConfig.TranslateModelId,
N: 1,
Quality: "standard",
Size: "1024x1024",
Style: "vivid",
Power: h.App.SysConfig.DallPower,
}
res = h.DB.Create(&job)
if res.Error != nil {
resp.ERROR(c, "创建 DALL-E 绘图任务失败："+res.Error.Error())
job := model.DallJob{
UserId: user.Id,
Prompt: prompt,
Power: h.App.SysConfig.DallPower,
TaskInfo: utils.JsonEncode(task),
}
err := h.DB.Create(&job).Error
if err != nil {
resp.ERROR(c, "创建 DALL-E 绘图任务失败："+err.Error())
return
}
content, err := h.dallService.Image(types.DallTask{
Id: job.Id,
UserId: user.Id,
Prompt: job.Prompt,
N: 1,
Quality: "standard",
Size: "1024x1024",
Style: "vivid",
Power: job.Power,
}, true)
task.Id = job.Id
content, err := h.dallService.Image(task, true)
if err != nil {
resp.ERROR(c, "任务执行失败："+err.Error())
return
}
// 扣减算力
err = h.userService.DecreasePower(int(user.Id), job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: task.ModelName,
Remark: fmt.Sprintf("绘画提示词：%s", utils.CutWords(job.Prompt, 10)),
})
if err != nil {
resp.ERROR(c, "扣减算力失败："+err.Error())
return
}
resp.SUCCESS(c, content)
}

1
api/main.go
View File

@@ -474,6 +474,7 @@ func main() {
group.GET("imgWall", h.ImgWall)
group.GET("remove", h.Remove)
group.GET("publish", h.Publish)
group.GET("models", h.GetModels)
}),
fx.Provide(handler.NewSunoHandler),
fx.Invoke(func(s *core.AppServer, h *handler.SunoHandler) {

85
api/service/dalle/service.go
View File

@@ -8,7 +8,6 @@ package dalle
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"errors"
"fmt"
"geekai/core/types"
logger2 "geekai/logger"
@@ -17,9 +16,11 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"io"
"time"
"github.com/go-redis/redis/v8"
"github.com/imroc/req/v3"
"gorm.io/gorm"
)
@@ -100,17 +101,18 @@ func (s *Service) Run() {
type imgReq struct {
Model string `json:"model"`
Prompt string `json:"prompt"`
N int `json:"n"`
Size string `json:"size"`
Quality string `json:"quality"`
Style string `json:"style"`
N int `json:"n,omitempty"`
Size string `json:"size,omitempty"`
Quality string `json:"quality,omitempty"`
Style string `json:"style,omitempty"`
}
type imgRes struct {
Created int64 `json:"created"`
Data []struct {
RevisedPrompt string `json:"revised_prompt"`
Url string `json:"url"`
RevisedPrompt string `json:"revised_prompt,omitempty"`
Url string `json:"url,omitempty"`
B64Json string `json:"b64_json,omitempty"`
} `json:"data"`
}
@@ -135,29 +137,20 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
}
}
var user model.User
s.db.Where("id", task.UserId).First(&user)
if user.Power < task.Power {
return "", errors.New("insufficient of power")
}
// 扣减算力
err := s.userService.DecreasePower(int(user.Id), task.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "dall-e-3",
Remark: fmt.Sprintf("绘画提示词：%s", utils.CutWords(task.Prompt, 10)),
})
if err != nil {
return "", fmt.Errorf("error with decrease power: %v", err)
}
var chatModel model.ChatModel
s.db.Where("id = ?", task.ModelId).First(&chatModel)
// get image generation API KEY
var apiKey model.ApiKey
err = s.db.Where("type", "dalle").
Where("enabled", true).
Order("last_used_at ASC").First(&apiKey).Error
session := s.db.Where("enabled", true)
if chatModel.KeyId > 0 {
session = session.Where("id = ?", chatModel.KeyId)
} else {
session = session.Where("type = ?", "dalle")
}
err := session.Order("last_used_at ASC").First(&apiKey).Error
if err != nil {
return "", fmt.Errorf("no available DALL-E api key: %v", err)
return "", fmt.Errorf("no available Image Generation api key: %v", err)
}
var res imgRes
@@ -167,7 +160,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
}
apiURL := fmt.Sprintf("%s/v1/images/generations", apiKey.ApiURL)
reqBody := imgReq{
Model: "dall-e-3",
Model: chatModel.Value,
Prompt: prompt,
N: 1,
Size: task.Size,
@@ -182,20 +175,39 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
SetSuccessResult(&res).
Post(apiURL)
if err != nil {
logger.Errorf("error with send request: %v", err)
return "", fmt.Errorf("error with send request: %v", err)
}
if r.IsErrorState() {
logger.Errorf("error with send request, status: %s, %+v", r.Status, errRes.Error)
return "", fmt.Errorf("error with send request, status: %s, %+v", r.Status, errRes.Error)
}
all, _ := io.ReadAll(r.Body)
logger.Debugf("response: %+v", string(all))
// update the api key last use time
s.db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
// update task progress
err = s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
var imgURL string
var data = map[string]interface{}{
"progress": 100,
"org_url": res.Data[0].Url,
"prompt": prompt,
}).Error
}
// 如果返回的是base64，则需要上传到oss
if res.Data[0].B64Json != "" {
imgURL, err = s.uploadManager.GetUploadHandler().PutBase64(res.Data[0].B64Json)
if err != nil {
return "", fmt.Errorf("error with upload image: %v", err)
}
logger.Infof("upload image to oss: %s", imgURL)
data["img_url"] = imgURL
} else {
imgURL = res.Data[0].Url
}
data["org_url"] = imgURL
// update task progress
err = s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(data).Error
if err != nil {
return "", fmt.Errorf("err with update database: %v", err)
}
@@ -252,9 +264,14 @@ func (s *Service) CheckTaskStatus() {
// 找出失败的任务，并恢复其扣减算力
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
for _, job := range jobs {
err := s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
var task types.DallTask
err := utils.JsonDecode(job.TaskInfo, &task)
if err != nil {
continue
}
err = s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "dall-e-3",
Model: task.ModelName,
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
})
if err != nil {

1
api/store/model/chat_model.go
View File

@@ -12,4 +12,5 @@ type ChatModel struct {
MaxContext int // 最大上下文长度
Temperature float32 // 模型温度
KeyId int // 绑定 API KEY ID
Type string // 模型类型
}

1
api/store/vo/chat_model.go
View File

@@ -13,4 +13,5 @@ type ChatModel struct {
Temperature float32 `json:"temperature"` // 模型温度
KeyId int `json:"key_id,omitempty"`
KeyName string `json:"key_name"`
Type string `json:"type"`
}