Generate default avatar fallbacks and label images for Gemini alignment (#278)

Author: yetibot-mps[bot]

src/yetibot/core/util/image_input.clj

@@ -6,12 +6,25 @@
 (def ^:private image-url-pattern
   #"(https?://\S+\.(?:jpg|jpeg|png|gif|webp)(?:[?\#]\S*)?|https?://(?:cdn\.discordapp\.com|media\.discordapp\.net|i\.imgur\.com)/\S+)")
 
+(defn- parse-id [id]
+  (try
+    (BigInteger. (str id))
+    (catch Exception _
+      BigInteger/ZERO)))
+
+(defn- default-avatar-url [id]
+  (let [id-val (parse-id id)
+        idx (if (pos? id-val)
+              (mod (.shiftRight id-val 22) 6)
+              0)]
+    (format "https://cdn.discordapp.com/embed/avatars/%d.png" (int idx))))
+
 (defn- discord-avatar-url [{:keys [id avatar]}]
   (cond
     (= id "269292446041636866") "https://i.imgflip.com/4/9omh8s.jpg"
     (and id avatar) (let [ext (if (str/starts-with? avatar "a_") "gif" "png")]
                       (format "https://cdn.discordapp.com/avatars/%s/%s.%s?size=256" id avatar ext))
-    :else nil))
+    :else (default-avatar-url id)))
 
 (defn- entity-index [prompt entity]
   (if (:is-me? entity)
@@ -48,12 +61,22 @@
                     (str/replace (str "<@" id ">") (str "@" username))))))
           prompt entities))
 
+(defn- build-labels-prefix [entities]
+  (if (seq entities)
+    (let [labels (map-indexed (fn [idx entity]
+                                (format "Image %d is @%s" (inc idx) (:username entity)))
+                              entities)]
+      (str "(" (str/join ". " labels) ") "))
+    ""))
+
 (defn- extract-mention-avatars [prompt raw-event]
   (let [entities (get-entities prompt raw-event)
         sorted-entities (sort-by #(entity-index prompt %) entities)]
     (if (seq sorted-entities)
-      {:urls (vec (keep discord-avatar-url sorted-entities))
-       :prompt (replace-entities prompt sorted-entities)}
+      (let [prefix (build-labels-prefix sorted-entities)
+            replaced-prompt (replace-entities prompt sorted-entities)]
+        {:urls (vec (keep discord-avatar-url sorted-entities))
+         :prompt (str prefix replaced-prompt)})
       {:urls [] :prompt prompt})))
 
 (defn- image-attachment? [{:keys [content-type content_type filename]}]

test/yetibot/core/test/util/image_input.clj

@@ -10,22 +10,24 @@
     (img-input/extract-images
       "<@123456789>"
       {:raw-event {:mentions [{:id "123456789" :avatar "avatar123" :username "alice"}]}})
-    => {:prompt "@alice" :image-urls ["https://cdn.discordapp.com/avatars/123456789/avatar123.png?size=256"]})
+    => {:prompt "(Image 1 is @alice) @alice"
+        :image-urls ["https://cdn.discordapp.com/avatars/123456789/avatar123.png?size=256"]})
 
   (fact
     "a hardcoded user mention extracts the custom mario death meme avatar"
     (img-input/extract-images
       "<@269292446041636866>"
       {:raw-event {:mentions [{:id "269292446041636866" :avatar "someavatar" :username "bob"}]}})
-    => {:prompt "@bob" :image-urls ["https://i.imgflip.com/4/9omh8s.jpg"]})
+    => {:prompt "(Image 1 is @bob) @bob"
+        :image-urls ["https://i.imgflip.com/4/9omh8s.jpg"]})
 
   (fact
     "user mentions are extracted and sorted in typing order, not snowflake order"
     (img-input/extract-images
       "draw <@333333333> and <@111111111>"
       {:raw-event {:mentions [{:id "111111111" :avatar "av1" :username "alice"}
                              {:id "333333333" :avatar "av3" :username "charlie"}]}})
-    => {:prompt "draw @charlie and @alice"
+    => {:prompt "(Image 1 is @charlie. Image 2 is @alice) draw @charlie and @alice"
         :image-urls ["https://cdn.discordapp.com/avatars/333333333/av3.png?size=256"
                      "https://cdn.discordapp.com/avatars/111111111/av1.png?size=256"]})
 
@@ -34,15 +36,15 @@
     (img-input/extract-images
       "draw me as a wizard"
       {:raw-event {:author {:id "222222222" :avatar "av2" :username "bob"}}})
-    => {:prompt "draw @bob as a wizard"
+    => {:prompt "(Image 1 is @bob) draw @bob as a wizard"
         :image-urls ["https://cdn.discordapp.com/avatars/222222222/av2.png?size=256"]})
 
   (fact
     "extracting speaking user's avatar is case-insensitive for 'me'"
     (img-input/extract-images
       "draw ME flying"
       {:raw-event {:author {:id "222222222" :avatar "av2" :username "bob"}}})
-    => {:prompt "draw @bob flying"
+    => {:prompt "(Image 1 is @bob) draw @bob flying"
         :image-urls ["https://cdn.discordapp.com/avatars/222222222/av2.png?size=256"]})
 
   (fact
@@ -59,6 +61,14 @@
       "draw me and <@111111111>"
       {:raw-event {:author {:id "222222222" :avatar "av2" :username "bob"}
                    :mentions [{:id "111111111" :avatar "av1" :username "alice"}]}})
-    => {:prompt "draw @bob and @alice"
+    => {:prompt "(Image 1 is @bob. Image 2 is @alice) draw @bob and @alice"
         :image-urls ["https://cdn.discordapp.com/avatars/222222222/av2.png?size=256"
-                     "https://cdn.discordapp.com/avatars/111111111/av1.png?size=256"]}))
+                     "https://cdn.discordapp.com/avatars/111111111/av1.png?size=256"]})
+
+  (fact
+    "a user without an avatar gets the default discord avatar URL"
+    (img-input/extract-images
+      "<@123456789>"
+      {:raw-event {:mentions [{:id "123456789" :username "alice"}]}})
+    => {:prompt "(Image 1 is @alice) @alice"
+        :image-urls ["https://cdn.discordapp.com/embed/avatars/5.png"]}))