feat: integrate the AI assistant.

msgcache.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/utils"
+	"github.com/dgraph-io/ristretto"
+	"github.com/eko/gocache/lib/v4/cache"
+	gocache_lib "github.com/eko/gocache/lib/v4/store"
+	ristretto_store "github.com/eko/gocache/store/ristretto/v4"
+	"github.com/go-errors/errors"
+	tele "gopkg.in/telebot.v3"
+)
+
+var (
+	msgCache *cache.Cache[tele.Message]
+	ctxVoid  = context.Background()
+)
+
+func initMsgCache() error {
+	ristrettoCache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 100_000,
+		MaxCost:     20 << 20, // 20 MiB
+		BufferItems: 64,
+	})
+	if err != nil {
+		return err
+	}
+	ristrettoStore := ristretto_store.NewRistretto(ristrettoCache)
+
+	msgCache = cache.New[tele.Message](ristrettoStore)
+	return nil
+}
+
+func cacheMessage(msg *tele.Message) error {
+	if msg == nil || msg.Chat == nil {
+		return errors.New("failed to get message ID")
+	}
+
+	msgId := fmt.Sprintf("%d:%d", msg.Chat.ID, msg.ID)
+	return msgCache.Set(ctxVoid, msgId, *msg)
+}
+
+func getCachedMessageById(id string) (*tele.Message, error) {
+	msg, err := msgCache.Get(ctxVoid, id)
+	if err != nil {
+		if _, ok := err.(*gocache_lib.NotFound); ok {
+			return nil, nil
+		}
+		return nil, errors.Wrap(err, 0)
+	}
+	return &msg, err
+}
+
+func getCachedMessage(msg *tele.Message) (*tele.Message, error) {
+	if msg == nil || msg.Chat == nil {
+		return nil, errors.New("failed to get message ID")
+	}
+
+	msgId := fmt.Sprintf("%d:%d", msg.Chat.ID, msg.ID)
+	return getCachedMessageById(msgId)
+}
+
+func getCachedThread(msg *tele.Message) ([]*tele.Message, error) {
+	if msg == nil {
+		return nil, errors.New("empty message given")
+	}
+
+	threadR := []*tele.Message{msg}
+	currentMsg := msg
+	for {
+		if currentMsg.ReplyTo == nil {
+			break
+		}
+
+		parentMsg, err := getCachedMessage(currentMsg.ReplyTo)
+		if err != nil {
+			return nil, err
+		}
+		if parentMsg == nil {
+			parentMsg = currentMsg.ReplyTo
+		}
+		threadR = append(threadR, parentMsg)
+		currentMsg = parentMsg
+	}
+
+	return utils.Reverse(threadR), nil
+}