Remove query string from URLs

Closes #56
igolaizola · Sep 26, 2023 · 6e9f31e · 6e9f31e
1 parent faf449b
commit 6e9f31e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 9 deletions.
diff --git a/pkg/ai/bluewillow/bluewillow.go b/pkg/ai/bluewillow/bluewillow.go
@@ -108,7 +108,7 @@ func New(client *discord.Client, cfg *Config) (ai.Client, error) {
 				}
 
 				// Attachment based message
-				cacheID = msg.Attachments[0].URL
+				cacheID = cleanURL(msg.Attachments[0].URL)
 
 				// Ignore message already in the cache
 				c.lck.Lock()
@@ -406,7 +406,7 @@ func (c *Client) Imagine(ctx context.Context, prompt string) (*ai.Preview, error
 		return nil, fmt.Errorf("bluewillow: message has no image ids")
 	}
 	return &ai.Preview{
-		URL:            preview.Attachments[0].URL,
+		URL:            cleanURL(preview.Attachments[0].URL),
 		Prompt:         prompt,
 		ResponsePrompt: responsePrompt,
 		MessageID:      preview.ID,
@@ -447,7 +447,7 @@ func (c *Client) Upscale(ctx context.Context, preview *ai.Preview, index int) (s
 	if err != nil {
 		return "", fmt.Errorf("bluewillow: couldn't receive links message: %w", err)
 	}
-	return msg.Attachments[0].URL, nil
+	return cleanURL(msg.Attachments[0].URL), nil
 }
 
 func (c *Client) Variation(ctx context.Context, preview *ai.Preview, index int) (*ai.Preview, error) {
@@ -503,7 +503,7 @@ func (c *Client) Variation(ctx context.Context, preview *ai.Preview, index int)
 		return nil, fmt.Errorf("bluewillow: message has no image ids")
 	}
 	return &ai.Preview{
-		URL:            msg.Attachments[0].URL,
+		URL:            cleanURL(msg.Attachments[0].URL),
 		Prompt:         preview.Prompt,
 		ResponsePrompt: preview.ResponsePrompt,
 		MessageID:      msg.ID,
@@ -522,3 +522,11 @@ var linkWrappedRegex = regexp.MustCompile(`<https?://[^\s]+>`)
 func fromResponsePrompt(s string) string {
 	return linkWrappedRegex.ReplaceAllString(s, "<LINK>")
 }
+
+func cleanURL(u string) string {
+	// Remove query string
+	if i := strings.Index(u, "?"); i >= 0 {
+		u = u[:i]
+	}
+	return u
+}
diff --git a/pkg/ai/midjourney/midjourney.go b/pkg/ai/midjourney/midjourney.go
@@ -143,7 +143,7 @@ func New(client *discord.Client, cfg *Config) (ai.Client, error) {
 				}
 
 				// Attachment based message
-				cacheID = msg.Attachments[0].URL
+				cacheID = cleanURL(msg.Attachments[0].URL)
 
 				// Ignore message already in the cache
 				c.lck.Lock()
@@ -625,7 +625,7 @@ func (c *Client) Imagine(ctx context.Context, prompt string) (*ai.Preview, error
 		return nil, fmt.Errorf("midjourney: message has no image ids")
 	}
 	return &ai.Preview{
-		URL:            preview.Attachments[0].URL,
+		URL:            cleanURL(preview.Attachments[0].URL),
 		Prompt:         prompt,
 		ResponsePrompt: responsePrompt,
 		MessageID:      preview.ID,
@@ -670,7 +670,7 @@ func (c *Client) Upscale(ctx context.Context, preview *ai.Preview, index int) (s
 	if err != nil {
 		return "", fmt.Errorf("midjourney: couldn't receive links message: %w", err)
 	}
-	return msg.Attachments[0].URL, nil
+	return cleanURL(msg.Attachments[0].URL), nil
 }
 
 func (c *Client) Variation(ctx context.Context, preview *ai.Preview, index int) (*ai.Preview, error) {
@@ -730,7 +730,7 @@ func (c *Client) Variation(ctx context.Context, preview *ai.Preview, index int)
 		return nil, fmt.Errorf("midjourney: message has no image ids")
 	}
 	return &ai.Preview{
-		URL:            msg.Attachments[0].URL,
+		URL:            cleanURL(msg.Attachments[0].URL),
 		Prompt:         preview.Prompt,
 		ResponsePrompt: preview.ResponsePrompt,
 		MessageID:      msg.ID,
@@ -758,7 +758,7 @@ func (c *Client) checkAction(msg *discord.Message) (bool, error) {
 	if msg.Embeds[0].Image == nil {
 		return false, fmt.Errorf("midjourney: missing image in embed")
 	}
-	image := msg.Embeds[0].Image.URL
+	image := cleanURL(msg.Embeds[0].Image.URL)
 	if image == "" {
 		return false, fmt.Errorf("midjourney: missing image url in embed")
 	}
@@ -830,3 +830,11 @@ func replaceLinks(s string) string {
 	s = linkWrappedRegex.ReplaceAllString(s, "<LINK>")
 	return linkRegex.ReplaceAllString(s, "<LINK>")
 }
+
+func cleanURL(u string) string {
+	// Remove query string
+	if i := strings.Index(u, "?"); i >= 0 {
+		u = u[:i]
+	}
+	return u
+}