from logging import warn import aiohttp import asyncio import itertools import json import meilisearch import os from aiobotocore.session import get_session as aiobotocore_get_session from db.attachment import Attachment from db.meme import Meme from db.migrate import migrate_db from dotenv import dotenv_values, load_dotenv from google.cloud import vision as google_vision from google.cloud.vision_v1 import types from io import BytesIO from PIL import Image from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker, Session from types_aiobotocore_s3.client import S3Client from urllib3.util import parse_url from urllib3.util import Url from uuid import uuid4 # need to change the env for the google api calls load_dotenv() config = dotenv_values(".env") MEILI_INDEX = 'memes' MEILI_URL = 'http://localhost:7700' CHANNEL_IDS = [ 367936426962190336, # TPM shitposts 1095060717117374687, # aella memes ] STORAGE_BUCKET = config['STORAGE_BUCKET'] assert STORAGE_BUCKET headers = { 'Authorization': config['DISCORD_AUTH'], } def detect_text(url, vision): request = { 'image': { 'source': { 'image_uri': url } }, 'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)], } response: types.AnnotateImageResponse = vision.annotate_image(request=request) vision_json = json.loads(types.AnnotateImageResponse.to_json(response)) return request def filter_message_img_urls(messages): url_lists = map( lambda msg: list(map(lambda attachment: (parse_url(attachment['url']), msg), msg['attachments'])) , messages) return [item for group in url_lists for item in group] async def process_channel( http_session: aiohttp.ClientSession, storage: S3Client, vision: google_vision.ImageAnnotatorClient, db_session_maker: sessionmaker[Session], channel_id: int ): async def fetch_channel_messages(channel_id): url = f"https://discord.com/api/v9/channels/{channel_id}/messages" async with http_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 http_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 http_session.get(url, headers=headers) as res: return await res.json() async def download_attachment(attachment: Attachment): async with http_session.get(attachment.url) as response: if response.status == 200: (filename, ext) = attachment.file_details() object_name = attachment.object_name() print(f"uploading object_name={object_name} filename={filename} ext={ext}") res_bytes = await response.read() # make a webp conversion for obvious image types if attachment.convert_to_webp(): # have to read the whole img before converting to webp image = Image.open(BytesIO(res_bytes)) webp_output = BytesIO() image.save(webp_output, format='WebP') webp_output.seek(0) return await storage.put_object(Bucket=STORAGE_BUCKET, Key=object_name, Body=webp_output) return await storage.put_object(Bucket=STORAGE_BUCKET, Key=object_name, Body=res_bytes) def create_attachments(message, channel, guild): urls = [parse_url(attachment['url']) for attachment in message['attachments']] def generate_attachment(url: Url): if not url.path: return None (filename, ext) = os.path.splitext(os.path.basename(url.path)) extensions = ['jpg', 'png', 'jpeg', 'webp', 'mov', 'gif', 'mp4'] if ext.lower()[1:] not in extensions: return None return Attachment( id=uuid4(), msg_id=int(message['id']), message=message, guild=guild, url=url.url, channel=channel ) return [generate_attachment(url) for url in urls] 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) # at this point we have the channel messages, the guild (server) and channel info # first thing we have to do is parse out the urls in the msg attachments # once we have those we can generate where we should put them in our storage # should only persist them later once the processing is done attachment_groups = [create_attachments(message, channel, guild) for message in messages] # super wack flatmap and filter attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups))) db_session = db_session_maker() # filter attachments that already exist in the db # dirty way is just check if msg_id exists in db, don't care about specific attachments # ok now we will try and upload each attachment # for the ones that succeed we will add them to our session to be stored async def upload_and_persist(attachment: Attachment): try: await download_attachment(attachment=attachment) # this is kinda dodge, the exception will catch if we fail to add to session db_session.add(attachment) except Exception as e: warn(e) warn(f'deleting attachment upload {attachment.object_name()}') await storage.delete_object(Bucket=STORAGE_BUCKET, Key=attachment.object_name()) attachments = [upload_and_persist(attachment=a) for a in attachments] await asyncio.gather(*attachments) db_session.commit() # NOW # ------------------- # insert attachment # upload to s3 # fetch vision for relevant attachments # create meme from attachment and vision # insert meme # add documents to meili # tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages] # await asyncio.gather(*tasks) # [attachment.detect_text(vision) for attachment in attachments] # [attachment.persist() for attachment in attachments] # client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY']) # # have to vars the dataclass out for meili # documents = [vars(attachment.to_meme()) for attachment in attachments] # client.index(MEILI_INDEX).add_documents(documents) async def check_bucket_exists(storage, bucket_name): try: await storage.head_bucket(Bucket=bucket_name) return True except storage.exceptions.ClientError as e: error_code = e.response.get('Error', {}).get('Code') if error_code == '404': return False # Bucket does not exist else: raise # Reraise other exceptions async def main(): assert config['STORAGE_URL'] assert config['POSTGRES_URL'] vision = google_vision.ImageAnnotatorClient() # db setup db_url = config['POSTGRES_URL'] db_engine = create_engine(db_url, echo=True) migrate_db(db_engine, db_url) db_session_maker = sessionmaker(bind=db_engine) aioboto = aiobotocore_get_session() async with aioboto.create_client( 's3', endpoint_url=config['STORAGE_URL'], aws_access_key_id=config['STORAGE_ACCESS_KEY'], aws_secret_access_key=config['STORAGE_SECRET_KEY'], ) as storage, aiohttp.ClientSession() as http_session: storage: S3Client # Make the bucket if it doesn't exist. exists = await check_bucket_exists(storage, STORAGE_BUCKET) if not exists: await storage.create_bucket(Bucket=STORAGE_BUCKET) print("Created bucket", STORAGE_BUCKET) else: print("Bucket", STORAGE_BUCKET, "already exists") channels = [process_channel(http_session, storage, vision, db_session_maker, channel_id) for channel_id in CHANNEL_IDS] memes = await asyncio.gather(*channels) asyncio.run(main()) # TODO # [ ] Backfill last N images from a channel # [ ] Set oldest_img_id for channel # logger