osmarks/autobotrobot
1
0
Fork 0
mirror of https://github.com/osmarks/autobotrobot synced 2025-04-04 22:36:56 +00:00
Code Issues Releases Wiki Activity

Meme search API changes

Browse Source
This commit is contained in:
osmarks 2024-05-21 15:26:28 +01:00
parent 6e30123e56
commit d64e7895f0
2 changed files with 22 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/commands.py
View File

@ -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):

21
src/util.py
View File

@ -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"]
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]