From f75db108a7d0e96b28f072a3b867c1261858a01b Mon Sep 17 00:00:00 2001 From: atridadl Date: Wed, 27 Dec 2023 12:47:14 -0700 Subject: [PATCH] Cooldown Logic is now distributed --- .env.example | 8 +++++--- go.mod | 7 ++++++- go.sum | 6 ++++++ lib/redis.go | 40 ++++++++++++++++++++++++++++++++++++++++ main.go | 32 ++++++++++++++++++++++++++++++++ 5 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 lib/redis.go diff --git a/.env.example b/.env.example index bf374c0..0a6ce53 100644 --- a/.env.example +++ b/.env.example @@ -1,4 +1,6 @@ # Tokens -DISCORD_TOKEN= -OPENAI_API_KEY= -REPLICATE_API_TOKEN= \ No newline at end of file +DISCORD_TOKEN="" +OPENAI_API_KEY="" +REPLICATE_API_TOKEN="" +REDIS_HOST="" +REDIS_PASSWORD="" \ No newline at end of file diff --git a/go.mod b/go.mod index eaac0b0..a42154f 100644 --- a/go.mod +++ b/go.mod @@ -4,12 +4,17 @@ 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/cespare/xxhash/v2 v2.2.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + 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/redis/go-redis/v9 v9.3.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 83b1c59..af85f56 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,7 @@ +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/diamondburned/arikawa/v3 v3.3.4 h1:UXOjM7PRlWLJ8kVAydX/VetqV7W4/d4xU92JRy3SpU4= github.com/diamondburned/arikawa/v3 v3.3.4/go.mod h1:5KMSeB9R2Kzi6K4EcqMz7mwAFpAi5jglX/Veq0+MPOo= github.com/gorilla/schema v1.2.0 h1:YufUaxZYCKGFuAq3c96BOhjgd5nmXiOY9NGzF247Tsc= @@ -6,6 +10,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/redis/go-redis/v9 v9.3.1 h1:KqdY8U+3X6z+iACvumCNxnoluToB+9Me+TvyFa21Mds= +github.com/redis/go-redis/v9 v9.3.1/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M= 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= diff --git a/lib/redis.go b/lib/redis.go new file mode 100644 index 0000000..ec486c6 --- /dev/null +++ b/lib/redis.go @@ -0,0 +1,40 @@ +package lib + +import ( + "context" + "os" + "time" + + "github.com/redis/go-redis/v9" +) + +var ctx = context.Background() +var redis_host = os.Getenv("REDIS_HOST") +var redis_password = os.Getenv("REDIS_PASSWORD") + +func SetCache(key string, value string, ttlMinutes int) bool { + rdb := redis.NewClient(&redis.Options{ + Addr: redis_host, + Password: redis_password, + DB: 0, + }) + + err := rdb.Set(ctx, key, value, time.Minute*time.Duration(ttlMinutes)).Err() + + return err != nil +} + +func GetCache(key string) string { + rdb := redis.NewClient(&redis.Options{ + Addr: redis_host, + Password: redis_password, + DB: 0, + }) + + val, err := rdb.Get(ctx, key).Result() + if err != nil { + return "nil" + } + + return val +} diff --git a/main.go b/main.go index 4f44dda..8e17cd6 100644 --- a/main.go +++ b/main.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "fmt" + "himbot/lib" "io" "log" "net/http" @@ -122,12 +123,23 @@ func newHandler(s *state.State) *handler { } func (h *handler) cmdPing(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + // Command Logic return &api.InteractionResponseData{ Content: option.NewNullableString("Pong!"), } } func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + // Cooldown Logic + cachedVal := lib.GetCache(data.Event.User.ID.String() + ":" + "ask") + if cachedVal != "nil" { + return &api.InteractionResponseData{ + Content: option.NewNullableString("Please wait for the cooldown!"), + } + } + lib.SetCache(data.Event.User.ID.String()+":"+"ask", data.Event.User.ID.String()+":"+"ask", 1) + + // Command Logic var options struct { Arg string `discord:"prompt"` } @@ -166,6 +178,16 @@ func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.In } func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + // Cooldown Logic + cachedVal := lib.GetCache(data.Event.User.ID.String() + ":" + "pic") + if cachedVal != "nil" { + return &api.InteractionResponseData{ + Content: option.NewNullableString("Please wait for the cooldown!"), + } + } + lib.SetCache(data.Event.User.ID.String()+":"+"pic", data.Event.User.ID.String()+":"+"pic", 1) + + // Command Logic var options struct { Prompt string `discord:"prompt"` } @@ -234,6 +256,16 @@ func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.In } func (h *handler) cmdHDPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData { + // Cooldown Logic + cachedVal := lib.GetCache(data.Event.User.ID.String() + ":" + "hdpic") + if cachedVal != "nil" { + return &api.InteractionResponseData{ + Content: option.NewNullableString("Please wait for the cooldown!"), + } + } + lib.SetCache(data.Event.User.ID.String()+":"+"hdpic", data.Event.User.ID.String()+":"+"hdpic", 10) + + // Command Logic var options struct { Prompt string `discord:"prompt"` }