mem0/embedchain/loaders/discord.py

import hashlib
import logging
import os

from embedchain.helpers.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader

logger = logging.getLogger(__name__)


@register_deserializable
class DiscordLoader(BaseLoader):
    """
    Load data from a Discord Channel ID.
    """

    def __init__(self):
        if not os.environ.get("DISCORD_TOKEN"):
            raise ValueError("DISCORD_TOKEN is not set")

        self.token = os.environ.get("DISCORD_TOKEN")

    @staticmethod
    def _format_message(message):
        return {
            "message_id": message.id,
            "content": message.content,
            "author": {
                "id": message.author.id,
                "name": message.author.name,
                "discriminator": message.author.discriminator,
            },
            "created_at": message.created_at.isoformat(),
            "attachments": [
                {
                    "id": attachment.id,
                    "filename": attachment.filename,
                    "size": attachment.size,
                    "url": attachment.url,
                    "proxy_url": attachment.proxy_url,
                    "height": attachment.height,
                    "width": attachment.width,
                }
                for attachment in message.attachments
            ],
            "embeds": [
                {
                    "title": embed.title,
                    "type": embed.type,
                    "description": embed.description,
                    "url": embed.url,
                    "timestamp": embed.timestamp.isoformat(),
                    "color": embed.color,
                    "footer": {
                        "text": embed.footer.text,
                        "icon_url": embed.footer.icon_url,
                        "proxy_icon_url": embed.footer.proxy_icon_url,
                    },
                    "image": {
                        "url": embed.image.url,
                        "proxy_url": embed.image.proxy_url,
                        "height": embed.image.height,
                        "width": embed.image.width,
                    },
                    "thumbnail": {
                        "url": embed.thumbnail.url,
                        "proxy_url": embed.thumbnail.proxy_url,
                        "height": embed.thumbnail.height,
                        "width": embed.thumbnail.width,
                    },
                    "video": {
                        "url": embed.video.url,
                        "height": embed.video.height,
                        "width": embed.video.width,
                    },
                    "provider": {
                        "name": embed.provider.name,
                        "url": embed.provider.url,
                    },
                    "author": {
                        "name": embed.author.name,
                        "url": embed.author.url,
                        "icon_url": embed.author.icon_url,
                        "proxy_icon_url": embed.author.proxy_icon_url,
                    },
                    "fields": [
                        {
                            "name": field.name,
                            "value": field.value,
                            "inline": field.inline,
                        }
                        for field in embed.fields
                    ],
                }
                for embed in message.embeds
            ],
        }

    def load_data(self, channel_id: str):
        """Load data from a Discord Channel ID."""
        import discord

        messages = []

        class DiscordClient(discord.Client):
            async def on_ready(self) -> None:
                logger.info("Logged on as {0}!".format(self.user))
                try:
                    channel = self.get_channel(int(channel_id))
                    if not isinstance(channel, discord.TextChannel):
                        raise ValueError(
                            f"Channel {channel_id} is not a text channel. " "Only text channels are supported for now."
                        )
                    threads = {}

                    for thread in channel.threads:
                        threads[thread.id] = thread

                    async for message in channel.history(limit=None):
                        messages.append(DiscordLoader._format_message(message))
                        if message.id in threads:
                            async for thread_message in threads[message.id].history(limit=None):
                                messages.append(DiscordLoader._format_message(thread_message))

                except Exception as e:
                    logger.error(e)
                    await self.close()
                finally:
                    await self.close()

        intents = discord.Intents.default()
        intents.message_content = True
        client = DiscordClient(intents=intents)
        client.run(self.token)

        metadata = {
            "url": channel_id,
        }

        messages = str(messages)

        doc_id = hashlib.sha256((messages + channel_id).encode()).hexdigest()

        return {
            "doc_id": doc_id,
            "data": [
                {
                    "content": messages,
                    "meta_data": metadata,
                }
            ],
        }