feat: add worker (#4)

* chore: add blocked user and allowed user

* chore: refine task

* chore: add worker

* chore: add worker

Author: nduyhai

go.mod

@@ -12,6 +12,7 @@ require (
 )
 
 require (
+	github.com/alitto/pond/v2 v2.5.0 // indirect
 	github.com/benbjohnson/clock v1.3.5 // indirect
 	github.com/bytedance/gopkg v0.1.3 // indirect
 	github.com/bytedance/sonic v1.14.1 // indirect

go.sum

@@ -1,3 +1,5 @@
+github.com/alitto/pond/v2 v2.5.0 h1:vPzS5GnvSDRhWQidmj2djHllOmjFExVFbDGCw1jdqDw=
+github.com/alitto/pond/v2 v2.5.0/go.mod h1:xkjYEgQ05RSpWdfSd1nM3OVv7TBhLdy7rMp3+2Nq+yE=
 github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o=
 github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M=

internal/adapters/llm/openai/openai.go

@@ -3,7 +3,6 @@ package openai
 import (
 	"context"
 	"log/slog"
-	"time"
 
 	"github.com/nduyhai/valjean/internal/app/entities"
 	goopenai "github.com/openai/openai-go/v3"
@@ -19,27 +18,32 @@ func NewClient(ai goopenai.Client, logger *slog.Logger) *Client {
 }
 
 func (c *Client) Evaluate(ctx context.Context, in entities.EvalInput) (entities.EvalOutput, error) {
-	ctxTimeout, cancelFunc := context.WithTimeout(ctx, 10*time.Second)
-	defer cancelFunc()
 
-	var messages []goopenai.ChatCompletionMessageParamUnion
-
-	for _, contextText := range in.ContextSnips {
-		if contextText != "" {
-			messages = append(messages, goopenai.UserMessage(contextText))
+	// Combine context snippets with the main message
+	var fullMessage string
+	if len(in.ContextSnips) > 0 {
+		fullMessage = "Context:\n"
+		for _, contextText := range in.ContextSnips {
+			if contextText != "" {
+				fullMessage += contextText + "\n"
+			}
 		}
+		fullMessage += "\nUser message:\n" + in.Text
+	} else {
+		fullMessage = in.Text
 	}
 
-	messages = append(messages, goopenai.UserMessage(in.Text))
+	messages := []goopenai.ChatCompletionMessageParamUnion{
+		goopenai.UserMessage(fullMessage),
+	}
 
-	chatCompletion, err := c.ai.Chat.Completions.New(ctxTimeout, goopenai.ChatCompletionNewParams{
+	chatCompletion, err := c.ai.Chat.Completions.New(ctx, goopenai.ChatCompletionNewParams{
 		Messages: messages,
 		Model:    goopenai.ChatModelGPT4o,
 	})
-	msg := "Please try again"
 
+	msg := "Please try again"
 	if err != nil {
-		c.logger.Error("failed to evaluate", slog.Any("error", err))
 		return entities.EvalOutput{
 			Summary: msg,
 		}, err

internal/adapters/worker/memory.go

@@ -0,0 +1,20 @@
+package worker
+
+import "github.com/alitto/pond/v2"
+
+type Memory struct {
+	pool pond.Pool
+}
+
+func NewMemory() *Memory {
+	pool := pond.NewPool(100)
+	return &Memory{pool: pool}
+}
+
+func (m *Memory) Submit(f func()) {
+	m.pool.Submit(f)
+}
+
+func (m *Memory) Shutdown() {
+	m.pool.StopAndWait()
+}

internal/app/service/moderation.go

@@ -2,12 +2,15 @@ package service
 
 import (
 	"context"
+	"slices"
 	"strings"
 
 	"github.com/nduyhai/valjean/internal/app/entities"
 	"github.com/nduyhai/valjean/internal/infra/config"
 )
 
+const PrivateChat = "private"
+
 type Moderation interface {
 	Allowed(ctx context.Context, in entities.EvalInput) bool // ok, reason
 }
@@ -22,6 +25,15 @@ func NewModeration(config config.Config) Moderation {
 
 func (m *moderation) Allowed(ctx context.Context, in entities.EvalInput) bool {
 	text := strings.TrimSpace(in.Text)
+
+	if len(m.telegram.BlockedUsers) > 0 && slices.Contains(m.telegram.BlockedUsers, in.UserHandle) {
+		return false
+	}
+
+	if len(m.telegram.AllowedUsers) > 0 && !slices.Contains(m.telegram.AllowedUsers, in.UserHandle) {
+		return false
+	}
+
 	if m.telegram.Prefix != "" && strings.HasPrefix(text, m.telegram.Prefix) {
 		return true
 	}
@@ -34,7 +46,7 @@ func (m *moderation) Allowed(ctx context.Context, in entities.EvalInput) bool {
 		return true
 	}
 
-	if in.ChatType == "private" {
+	if in.ChatType == PrivateChat {
 		return true
 	}
 

internal/app/usecase/evaluate.go

@@ -18,13 +18,23 @@ type EvaluateUseCase struct {
 	moderation    service.Moderation
 	rateLimiter   ports.RateLimiter
 	eventProducer ports.EventProducer
+	worker        ports.Worker
 	telegram      config.Telegram
 	cooldown      time.Duration
 	logger        *slog.Logger
 }
 
-func NewEvaluateUseCase(evaluator ports.Evaluator, moderation service.Moderation, rateLimiter ports.RateLimiter, eventProducer ports.EventProducer, config config.Config, logger *slog.Logger) *EvaluateUseCase {
-	return &EvaluateUseCase{evaluator: evaluator, moderation: moderation, rateLimiter: rateLimiter, eventProducer: eventProducer, telegram: config.Telegram, cooldown: 2 * time.Second, logger: logger}
+func NewEvaluateUseCase(evaluator ports.Evaluator, moderation service.Moderation, rateLimiter ports.RateLimiter, eventProducer ports.EventProducer, worker ports.Worker, config config.Config, logger *slog.Logger) *EvaluateUseCase {
+	return &EvaluateUseCase{
+		evaluator:     evaluator,
+		moderation:    moderation,
+		rateLimiter:   rateLimiter,
+		eventProducer: eventProducer,
+		worker:        worker,
+		telegram:      config.Telegram,
+		cooldown:      60 * time.Second,
+		logger:        logger,
+	}
 
 }
 
@@ -38,11 +48,23 @@ func (e *EvaluateUseCase) Handle(ctx context.Context, in entities.EvalInput) err
 
 		return errors.New("cooling down—try again in a moment")
 	}
+
+	e.worker.Submit(func() {
+		e.process(in)
+	})
+
+	return nil
+}
+
+func (e *EvaluateUseCase) process(in entities.EvalInput) {
+	ctx, cancelFunc := context.WithTimeout(context.Background(), e.cooldown)
+	defer cancelFunc()
+
 	// moderation
 	allowed := e.moderation.Allowed(ctx, in)
 	if !allowed {
 		e.logger.Warn("message skipped")
-		return errors.New("message skipped")
+		return
 	}
 	out, err := e.evaluator.Evaluate(ctx, in)
 	if err != nil || out.Summary == "" {
@@ -51,12 +73,10 @@ func (e *EvaluateUseCase) Handle(ctx context.Context, in entities.EvalInput) err
 
 		e.sendMsg(ctx, in, "i couldn’t evaluate that right now")
 
-		return errors.New("i couldn’t evaluate that right now")
+		return
 	}
 
 	e.sendMsg(ctx, in, out.Summary)
-
-	return nil
 }
 
 func (e *EvaluateUseCase) sendMsg(ctx context.Context, in entities.EvalInput, replyMsg string) {

internal/infra/config/config.go

@@ -13,10 +13,12 @@ type HTTPConfig struct {
 }
 
 type Telegram struct {
-	Prefix        string `env:"TELEGRAM_PREFIX" envDefault:"!eval"` // "!eval "
-	BotUsername   string `env:"TELEGRAM_BOT_USERNAME" envDefault:"valjean"`
-	Token         string `env:"TELEGRAM_BOT_TOKEN"`
-	WebhookSecret string `env:"TELEGRAM_WEBHOOK_SECRET"`
+	Prefix        string   `env:"TELEGRAM_PREFIX" envDefault:"!eval"` // "!eval "
+	BotUsername   string   `env:"TELEGRAM_BOT_USERNAME" envDefault:"valjean"`
+	Token         string   `env:"TELEGRAM_BOT_TOKEN"`
+	WebhookSecret string   `env:"TELEGRAM_WEBHOOK_SECRET"`
+	BlockedUsers  []string `env:"TELEGRAM_BLOCKED_USERS"`
+	AllowedUsers  []string `env:"TELEGRAM_ALLOWED_USERS"`
 }
 
 type OpenAI struct {

internal/infra/fxmodules/modules.go

@@ -8,6 +8,7 @@ import (
 	"github.com/nduyhai/valjean/internal/adapters/limiter"
 	"github.com/nduyhai/valjean/internal/adapters/llm/openai"
 	"github.com/nduyhai/valjean/internal/adapters/producer"
+	"github.com/nduyhai/valjean/internal/adapters/worker"
 	"github.com/nduyhai/valjean/internal/app/service"
 	"github.com/nduyhai/valjean/internal/app/usecase"
 	"github.com/nduyhai/valjean/internal/infra/config"
@@ -43,6 +44,7 @@ var LimiterModule = fx.Module("adapters",
 		fx.Annotate(openai.NewClient, fx.As(new(ports.Evaluator))),
 		NewTelegramSdk,
 		fx.Annotate(producer.NewTelegram, fx.As(new(ports.EventProducer))),
+		fx.Annotate(worker.NewMemory, fx.As(new(ports.Worker))),
 	),
 )
 

internal/infra/fxmodules/server.go

@@ -8,6 +8,7 @@ import (
 	"github.com/nduyhai/valjean/internal/adapters/http"
 	"github.com/nduyhai/valjean/internal/infra/config"
 	"github.com/nduyhai/valjean/internal/infra/httpserver"
+	"github.com/nduyhai/valjean/internal/ports"
 	"go.uber.org/fx"
 )
 
@@ -58,6 +59,7 @@ var ServerModule = fx.Options(
 type ServerParams struct {
 	fx.In
 	HTTPServer *httpserver.Server
+	Worker     ports.Worker
 }
 
 func ServerLifecycle(lc fx.Lifecycle, servers ServerParams) {
@@ -79,6 +81,8 @@ func ServerLifecycle(lc fx.Lifecycle, servers ServerParams) {
 				// Log error
 			}
 
+			servers.Worker.Shutdown()
+
 			return nil
 		},
 	})

internal/ports/worker.go

@@ -0,0 +1,6 @@
+package ports
+
+type Worker interface {
+	Submit(func())
+	Shutdown()
+}