diff --git a/src/commands.py b/src/commands.py index f9f8585..52243cc 100644 --- a/src/commands.py +++ b/src/commands.py @@ -232,15 +232,14 @@ AutoBotRobot is operated by gollark/osmarks. 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 + "terms": [{"text": query, "weight": 1}], + "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 ] + o_files = [ discord.File(Path(util.config["memetics"]["memes_local"]) / Path(util.config["memetics"]["meme_base"]) / m[1]) for m in results["matches"] ] 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 ] + o_files = [ discord.File(Path(util.config["memetics"]["memes_local"]) / util.meme_thumbnail(results, m)) for m in results["matches"] ] await ctx.send(files=o_files) def setup(bot): diff --git a/src/util.py b/src/util.py index bf4bc5a..d098a59 100644 --- a/src/util.py +++ b/src/util.py @@ -17,6 +17,7 @@ import pytz import collections import aiohttp import string +from pathlib import Path config = {} @@ -341,12 +342,26 @@ async def generate(response: aiohttp.ClientSession, prompt): async with response.post(config["ai"]["llm_backend"], json={ "prompt": prompt, "max_tokens": 200, - "stop": ["\n"] + "stop": ["\n"], + "client": "abr" }) as res: data = await res.json() 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 +TARGET_FORMAT = "jpegh" +def meme_thumbnail(results, result): + try: + format_id = results["formats"].index(TARGET_FORMAT) + except ValueError: + format_id = None + + if not format_id: + return Path(config["memetics"]["meme_base"]) / result[1] + else: + format_id = 1 << format_id + if result[3] & format_id != 0: + return Path(config["memetics"]["thumb_base"]) / f"{result[2]}{TARGET_FORMAT}.{results['extensions'][TARGET_FORMAT]}" + else: + return Path(config["memetics"]["meme_base"]) / result[1] \ No newline at end of file