docs_site_loader.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. import hashlib
  2. import logging
  3. from urllib.parse import urljoin, urlparse
  4. import requests
  5. try:
  6. from bs4 import BeautifulSoup
  7. except ImportError:
  8. raise ImportError(
  9. 'DocsSite requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
  10. ) from None
  11. from embedchain.helpers.json_serializable import register_deserializable
  12. from embedchain.loaders.base_loader import BaseLoader
  13. @register_deserializable
  14. class DocsSiteLoader(BaseLoader):
  15. def __init__(self):
  16. self.visited_links = set()
  17. def _get_child_links_recursive(self, url):
  18. parsed_url = urlparse(url)
  19. base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
  20. current_path = parsed_url.path
  21. response = requests.get(url)
  22. if response.status_code != 200:
  23. logging.info(f"Failed to fetch the website: {response.status_code}")
  24. return
  25. soup = BeautifulSoup(response.text, "html.parser")
  26. all_links = [link.get("href") for link in soup.find_all("a")]
  27. child_links = [link for link in all_links if link and link.startswith(current_path) and link != current_path]
  28. absolute_paths = [urljoin(base_url, link) for link in child_links]
  29. for link in absolute_paths:
  30. if link not in self.visited_links:
  31. self.visited_links.add(link)
  32. self._get_child_links_recursive(link)
  33. def _get_all_urls(self, url):
  34. self.visited_links = set()
  35. self._get_child_links_recursive(url)
  36. urls = [link for link in self.visited_links if urlparse(link).netloc == urlparse(url).netloc]
  37. return urls
  38. def _load_data_from_url(self, url):
  39. response = requests.get(url)
  40. if response.status_code != 200:
  41. logging.info(f"Failed to fetch the website: {response.status_code}")
  42. return []
  43. soup = BeautifulSoup(response.content, "html.parser")
  44. selectors = [
  45. "article.bd-article",
  46. 'article[role="main"]',
  47. "div.md-content",
  48. 'div[role="main"]',
  49. "div.container",
  50. "div.section",
  51. "article",
  52. "main",
  53. ]
  54. output = []
  55. for selector in selectors:
  56. element = soup.select_one(selector)
  57. if element:
  58. content = element.prettify()
  59. break
  60. else:
  61. content = soup.get_text()
  62. soup = BeautifulSoup(content, "html.parser")
  63. ignored_tags = [
  64. "nav",
  65. "aside",
  66. "form",
  67. "header",
  68. "noscript",
  69. "svg",
  70. "canvas",
  71. "footer",
  72. "script",
  73. "style",
  74. ]
  75. for tag in soup(ignored_tags):
  76. tag.decompose()
  77. content = " ".join(soup.stripped_strings)
  78. output.append(
  79. {
  80. "content": content,
  81. "meta_data": {"url": url},
  82. }
  83. )
  84. return output
  85. def load_data(self, url):
  86. all_urls = self._get_all_urls(url)
  87. output = []
  88. for u in all_urls:
  89. output.extend(self._load_data_from_url(u))
  90. doc_id = hashlib.sha256((" ".join(all_urls) + url).encode()).hexdigest()
  91. return {
  92. "doc_id": doc_id,
  93. "data": output,
  94. }