Merge branch 'yangjian102621:main' into main

api/handler/admin/function_handler.go

@@ -6,7 +6,11 @@ import (
 	"chatplus/handler"
 	"chatplus/store/model"
 	"chatplus/store/vo"
+	"chatplus/utils"
 	"chatplus/utils/resp"
+
+	"github.com/golang-jwt/jwt/v5"
+
 	"github.com/gin-gonic/gin"
 	"gorm.io/gorm"
 )
@@ -29,13 +33,65 @@ func (h *FunctionHandler) Save(c *gin.Context) {
 		return
 	}
 
-	logger.Info(data)
+	var f = model.Function{
+		Id:          data.Id,
+		Name:        data.Name,
+		Label:       data.Label,
+		Description: data.Description,
+		Parameters:  utils.JsonEncode(data.Parameters),
+		Required:    utils.JsonEncode(data.Required),
+		Action:      data.Action,
+		Token:       data.Token,
+		Enabled:     data.Enabled,
+	}
+
+	res := h.db.Save(&f)
+	if res.Error != nil {
+		resp.ERROR(c, "error with save data:"+res.Error.Error())
+		return
+	}
+	data.Id = f.Id
+	resp.SUCCESS(c, data)
+}
+
+func (h *FunctionHandler) Set(c *gin.Context) {
+	var data struct {
+		Id    uint        `json:"id"`
+		Filed string      `json:"filed"`
+		Value interface{} `json:"value"`
+	}
+
+	if err := c.ShouldBindJSON(&data); err != nil {
+		resp.ERROR(c, types.InvalidArgs)
+		return
+	}
+
+	res := h.db.Model(&model.Function{}).Where("id = ?", data.Id).Update(data.Filed, data.Value)
+	if res.Error != nil {
+		resp.ERROR(c, "更新数据库失败！")
+		return
+	}
 	resp.SUCCESS(c)
 }
 
 func (h *FunctionHandler) List(c *gin.Context) {
+	var items []model.Function
+	res := h.db.Find(&items)
+	if res.Error != nil {
+		resp.ERROR(c, "No data found")
+		return
+	}
 
-	resp.SUCCESS(c)
+	functions := make([]vo.Function, 0)
+	for _, v := range items {
+		var f vo.Function
+		err := utils.CopyObject(v, &f)
+		if err != nil {
+			continue
+		}
+		functions = append(functions, f)
+	}
+	resp.SUCCESS(c, functions)
 }
 
 func (h *FunctionHandler) Remove(c *gin.Context) {
@@ -50,3 +106,20 @@ func (h *FunctionHandler) Remove(c *gin.Context) {
 	}
 	resp.SUCCESS(c)
 }
+
+// GenToken generate function api access token
+func (h *FunctionHandler) GenToken(c *gin.Context) {
+	// 创建 token
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+		"user_id": 0,
+		"expired": 0,
+	})
+	tokenString, err := token.SignedString([]byte(h.App.Config.Session.SecretKey))
+	if err != nil {
+		logger.Error("error with generate token", err)
+		resp.ERROR(c)
+		return
+	}
+
+	resp.SUCCESS(c, tokenString)
+}

api/handler/chatimpl/chat_handler.go

@@ -442,7 +442,7 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, platf
 	} else {
 		client = http.DefaultClient
 	}
-	logger.Infof("Sending %s request, KEY: %s, PROXY: %s, Model: %s", platform, *apiKey, proxyURL, req.Model)
+	logger.Infof("Sending %s request, ApiURL:%s, PROXY: %s, Model: %s", platform, apiURL, proxyURL, req.Model)
 	switch platform {
 	case types.Azure:
 		request.Header.Set("api-key", *apiKey)
@@ -452,7 +452,6 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, platf
 		if err != nil {
 			return nil, err
 		}
-		logger.Info(token)
 		request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
 		break
 	case types.Baidu:

api/handler/chatimpl/xunfei_handler.go

@@ -50,7 +50,7 @@ type xunFeiResp struct {
 }
 
 var Model2URL = map[string]string{
-	"generalv1": "1.1",
+	"general":   "v1.1",
 	"generalv2": "v2.1",
 	"generalv3": "v3.1",
 }

api/handler/mj_handler.go

@@ -13,7 +13,9 @@ import (
 	"encoding/base64"
 	"fmt"
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
+	"net/http"
 	"strings"
 	"time"
 )
@@ -58,6 +60,27 @@ func (h *MidJourneyHandler) preCheck(c *gin.Context) bool {
 
 }
 
+// Client WebSocket 客户端，用于通知任务状态变更
+func (h *MidJourneyHandler) Client(c *gin.Context) {
+	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error(err)
+		c.Abort()
+		return
+	}
+
+	userId := h.GetInt(c, "user_id", 0)
+	if userId == 0 {
+		logger.Info("Invalid user ID")
+		c.Abort()
+		return
+	}
+
+	client := types.NewWsClient(ws)
+	h.pool.Clients.Put(uint(userId), client)
+	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
+}
+
 // Image 创建一个绘画任务
 func (h *MidJourneyHandler) Image(c *gin.Context) {
 	var data struct {
@@ -147,6 +170,9 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
 		UserId:    userId,
 	})
 
+	client := h.pool.Clients.Get(uint(job.UserId))
+	_ = client.Send([]byte("Task Updated"))
+
 	// update user's img calls
 	h.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
 	resp.SUCCESS(c)
@@ -205,6 +231,10 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
 		MessageId:   data.MessageId,
 		MessageHash: data.MessageHash,
 	})
+
+	client := h.pool.Clients.Get(uint(job.UserId))
+	_ = client.Send([]byte("Task Updated"))
+
 	resp.SUCCESS(c)
 }
 
@@ -226,6 +256,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
 
 	job := model.MidJourneyJob{
 		Type:        types.TaskVariation.String(),
+		ChannelId:   data.ChannelId,
 		ReferenceId: data.MessageId,
 		UserId:      userId,
 		TaskId:      data.TaskId,
@@ -250,6 +281,9 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
 		MessageHash: data.MessageHash,
 	})
 
+	client := h.pool.Clients.Get(uint(job.UserId))
+	_ = client.Send([]byte("Task Updated"))
+
 	// update user's img calls
 	h.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
 	resp.SUCCESS(c)
@@ -320,6 +354,7 @@ func (h *MidJourneyHandler) JobList(c *gin.Context) {
 func (h *MidJourneyHandler) Remove(c *gin.Context) {
 	var data struct {
 		Id     uint   `json:"id"`
+		UserId uint   `json:"user_id"`
 		ImgURL string `json:"img_url"`
 	}
 	if err := c.ShouldBindJSON(&data); err != nil {
@@ -340,5 +375,8 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
 		logger.Error("remove image failed: ", err)
 	}
 
+	client := h.pool.Clients.Get(data.UserId)
+	_ = client.Send([]byte("Task Updated"))
+
 	resp.SUCCESS(c)
 }