diff --git a/bot/bot.go b/bot/bot.go index e75dc17..9e1f17d 100644 --- a/bot/bot.go +++ b/bot/bot.go @@ -10,6 +10,7 @@ import ( "strings" "telegram-ollama-reply-bot/extractor" "telegram-ollama-reply-bot/llm" + "telegram-ollama-reply-bot/stats" ) var ( @@ -18,12 +19,11 @@ var ( ErrHandlerInit = errors.New("cannot initialize handler") ) -const contextUserKey = "user" - type Bot struct { api *telego.Bot llm *llm.LlmConnector extractor *extractor.Extractor + stats *stats.Stats } func NewBot(api *telego.Bot, llm *llm.LlmConnector, extractor *extractor.Extractor) *Bot { @@ -31,6 +31,7 @@ func NewBot(api *telego.Bot, llm *llm.LlmConnector, extractor *extractor.Extract api: api, llm: llm, extractor: extractor, + stats: stats.NewStats(), } } @@ -66,6 +67,10 @@ func (b *Bot) Run() error { defer bh.Stop() defer b.api.StopLongPolling() + // Middlewares + bh.Use(b.chatTypeStatsCounter) + + // Handlers bh.Handle(b.startHandler, th.CommandEqual("start")) bh.Handle(b.heyHandler, th.CommandEqual("hey")) bh.Handle(b.summarizeHandler, th.CommandEqual("summarize")) @@ -79,6 +84,8 @@ func (b *Bot) Run() error { func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) { slog.Info("/hey") + b.stats.HeyRequest() + chatID := tu.ID(update.Message.Chat.ID) b.sendTyping(chatID) @@ -112,6 +119,8 @@ func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) { func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) { slog.Info("/summarize", update.Message.Text) + b.stats.SummarizeRequest() + chatID := tu.ID(update.Message.Chat.ID) b.sendTyping(chatID) @@ -208,6 +217,25 @@ func (b *Bot) startHandler(bot *telego.Bot, update telego.Update) { } } +func (b *Bot) statsHandler(bot *telego.Bot, update telego.Update) { + slog.Info("/stats") + + chatID := tu.ID(update.Message.Chat.ID) + + b.sendTyping(chatID) + + _, err := bot.SendMessage(b.reply(update.Message, tu.Message( + chatID, + "Current bot stats:\r\n"+ + "```json\r\n"+ + b.stats.String()+"\r\n"+ + "```", + )).WithParseMode("Markdown")) + if err != nil { + slog.Error("Cannot send a message", err) + } +} + func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams { return newMessage.WithReplyParameters(&telego.ReplyParameters{ MessageID: originalMessage.MessageID, diff --git a/bot/middleware.go b/bot/middleware.go new file mode 100644 index 0000000..e3a7467 --- /dev/null +++ b/bot/middleware.go @@ -0,0 +1,24 @@ +package bot + +import ( + "github.com/mymmrac/telego" + "github.com/mymmrac/telego/telegohandler" +) + +func (b *Bot) chatTypeStatsCounter(bot *telego.Bot, update telego.Update, next telegohandler.Handler) { + message := update.Message + + if message == nil { + next(bot, update) + } + + switch message.Chat.Type { + case telego.ChatTypeGroup, telego.ChatTypeSupergroup: + b.stats.GroupRequest() + case telego.ChatTypePrivate: + b.stats.PrivateRequest() + b.stats.PrivateRequest() + } + + next(bot, update) +} diff --git a/stats/stats.go b/stats/stats.go new file mode 100644 index 0000000..44f8eda --- /dev/null +++ b/stats/stats.go @@ -0,0 +1,84 @@ +package stats + +import ( + "encoding/json" + "sync" + "time" +) + +type Stats struct { + mu sync.Mutex + + RunningSince time.Time + + GroupRequests uint64 + PrivateRequests uint64 + + HeyRequests uint64 + SummarizeRequests uint64 +} + +func NewStats() *Stats { + return &Stats{ + RunningSince: time.Now(), + + GroupRequests: 0, + PrivateRequests: 0, + + HeyRequests: 0, + SummarizeRequests: 0, + } +} + +func (s *Stats) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Uptime string `json:"uptime"` + + GroupRequests uint64 `json:"group_requests"` + PrivateRequests uint64 `json:"private_requests"` + + HeyRequests uint64 `json:"hey_requests"` + SummarizeRequests uint64 `json:"summarize_requests"` + }{ + Uptime: time.Now().Sub(s.RunningSince).String(), + + GroupRequests: s.GroupRequests, + PrivateRequests: s.PrivateRequests, + + HeyRequests: s.HeyRequests, + SummarizeRequests: s.SummarizeRequests, + }) +} + +func (s *Stats) String() string { + data, err := json.MarshalIndent(s, "", " ") + if err != nil { + return "{\"error\": \"cannot serialize stats\"}" + } + + return string(data) +} + +func (s *Stats) GroupRequest() { + s.mu.Lock() + defer s.mu.Unlock() + s.GroupRequests++ +} + +func (s *Stats) PrivateRequest() { + s.mu.Lock() + defer s.mu.Unlock() + s.PrivateRequests++ +} + +func (s *Stats) HeyRequest() { + s.mu.Lock() + defer s.mu.Unlock() + s.HeyRequests++ +} + +func (s *Stats) SummarizeRequest() { + s.mu.Lock() + defer s.mu.Unlock() + s.SummarizeRequests++ +}