meili

main.py

@@ -0,0 +1,155 @@
+import os
+import asyncio
+import aiohttp
+import json
+import meilisearch
+
+from dataclasses import dataclass
+from typing import Dict, Any
+
+from dotenv import dotenv_values, load_dotenv
+from google.cloud import vision as google_vision
+from google.protobuf.json_format import MessageToDict, MessageToJson
+from google.cloud.vision_v1 import types
+
+# need to change the env for the google api calls
+load_dotenv()
+config = dotenv_values(".env")
+
+CHANNEL_IDS = [
+    367936426962190336, # TPM shitposts
+    1095060717117374687, # aella memes
+]
+
+headers = {
+    'Authorization': config['DISCORD_AUTH'],
+}
+
+# TODO
+# [ ] Backfill last N images from a channel
+# [ ] Set oldest_img_id for channel
+# [ ] save img as json + imagefile
+# [ ] check extensions (ignore mp4 maybe)
+
+@dataclass
+class MemeMessage:
+    id: int
+    url: str
+    guild_name: str
+    channel_name: str
+
+    text: str | None
+    vision_json: str | None
+
+    # maybe don't need these
+    # kind of extra
+    message_json: Dict[str, Any]
+    guild_json: Dict[str, Any]
+    channel_json: Dict[str, Any]
+
+    async def save_img(self, session: aiohttp.ClientSession):
+        output_dir = self.output_dir()
+
+        # there's a big param of some kind after the image url that we remove
+        url_filename = os.path.basename(self.url).split('?')[0]
+        (filename, ext) = os.path.splitext(url_filename)
+        print(f"filename={filename} ext={ext}")
+        os.makedirs(output_dir, exist_ok=True);
+
+        async with session.get(self.url) as response:
+            if response.status == 200:
+                with open(f'{output_dir}/{self.id}{ext}', 'wb') as f:
+                    f.write(await response.read())
+
+    def output_dir(self):
+        return f'./data/{self.guild_name}/{self.channel_name}'
+
+    def annotate_image_request(self, vision):
+        request = {
+            'image': { 'source': { 'image_uri': self.url } },
+            'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
+        }
+
+        response: types.AnnotateImageResponse = vision.annotate_image(request=request)
+        self.vision_json = json.loads(types.AnnotateImageResponse.to_json(response))
+
+        if len(response.text_annotations) > 0:
+            self.text = response.text_annotations[0].description
+
+        return request
+
+    def persist(self):
+        with open(f'{self.output_dir()}/{self.id}.json', 'w') as f:
+            json.dump(vars(self), f, indent=4)
+
+    def to_meili_document(self):
+        return {
+            'id': self.id,
+            'text': self.text,
+            'url': self.url,
+        }
+
+def filter_message_img_urls(messages):
+    url_lists = map(
+        lambda msg: list(map(lambda attachment: (attachment['url'], msg), msg['attachments']))
+    , messages)
+    return [item for group in url_lists for item in group]
+
+async def download_channel_imgs(session, vision, channel_id):
+    async def fetch_channel_messages(channel_id):
+        url = f"https://discord.com/api/v9/channels/{channel_id}/messages"
+        async with session.get(url, headers=headers) as res:
+            return await res.json()
+
+    async def fetch_guild_details(guild_id):
+        url = f"https://discord.com/api/v9/guilds/{guild_id}"
+        async with session.get(url, headers=headers) as res:
+            data = await res.json()
+
+            keys_to_remove = ['emojis', 'stickers', 'roles']
+            [data.pop(key, None) for key in keys_to_remove]
+
+            return data
+
+    async def fetch_channel(channel_id):
+        url = f"https://discord.com/api/v9/channels/{channel_id}"
+        async with session.get(url, headers=headers) as res:
+            return await res.json()
+
+    channel = await fetch_channel(channel_id)
+    guild = await fetch_guild_details(channel['guild_id'])
+    print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}")
+    messages = await fetch_channel_messages(channel_id)
+    urls = filter_message_img_urls(messages)
+
+    memes = list(map(
+        lambda msg: MemeMessage(
+            id=msg[1]['id'],
+            url=msg[0],
+            guild_name=guild['name'],
+            channel_name=channel['name'],
+            guild_json=guild,
+            channel_json=channel,
+            message_json=msg[1],
+            vision_json=None,
+            text=None,
+        )
+    , urls))
+
+    tasks = [meme.save_img(session) for meme in memes]
+    await asyncio.gather(*tasks)
+
+    [meme.annotate_image_request(vision) for meme in memes]
+    [meme.persist() for meme in memes]
+
+    client = meilisearch.Client('http://localhost:7700', 'aSampleMasterKey')
+    documents = [meme.to_meili_document() for meme in memes]
+    client.index('memes').add_documents(documents)
+
+async def main():
+    vision = google_vision.ImageAnnotatorClient()
+    async with aiohttp.ClientSession() as session:
+        channels = [download_channel_imgs(session, vision, channel_id) for channel_id in CHANNEL_IDS]
+        memes = await asyncio.gather(*channels)
+
+asyncio.run(main())