Source code for langchain_community.document_loaders.chatgpt

import datetime
import json
from typing import List

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader


def concatenate_rows(message: dict, title: str) -> str:
    """将消息信息以可读格式组合在一起，准备好供使用。
参数：
    message: 要连接的消息
    title: 对话框的标题

返回：
    连接后的消息
"""
    if not message:
        return ""

    sender = message["author"]["role"] if message["author"] else "unknown"
    text = message["content"]["parts"][0]
    date = datetime.datetime.fromtimestamp(message["create_time"]).strftime(
        "%Y-%m-%d %H:%M:%S"
    )
    return f"{title} - {sender} on {date}: {text}\n\n"


class ChatGPTLoader(BaseLoader):
    """从导出的`ChatGPT`数据加载对话。"""

    def __init__(self, log_file: str, num_logs: int = -1):
        """初始化一个类对象。

参数：
    log_file：日志文件的路径
    num_logs：要加载的日志数量。如果为0，则加载所有日志。
"""
        self.log_file = log_file
        self.num_logs = num_logs

    def load(self) -> List[Document]:
        with open(self.log_file, encoding="utf8") as f:
            data = json.load(f)[: self.num_logs] if self.num_logs else json.load(f)

        documents = []
        for d in data:
            title = d["title"]
            messages = d["mapping"]
            text = "".join(
                [
                    concatenate_rows(messages[key]["message"], title)
                    for idx, key in enumerate(messages)
                    if not (
                        idx == 0
                        and messages[key]["message"]["author"]["role"] == "system"
                    )
                ]
            )
            metadata = {"source": str(self.log_file)}
            documents.append(Document(page_content=text, metadata=metadata))

        return documents