import os import asyncio import json import itertools from urllib3.util import parse_url from pathlib import Path from io import BytesIO from dataclasses import dataclass from uuid import uuid4, UUID import aiohttp from aiobotocore.session import get_session as aiobotocore_get_session from types_aiobotocore_s3.client import S3Client from PIL import Image 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 import meilisearch from urllib3.util import Url, url from sqlalchemy import create_engine from db.migrate import migrate_db from db.attachment import Attachment from db.meme import Meme # need to change the env for the google api calls load_dotenv() config = dotenv_values(".env") DATA_DIR = './data/memestream' MEILI_INDEX = 'memes' MEILI_URL = 'http://localhost:7700' 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 MemeDocument: id: UUID url: Url text: str tags: list[str] 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 download_channel_attachments(session, storage, 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() async def upload_object(object_name, body): storage_bucket = config['STORAGE_BUCKET'] if not storage_bucket: raise RuntimeError(f"config missing STORAGE_BUCKET") return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=body) async def download_attachment( session: aiohttp.ClientSession, storage: S3Client, attachment: Attachment ): async with session.get(attachment.url) as response: if response.status == 200: # make a webp conversion for obvious image types res_bytes = await response.read() (filename, ext) = attachment.file_details() object_name = attachment.object_name() print(f"uploading object_name={object_name} filename={filename} ext={ext}") if attachment.convert_to_webp(): image = Image.open(BytesIO(res_bytes)) webp_output = BytesIO() image.save(webp_output, format='WebP') webp_output.seek(0) return await upload_object(object_name=object_name, body=webp_output) return await upload_object(object_name=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))) # 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) 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: storage: S3Client storage_bucket = config['STORAGE_BUCKET'] assert storage_bucket # 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") async with aiohttp.ClientSession() as session: channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS] memes = await asyncio.gather(*channels) asyncio.run(main())