✨ メッセージからFileのUUIDをパースし、Base64に変換し、OpenAIに画像認識できる機能を追加

internal/bot/bot.go

@@ -8,14 +8,15 @@ import (
 	traqwsbot "github.com/traPtitech/traq-ws-bot"
 	"log"
 	"os"
+	"strings"
 )
 
 var (
 	Bot  *traqwsbot.Bot
 	Info *traq.MyUserDetail
 )
 
-func init() {
+func InitBot() {
 	token := GetToken()
 
 	bot, err := traqwsbot.NewBot(&traqwsbot.Options{
@@ -29,9 +30,20 @@ func init() {
 	if err != nil || res.StatusCode != 200 {
 		log.Fatalf("error: 自分の情報を取得できませんでした: %v", err)
 	}
+
+	Bot = bot
 	Info = botInfo
 }
 
+func RemoveFirstBotId(input string) string {
+	BotId := Info.Id
+	index := strings.Index(input, BotId)
+	if index == -1 {
+		return input
+	}
+	return input[:index] + input[index+len(BotId):]
+}
+
 func GetToken() (token string) {
 	token, exist := os.LookupEnv("BOT_TOKEN")
 	if !exist {
@@ -40,8 +52,8 @@ func GetToken() (token string) {
 	return token
 }
 
-func GetBot() (bot *traqwsbot.Bot) {
-	return bot
+func GetBot() *traqwsbot.Bot {
+	return Bot
 }
 
 func BotJoin(ChannelID string) error {

internal/bot/file.go

@@ -0,0 +1,112 @@
+package bot
+
+import (
+	"context"
+	"encoding/base64"
+	"github.com/traPtitech/go-traq"
+	"io"
+	"log"
+	"os"
+	"regexp"
+	"strings"
+)
+
+func GetFileMetadata(fileID string) *traq.FileInfo {
+	bot := GetBot()
+
+	fileInfo, _, err := bot.API().
+		FileApi.GetFileMeta(context.Background(), fileID).
+		Execute()
+	if err != nil {
+		log.Println(err)
+	}
+
+	return fileInfo
+}
+
+func GetFileData(fileID string) *os.File {
+	bot := GetBot()
+
+	fileData, _, err := bot.API().
+		FileApi.GetFile(context.Background(), fileID).
+		Execute()
+	if err != nil {
+		log.Println(err)
+	}
+
+	return fileData
+}
+
+func ConvertFileToBase64IfFileIsImage(fileID string) *string {
+	if !isImage(fileID) {
+		log.Println("Not an image")
+		return nil
+	}
+
+	fileData := GetFileData(fileID)
+	if fileData == nil {
+		log.Println("Failed to get file data")
+		return nil
+	}
+	defer fileData.Close()
+
+	base64Data, err := fileToBase64(fileData)
+	if err != nil {
+		log.Printf("Error reading file: %v\n", err)
+		return nil
+	}
+
+	return base64Data
+}
+
+// isImage は、MIMEタイプが画像であるかを判定する関数
+func isImage(fileID string) bool {
+	fileInfo := GetFileMetadata(fileID)
+	if fileInfo == nil {
+		log.Println("Failed to get file metadata")
+		return false
+	}
+
+	return strings.HasPrefix(fileInfo.Mime, "image/")
+}
+
+// fileToBase64 は、ファイルを読み込み、BASE64エンコードされた文字列に変換する関数
+func fileToBase64(file *os.File) (*string, error) {
+	data, err := io.ReadAll(file)
+	if err != nil {
+		return nil, err
+	}
+
+	base64Data := base64.StdEncoding.EncodeToString(data)
+	return &base64Data, nil
+}
+
+func ExtractFileUUIDs(text string) []string {
+	const pattern = `https://q\.trap\.jp/files/([a-fA-F0-9-]{36})`
+
+	re := regexp.MustCompile(pattern)
+	matches := re.FindAllStringSubmatch(text, -1)
+
+	var uuids []string
+	for _, match := range matches {
+		if len(match) > 1 {
+			uuids = append(uuids, match[1])
+		}
+	}
+
+	return uuids
+}
+
+func GetBase64ImagesFromMessage(message string) []string {
+	fileUUIDs := ExtractFileUUIDs(message)
+	var base64Images []string
+
+	for _, uuid := range fileUUIDs {
+		base64Data := ConvertFileToBase64IfFileIsImage(uuid)
+		if base64Data != nil {
+			base64Images = append(base64Images, *base64Data)
+		}
+	}
+
+	return base64Images
+}

internal/gpt/gpt.go

@@ -38,7 +38,7 @@ var SystemRoleMessage = "あなたはサークルである東京工業大学デ
 
 const GptSystemString = "FirstSystemMessageを変更しました。/gptsys showで確認できます。\nFirstSystemMessageとは、常に履歴の一番最初に入り、最初にgptに情報や状況を説明するのに使用する文字列です"
 
-func init() {
+func InitGPT() {
 	apiKey = getApiKey()
 }
 
@@ -108,8 +108,12 @@ func OpenAIStream(messages []Message, do func(string)) (responseMessage string,
 	return
 }
 
-func Chat(channelID, newMessage string) {
-	addMessageAsUser(newMessage)
+func Chat(channelID, newMessage string, imageBase64 []string) {
+	if len(imageBase64) >= 1 {
+		addImageAndTextAsUser(newMessage, imageBase64)
+	} else {
+		addMessageAsUser(newMessage)
+	}
 	updateSystemRoleMessage(SystemRoleMessage)
 	postMessage, err := bot.PostMessageWithErr(channelID, blobs[rand.Intn(len(blobs))]+":loading:")
 	if err != nil {
@@ -170,6 +174,33 @@ func addMessageAsUser(message string) {
 	})
 }
 
+func addImageAndTextAsUser(message string, imageDataBase64 []string) {
+	var parts []openai.ChatMessagePart
+
+	parts = append(parts, openai.ChatMessagePart{
+		Type: openai.ChatMessagePartTypeText,
+		Text: message,
+	})
+
+	for _, b64 := range imageDataBase64 {
+		imageURL := "data:image/jpeg;base64," + b64
+		imagePart := openai.ChatMessagePart{
+			Type: openai.ChatMessagePartTypeImageURL,
+			ImageURL: &openai.ChatMessageImageURL{
+				URL: imageURL,
+			},
+		}
+		parts = append(parts, imagePart)
+	}
+
+	messagePart := Message{
+		Role:         "user",
+		MultiContent: parts,
+	}
+
+	Messages = append(Messages, messagePart)
+}
+
 func addMessageAsAssistant(message string) {
 	Messages = append(Messages, Message{
 		Role:    openai.ChatMessageRoleAssistant,

internal/handler/OnMessageCreated.go

@@ -1,6 +1,8 @@
 package handler
 
 import (
+	"fmt"
+	"github.com/pikachu0310/BOT_GPT/internal/bot"
 	"github.com/pikachu0310/BOT_GPT/internal/gpt"
 	"github.com/traPtitech/traq-ws-bot/payload"
 	"log"
@@ -18,6 +20,16 @@ func (h *Handler) MessageReceived() func(p *payload.MessageCreated) {
 			return
 		}
 
-		gpt.Chat(p.Message.ChannelID, p.Message.PlainText)
+		plainTextWithoutMention := bot.RemoveFirstBotId(p.Message.PlainText)
+
+		// if first 5 text = debug
+		if len(plainTextWithoutMention) >= 5 && plainTextWithoutMention[:5] == "debug" {
+			fmt.Printf("embed: %#v\n", p.Message.Embedded)
+			return
+		}
+
+		imagesBase64 := bot.GetBase64ImagesFromMessage(p.Message.Text)
+
+		gpt.Chat(p.Message.ChannelID, plainTextWithoutMention, imagesBase64)
 	}
 }

main.go

@@ -5,6 +5,7 @@ import (
 	"github.com/jmoiron/sqlx"
 	"github.com/joho/godotenv"
 	"github.com/pikachu0310/BOT_GPT/internal/bot"
+	"github.com/pikachu0310/BOT_GPT/internal/gpt"
 	"github.com/pikachu0310/BOT_GPT/internal/handler"
 	"github.com/pikachu0310/BOT_GPT/internal/pkg/config"
 	"github.com/pikachu0310/BOT_GPT/internal/repository"
@@ -14,9 +15,12 @@ import (
 func main() {
 	err := godotenv.Load(".env")
 	if err != nil {
-		fmt.Printf("error: .env is not exist: %v", err)
+		fmt.Printf(".env is not exist: %v", err)
 	}
 
+	bot.InitBot()
+	gpt.InitGPT()
+
 	// connect to database
 	db, err := sqlx.Connect("mysql", config.MySQL().FormatDSN())
 	if err != nil {
@@ -33,4 +37,8 @@ func main() {
 	// setup bot
 	traQBot := bot.GetBot()
 	traQBot.OnMessageCreated(h.MessageReceived())
+
+	if err := traQBot.Start(); err != nil {
+		panic(err)
+	}
 }