Merge pull request #576 from innogames/md_dalle_buttons

Dall-E: 2nd button + add "quality" option
innogames · Jun 13, 2024 · d4b15e6 · d4b15e6
2 parents 3481901 + 27b20ec
commit d4b15e6
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 15 deletions.
diff --git a/bot/util/delay_test.go b/bot/util/delay_test.go
@@ -12,17 +12,20 @@ func TestIncreasingDelay(t *testing.T) {
 	subject.randomAdd = time.Duration(0)
 	var actual time.Duration
 
+	// fuzzy comparison as we are working with exact time
+	delta := float64(time.Millisecond * 1)
+
 	// just started -> should return min
 	actual = subject.GetNextDelay()
-	assert.Equal(t, 10, int(actual.Seconds()))
+	assert.InDelta(t, 10, int(actual.Seconds()), delta)
 
 	// 6h = 10 + 1/4 * 50
 	subject.startedAt = time.Now().Add(-time.Hour * 6)
 	actual = subject.GetNextDelay()
-	assert.Equal(t, time.Second*22+time.Millisecond*500, actual)
+	assert.InDelta(t, time.Second*22+time.Millisecond*500, actual, delta)
 
 	// > 1d -> max
 	subject.startedAt = time.Now().Add(-time.Hour * 40)
 	actual = subject.GetNextDelay()
-	assert.Equal(t, time.Second*60, actual)
+	assert.InDelta(t, time.Second*60, actual, delta)
 }
diff --git a/command/openai/api.go b/command/openai/api.go
@@ -106,10 +106,11 @@ type ChatChoice struct {
 	  }
 */
 type DalleRequest struct {
-	Model  string `json:"model"`
-	Prompt string `json:"prompt"`
-	N      int    `json:"n"`
-	Size   string `json:"size"`
+	Model   string `json:"model"`
+	Quality string `json:"quality,omitempty"`
+	Prompt  string `json:"prompt"`
+	N       int    `json:"n"`
+	Size    string `json:"size"`
 }
 
 /*

diff --git a/command/openai/config.go b/command/openai/config.go
@@ -32,6 +32,7 @@ type Config struct {
 	DalleModel          string `mapstructure:"dalle_model"`
 	DalleImageSize      string `mapstructure:"dalle_image_size"`
 	DalleNumberOfImages int    `mapstructure:"dalle_number_of_images"`
+	DalleQuality        string `mapstructure:"dalle_quality"`
 }
 
 // IsEnabled checks if token is set

diff --git a/command/openai/dalle.go b/command/openai/dalle.go
@@ -36,7 +36,7 @@ func (c *openaiCommand) dalleGenerateImage(match matcher.Result, message msg.Mes
 
 		startTime := time.Now()
 		for _, image := range images {
-			err := c.sendImageInSlack(image, message)
+			err := c.sendImageInSlack(image, message, prompt)
 			if err != nil {
 				c.ReplyError(
 					message,
@@ -48,7 +48,7 @@ func (c *openaiCommand) dalleGenerateImage(match matcher.Result, message msg.Mes
 	}()
 }
 
-func (c *openaiCommand) sendImageInSlack(image DalleResponseImage, message msg.Message) error {
+func (c *openaiCommand) sendImageInSlack(image DalleResponseImage, message msg.Message, prompt string) error {
 	req, err := http.NewRequest("GET", image.URL, nil)
 	if err != nil {
 		return err
@@ -73,7 +73,8 @@ func (c *openaiCommand) sendImageInSlack(image DalleResponseImage, message msg.M
 		[]slack.Block{
 			slack.NewActionBlock(
 				"",
-				client.GetInteractionButton("dalle", "Regenerate", fmt.Sprintf("dall-e %s", image.RevisedPrompt)),
+				client.GetInteractionButton("dalle_original", "New (original prompt)", "dall-e "+prompt),
+				client.GetInteractionButton("dalle_advanced", "New (advanced prompt)", "dall-e "+image.RevisedPrompt),
 			),
 		},
 		slack.MsgOptionTS(message.Timestamp),
@@ -84,10 +85,11 @@ func (c *openaiCommand) sendImageInSlack(image DalleResponseImage, message msg.M
 
 func generateImages(cfg Config, prompt string) ([]DalleResponseImage, error) {
 	jsonData, _ := json.Marshal(DalleRequest{
-		Model:  cfg.DalleModel,
-		Size:   cfg.DalleImageSize,
-		N:      cfg.DalleNumberOfImages,
-		Prompt: prompt,
+		Model:   cfg.DalleModel,
+		Size:    cfg.DalleImageSize,
+		N:       cfg.DalleNumberOfImages,
+		Quality: cfg.DalleQuality,
+		Prompt:  prompt,
 	})
 
 	start := time.Now()

diff --git a/command/openai/dalle_test.go b/command/openai/dalle_test.go
@@ -104,7 +104,7 @@ func TestDalle(t *testing.T) {
 			t,
 			slackClient,
 			message,
-			`[{"type":"actions","elements":[{"type":"button","text":{"type":"plain_text","text":"Regenerate","emoji":true},"action_id":"dalle","value":"dall-e revised prompt 1234"}]}]`,
+			`[{"type":"actions","elements":[{"type":"button","text":{"type":"plain_text","text":"New (original prompt)","emoji":true},"action_id":"dalle_original","value":"dall-e a nice cat"},{"type":"button","text":{"type":"plain_text","text":"New (advanced prompt)","emoji":true},"action_id":"dalle_advanced","value":"dall-e revised prompt 1234"}]}]`,
 		)
 
 		slackClient.On(