From 6e30123e56a5592a728e93b427f7ae5002a843e3 Mon Sep 17 00:00:00 2001
From: osmarks <osmarks@protonmail.com>
Date: Sun, 12 May 2024 17:44:00 +0100
Subject: [PATCH] Meme Search Engine frontend

---
 requirements.txt |  6 +++---
 src/commands.py  | 18 ++++++++++++++++++
 src/util.py      | 12 +++++++++---
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 2d1f4e3..e3ec73a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,9 +1,9 @@
 # TODO
 pytio==0.3.1
-aiohttp==3.8.3
-aiosqlite==0.17.0
+aiohttp==3.9.3
+aiosqlite==0.19.0
 nextcord==2.3.2
-numpy==1.23
+numpy==1.26
 prometheus-async==19.2.0
 prometheus-client==0.15.0
 pydot==1.4.2
diff --git a/src/commands.py b/src/commands.py
index f94708d..f9f8585 100644
--- a/src/commands.py
+++ b/src/commands.py
@@ -6,7 +6,9 @@ import re
 import aiohttp
 import subprocess
 import discord.ext.commands as commands
+import discord
 from datetime import datetime
+from pathlib import Path
 
 import tio
 import util
@@ -225,5 +227,21 @@ AutoBotRobot is operated by gollark/osmarks.
         generation = await util.generate(self.session, util.config["ai"]["prompt_start"] + "".join(prompt))
         if generation.strip(): await ctx.send(generation.strip())
 
+    @commands.command(help="Search meme library.", aliases=["memes"])
+    async def meme(self, ctx, *, query=None):
+        search_many = ctx.invoked_with == "memes"
+        raw_memes = await util.user_config_lookup(ctx, "enable_raw_memes") == "true"
+        async with self.session.post(util.config["memetics"]["meme_search_backend"], json={
+            "text": [[query, 1]],
+            "top_k": 4 if search_many else 1
+        }) as res:
+            results = await res.json()
+        files = [ x["file"] for x in results ]
+        if raw_memes:
+            o_files = [ discord.File(Path(util.config["memetics"]["memes_local"]) / Path(util.config["memetics"]["meme_base"]) / f) for f in files ]
+        else:
+            o_files = [ discord.File(Path(util.config["memetics"]["memes_local"]) / Path(util.config["memetics"]["thumb_base"]) / util.meme_thumbnail(f)) for f in files ]
+        await ctx.send(files=o_files)
+
 def setup(bot):
     bot.add_cog(GeneralCommands(bot))
diff --git a/src/util.py b/src/util.py
index 346d29e..bf4bc5a 100644
--- a/src/util.py
+++ b/src/util.py
@@ -16,6 +16,7 @@ import math
 import pytz
 import collections
 import aiohttp
+import string
 
 config = {}
 
@@ -36,7 +37,7 @@ prefixes = {
     # highly dubiously useful unofficial prefixes
     "R": 27, "r": -27, "Q": 30, "q": -30, "X": 27, "x": -27, "W": 30, "w": -30
 }
-number = "(-?[0-9]+(?:\.[0-9]+)?)(" + "|".join(prefixes.keys()) + ")?"
+number = "(-?[0-9]+(?:\\.[0-9]+)?)(" + "|".join(prefixes.keys()) + ")?"
 
 time_units = (
     ("galacticyears", "cosmicyears", "gy", "[Cc]y"),
@@ -140,7 +141,7 @@ def format_timedelta(from_, to):
         if x != 0: out += str(x) + short
     return "0s" if out == "" else out
 
-CODEBLOCK_REGEX = "^[^`]*```[a-zA-Z0-9_\-+]*\n(.+)```$"
+CODEBLOCK_REGEX = "^[^`]*```[a-zA-Z0-9_\\-+]*\n(.+)```$"
 CODELINE_REGEX = "^[^`]*`(.*)`$"
 def extract_codeblock(s):
     match1 = re.match(CODEBLOCK_REGEX, s, flags=re.DOTALL)
@@ -343,4 +344,9 @@ async def generate(response: aiohttp.ClientSession, prompt):
         "stop": ["\n"]
     }) as res:
         data = await res.json()
-        return data["choices"][0]["text"]
\ No newline at end of file
+        return data["choices"][0]["text"]
+
+filesafe_charset = string.ascii_letters + string.digits + "-"
+
+def meme_thumbnail(original):
+    return ''.join([ i if i in filesafe_charset else '_' for i in (config["memetics"]["meme_base"] + "/" + original) ]) + "." + config["memetics"]["target_format"] + config["memetics"]["target_format_ext"]
\ No newline at end of file