Implementing simple stats counter.
This commit is contained in:
parent
2272570377
commit
2ed9c2f6dc
32
bot/bot.go
32
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,
|
||||
|
|
24
bot/middleware.go
Normal file
24
bot/middleware.go
Normal 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
84
stats/stats.go
Normal 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++
|
||||
}
|
Loading…
Reference in a new issue