mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-09-18 01:06:39 +08:00
62 lines
1.4 KiB
Go
62 lines
1.4 KiB
Go
package function
|
|
|
|
import (
|
|
"chatplus/core/types"
|
|
"chatplus/utils"
|
|
"errors"
|
|
"fmt"
|
|
"github.com/imroc/req/v3"
|
|
"time"
|
|
)
|
|
|
|
// AI 绘画函数
|
|
|
|
type FuncMidJourney struct {
|
|
name string
|
|
config types.ChatPlusExtConfig
|
|
client *req.Client
|
|
}
|
|
|
|
func NewMidJourneyFunc(config types.ChatPlusExtConfig) FuncMidJourney {
|
|
return FuncMidJourney{
|
|
name: "MidJourney AI 绘画",
|
|
config: config,
|
|
client: req.C().SetTimeout(10 * time.Second)}
|
|
}
|
|
|
|
func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
|
|
if f.config.Token == "" {
|
|
return "", errors.New("无效的 API Token")
|
|
}
|
|
|
|
logger.Infof("MJ 绘画参数：%+v", params)
|
|
prompt := utils.InterfaceToString(params["prompt"])
|
|
if !utils.IsEmptyValue(params["ar"]) {
|
|
prompt = prompt + fmt.Sprintf(" --ar %v", params["ar"])
|
|
delete(params, "ar")
|
|
}
|
|
prompt = prompt + " --niji 5"
|
|
url := fmt.Sprintf("%s/api/mj/image", f.config.ApiURL)
|
|
var res types.BizVo
|
|
r, err := f.client.R().
|
|
SetHeader("Authorization", f.config.Token).
|
|
SetHeader("Content-Type", "application/json").
|
|
SetBody(params).
|
|
SetSuccessResult(&res).Post(url)
|
|
if err != nil || r.IsErrorState() {
|
|
return "", fmt.Errorf("%v%v", r.String(), err)
|
|
}
|
|
|
|
if res.Code != types.Success {
|
|
return "", errors.New(res.Message)
|
|
}
|
|
|
|
return prompt, nil
|
|
}
|
|
|
|
func (f FuncMidJourney) Name() string {
|
|
return f.name
|
|
}
|
|
|
|
var _ Function = &FuncMidJourney{}
|