codekingpro/portable-devtools

1"""Loads .ipynb notebook files."""2 3import json4from pathlib import Path5from typing import Any, List, Union6 7from langchain_core.documents import Document8 9from langchain_community.document_loaders.base import BaseLoader10 11 12def concatenate_cells(13    cell: dict, include_outputs: bool, max_output_length: int, traceback: bool14) -> str:15    """Combine cells information in a readable format ready to be used.16 17    Args:18        cell: A dictionary19        include_outputs: Whether to include the outputs of the cell.20        max_output_length: Maximum length of the output to be displayed.21        traceback: Whether to return a traceback of the error.22 23    Returns:24        A string with the cell information.25 26    """27    cell_type = cell["cell_type"]28    source = cell["source"]29    if include_outputs:30        try:31            output = cell["outputs"]32        except KeyError:33            pass34 35    if include_outputs and cell_type == "code" and output:36        if "ename" in output[0].keys():37            error_name = output[0]["ename"]38            error_value = output[0]["evalue"]39            if traceback:40                traceback = output[0]["traceback"]41                return (42                    f"'{cell_type}' cell: '{source}'\n, gives error '{error_name}',"43                    f" with description '{error_value}'\n"44                    f"and traceback '{traceback}'\n\n"45                )46            else:47                return (48                    f"'{cell_type}' cell: '{source}'\n, gives error '{error_name}',"49                    f"with description '{error_value}'\n\n"50                )51        elif output[0]["output_type"] == "stream":52            output = output[0]["text"]53            min_output = min(max_output_length, len(output))54            return (55                f"'{cell_type}' cell: '{source}'\n with "56                f"output: '{output[:min_output]}'\n\n"57            )58    else:59        return f"'{cell_type}' cell: '{source}'\n\n"60 61    return ""62 63 64def remove_newlines(x: Any) -> Any:65    """Recursively remove newlines, no matter the data structure they are stored in."""66 67    if isinstance(x, str):68        return x.replace("\n", "")69    elif isinstance(x, list):70        return [remove_newlines(elem) for elem in x]71    elif isinstance(x, dict):72        return {k: remove_newlines(v) for (k, v) in x.items()}73    else:74        return x75 76 77class NotebookLoader(BaseLoader):78    """Load `Jupyter notebook` (.ipynb) files."""79 80    def __init__(81        self,82        path: Union[str, Path],83        include_outputs: bool = False,84        max_output_length: int = 10,85        remove_newline: bool = False,86        traceback: bool = False,87    ):88        """Initialize with a path.89 90        Args:91            path: The path to load the notebook from.92            include_outputs: Whether to include the outputs of the cell.93                Defaults to False.94            max_output_length: Maximum length of the output to be displayed.95                Defaults to 10.96            remove_newline: Whether to remove newlines from the notebook.97                Defaults to False.98            traceback: Whether to return a traceback of the error.99                Defaults to False.100        """101        self.file_path = path102        self.include_outputs = include_outputs103        self.max_output_length = max_output_length104        self.remove_newline = remove_newline105        self.traceback = traceback106 107    def load(108        self,109    ) -> List[Document]:110        """Load documents."""111        p = Path(self.file_path)112 113        with open(p, encoding="utf8") as f:114            d = json.load(f)115 116        filtered_data = [117            {k: v for (k, v) in cell.items() if k in ["cell_type", "source", "outputs"]}118            for cell in d["cells"]119        ]120 121        if self.remove_newline:122            filtered_data = list(map(remove_newlines, filtered_data))123 124        text = "".join(125            list(126                map(127                    lambda x: concatenate_cells(128                        x, self.include_outputs, self.max_output_length, self.traceback129                    ),130                    filtered_data,131                )132            )133        )134 135        metadata = {"source": str(p)}136 137        return [Document(page_content=text, metadata=metadata)]138