import os import asyncio import json from urllib3.util import parse_url from pathlib import Path from io import BytesIO from dataclasses import dataclass from uuid import uuid4, UUID from typing import Dict, Any, List 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 # 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 download_attachment( session: aiohttp.ClientSession, storage: S3Client, url: Url, message, ): # there's a big param of some kind after the image url that we remove if not url.path: raise RuntimeError(f"url missing path url={url}") storage_bucket = config['UPLOAD_BUCKET'] if not storage_bucket: raise RuntimeError(f"config missing UPLOAD_BUCKET") (filename, ext) = os.path.splitext(os.path.basename(url.path)) # TODO # FILTER WANTED TYPES id = uuid4() async with session.get(url.url) as response: if response.status == 200: # make a webp conversion for obvious image types convert_extensions = ['jpg', 'png', 'jpeg', 'webp'] res_bytes = await response.read() object_name: str # trim the . if ext.lower()[1:] in convert_extensions: image = Image.open(BytesIO(res_bytes)) webp_output = BytesIO() image.save(webp_output, format='WebP') webp_output.seek(0) object_name = f"{id}.webp" print(f"uploading object_name={object_name} filename={filename} ext={ext}") return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=webp_output) object_name = f"{id}{ext}" print(f"uploading object_name={object_name} filename={filename} ext={ext}") return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=res_bytes) # add discord id to server processed msgs list # if there then nothing to do # return MemeDocument( # id=id, # url=self.upload_url, # text=self.text or '', # tags=[self.guild_name, self.channel_name], # ) 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_and_messages = filter_message_img_urls(messages) 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) 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['UPLOAD_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())