codekingpro/portable-devtools

1from __future__ import annotations2 3from typing import TYPE_CHECKING, List4 5from langchain_core.documents import Document6 7from langchain_community.document_loaders.base import BaseLoader8 9if TYPE_CHECKING:10    import pandas as pd11 12 13class DiscordChatLoader(BaseLoader):14    """Load `Discord` chat logs."""15 16    def __init__(self, chat_log: pd.DataFrame, user_id_col: str = "ID"):17        """Initialize with a Pandas DataFrame containing chat logs.18 19        Args:20            chat_log: Pandas DataFrame containing chat logs.21            user_id_col: Name of the column containing the user ID. Defaults to "ID".22        """23        if not isinstance(chat_log, pd.DataFrame):24            raise ValueError(25                f"Expected chat_log to be a pd.DataFrame, got {type(chat_log)}"26            )27        self.chat_log = chat_log28        self.user_id_col = user_id_col29 30    def load(self) -> List[Document]:31        """Load all chat messages."""32        result = []33        for _, row in self.chat_log.iterrows():34            user_id = row[self.user_id_col]35            metadata = row.to_dict()36            metadata.pop(self.user_id_col)37            result.append(Document(page_content=user_id, metadata=metadata))38        return result39