codekingpro/portable-devtools

1import datetime2import json3from typing import List4 5from langchain_core.documents import Document6 7from langchain_community.document_loaders.base import BaseLoader8 9 10def concatenate_rows(message: dict, title: str) -> str:11    """12    Combine message information in a readable format ready to be used.13    Args:14        message: Message to be concatenated15        title: Title of the conversation16 17    Returns:18        Concatenated message19    """20    if not message:21        return ""22 23    sender = message["author"]["role"] if message["author"] else "unknown"24    text = message["content"]["parts"][0]25    date = datetime.datetime.fromtimestamp(message["create_time"]).strftime(26        "%Y-%m-%d %H:%M:%S"27    )28    return f"{title} - {sender} on {date}: {text}\n\n"29 30 31class ChatGPTLoader(BaseLoader):32    """Load conversations from exported `ChatGPT` data."""33 34    def __init__(self, log_file: str, num_logs: int = -1):35        """Initialize a class object.36 37        Args:38            log_file: Path to the log file39            num_logs: Number of logs to load. If 0, load all logs.40        """41        self.log_file = log_file42        self.num_logs = num_logs43 44    def load(self) -> List[Document]:45        with open(self.log_file, encoding="utf8") as f:46            data = json.load(f)[: self.num_logs] if self.num_logs else json.load(f)47 48        documents = []49        for d in data:50            title = d["title"]51            messages = d["mapping"]52            text = "".join(53                [54                    concatenate_rows(messages[key]["message"], title)55                    for idx, key in enumerate(messages)56                    if not (57                        idx == 058                        and messages[key]["message"]["author"]["role"] == "system"59                    )60                ]61            )62            metadata = {"source": str(self.log_file)}63            documents.append(Document(page_content=text, metadata=metadata))64 65        return documents66