add custom readthedocs loader

config.py

@@ -22,7 +22,7 @@ def get_index_names():
     keys = [
         k for k in [
             k.strip().lower()
-            for k in os.environ["INDEX_NAMES"].split(",")
+            for k in os.environ.get("INDEX_NAMES", "").split(",")
         ]
         if k
     ]

loaders/__init__.py

@@ -1,9 +1,11 @@
 from .wikipage import WikiPageLoader
 from .github_issue import GithubIssueLoader
+from .rtdhtmlpage import RTDHtmlPageLoader
 
 LOADERS = {
     "wikipage": WikiPageLoader,
-    "github_issue": GithubIssueLoader
+    "github_issue": GithubIssueLoader,
+    "rtdhtmlpage": RTDHtmlPageLoader,
 }
 LOADER_NAMES = tuple(LOADERS.keys())
 

loaders/rtdhtmlpage.py

@@ -0,0 +1,77 @@
+from datetime import datetime
+from pathlib import Path
+from typing import Iterator
+
+from langchain.docstore.document import Document
+from langchain.document_loaders import ReadTheDocsLoader
+
+
+class RTDHtmlPageLoader(ReadTheDocsLoader):
+    """directory path for readthedocs documents
+
+    $ wget -r -np -A.html https://docs.djangoproject.com/en/4.2/
+    $ python store.py -l rtdhtmlpage django ./docs.djangoproject.com/
+    """
+    def __init__(self, index: str, inputfile: Path, *args, **kwargs):
+        self.index = index
+        kwargs["custom_html_tag"] = ("div", {"id": "docs-content"})
+        super().__init__(inputfile, *args, **kwargs)
+
+    def _my_clean_data(self, data: str) -> str:
+        from bs4 import BeautifulSoup
+
+        soup = BeautifulSoup(data, **self.bs_kwargs)
+
+        # default tags
+        html_tags = [
+            ("div", {"role": "main"}),
+            ("main", {"id": "main-content"}),
+        ]
+
+        if self.custom_html_tag is not None:
+            html_tags.append(self.custom_html_tag)
+
+        text = None
+
+        # reversed order. check the custom one first
+        for tag, attrs in html_tags[::-1]:
+            text = soup.find(tag, attrs)
+            # if found, break
+            if text is not None:
+                break
+
+        if text is not None:
+            title = "".join(t.text for t in text.find("h1") if t.name!="a")
+            text = text.get_text()
+        else:
+            text = ""
+            title = ""
+
+        # trim empty lines
+        text = "\n".join([t for t in text.split("\n") if t])
+
+        return text, title
+
+    def lazy_load(self) -> Iterator[Document]:
+        """Load documents."""
+        for p in self.file_path.rglob("*"):
+            if p.is_dir():
+                continue
+            with open(p, encoding=self.encoding, errors=self.errors) as f:
+                text, title = self._my_clean_data(f.read())
+
+            metadata = {
+                "title": title,
+                "ctime": int(datetime.now().timestamp()),
+                "user": "rtd",
+                "type": "rtd",
+                "url": f"https://{str(p)}",
+                "index": self.index,
+                "id": str(p),
+            }
+            # print(metadata)
+            yield Document(page_content=text, metadata=metadata)
+
+
+    def load(self) -> list[Document]:
+        return list(self.lazy_load())

requirements.txt

@@ -9,4 +9,5 @@ sentence_transformers
 streamlit
 python-dateutil
 openai
-tqdm
+tqdm
+beautifulsoup4

store.py

@@ -56,7 +56,7 @@ def store(texts):
 def get_parser():
     p = argparse.ArgumentParser()
     p.add_argument("index", type=str)
-    p.add_argument("inputfile", metavar="INPUTFILE", type=argparse.FileType("rt"))
+    p.add_argument("inputfile", metavar="INPUTFILE", type=str)
     p.add_argument("-l", "--loader", type=str, choices=LOADER_NAMES, required=True)
     return p
 
@@ -65,13 +65,14 @@ def main():
     """
     $ python store.py --loader wikipage "index" "FILE_PATH"
     $ python store.py -l wikipage wiki data/wiki.json
+    $ python store.py -l rtdhtmlpage django ./docs.djangoproject.com/
     """
     p = get_parser()
     args = p.parse_args()
     loader = get_loader(
         args.loader,
         index=args.index,
-        inputfile=Path(args.inputfile.name),
+        inputfile=Path(args.inputfile),
     )
 
     docs = loader.load()