From c96df72cd09e4058d2e4ac1e4eda3a8561e1f456 Mon Sep 17 00:00:00 2001 From: Taranjeet Singh Date: Wed, 9 Aug 2023 16:14:02 -0700 Subject: [PATCH] Fix: lazy load Notion loader (#419) --- embedchain/data_formatter/data_formatter.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/embedchain/data_formatter/data_formatter.py b/embedchain/data_formatter/data_formatter.py index f478ca55..f8834828 100644 --- a/embedchain/data_formatter/data_formatter.py +++ b/embedchain/data_formatter/data_formatter.py @@ -11,7 +11,6 @@ from embedchain.loaders.docs_site_loader import DocsSiteLoader from embedchain.loaders.docx_file import DocxFileLoader from embedchain.loaders.local_qna_pair import LocalQnaPairLoader from embedchain.loaders.local_text import LocalTextLoader -from embedchain.loaders.notion import NotionLoader from embedchain.loaders.pdf_file import PdfFileLoader from embedchain.loaders.sitemap import SitemapLoader from embedchain.loaders.web_page import WebPageLoader @@ -46,10 +45,16 @@ class DataFormatter: "docx": DocxFileLoader(), "sitemap": SitemapLoader(), "docs_site": DocsSiteLoader(), - "notion": NotionLoader(), } + lazy_loaders = ("notion", ) if data_type in loaders: return loaders[data_type] + elif data_type in lazy_loaders: + if data_type == "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}")