Implementing simple stats counter.
All checks were successful
continuous-integration/drone/push Build is passing
continuous-integration/drone/pr Build is passing

This commit is contained in:
Alexey Skobkin 2024-03-11 23:15:27 +03:00
parent 2272570377
commit 2ed9c2f6dc
No known key found for this signature in database
GPG key ID: 5D5CEF6F221278E7
3 changed files with 138 additions and 2 deletions

View file

@ -10,6 +10,7 @@ import (
"strings" "strings"
"telegram-ollama-reply-bot/extractor" "telegram-ollama-reply-bot/extractor"
"telegram-ollama-reply-bot/llm" "telegram-ollama-reply-bot/llm"
"telegram-ollama-reply-bot/stats"
) )
var ( var (
@ -18,12 +19,11 @@ var (
ErrHandlerInit = errors.New("cannot initialize handler") ErrHandlerInit = errors.New("cannot initialize handler")
) )
const contextUserKey = "user"
type Bot struct { type Bot struct {
api *telego.Bot api *telego.Bot
llm *llm.LlmConnector llm *llm.LlmConnector
extractor *extractor.Extractor extractor *extractor.Extractor
stats *stats.Stats
} }
func NewBot(api *telego.Bot, llm *llm.LlmConnector, extractor *extractor.Extractor) *Bot { 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, api: api,
llm: llm, llm: llm,
extractor: extractor, extractor: extractor,
stats: stats.NewStats(),
} }
} }
@ -66,6 +67,10 @@ func (b *Bot) Run() error {
defer bh.Stop() defer bh.Stop()
defer b.api.StopLongPolling() defer b.api.StopLongPolling()
// Middlewares
bh.Use(b.chatTypeStatsCounter)
// Handlers
bh.Handle(b.startHandler, th.CommandEqual("start")) bh.Handle(b.startHandler, th.CommandEqual("start"))
bh.Handle(b.heyHandler, th.CommandEqual("hey")) bh.Handle(b.heyHandler, th.CommandEqual("hey"))
bh.Handle(b.summarizeHandler, th.CommandEqual("summarize")) 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) { func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) {
slog.Info("/hey") slog.Info("/hey")
b.stats.HeyRequest()
chatID := tu.ID(update.Message.Chat.ID) chatID := tu.ID(update.Message.Chat.ID)
b.sendTyping(chatID) 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) { func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) {
slog.Info("/summarize", update.Message.Text) slog.Info("/summarize", update.Message.Text)
b.stats.SummarizeRequest()
chatID := tu.ID(update.Message.Chat.ID) chatID := tu.ID(update.Message.Chat.ID)
b.sendTyping(chatID) 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 { func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams {
return newMessage.WithReplyParameters(&telego.ReplyParameters{ return newMessage.WithReplyParameters(&telego.ReplyParameters{
MessageID: originalMessage.MessageID, 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++
}