Repository: mem0ai/mem0
File: embedchain/data_formatter/data_formatter.py (Python)

Code before:
    def _get_loader(self, data_type: DataType, config: LoaderConfig) -> BaseLoader:
        """
        Returns the appropriate data loader for the given data type.

        :param data_type: The type of the data to load.
        :type data_type: DataType
        :param config: Config to initialize the loader with.
        :type config: LoaderConfig
        :raises ValueError: If an unsupported data type is provided.
        :return: The loader for the given data type.
        :rtype: BaseLoader
        """
        loaders = {
            DataType.YOUTUBE_VIDEO: YoutubeVideoLoader,
            DataType.PDF_FILE: PdfFileLoader,
            DataType.WEB_PAGE: WebPageLoader,
            DataType.QNA_PAIR: LocalQnaPairLoader,
            DataType.TEXT: LocalTextLoader,
            DataType.DOCX: DocxFileLoader,
            DataType.SITEMAP: SitemapLoader,
            DataType.XML: XmlLoader,
            DataType.DOCS_SITE: DocsSiteLoader,
            DataType.CSV: CsvLoader,
            DataType.MDX: MdxLoader,
            DataType.IMAGES: ImagesLoader,
            DataType.JSON: JsonLoader,
        }
        lazy_loaders = {DataType.NOTION}
        if data_type in loaders:
            loader_class: type = loaders[data_type]
            loader: BaseLoader = loader_class()
            return loader
        elif data_type in lazy_loaders:
            if data_type == DataType.NOTION:
                from embedchain.loaders.notion import NotionLoader

                return NotionLoader()
            else:
                raise ValueError(f"Unsupported data type: {data_type}")
        else:
            raise ValueError(f"Unsupported data type: {data_type}")

    def _get_chunker(self, data_type: DataType, config: ChunkerConfig) -> BaseChunker:
        """Returns the appropriate chunker for the given data type.

        :param data_type: The type of the data to chunk.
        :type data_type: DataType
        :param config: Config to initialize the chunker with.
        :type config: ChunkerConfig
        :raises ValueError: If an unsupported data type is provided.
        :return: The chunker for the given data type.

Review comment (suggestion):
```suggestion
            DataType.JSON: JSONLoader,
```

Code after:
    def _get_loader(self, data_type: DataType, config: LoaderConfig) -> BaseLoader:
        """
        Returns the appropriate data loader for the given data type.

        :param data_type: The type of the data to load.
        :type data_type: DataType
        :param config: Config to initialize the loader with.
        :type config: LoaderConfig
        :raises ValueError: If an unsupported data type is provided.
        :return: The loader for the given data type.
        :rtype: BaseLoader
        """
        loaders = {
            DataType.YOUTUBE_VIDEO: YoutubeVideoLoader,
            DataType.PDF_FILE: PdfFileLoader,
            DataType.WEB_PAGE: WebPageLoader,
            DataType.QNA_PAIR: LocalQnaPairLoader,
            DataType.TEXT: LocalTextLoader,
            DataType.DOCX: DocxFileLoader,
            DataType.SITEMAP: SitemapLoader,
            DataType.XML: XmlLoader,
            DataType.DOCS_SITE: DocsSiteLoader,
            DataType.CSV: CsvLoader,
            DataType.MDX: MdxLoader,
            DataType.IMAGES: ImagesLoader,
            DataType.JSON: JSONLoader,
        }
        lazy_loaders = {DataType.NOTION}
        if data_type in loaders:
            loader_class: type = loaders[data_type]
            loader: BaseLoader = loader_class()
            return loader
        elif data_type in lazy_loaders:
            if data_type == DataType.NOTION:
                from embedchain.loaders.notion import NotionLoader

                return NotionLoader()
            else:
                raise ValueError(f"Unsupported data type: {data_type}")
        else:
            raise ValueError(f"Unsupported data type: {data_type}")

    def _get_chunker(self, data_type: DataType, config: ChunkerConfig) -> BaseChunker:
        """Returns the appropriate chunker for the given data type.

        :param data_type: The type of the data to chunk.
        :type data_type: DataType
        :param config: Config to initialize the chunker with.
        :type config: ChunkerConfig
        :raises ValueError: If an unsupported data type is provided.
        :return: The chunker for the given data type.