cortexdb-connectors 0.2.0__py3-none-any.whl

cortexdb_connectors/discord/__init__.py

@@ -0,0 +1,646 @@
+"""
+Discord connector for CortexDB.
+
+Ingests messages, thread messages, and forum posts from Discord guilds into
+CortexDB as Episodes.  Uses ``aiohttp`` for lightweight REST API access
+against the Discord API v10 and respects per-route rate limits.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+from cortexdb_connectors.base import (
+    Actor,
+    CortexConnector,
+    EntityRef,
+    Episode,
+    EpisodeType,
+    RawEvent,
+    Source,
+    Visibility,
+    VisibilityLevel,
+)
+
+logger = logging.getLogger(__name__)
+
+_SOURCE = Source(system="discord")
+
+_API_BASE = "https://discord.com/api/v10"
+
+# Per-route rate limit safety delay.
+_RATE_LIMIT_DELAY = 0.25
+
+# Number of messages to fetch per request (max 100 per Discord docs).
+_MESSAGE_LIMIT = 100
+
+# Discord epoch: 2015-01-01T00:00:00Z in milliseconds.
+_DISCORD_EPOCH = 1420070400000
+
+
+class DiscordConnector(CortexConnector):
+    """Connector that syncs Discord guild messages, threads, and forum
+    posts into CortexDB.
+
+    Parameters
+    ----------
+    cortex_url:
+        Base URL of the CortexDB API.
+    cortex_api_key:
+        Bearer token for CortexDB authentication.
+    bot_token:
+        Discord bot token.
+    guild_ids:
+        List of guild (server) IDs to sync.
+    channel_ids:
+        Explicit list of channel IDs to sync.  If empty, all text
+        channels in the configured guilds are discovered.
+    tenant_id:
+        CortexDB tenant identifier.
+    namespace:
+        CortexDB namespace for ingested episodes.
+    backfill_days:
+        Number of days of history to backfill on first sync.
+    include_bots:
+        Whether to include messages from bot accounts.
+    """
+
+    connector_name: str = "discord"
+
+    def __init__(
+        self,
+        cortex_url: str,
+        cortex_api_key: str,
+        bot_token: str,
+        guild_ids: list[str] | None = None,
+        channel_ids: list[str] | None = None,
+        tenant_id: str = "default",
+        namespace: str = "discord",
+        backfill_days: int = 30,
+        include_bots: bool = False,
+    ) -> None:
+        super().__init__(cortex_url, cortex_api_key, tenant_id)
+        self.bot_token = bot_token
+        self.guild_ids = guild_ids or []
+        self.channel_ids = channel_ids or []
+        self.namespace = namespace
+        self.backfill_days = backfill_days
+        self.include_bots = include_bots
+        self._guild_cache: dict[str, dict[str, Any]] = {}
+        self._channel_cache: dict[str, dict[str, Any]] = {}
+
+    # -- CortexConnector interface -------------------------------------------
+
+    async def fetch_events(
+        self, since: datetime | None = None
+    ) -> list[RawEvent]:
+        """Fetch messages from configured Discord guilds and channels.
+
+        If *since* is provided, only messages newer than that timestamp are
+        returned.  Uses ``before`` message-ID pagination.
+        """
+        import aiohttp
+
+        cutoff = since or datetime.now(timezone.utc) - timedelta(
+            days=self.backfill_days
+        )
+        events: list[RawEvent] = []
+
+        headers = {
+            "Authorization": f"Bot {self.bot_token}",
+            "Content-Type": "application/json",
+        }
+
+        async with aiohttp.ClientSession(headers=headers) as session:
+            # Resolve channels to sync.
+            channels = await self._resolve_channels(session)
+
+            for channel in channels:
+                channel_id = channel["id"]
+                channel_type = channel.get("type", 0)
+
+                # Text channels (0), announcement (5), voice text (2 is
+                # voice but may have text), forum (15).
+                if channel_type == 15:
+                    # Forum channel: fetch threads as documents.
+                    thread_events = await self._fetch_forum_threads(
+                        session, channel, cutoff
+                    )
+                    events.extend(thread_events)
+                elif channel_type in (0, 5, 11, 12):
+                    # Regular text / announcement / public/private thread.
+                    msg_events = await self._fetch_channel_messages(
+                        session, channel, cutoff
+                    )
+                    events.extend(msg_events)
+
+            # Fetch active threads in each guild.
+            for guild_id in self.guild_ids:
+                thread_events = await self._fetch_active_threads(
+                    session, guild_id, cutoff
+                )
+                events.extend(thread_events)
+
+        return events
+
+    def normalize(self, raw: RawEvent) -> Episode:
+        """Convert a Discord ``RawEvent`` into a CortexDB ``Episode``."""
+        data = raw.data
+        event_kind = data.get("_event_kind", "message")
+
+        if event_kind == "forum_post":
+            return self._normalize_forum_post(raw)
+        return self._normalize_message(raw)
+
+    def map_permissions(self, raw: RawEvent) -> Visibility:
+        """Map Discord channel permissions to CortexDB visibility.
+
+        DMs are ``Private``, private threads and channels with restricted
+        overwrites are ``Restricted``, and public channels are
+        ``Organization``.
+        """
+        perms = raw.permissions
+        channel_type = perms.get("channel_type", 0)
+
+        # DM channels.
+        if channel_type in (1, 3):
+            return Visibility(level=VisibilityLevel.PRIVATE)
+
+        # Private threads.
+        if channel_type == 12:
+            members = tuple(perms.get("thread_member_ids", []))
+            return Visibility(
+                level=VisibilityLevel.RESTRICTED,
+                allowed_principals=members,
+            )
+
+        # Channels with specific role restrictions.
+        if perms.get("is_restricted"):
+            roles = tuple(perms.get("allowed_role_ids", []))
+            return Visibility(
+                level=VisibilityLevel.RESTRICTED,
+                allowed_principals=roles,
+            )
+
+        return Visibility(level=VisibilityLevel.ORGANIZATION)
+
+    # -- message normalization -----------------------------------------------
+
+    def _normalize_message(self, raw: RawEvent) -> Episode:
+        """Normalize a Discord message into an Episode."""
+        data = raw.data
+        author = data.get("author", {})
+        actor = Actor(
+            id=author.get("id", "unknown"),
+            name=author.get("username", "Unknown"),
+            avatar_url=(
+                f"https://cdn.discordapp.com/avatars/"
+                f"{author.get('id', '')}/{author.get('avatar', '')}.png"
+                if author.get("avatar")
+                else None
+            ),
+        )
+
+        channel_id = data.get("channel_id", "")
+        channel_name = data.get("_channel_name", channel_id)
+        guild_id = data.get("guild_id", "")
+        guild_name = data.get("_guild_name", guild_id)
+
+        content = data.get("content", "")
+
+        # Append attachment descriptions.
+        attachments = data.get("attachments", [])
+        if attachments:
+            att_desc = ", ".join(
+                a.get("filename", "attachment") for a in attachments
+            )
+            content += f"\n[Attachments: {att_desc}]"
+
+        # Append embed titles.
+        embeds = data.get("embeds", [])
+        if embeds:
+            embed_desc = ", ".join(
+                e.get("title", "embed") for e in embeds if e.get("title")
+            )
+            if embed_desc:
+                content += f"\n[Embeds: {embed_desc}]"
+
+        entities = self._extract_message_entities(data)
+
+        # Thread context.
+        thread_id_val: str | None = None
+        parent_id_val: str | None = None
+        if data.get("_thread_id"):
+            thread_id_val = f"discord:{data['_thread_id']}"
+            parent_id_val = f"discord:{channel_id}"
+        elif data.get("message_reference"):
+            ref = data["message_reference"]
+            parent_id_val = f"discord:{ref.get('message_id', '')}"
+
+        # Reaction metadata.
+        reactions = data.get("reactions", [])
+        reaction_summary = [
+            {
+                "emoji": r.get("emoji", {}).get("name", ""),
+                "count": r.get("count", 0),
+            }
+            for r in reactions
+        ]
+
+        return Episode(
+            actor=actor,
+            source=_SOURCE,
+            episode_type=EpisodeType.MESSAGE,
+            occurred_at=raw.timestamp,
+            content=content,
+            raw_payload=data,
+            tenant_id=self.tenant_id,
+            namespace=self.namespace,
+            entities=entities,
+            tags=["discord", guild_name.lower().replace(" ", "_")],
+            metadata={
+                "message_id": data.get("id", ""),
+                "channel_id": channel_id,
+                "channel_name": channel_name,
+                "guild_id": guild_id,
+                "guild_name": guild_name,
+                "attachment_count": len(attachments),
+                "embed_count": len(embeds),
+                "reactions": reaction_summary,
+                "mention_count": len(data.get("mentions", [])),
+                "pinned": data.get("pinned", False),
+            },
+            thread_id=thread_id_val,
+            parent_id=parent_id_val,
+            idempotency_key=f"discord:{data.get('id', raw.external_id)}",
+        )
+
+    def _normalize_forum_post(self, raw: RawEvent) -> Episode:
+        """Normalize a Discord forum post (thread starter) into an Episode."""
+        data = raw.data
+        author = data.get("author", {})
+        actor = Actor(
+            id=author.get("id", "unknown"),
+            name=author.get("username", "Unknown"),
+            avatar_url=(
+                f"https://cdn.discordapp.com/avatars/"
+                f"{author.get('id', '')}/{author.get('avatar', '')}.png"
+                if author.get("avatar")
+                else None
+            ),
+        )
+
+        thread_name = data.get("_thread_name", "")
+        content = data.get("content", "")
+        if thread_name:
+            content = f"[{thread_name}]\n\n{content}"
+
+        entities = self._extract_message_entities(data)
+
+        return Episode(
+            actor=actor,
+            source=_SOURCE,
+            episode_type=EpisodeType.DOCUMENT,
+            occurred_at=raw.timestamp,
+            content=content,
+            raw_payload=data,
+            tenant_id=self.tenant_id,
+            namespace=self.namespace,
+            entities=entities,
+            tags=["discord", "forum_post"],
+            metadata={
+                "message_id": data.get("id", ""),
+                "thread_id": data.get("_thread_id", ""),
+                "thread_name": thread_name,
+                "channel_id": data.get("channel_id", ""),
+                "guild_id": data.get("guild_id", ""),
+                "applied_tags": data.get("_applied_tags", []),
+            },
+            thread_id=f"discord:{data.get('_thread_id', '')}",
+            idempotency_key=f"discord:forum:{raw.external_id}",
+        )
+
+    # -- fetchers ------------------------------------------------------------
+
+    async def _api_get(
+        self, session: Any, path: str, params: dict[str, Any] | None = None
+    ) -> Any:
+        """Make a GET request to the Discord API with rate-limit handling."""
+        url = f"{_API_BASE}{path}"
+
+        for attempt in range(5):
+            async with session.get(url, params=params) as resp:
+                if resp.status == 429:
+                    retry_data = await resp.json()
+                    retry_after = retry_data.get("retry_after", 1.0)
+                    logger.warning(
+                        "Discord rate-limited on %s; sleeping %.1f s",
+                        path,
+                        retry_after,
+                    )
+                    await asyncio.sleep(retry_after)
+                    continue
+
+                if resp.status == 404:
+                    logger.debug("Discord 404 for %s", path)
+                    return None
+
+                if resp.status == 403:
+                    logger.warning("Discord 403 forbidden for %s", path)
+                    return None
+
+                resp.raise_for_status()
+                await asyncio.sleep(_RATE_LIMIT_DELAY)
+                return await resp.json()
+
+        logger.error("Exhausted retries for %s", path)
+        return None
+
+    async def _resolve_channels(
+        self, session: Any
+    ) -> list[dict[str, Any]]:
+        """Resolve the list of channels to sync."""
+        if self.channel_ids:
+            channels: list[dict[str, Any]] = []
+            for cid in self.channel_ids:
+                ch = await self._api_get(session, f"/channels/{cid}")
+                if ch:
+                    self._channel_cache[cid] = ch
+                    channels.append(ch)
+            return channels
+
+        # Discover text channels from configured guilds.
+        channels = []
+        for guild_id in self.guild_ids:
+            guild = await self._api_get(session, f"/guilds/{guild_id}")
+            if guild:
+                self._guild_cache[guild_id] = guild
+
+            guild_channels = await self._api_get(
+                session, f"/guilds/{guild_id}/channels"
+            )
+            if not guild_channels:
+                continue
+
+            for ch in guild_channels:
+                ch_type = ch.get("type", 0)
+                # 0=text, 5=announcement, 15=forum.
+                if ch_type in (0, 5, 15):
+                    self._channel_cache[ch["id"]] = ch
+                    channels.append(ch)
+
+        return channels
+
+    async def _fetch_channel_messages(
+        self,
+        session: Any,
+        channel: dict[str, Any],
+        cutoff: datetime,
+    ) -> list[RawEvent]:
+        """Fetch messages from a text channel newer than *cutoff*."""
+        channel_id = channel["id"]
+        channel_name = channel.get("name", channel_id)
+        guild_id = channel.get("guild_id", "")
+        guild_name = self._guild_cache.get(guild_id, {}).get(
+            "name", guild_id
+        )
+        channel_type = channel.get("type", 0)
+
+        events: list[RawEvent] = []
+        # Convert cutoff to a Discord snowflake for efficient pagination.
+        cutoff_snowflake = self._datetime_to_snowflake(cutoff)
+        before: str | None = None
+
+        while True:
+            params: dict[str, Any] = {"limit": _MESSAGE_LIMIT}
+            if before:
+                params["before"] = before
+            else:
+                # Start from the newest message.
+                pass
+
+            messages = await self._api_get(
+                session, f"/channels/{channel_id}/messages", params=params
+            )
+            if not messages:
+                break
+
+            reached_cutoff = False
+            for msg in messages:
+                msg_id = msg.get("id", "")
+
+                # Check if we've gone past the cutoff.
+                if int(msg_id) < cutoff_snowflake:
+                    reached_cutoff = True
+                    break
+
+                # Skip bot messages unless configured to include them.
+                author = msg.get("author", {})
+                if author.get("bot") and not self.include_bots:
+                    continue
+
+                timestamp = self._parse_discord_timestamp(
+                    msg.get("timestamp", "")
+                )
+
+                events.append(
+                    RawEvent(
+                        source="discord",
+                        external_id=msg_id,
+                        timestamp=timestamp,
+                        data={
+                            **msg,
+                            "_channel_name": channel_name,
+                            "_guild_name": guild_name,
+                        },
+                        permissions={
+                            "channel_type": channel_type,
+                            "guild_id": guild_id,
+                            "is_restricted": False,
+                        },
+                    )
+                )
+
+            if reached_cutoff or len(messages) < _MESSAGE_LIMIT:
+                break
+
+            before = messages[-1]["id"]
+
+        return events
+
+    async def _fetch_forum_threads(
+        self,
+        session: Any,
+        channel: dict[str, Any],
+        cutoff: datetime,
+    ) -> list[RawEvent]:
+        """Fetch forum threads and their starter messages."""
+        channel_id = channel["id"]
+        guild_id = channel.get("guild_id", "")
+        events: list[RawEvent] = []
+
+        # Fetch archived threads.
+        for archive_type in ("public", "private"):
+            path = (
+                f"/channels/{channel_id}/threads/archived/{archive_type}"
+            )
+            data = await self._api_get(session, path)
+            if not data:
+                continue
+
+            for thread in data.get("threads", []):
+                thread_events = await self._ingest_thread(
+                    session, thread, channel, cutoff, is_forum=True
+                )
+                events.extend(thread_events)
+
+        return events
+
+    async def _fetch_active_threads(
+        self, session: Any, guild_id: str, cutoff: datetime
+    ) -> list[RawEvent]:
+        """Fetch active threads in a guild."""
+        events: list[RawEvent] = []
+        data = await self._api_get(
+            session, f"/guilds/{guild_id}/threads/active"
+        )
+        if not data:
+            return events
+
+        for thread in data.get("threads", []):
+            parent_id = thread.get("parent_id", "")
+            parent_channel = self._channel_cache.get(parent_id, {})
+            thread_events = await self._ingest_thread(
+                session,
+                thread,
+                parent_channel or {"id": parent_id, "guild_id": guild_id},
+                cutoff,
+                is_forum=(parent_channel.get("type") == 15),
+            )
+            events.extend(thread_events)
+
+        return events
+
+    async def _ingest_thread(
+        self,
+        session: Any,
+        thread: dict[str, Any],
+        parent_channel: dict[str, Any],
+        cutoff: datetime,
+        is_forum: bool = False,
+    ) -> list[RawEvent]:
+        """Fetch messages inside a thread / forum post."""
+        thread_id = thread["id"]
+        thread_name = thread.get("name", "")
+        guild_id = parent_channel.get("guild_id", "")
+        guild_name = self._guild_cache.get(guild_id, {}).get(
+            "name", guild_id
+        )
+        channel_type = thread.get("type", 11)
+        applied_tags = thread.get("applied_tags", [])
+
+        fake_channel = {
+            "id": thread_id,
+            "name": thread_name,
+            "guild_id": guild_id,
+            "type": channel_type,
+        }
+        events = await self._fetch_channel_messages(
+            session, fake_channel, cutoff
+        )
+
+        # Mark events with thread context.
+        enriched: list[RawEvent] = []
+        for raw in events:
+            raw.data["_thread_id"] = thread_id
+            raw.data["_thread_name"] = thread_name
+            raw.data["_guild_name"] = guild_name
+            raw.data["_applied_tags"] = applied_tags
+
+            # Mark the first message of a forum thread as a forum post.
+            if is_forum and raw.external_id == thread_id:
+                raw.data["_event_kind"] = "forum_post"
+            enriched.append(raw)
+
+        return enriched
+
+    # -- helpers -------------------------------------------------------------
+
+    @staticmethod
+    def _parse_discord_timestamp(ts: str) -> datetime:
+        """Parse an ISO-8601 timestamp from the Discord API."""
+        if not ts:
+            return datetime.now(timezone.utc)
+        try:
+            return datetime.fromisoformat(ts.replace("Z", "+00:00"))
+        except (ValueError, TypeError):
+            return datetime.now(timezone.utc)
+
+    @staticmethod
+    def _datetime_to_snowflake(dt: datetime) -> int:
+        """Convert a datetime to a Discord snowflake ID for pagination."""
+        ms = int(dt.timestamp() * 1000) - _DISCORD_EPOCH
+        if ms < 0:
+            ms = 0
+        return ms << 22
+
+    @staticmethod
+    def _extract_message_entities(data: dict[str, Any]) -> list[EntityRef]:
+        """Extract entity references from a Discord message."""
+        entities: list[EntityRef] = []
+
+        guild_id = data.get("guild_id", "")
+        if guild_id:
+            entities.append(
+                EntityRef(
+                    entity_type="guild",
+                    entity_id=guild_id,
+                    display_name=data.get("_guild_name", guild_id),
+                )
+            )
+
+        channel_id = data.get("channel_id", "")
+        if channel_id:
+            entities.append(
+                EntityRef(
+                    entity_type="channel",
+                    entity_id=channel_id,
+                    display_name=data.get("_channel_name", channel_id),
+                )
+            )
+
+        author = data.get("author", {})
+        if author.get("id"):
+            entities.append(
+                EntityRef(
+                    entity_type="user",
+                    entity_id=author["id"],
+                    display_name=author.get("username", ""),
+                )
+            )
+
+        # Mentioned users.
+        for mention in data.get("mentions", []):
+            if mention.get("id") and mention["id"] != author.get("id"):
+                entities.append(
+                    EntityRef(
+                        entity_type="user",
+                        entity_id=mention["id"],
+                        display_name=mention.get("username", ""),
+                    )
+                )
+
+        # Mentioned roles.
+        for role_id in data.get("mention_roles", []):
+            entities.append(
+                EntityRef(
+                    entity_type="role",
+                    entity_id=str(role_id),
+                    display_name="",
+                )
+            )
+
+        return entities