Implementing simple stats counter #6

Merged
skobkin merged 1 commit from feature_4_configuration into main 2024-03-11 20:17:14 +00:00
3 changed files with 138 additions and 2 deletions

View file

@ -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,

24
bot/middleware.go Normal file
View file

@ -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)
}

84
stats/stats.go Normal file
View file

@ -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++
}