From 62706e830938e723b103d5a19a1660152d14f4a0 Mon Sep 17 00:00:00 2001
From: Jack Merrill <me@jackmerrill.com>
Date: Fri, 18 Aug 2023 23:01:43 -0500
Subject: [PATCH] feat: dalle

---
 internal/commands/fun/dalle.go | 77 ++++++++++++++++++++++++++++++++++
 main.go                        |  6 ++-
 2 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 internal/commands/fun/dalle.go

diff --git a/internal/commands/fun/dalle.go b/internal/commands/fun/dalle.go
new file mode 100644
index 0000000..b443edc
--- /dev/null
+++ b/internal/commands/fun/dalle.go
@@ -0,0 +1,77 @@
+package fun
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/jackmerrill/hampbot/internal/utils/config"
+	"github.com/jackmerrill/hampbot/internal/utils/embed"
+	"github.com/sashabaranov/go-openai"
+	"github.com/zekroTJA/shireikan"
+)
+
+type Dalle struct {
+}
+
+func (c *Dalle) GetInvokes() []string {
+	return []string{"dalle", "ai-art"}
+}
+
+func (c *Dalle) GetDescription() string {
+	return "Generate AI art with DALL-E"
+}
+
+func (c *Dalle) GetHelp() string {
+	return "`dalle [prompt]` - Generate AI art with DALL-E"
+}
+
+func (c *Dalle) GetGroup() string {
+	return config.GroupFun
+}
+
+func (c *Dalle) GetDomainName() string {
+	return "hamp.fun.dalle"
+}
+
+func (c *Dalle) GetSubPermissionRules() []shireikan.SubPermission {
+	return nil
+}
+func (c *Dalle) IsExecutableInDMChannels() bool {
+	return true
+}
+
+func (c *Dalle) Exec(ctx shireikan.Context) error {
+	openaiToken := os.Getenv("OPENAI_TOKEN")
+
+	if openaiToken == "" {
+		ctx.ReplyEmbed(embed.NewErrorEmbed(ctx).SetDescription("No OpenAI token set.").MessageEmbed)
+		return fmt.Errorf("no openai token set")
+	}
+
+	prompt := strings.TrimPrefix(ctx.GetMessage().Content, fmt.Sprintf("%sdalle ", config.BotPrefix))
+
+	client := openai.NewClient(openaiToken)
+
+	ctx.GetSession().ChannelTyping(ctx.GetChannel().ID)
+
+	resp, err := client.CreateImage(context.Background(), openai.ImageRequest{
+		Prompt:         prompt,
+		Size:           openai.CreateImageSize1024x1024,
+		ResponseFormat: openai.CreateImageResponseFormatURL,
+		N:              1,
+	})
+
+	if err != nil {
+		ctx.ReplyEmbed(embed.NewErrorEmbed(ctx).SetDescription("Failed to generate image.").MessageEmbed)
+		return err
+	}
+
+	ctx.ReplyEmbed(embed.NewSuccessEmbed(ctx).
+		SetTitle("DALL-E").
+		SetDescription(prompt).
+		SetImage(resp.Data[0].URL).
+		MessageEmbed)
+	return nil
+}
diff --git a/main.go b/main.go
index 3ee2a15..1f940d9 100644
--- a/main.go
+++ b/main.go
@@ -11,6 +11,7 @@ import (
 	"github.com/jackmerrill/hampbot/internal/commands/fun"
 	studentlife "github.com/jackmerrill/hampbot/internal/commands/studentlife"
 	util "github.com/jackmerrill/hampbot/internal/commands/util"
+	"github.com/jackmerrill/hampbot/internal/utils/config"
 	"github.com/zekroTJA/shireikan"
 
 	"github.com/charmbracelet/log"
@@ -38,7 +39,7 @@ func main() {
 	}()
 
 	handler := shireikan.New(&shireikan.Config{
-		GeneralPrefix:         ">",
+		GeneralPrefix:         config.BotPrefix,
 		AllowBots:             false,
 		AllowDM:               true,
 		ExecuteOnEdit:         true,
@@ -69,6 +70,9 @@ func main() {
 	handler.Register(&fun.XKCD{})
 	log.Debug("Registered xkcd command")
 
+	handler.Register(&fun.Dalle{})
+	log.Debug("Registered dalle command")
+
 	log.Info("Registered all commands")
 
 	handler.Setup(session)