From ce9233c258a6769ddf4e640dd49623c5c7b0f37d Mon Sep 17 00:00:00 2001
From: Timothy Carambat <rambat1010@gmail.com>
Date: Mon, 11 Dec 2023 14:40:33 -0800
Subject: [PATCH] feat: enable HTML uploads from UI (#422)

resolves #418
---
 collector/scripts/watch/convert/as_html.py | 42 ++++++++++++++++++++++
 collector/scripts/watch/filetypes.py       |  3 ++
 2 files changed, 45 insertions(+)
 create mode 100644 collector/scripts/watch/convert/as_html.py

diff --git a/collector/scripts/watch/convert/as_html.py b/collector/scripts/watch/convert/as_html.py
new file mode 100644
index 00000000..a70ddb56
--- /dev/null
+++ b/collector/scripts/watch/convert/as_html.py
@@ -0,0 +1,42 @@
+import os, re
+from slugify import slugify
+from langchain.document_loaders import BSHTMLLoader
+from ..utils import guid, file_creation_time, write_to_server_documents, move_source
+from ...utils import tokenize
+
+# Process all html-related documents.
+def as_html(**kwargs):
+  parent_dir = kwargs.get('directory', 'hotdir')
+  filename = kwargs.get('filename')
+  ext = kwargs.get('ext', '.html')
+  remove = kwargs.get('remove_on_complete', False)
+  fullpath = f"{parent_dir}/{filename}{ext}"
+
+  loader = BSHTMLLoader(fullpath)
+  document = loader.load()[0]
+  content = re.sub(r"\n+", "\n", document.page_content)
+
+  if len(content) == 0:
+    print(f"Resulting text content was empty for {filename}{ext}.")
+    return(False, f"No text content found in {filename}{ext}")
+
+  print(f"-- Working {fullpath} --")
+  data = {
+    'id': guid(),
+    'url': "file://"+os.path.abspath(f"{parent_dir}/processed/{filename}{ext}"),
+    'title':  document.metadata.get('title', f"{filename}{ext}"),
+    'docAuthor': 'Unknown', # TODO: Find a better author
+    'description': 'Unknown', # TODO: Find a better description
+    'docSource': 'an HTML file uploaded by the user.',
+    'chunkSource': f"{filename}{ext}",
+    'published': file_creation_time(fullpath),
+    'wordCount': len(content),
+    'pageContent': content,
+    'token_count_estimate': len(tokenize(content))
+  }
+
+  write_to_server_documents(data, f"{slugify(filename)}-{data.get('id')}")
+  move_source(parent_dir, f"{filename}{ext}", remove=remove)
+
+  print(f"[SUCCESS]: {filename}{ext} converted & ready for embedding.\n")
+  return(True, None)
diff --git a/collector/scripts/watch/filetypes.py b/collector/scripts/watch/filetypes.py
index 809c8207..ebc11a6e 100644
--- a/collector/scripts/watch/filetypes.py
+++ b/collector/scripts/watch/filetypes.py
@@ -3,6 +3,7 @@ from .convert.as_markdown import as_markdown
 from .convert.as_pdf import as_pdf
 from .convert.as_docx import as_docx, as_odt
 from .convert.as_mbox import as_mbox
+from .convert.as_html import as_html
 
 FILETYPES = {
     '.txt': as_text,
@@ -11,10 +12,12 @@ FILETYPES = {
     '.docx': as_docx,
     '.odt': as_odt,
     '.mbox': as_mbox, 
+    '.html': as_html, 
 }
 
 ACCEPTED_MIMES = {
     'text/plain': ['.txt', '.md'],
+    'text/html': ['.html'], 
     'application/vnd.openxmlformats-officedocument.wordprocessingml.document': ['.docx'],
     'application/vnd.oasis.opendocument.text': ['.odt'],
     'application/pdf': ['.pdf'],