From 94ed7d4692acafce7583063cbda5f6fa3bb57123 Mon Sep 17 00:00:00 2001
From: gohumble <55599638+gohumble@users.noreply.github.com>
Date: Wed, 31 Aug 2022 20:09:55 +0200
Subject: [PATCH] dalle admin api. Added configurable worker (#410)

---
 config.yaml.example           |  3 ++-
 internal/api/admin/dalle.go   | 14 ++++++++++++++
 internal/config.go            |  1 +
 internal/dalle/client.go      |  9 +++++++++
 internal/telegram/generate.go |  9 +++++++--
 main.go                       |  2 ++
 6 files changed, 35 insertions(+), 3 deletions(-)
 create mode 100644 internal/api/admin/dalle.go

diff --git a/config.yaml.example b/config.yaml.example
index 6713a10..7bfa553 100644
--- a/config.yaml.example
+++ b/config.yaml.example
@@ -27,4 +27,5 @@ database:
   groupsdb_path: "data/groups.db"
 generate:
   dalle_key: "asd"
-  dalle_price: 1000
\ No newline at end of file
+  dalle_price: 1000
+  worker: 2
\ No newline at end of file
diff --git a/internal/api/admin/dalle.go b/internal/api/admin/dalle.go
new file mode 100644
index 0000000..385b0a1
--- /dev/null
+++ b/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
+}
diff --git a/internal/config.go b/internal/config.go
index c6c1034..7cdc5df 100644
--- a/internal/config.go
+++ b/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 {
diff --git a/internal/dalle/client.go b/internal/dalle/client.go
index 42a7ca5..d22b384 100644
--- a/internal/dalle/client.go
+++ b/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)
diff --git a/internal/telegram/generate.go b/internal/telegram/generate.go
index d9ba5c9..1626927 100644
--- a/internal/telegram/generate.go
+++ b/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)
diff --git a/main.go b/main.go
index 0aa1cf7..8b755b6 100644
--- a/main.go
+++ b/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)
 
 }