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

api/handler/admin/chat_model_handler.go

@@ -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)

api/handler/chat_model_handler.go

@@ -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)

api/handler/chat_openai_handler.go

@@ -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)

api/handler/dalle_handler.go

@@ -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)
+}

api/handler/function_handler.go

@@ -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)
 }