dalle admin api. Added configurable worker (#410)

config.yaml.example

@@ -27,4 +27,5 @@ database:
   groupsdb_path: "data/groups.db"
 generate:
   dalle_key: "asd"
-  dalle_price: 1000
+  dalle_price: 1000
+  worker: 2

internal/api/admin/dalle.go

@@ -0,0 +1,14 @@
+package admin
+
+import (
+	"github.com/LightningTipBot/LightningTipBot/internal/dalle"
+	"net/http"
+)
+
+func (s Service) DisableDalle(w http.ResponseWriter, r *http.Request) {
+	dalle.Enabled = false
+}
+
+func (s Service) EnableDalle(w http.ResponseWriter, r *http.Request) {
+	dalle.Enabled = true
+}

internal/config.go

@@ -20,6 +20,7 @@ var Configuration = struct {
 type GenerateConfiguration struct {
 	DalleKey   string `yaml:"dalle_key"`
 	DallePrice int64  `yaml:"dalle_price"`
+	Worker     int    `yaml:"worker"`
 }
 
 type SocksConfiguration struct {

internal/dalle/client.go

@@ -2,9 +2,18 @@ package dalle
 
 import (
 	"context"
+	"github.com/LightningTipBot/LightningTipBot/internal"
 	"io"
 )
 
+var Enabled bool
+
+func init() {
+	if internal.Configuration.Generate.DalleKey != "" {
+		Enabled = true
+	}
+}
+
 type Client interface {
 	Generate(ctx context.Context, prompt string) (*Task, error)
 	ListTasks(ctx context.Context, req *ListTasksRequest) (*ListTasksResponse, error)

internal/telegram/generate.go

@@ -22,7 +22,7 @@ import (
 // generateImages is called when the user enters /generate or /generate <prompt>
 // asks the user for a prompt if not given
 func (bot *TipBot) generateImages(ctx intercept.Context) (intercept.Context, error) {
-	if internal.Configuration.Generate.DalleKey == "" {
+	if !dalle.Enabled {
 		bot.trySendMessage(ctx.Message().Sender, "🤖💤 Dalle image generation is currently not available. Please try again later.")
 		return ctx, nil
 	}
@@ -102,9 +102,14 @@ func (bot *TipBot) confirmGenerateImages(ctx intercept.Context) (intercept.Conte
 }
 
 var jobChan chan func(workerId int)
-var workers = 2
+var workers = internal.Configuration.Generate.Worker
 
 func init() {
+	if workers == 0 {
+		log.Printf("Dalle is disabled. No worker started.")
+		return
+	}
+	log.Printf("Starting Dalle image generation. Worker: %d, Price: %d sat", workers, internal.Configuration.Generate.DallePrice)
 	jobChan = make(chan func(workerId int), workers)
 	for i := 0; i < workers; i++ {
 		go worker(jobChan, i)

main.go

@@ -81,6 +81,8 @@ func startApiServer(bot *telegram.TipBot) {
 	internalAdminServer.AppendRoute("/mutex/unlock/{id}", mutex.UnlockHTTP)
 	internalAdminServer.AppendRoute("/admin/ban/{id}", adminService.BanUser)
 	internalAdminServer.AppendRoute("/admin/unban/{id}", adminService.UnbanUser)
+	internalAdminServer.AppendRoute("/admin/dalle/enable", adminService.EnableDalle)
+	internalAdminServer.AppendRoute("/admin/dalle/disable", adminService.DisableDalle)
 	internalAdminServer.PathPrefix("/debug/pprof/", http.DefaultServeMux)
 
 }