[Feature] Improve GitHub loader

embedchain/loaders/github.py

@@ -2,6 +2,7 @@
 import hashlib
 import logging
 import os
+from tqdm import tqdm
 
 from embedchain.loaders.base_loader import BaseLoader
 from embedchain.loaders.json import JSONLoader
@@ -53,14 +54,26 @@ def _load_file(file_path: str):
 
             return data.get("data", [])
 
+
+        def _is_file_empty(file_path):
+            return os.path.getsize(file_path) == 0
+
+
+        def _is_whitelisted(file_path):
+            whitelisted_extensions = ['md', 'txt', 'html', 'json', 'py', 'js', 'jsx', 'ts', 'tsx', 'mdx', 'rst']
+            _, file_extension = os.path.splitext(file_path)
+            return file_extension[1:] in whitelisted_extensions
+
+
         def _add_repo_files(repo_path: str):
             with concurrent.futures.ThreadPoolExecutor() as executor:
                 future_to_file = {
                     executor.submit(_load_file, os.path.join(root, filename)): os.path.join(root, filename)
                     for root, _, files in os.walk(repo_path)
                     for filename in files
+                    if (_is_whitelisted(os.path.join(root, filename)) and not _is_file_empty(os.path.join(root, filename)))
                 }  # noqa: E501
-                for future in concurrent.futures.as_completed(future_to_file):
+                for future in tqdm(concurrent.futures.as_completed(future_to_file), total=len(future_to_file)):
                     file = future_to_file[future]
                     try:
                         results = future.result()
@@ -70,6 +83,7 @@ def _add_repo_files(repo_path: str):
                     except Exception as e:
                         logging.warn(f"Failed to process {file}: {e}")
 
+
         source_hash = hashlib.sha256(repo_url.encode()).hexdigest()
         repo_path = f"/tmp/{source_hash}"
         _fetch_or_clone_repo(repo_url=repo_url, local_path=repo_path)

embedchain/utils.py

@@ -215,6 +215,10 @@ def is_openapi_yaml(yaml_content):
         if url.path.endswith(".csv"):
             logging.debug(f"Source of `{formatted_source}` detected as `csv`.")
             return DataType.CSV
+
+        if url.path.endswith(".mdx") or url.path.endswith(".md"):
+            logging.debug(f"Source of `{formatted_source}` detected as `mdx`.")
+            return DataType.MDX
 
         if url.path.endswith(".docx"):
             logging.debug(f"Source of `{formatted_source}` detected as `docx`.")
@@ -291,6 +295,10 @@ def is_openapi_yaml(yaml_content):
         if source.endswith(".xml"):
             logging.debug(f"Source of `{formatted_source}` detected as `xml`.")
             return DataType.XML
+
+        if source.endswith(".mdx") or source.endswith(".md"):
+            logging.debug(f"Source of `{formatted_source}` detected as `mdx`.")
+            return DataType.MDX
 
         if source.endswith(".yaml"):
             with open(source, "r") as file: