web_page.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. import hashlib
  2. import logging
  3. import requests
  4. from bs4 import BeautifulSoup
  5. from embedchain.helper.json_serializable import register_deserializable
  6. from embedchain.loaders.base_loader import BaseLoader
  7. from embedchain.utils import clean_string
  8. @register_deserializable
  9. class WebPageLoader(BaseLoader):
  10. def load_data(self, url):
  11. """Load data from a web page."""
  12. response = requests.get(url)
  13. data = response.content
  14. content = self._get_clean_content(data, url)
  15. meta_data = {
  16. "url": url,
  17. }
  18. doc_id = hashlib.sha256((content + url).encode()).hexdigest()
  19. return {
  20. "doc_id": doc_id,
  21. "data": [
  22. {
  23. "content": content,
  24. "meta_data": meta_data,
  25. }
  26. ],
  27. }
  28. def _get_clean_content(self, html, url) -> str:
  29. soup = BeautifulSoup(html, "html.parser")
  30. original_size = len(str(soup.get_text()))
  31. tags_to_exclude = [
  32. "nav",
  33. "aside",
  34. "form",
  35. "header",
  36. "noscript",
  37. "svg",
  38. "canvas",
  39. "footer",
  40. "script",
  41. "style",
  42. ]
  43. for tag in soup(tags_to_exclude):
  44. tag.decompose()
  45. ids_to_exclude = ["sidebar", "main-navigation", "menu-main-menu"]
  46. for id in ids_to_exclude:
  47. tags = soup.find_all(id=id)
  48. for tag in tags:
  49. tag.decompose()
  50. classes_to_exclude = [
  51. "elementor-location-header",
  52. "navbar-header",
  53. "nav",
  54. "header-sidebar-wrapper",
  55. "blog-sidebar-wrapper",
  56. "related-posts",
  57. ]
  58. for class_name in classes_to_exclude:
  59. tags = soup.find_all(class_=class_name)
  60. for tag in tags:
  61. tag.decompose()
  62. content = soup.get_text()
  63. content = clean_string(content)
  64. cleaned_size = len(content)
  65. if original_size != 0:
  66. logging.info(
  67. f"[{url}] Cleaned page size: {cleaned_size} characters, down from {original_size} (shrunk: {original_size-cleaned_size} chars, {round((1-(cleaned_size/original_size)) * 100, 2)}%)" # noqa:E501
  68. )
  69. return content