memestream bb

main.py

@@ -1,24 +1,35 @@
 import os
 import asyncio
-import aiohttp
 import json
-import meilisearch
-
+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
-from PIL import Image
+from minio import Minio
+import meilisearch
+from urllib3.util import Url, url
+
+import uploader
 
 # need to change the env for the google api calls
 load_dotenv()
 config = dotenv_values(".env")
 
-DATA_DIR = './data/memesearch'
+DATA_DIR = './data/memestream'
 MEILI_INDEX = 'memes'
 MEILI_URL = 'http://localhost:7700'
 
@@ -39,57 +50,31 @@ headers = {
 
 @dataclass
 class MemeDocument:
-    id: int
-    url: str
-    text: str | None
+    id: UUID
+    url: Url
+    text: str
     tags: List[str]
 
 @dataclass
-class MemeMessage:
-    id: int
+class DiscordAttachmentMessage:
+    id: UUID
+    msg_id: int
     url: str
+    upload_url: str
     guild_name: str
     channel_name: str
-
     vision_json: str | None
-    text: 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]
 
-    def output_dir(self):
+    def document_dir(self):
         return f'{DATA_DIR}/guilds/{self.guild_name}/{self.channel_name}'
 
-    def img_dir(self):
+    def file_dir(self):
         return f"{DATA_DIR}/imgs/"
 
-    async def save_img(self, session: aiohttp.ClientSession):
-        webp_output = f'{DATA_DIR}/imgs/'
-
-        # 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(self.output_dir(), exist_ok=True);
-        os.makedirs(self.img_dir(), exist_ok=True);
-
-        async with session.get(self.url) as response:
-            if response.status == 200:
-                res_bytes = await response.read()
-
-                with open(f'{self.output_dir()}/{self.id}{ext}', 'wb') as f:
-                    f.write(res_bytes)
-
-                convert_extensions = ['jpg', 'png', 'jpeg', 'webp']
-                if ext.lower() in convert_extensions:
-                    image = Image.open(BytesIO(res_bytes))
-                    image.save(f"{self.img_dir()}/{self.id}.webp", format='WebP')
-
-    def annotate_image_request(self, vision):
+    def detect_text(self, vision):
         request = {
             'image': { 'source': { 'image_uri': self.url } },
             'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
@@ -104,24 +89,16 @@ class MemeMessage:
         return request
 
     def persist(self):
-        with open(f'{self.output_dir()}/{self.id}.json', 'w') as f:
+        with open(f'{self.document_dir()}/{self.id}.json', 'w') as f:
             json.dump(vars(self), f, indent=4)
 
-    def to_document(self):
-        return MemeDocument(
-            id=self.id,
-            url=f'{DATA_DIR}/{self.id}.webp',
-            text=self.text,
-            tags=[self.guild_name, self.channel_name],
-        )
-
 def filter_message_img_urls(messages):
     url_lists = map(
-        lambda msg: list(map(lambda attachment: (attachment['url'], msg), msg['attachments']))
+        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_imgs(session, vision, channel_id):
+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:
@@ -142,42 +119,122 @@ async def download_channel_imgs(session, vision, 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 = filter_message_img_urls(messages)
+    urls_and_messages = 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]
+    tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages]
     await asyncio.gather(*tasks)
 
-    [meme.annotate_image_request(vision) for meme in memes]
-    [meme.persist() for meme in memes]
+    # [attachment.detect_text(vision) for attachment in attachments]
+    # [attachment.persist() for attachment in attachments]
 
-    client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY'])
+    # client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY'])
 
-    # have to vars the dataclass out for meili
-    documents = [vars(meme.to_document()) for meme in memes]
-    client.index(MEILI_INDEX).add_documents(documents)
+    # # 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['MINIO_URL']
+
     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)
+
+    # storage = Minio(config['MINIO_URL'],
+    #     access_key=config['MINIO_ACCESS_KEY'],
+    #     secret_key=config['MINIO_SECRET_KEY'],
+    # )
+
+    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())