diff --git a/.DS_Store b/.DS_Store index 05d6da1..218ad48 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/.github/workflows/fly.yml b/.github/workflows/fly.yml new file mode 100644 index 0000000..0e6ae31 --- /dev/null +++ b/.github/workflows/fly.yml @@ -0,0 +1,15 @@ +name: Fly Deploy +on: + push: + branches: + - main +jobs: + deploy: + name: Deploy app + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: superfly/flyctl-actions/setup-flyctl@master + - run: flyctl deploy --remote-only + env: + FLY_API_TOKEN: ${{ secrets.FLY_API_TOKEN }} \ No newline at end of file diff --git a/.gitignore b/.gitignore index 00925dc..e10e601 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ .env.test.local .env.production.local .env +himbot \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..235a041 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,13 @@ +FROM golang:1.21.5 + +WORKDIR /app + +COPY go.mod . +COPY go.sum . +COPY main.go . + +RUN go mod download + +RUN go build . + +CMD [ "./himbot" ] \ No newline at end of file diff --git a/go.mod b/go.mod index d726ae7..eaac0b0 100644 --- a/go.mod +++ b/go.mod @@ -4,10 +4,13 @@ go 1.21.5 require github.com/diamondburned/arikawa/v3 v3.3.4 +require golang.org/x/sync v0.5.0 // indirect + require ( github.com/gorilla/schema v1.2.0 // indirect github.com/gorilla/websocket v1.4.2 // indirect github.com/joho/godotenv v1.5.1 + github.com/replicate/replicate-go v0.14.2 github.com/sashabaranov/go-openai v1.17.9 golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac // indirect ) diff --git a/go.sum b/go.sum index faeba57..83b1c59 100644 --- a/go.sum +++ b/go.sum @@ -6,6 +6,8 @@ github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0U github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= +github.com/replicate/replicate-go v0.14.2 h1:XgK+REvYrWs7qDeyugxHA93h31qBhEFk/3p1/p2w3W8= +github.com/replicate/replicate-go v0.14.2/go.mod h1:otIrl1vDmyjNhTzmVmp/mQU3Wt1+3387gFNEsAZq0ig= github.com/sashabaranov/go-openai v1.17.9 h1:QEoBiGKWW68W79YIfXWEFZ7l5cEgZBV4/Ow3uy+5hNY= github.com/sashabaranov/go-openai v1.17.9/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= @@ -19,6 +21,8 @@ golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.5.0 h1:60k92dhOjHxJkrqnwsfl8KuaHbn/5dl0lUPUklKo3qE= +golang.org/x/sync v0.5.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/main.go b/main.go index 24a02eb..6c3f8fa 100644 --- a/main.go +++ b/main.go @@ -1,13 +1,14 @@ package main import ( + "bytes" "context" "fmt" + "io" "log" - "math/rand" + "net/http" "os" "os/signal" - "time" "github.com/diamondburned/arikawa/v3/api" "github.com/diamondburned/arikawa/v3/api/cmdroute" @@ -15,7 +16,9 @@ import ( "github.com/diamondburned/arikawa/v3/gateway" "github.com/diamondburned/arikawa/v3/state" "github.com/diamondburned/arikawa/v3/utils/json/option" + "github.com/diamondburned/arikawa/v3/utils/sendpart" "github.com/joho/godotenv" + "github.com/replicate/replicate-go" openai "github.com/sashabaranov/go-openai" ) @@ -25,27 +28,45 @@ var commands = []api.CreateCommandData{ Description: "ping pong!", }, { - Name: "echo", - Description: "echo back the argument", + Name: "ask", + Description: "Ask Himbot!", Options: []discord.CommandOption{ &discord.StringOption{ - OptionName: "argument", - Description: "what's echoed back", + OptionName: "prompt", + Description: "The prompt to send to Himbot.", Required: true, }, }, }, { - Name: "thonk", - Description: "biiiig thonk", - }, - { - Name: "ask", - Description: "Ask Himbot!", + Name: "pic", + Description: "Generate an image using Stable Diffusion!", Options: []discord.CommandOption{ &discord.StringOption{ - OptionName: "argument", - Description: "the prompt", + OptionName: "prompt", + Description: "The prompt for the image generation.", + Required: true, + }, + }, + }, + { + Name: "hdpic", + Description: "Generate an image using DALL·E 3!", + Options: []discord.CommandOption{ + &discord.StringOption{ + OptionName: "prompt", + Description: "The prompt for the image generation.", + Required: true, + }, + }, + }, + { + Name: "hs", + Description: "This command was your nickname in highschool!", + Options: []discord.CommandOption{ + &discord.StringOption{ + OptionName: "nickname", + Description: "Your nickname in highschool.", Required: true, }, }, @@ -92,37 +113,23 @@ func newHandler(s *state.State) *handler { // Automatically defer handles if they're slow. h.Use(cmdroute.Deferrable(s, cmdroute.DeferOpts{})) h.AddFunc("ping", h.cmdPing) - h.AddFunc("echo", h.cmdEcho) - h.AddFunc("thonk", h.cmdThonk) h.AddFunc("ask", h.cmdAsk) + h.AddFunc("pic", h.cmdPic) + h.AddFunc("hdpic", h.cmdHDPic) + h.AddFunc("hs", h.cmdHS) return h } -func (h *handler) cmdPing(ctx context.Context, cmd cmdroute.CommandData) *api.InteractionResponseData { +func (h *handler) cmdPing(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { return &api.InteractionResponseData{ Content: option.NewNullableString("Pong!"), } } -func (h *handler) cmdEcho(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { - var options struct { - Arg string `discord:"argument"` - } - - if err := data.Options.Unmarshal(&options); err != nil { - return errorResponse(err) - } - - return &api.InteractionResponseData{ - Content: option.NewNullableString(options.Arg), - AllowedMentions: &api.AllowedMentions{}, // don't mention anyone - } -} - func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { var options struct { - Arg string `discord:"argument"` + Arg string `discord:"prompt"` } if err := data.Options.Unmarshal(&options); err != nil { @@ -158,10 +165,132 @@ func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.In } } -func (h *handler) cmdThonk(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { - time.Sleep(time.Duration(3+rand.Intn(5)) * time.Second) +func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + var options struct { + Prompt string `discord:"prompt"` + } + + if err := data.Options.Unmarshal(&options); err != nil { + return errorResponse(err) + } + + client, clientError := replicate.NewClient(replicate.WithTokenFromEnv()) + if clientError != nil { + return errorResponse(clientError) + } + if err := data.Options.Unmarshal(&options); err != nil { + return errorResponse(err) + } + + input := replicate.PredictionInput{ + "prompt": options.Prompt, + } + webhook := replicate.Webhook{ + URL: "https://example.com/webhook", + Events: []replicate.WebhookEventType{"start", "completed"}, + } + + prediction, predictionError := client.Run(context.Background(), "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b", input, &webhook) + + if predictionError != nil { + return errorResponse(predictionError) + } + + test, ok := prediction.([]interface{}) + + if !ok { + fmt.Println("prediction is not []interface{}") + } + + imgUrl, ok := test[0].(string) + + if !ok { + fmt.Println("prediction.Output[0] is not a string") + } + + imageRes, imageGetErr := http.Get(imgUrl) + if imageGetErr != nil { + log.Fatalln(imageGetErr) + } + + defer imageRes.Body.Close() + + imageBytes, imgReadErr := io.ReadAll(imageRes.Body) + if imgReadErr != nil { + log.Fatalln(imgReadErr) + } + + imageFile := bytes.NewBuffer(imageBytes) + + file := sendpart.File{ + Name: "image.png", + Reader: imageFile, + } + return &api.InteractionResponseData{ - Content: option.NewNullableString("https://tenor.com/view/thonk-thinking-sun-thonk-sun-thinking-sun-gif-14999983"), + Files: []sendpart.File{file}, + } +} + +func (h *handler) cmdHDPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + var options struct { + Prompt string `discord:"prompt"` + } + + if err := data.Options.Unmarshal(&options); err != nil { + return errorResponse(err) + } + + client := openai.NewClient(os.Getenv("OPENAI_API_KEY")) + + // Send the generation request to DALL·E 3 + resp, err := client.CreateImage(context.Background(), openai.ImageRequest{ + Prompt: options.Prompt, + Model: "dall-e-3", + Size: "1024x1024", + }) + if err != nil { + log.Printf("Image creation error: %v\n", err) + return errorResponse(fmt.Errorf("failed to generate image")) + } + + imageRes, err := http.Get(resp.Data[0].URL) + + if err != nil { + log.Fatalln(err) + } + + defer imageRes.Body.Close() + + imageBytes, err := io.ReadAll(imageRes.Body) + + if err != nil { + log.Fatalln(err) + } + + imageFile := bytes.NewBuffer(imageBytes) + + file := sendpart.File{ + Name: "image.png", + Reader: imageFile, + } + + return &api.InteractionResponseData{ + Files: []sendpart.File{file}, + } +} + +func (h *handler) cmdHS(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + var options struct { + Arg string `discord:"nickname"` + } + + if err := data.Options.Unmarshal(&options); err != nil { + return errorResponse(err) + } + + return &api.InteractionResponseData{ + Content: option.NewNullableString(options.Arg + " was " + data.Event.User.DisplayName + "'s nickname in highschool!"), } }