Merge pull request #85 from whatwewant/feat/support-image-generation

feat: support image generation (closes #86)

README.md

@@ -25,6 +25,11 @@
 
 ## 功能支持
 
+### 核心
+* [x] 支持文本生成
+* [x] 支持图片生成
+
+### 扩展
 * [x] 支持长对话，自动联系上下文
 * [x] 支持私人对话
 * [x] 支持群聊

commands/custom.go

@@ -0,0 +1,76 @@
+package commands
+
+import (
+	"net/http"
+
+	"github.com/go-zoox/chatbot-feishu"
+	chatgpt "github.com/go-zoox/chatgpt-client"
+	"github.com/go-zoox/chatgpt-for-chatbot-feishu/config"
+	"github.com/go-zoox/core-utils/fmt"
+	"github.com/go-zoox/feishu"
+	feishuEvent "github.com/go-zoox/feishu/event"
+	"github.com/go-zoox/fetch"
+	"github.com/go-zoox/logger"
+)
+
+func CreateCustomCommand(
+	feishuClient feishu.Client,
+	chatgptClient chatgpt.Client,
+	cfg *config.Config,
+) *chatbot.Command {
+	return &chatbot.Command{
+		ArgsLength: 1,
+		Handler: func(args []string, request *feishuEvent.EventRequest, reply chatbot.MessageReply) error {
+			if len(args) != 1 {
+				return fmt.Errorf("invalid args: %v", args)
+			}
+
+			question := args[0]
+			logger.Debugf("[custom command: %s, service: %s] question: %s", cfg.CustomCommand, cfg.CustomCommandService, question)
+
+			response, err := fetch.Post(cfg.CustomCommandService, &fetch.Config{
+				Headers: fetch.Headers{
+					"Content-Type": "application/json",
+					"Accept":       "application/json",
+					"User-Agent":   fmt.Sprintf("go-zoox_fetch/%s chatgpt-for-chatbot-feishu/%s", fetch.Version, cfg.Version),
+				},
+				Body: map[string]interface{}{
+					"question": args[0],
+				},
+			})
+			if err != nil {
+				logger.Errorf("failed to request from custom command service(%s)(1): %v", cfg.CustomCommandService, err)
+				if err2 := replyText(reply, fmt.Sprintf("failed to interact with command service(err: %v)", err)); err2 != nil {
+					return fmt.Errorf("failed to reply: %v", err)
+				}
+
+				return nil
+			}
+
+			if response.Status != http.StatusOK {
+				logger.Errorf("failed to request from custom command service(%s)(2): %d", cfg.CustomCommandService, response.Status)
+				if err := replyText(reply, fmt.Sprintf("failed to interact with command service (status: %d, response: %s)", response.Status, response.String())); err != nil {
+					return fmt.Errorf("failed to reply: %v", err)
+				}
+
+				return nil
+			}
+
+			answer := response.Get("answer").String()
+			if answer == "" {
+				logger.Error("failed to request from custom command service(%s): empty answer (response: %s)", cfg.CustomCommandService, response.String())
+				if err := replyText(reply, fmt.Sprintf("no answer found, unexpected response from custom command service(response: %s)", response.String())); err != nil {
+					return fmt.Errorf("failed to reply: %v", err)
+				}
+
+				return nil
+			}
+
+			if err := replyText(reply, answer); err != nil {
+				return fmt.Errorf("failed to reply: %v", err)
+			}
+
+			return nil
+		},
+	}
+}

commands/draw.go

@@ -0,0 +1,68 @@
+package commands
+
+import (
+	"github.com/go-zoox/chatbot-feishu"
+	chatgpt "github.com/go-zoox/chatgpt-client"
+	"github.com/go-zoox/core-utils/fmt"
+	"github.com/go-zoox/core-utils/strings"
+	"github.com/go-zoox/feishu"
+	feishuEvent "github.com/go-zoox/feishu/event"
+	feishuImage "github.com/go-zoox/feishu/image"
+	"github.com/go-zoox/fetch"
+	"github.com/go-zoox/fs"
+	"github.com/go-zoox/logger"
+	openaiclient "github.com/go-zoox/openai-client"
+)
+
+func CreateDrawCommand(
+	feishuClient feishu.Client,
+	chatgptClient chatgpt.Client,
+) *chatbot.Command {
+	return &chatbot.Command{
+		Handler: func(args []string, request *feishuEvent.EventRequest, reply chatbot.MessageReply) error {
+			prompt := strings.Join(args, " ")
+			if prompt == "" {
+				return replyText(reply, fmt.Sprintf("prompt is required (args: %s)", strings.Join(args, " ")))
+			}
+
+			logger.Infof("[draw]: %v", prompt)
+
+			logger.Infof("[draw]: request image generation ...")
+			response, err := chatgptClient.ImageGeneration(&openaiclient.ImageGenerationRequest{
+				Prompt: prompt,
+			})
+			if err != nil {
+				return replyText(reply, fmt.Sprintf("failed to request image generation: %v", err))
+			}
+
+			for _, image := range response.Data {
+				tmpFilePath := fs.TmpFilePath()
+
+				logger.Infof("[draw] download image from chatgpt: %v", image.URL)
+				_, err := fetch.Download(image.URL, tmpFilePath, &fetch.Config{})
+				if err != nil {
+					return replyText(reply, fmt.Sprintf("failed to download image: %v", err))
+				}
+
+				tmpFile, err := fs.Open(tmpFilePath)
+				if err != nil {
+					return replyText(reply, fmt.Sprintf("failed to open image: %v", err))
+				}
+
+				logger.Infof("[draw] upload image to feishu ...")
+				response, err := feishuClient.Image().Upload(&feishuImage.UploadRequest{
+					ImageType: "message",
+					Image:     tmpFile,
+				})
+				if err != nil {
+					return replyText(reply, fmt.Sprintf("failed to upload image: %v", err))
+				}
+
+				logger.Infof("[draw] reply image to feishu: %v", response.ImageKey)
+				replyImage(reply, response.ImageKey)
+			}
+
+			return nil
+		},
+	}
+}

commands/message.go

@@ -0,0 +1,166 @@
+package commands
+
+import (
+	"time"
+
+	"github.com/go-zoox/core-utils/regexp"
+
+	"github.com/go-zoox/chatbot-feishu"
+	chatgpt "github.com/go-zoox/chatgpt-client"
+	"github.com/go-zoox/chatgpt-for-chatbot-feishu/config"
+	"github.com/go-zoox/core-utils/fmt"
+	"github.com/go-zoox/core-utils/strings"
+	"github.com/go-zoox/feishu"
+	feishuEvent "github.com/go-zoox/feishu/event"
+	mc "github.com/go-zoox/feishu/message/content"
+	"github.com/go-zoox/logger"
+	"github.com/go-zoox/retry"
+)
+
+func CreateMessageCommand(
+	feishuClient feishu.Client,
+	chatgptClient chatgpt.Client,
+	cfg *config.Config,
+) *chatbot.Command {
+	return &chatbot.Command{
+		Handler: func(args []string, request *feishuEvent.EventRequest, reply func(content string, msgType ...string) error) (err error) {
+			text := strings.Join(args, " ")
+
+			// fmt.PrintJSON(request)
+			if cfg.BotInfo == nil {
+				logger.Infof("Trying to get bot info ...")
+				cfg.BotInfo, err = feishuClient.Bot().GetBotInfo()
+				if err != nil {
+					return fmt.Errorf("failed to get bot info: %v", err)
+				}
+			}
+
+			user, err := getUser(feishuClient, request, cfg)
+			if err != nil {
+				return fmt.Errorf("failed to get user: %v", err)
+			}
+
+			textMessage := strings.TrimSpace(text)
+			if textMessage == "" {
+				return nil
+			}
+
+			var question string
+			// group chat
+			if request.IsGroupChat() {
+				// @
+				if ok := regexp.Match("^@_user_1", textMessage); ok {
+					for _, metion := range request.Event.Message.Mentions {
+						if metion.Key == "@_user_1" && metion.ID.OpenID == cfg.BotInfo.OpenID {
+							question = textMessage[len("@_user_1"):]
+							question = strings.TrimSpace(question)
+							break
+						}
+					}
+				} else if ok := regexp.Match("^/chatgpt\\s+", textMessage); ok {
+					// command: /chatgpt
+					question = textMessage[len("/chatgpt "):]
+				}
+			} else if request.IsP2pChat() {
+				question = textMessage
+			}
+
+			question = strings.TrimSpace(question)
+			if question == "" {
+				logger.Infof("ignore empty question message")
+				return nil
+			}
+
+			// @TODO 离线服务
+			if !cfg.IsInService {
+				return replyText(reply, cfg.OfflineMessage)
+			}
+
+			go func() {
+				logger.Debugf("%s 问 ChatGPT：%s", user.User.Name, question)
+
+				var err error
+
+				conversation, err := chatgptClient.GetOrCreateConversation(request.ChatID(), &chatgpt.ConversationConfig{
+					MaxMessages: 50,
+					Model:       cfg.OpenAIModel,
+				})
+				if err != nil {
+					logger.Errorf("failed to get or create conversation by ChatID %s", request.ChatID())
+					return
+				}
+
+				if err := conversation.IsQuestionAsked(request.Event.Message.MessageID); err != nil {
+					logger.Warnf("duplicated event(id: %s): %v", request.Event.Message.MessageID, err)
+					return
+				}
+
+				var answer []byte
+				err = retry.Retry(func() error {
+
+					answer, err = conversation.Ask([]byte(question), &chatgpt.ConversationAskConfig{
+						ID:   request.Event.Message.MessageID,
+						User: user.User.Name,
+					})
+					if err != nil {
+						logger.Errorf("failed to request answer: %v", err)
+						return fmt.Errorf("failed to request answer: %v", err)
+					}
+
+					return nil
+				}, 5, 3*time.Second)
+				if err != nil {
+					logger.Errorf("failed to get answer: %v", err)
+					msgType, content, err := mc.
+						NewContent().
+						Text(&mc.ContentTypeText{
+							Text: "ChatGPT 繁忙，请稍后重试",
+						}).
+						Build()
+					if err != nil {
+						logger.Errorf("failed to build content: %v", err)
+						return
+					}
+					if err := reply(string(content), msgType); err != nil {
+						return
+					}
+					return
+				}
+
+				logger.Debugf("ChatGPT 答 %s：%s", user.User.Name, answer)
+
+				responseMessage := string(answer)
+				// if request.IsGroupChat() {
+				// 	responseMessage = fmt.Sprintf("%s\n-------------\n%s", question, answer)
+				// }
+
+				msgType, content, err := mc.
+					NewContent().
+					Post(&mc.ContentTypePost{
+						ZhCN: &mc.ContentTypePostBody{
+							Content: [][]mc.ContentTypePostBodyItem{
+								{
+									{
+										Tag:      "text",
+										UnEscape: true,
+										Text:     responseMessage,
+									},
+								},
+							},
+						},
+					}).
+					Build()
+				if err != nil {
+					logger.Errorf("failed to build content: %v", err)
+					return
+				}
+				if err := reply(string(content), msgType); err != nil {
+					logger.Errorf("failed to reply: %v", err)
+					return
+				}
+			}()
+
+			return nil
+		},
+	}
+}

commands/model.go

@@ -0,0 +1,60 @@
+package commands
+
+import (
+	"github.com/go-zoox/chatbot-feishu"
+	chatgpt "github.com/go-zoox/chatgpt-client"
+	"github.com/go-zoox/chatgpt-for-chatbot-feishu/config"
+	"github.com/go-zoox/core-utils/fmt"
+	"github.com/go-zoox/core-utils/strings"
+	"github.com/go-zoox/feishu"
+	feishuEvent "github.com/go-zoox/feishu/event"
+)
+
+func CreateModelCommand(
+	feishuClient feishu.Client,
+	chatgptClient chatgpt.Client,
+	cfg *config.Config,
+) *chatbot.Command {
+	return &chatbot.Command{
+		ArgsLength: 1,
+		Handler: func(args []string, request *feishuEvent.EventRequest, reply func(content string, msgType ...string) error) error {
+			if err := isAllowToDo(feishuClient, cfg, request, "model"); err != nil {
+				return err
+			}
+
+			if len(args) == 0 || args[0] == "" {
+				currentModel, err := chatgptClient.GetConversationModel(request.ChatID(), &chatgpt.ConversationConfig{
+					MaxMessages: 50,
+					Model:       cfg.OpenAIModel,
+				})
+				if err != nil {
+					return fmt.Errorf("failed to get model by conversation(%s)", request.ChatID())
+				}
+
+				if err := replyText(reply, fmt.Sprintf("当前模型：%s", currentModel)); err != nil {
+					return fmt.Errorf("failed to reply: %v", err)
+				}
+
+				return nil
+			}
+
+			model := args[0]
+			if model == "" {
+				return fmt.Errorf("model name is required (args: %s)", strings.Join(args, " "))
+			}
+
+			if err := chatgptClient.ChangeConversationModel(request.ChatID(), model, &chatgpt.ConversationConfig{
+				MaxMessages: 50,
+				Model:       cfg.OpenAIModel,
+			}); err != nil {
+				return fmt.Errorf("failed to set model(%s) for conversation(%s)", model, request.ChatID())
+			}
+
+			if err := replyText(reply, fmt.Sprintf("succeed to set model: %s", model)); err != nil {
+				return fmt.Errorf("failed to reply: %v", err)
+			}
+
+			return nil
+		},
+	}
+}