From fbeea444fc526418b545630cbd637b382ba92f5d Mon Sep 17 00:00:00 2001 From: Sid Sun Date: Thu, 16 May 2024 00:38:16 +0530 Subject: [PATCH] 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 --- Dockerfile | 3 +- pkg/bot/contract/contract.go | 7 ++-- pkg/bot/handlers/completion/completion.go | 25 ++++++-------- pkg/bot/router/router.go | 13 ++++--- pkg/bot/server.go | 42 +++++++++++++++++++++++ 5 files changed, 67 insertions(+), 23 deletions(-) diff --git a/Dockerfile b/Dockerfile index b9049b4..fd56f46 100644 --- a/Dockerfile +++ b/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"] diff --git a/pkg/bot/contract/contract.go b/pkg/bot/contract/contract.go index f1fa388..18bf9fd 100644 --- a/pkg/bot/contract/contract.go +++ b/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 { diff --git a/pkg/bot/handlers/completion/completion.go b/pkg/bot/handlers/completion/completion.go index 5bf4d61..a9689b9 100644 --- a/pkg/bot/handlers/completion/completion.go +++ b/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, "", " ") diff --git a/pkg/bot/router/router.go b/pkg/bot/router/router.go index 8a525f5..059432d 100644 --- a/pkg/bot/router/router.go +++ b/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,8 +27,13 @@ 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}, @@ -36,7 +41,7 @@ func New(cfg config.BotConfig) bot { if err != nil { panic(err) } - return bot{ + return &Bot{ bot: b, } } diff --git a/pkg/bot/server.go b/pkg/bot/server.go index ba627d7..0098701 100644 --- a/pkg/bot/server.go +++ 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") +}