refactor MessageLink: clean up unused prop and optimize insert message

we don't actually need the child or to recurse over every previous message to insert a new message
additionally, add load and dump message store to disk in JSON

Signed-off-by: Sid Sun <[email protected]>

Dockerfile

@@ -1,4 +1,4 @@
-FROM golang as builder
+FROM golang:1.21 as builder
 WORKDIR /app
 COPY go.mod go.sum ./
 RUN go mod download
@@ -8,5 +8,6 @@ RUN make build
 FROM alpine:latest
 RUN apk --no-cache add ca-certificates
 WORKDIR /app
+RUN mkdir /app/store
 COPY --from=builder /app/bin/ /app/
 CMD ["/app/openwebui-telegram"]

pkg/bot/contract/contract.go

@@ -1,10 +1,9 @@
 package contract
 
 type MessageLink struct {
-	Parent   int
-	Children []int
-	Text     string
-	From     string
+	Parent int
+	Text   string
+	From   string
 }
 
 type CompletionUpdate struct {

pkg/bot/handlers/completion/completion.go

@@ -71,6 +71,12 @@ func Handler(b *tele.Bot, isResend bool) tele.HandlerFunc {
 		var finalMessage *string
 		debounced := debounce.New(20 * time.Millisecond)
 		for completion := range updatesChan {
+			if finalMessage != nil && *finalMessage == completion.Message {
+				if completion.IsLast {
+					break
+				}
+				continue
+			}
 			finalMessage = &completion.Message
 			send := func() {
 				_, err := b.Edit(botMessage, completion.Message)
@@ -111,25 +117,16 @@ func addMessageToChain(m *tele.Message) *contract.MessageLink {
 	}
 
 	var parent int
-	if m.ReplyTo != nil {
+	// If this is a reply message, set the parent ID if it exists in the store
+	if m.ReplyTo != nil && store.ChatStore[m.Chat.ID][m.ReplyTo.ID] != nil {
 		parent = m.ReplyTo.ID
-		if store.ChatStore[m.Chat.ID][parent] == nil {
-			addMessageToChain(m.ReplyTo)
-		}
-		store.ChatStore[m.Chat.ID][parent].Children = append(store.ChatStore[m.Chat.ID][parent].Children, m.ID)
 	}
 
 	// fmt.Printf("Message: %+v\n", m)
-	if m.Sender == nil {
-		m.Sender = &tele.User{
-			Username: "unknown",
-		}
-	}
 	store.ChatStore[m.Chat.ID][m.ID] = &contract.MessageLink{
-		Parent:   parent,
-		Children: []int{},
-		Text:     m.Text,
-		From:     m.Sender.Username,
+		Parent: parent,
+		Text:   m.Text,
+		From:   m.Sender.Username,
 	}
 
 	// x, err := json.MarshalIndent(store.ChatStore, "", "  ")

pkg/bot/router/router.go

@@ -11,12 +11,12 @@ import (
 	tele "gopkg.in/telebot.v3"
 )
 
-type bot struct {
+type Bot struct {
 	bot *tele.Bot
 }
 
 // ListenAndServe starts listens on the update channel and handles routing the update to handlers
-func (b bot) Start() {
+func (b Bot) Start() {
 	store.BotUsername = b.bot.Me.Username
 	slog.Info("[StartBot] Started Bot", slog.String("bot_name", b.bot.Me.FirstName))
 	r := b.bot.Group()
@@ -27,16 +27,21 @@ func (b bot) Start() {
 	b.bot.Start()
 }
 
+func (b Bot) Stop() {
+	slog.Info("[StopBot] Stopping Bot")
+	b.bot.Stop()
+}
+
 // New returns a new instance of the router
-func New(cfg config.BotConfig) bot {
+func New(cfg config.BotConfig) *Bot {
 	b, err := tele.NewBot(tele.Settings{
 		Token:  cfg.Token(),
 		Poller: &tele.LongPoller{Timeout: 10 * time.Second},
 	})
 	if err != nil {
 		panic(err)
 	}
-	return bot{
+	return &Bot{
 		bot: b,
 	}
 }

pkg/bot/server.go

@@ -1,7 +1,10 @@
 package bot
 
 import (
+	"encoding/json"
 	"log/slog"
+	"os"
+	"os/signal"
 
 	"github.com/sid-sun/openwebui-bot/cmd/config"
 	"github.com/sid-sun/openwebui-bot/pkg/bot/router"
@@ -11,8 +14,47 @@ import (
 // StartBot starts the bot, inits all the requited submodules and routine for shutdown
 func StartBot(cfg config.Config) {
 	store.NewStore()
+	loadStore()
 	ch := router.New(cfg.Bot)
 
 	slog.Info("[StartBot] Starting Bot")
+	go dumpStore(ch)
 	ch.Start()
 }
+
+// Dump store data to disk as JSON on interrupt
+func dumpStore(ch *router.Bot) {
+	shutDown := make(chan os.Signal, 1)
+	signal.Notify(shutDown, os.Interrupt)
+	<-shutDown
+	slog.Info("[DumpStore] Dumping store data to disk")
+	// Implement store dumping logic here
+	x, err := json.MarshalIndent(store.ChatStore, "", "  ")
+	if err != nil {
+		slog.Error("[DumpStore] Error dumping store data to disk", slog.Any("error", err))
+		return
+	}
+	err = os.WriteFile("./store/chat_store.json", x, 0644)
+	if err != nil {
+		slog.Error("[DumpStore] Error writing store data to file", slog.Any("error", err))
+		return
+	}
+	slog.Info("[LoadStore] Dumped store data to disk")
+	ch.Stop()
+}
+
+// load data from JSON file on startup
+func loadStore() {
+	slog.Info("[LoadStore] Loading store data from disk")
+	data, err := os.ReadFile("./store/chat_store.json")
+	if err != nil {
+		slog.Error("[LoadStore] Error reading store data from file", slog.Any("error", err))
+		return
+	}
+	err = json.Unmarshal(data, &store.ChatStore)
+	if err != nil {
+		slog.Error("[LoadStore] Error unmarshaling store data from file", slog.Any("error", err))
+		return
+	}
+	slog.Info("[LoadStore] Loaded store data from disk")
+}