Source code for langchain_community.document_loaders.discord

from __future__ import annotations

from typing import TYPE_CHECKING, List

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader

if TYPE_CHECKING:
    import pandas as pd


class DiscordChatLoader(BaseLoader):
    """加载`Discord`聊天记录。"""

    def __init__(self, chat_log: pd.DataFrame, user_id_col: str = "ID"):
        """使用包含聊天记录的Pandas DataFrame 进行初始化。

参数:
    chat_log: 包含聊天记录的Pandas DataFrame。
    user_id_col: 包含用户ID的列的名称。默认为"ID"。
"""
        if not isinstance(chat_log, pd.DataFrame):
            raise ValueError(
                f"Expected chat_log to be a pd.DataFrame, got {type(chat_log)}"
            )
        self.chat_log = chat_log
        self.user_id_col = user_id_col

    def load(self) -> List[Document]:
        """加载所有聊天消息。"""
        result = []
        for _, row in self.chat_log.iterrows():
            user_id = row[self.user_id_col]
            metadata = row.to_dict()
            metadata.pop(self.user_id_col)
            result.append(Document(page_content=user_id, metadata=metadata))
        return result