sitemap.py 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. import concurrent.futures
  2. import hashlib
  3. import logging
  4. import os
  5. from urllib.parse import urlparse
  6. import requests
  7. from tqdm import tqdm
  8. try:
  9. from bs4 import BeautifulSoup
  10. from bs4.builder import ParserRejectedMarkup
  11. except ImportError:
  12. raise ImportError(
  13. 'Sitemap requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
  14. ) from None
  15. from embedchain.helpers.json_serializable import register_deserializable
  16. from embedchain.loaders.base_loader import BaseLoader
  17. from embedchain.loaders.web_page import WebPageLoader
  18. @register_deserializable
  19. class SitemapLoader(BaseLoader):
  20. """
  21. This method takes a sitemap URL or local file path as input and retrieves
  22. all the URLs to use the WebPageLoader to load content
  23. of each page.
  24. """
  25. def load_data(self, sitemap_source):
  26. output = []
  27. web_page_loader = WebPageLoader()
  28. if urlparse(sitemap_source).scheme in ("http", "https"):
  29. try:
  30. response = requests.get(sitemap_source)
  31. response.raise_for_status()
  32. soup = BeautifulSoup(response.text, "xml")
  33. except requests.RequestException as e:
  34. logging.error(f"Error fetching sitemap from URL: {e}")
  35. return
  36. elif os.path.isfile(sitemap_source):
  37. with open(sitemap_source, "r") as file:
  38. soup = BeautifulSoup(file, "xml")
  39. else:
  40. raise ValueError("Invalid sitemap source. Please provide a valid URL or local file path.")
  41. links = [link.text for link in soup.find_all("loc") if link.parent.name == "url"]
  42. if len(links) == 0:
  43. links = [link.text for link in soup.find_all("loc")]
  44. doc_id = hashlib.sha256((" ".join(links) + sitemap_source).encode()).hexdigest()
  45. def load_web_page(link):
  46. try:
  47. loader_data = web_page_loader.load_data(link)
  48. return loader_data.get("data")
  49. except ParserRejectedMarkup as e:
  50. logging.error(f"Failed to parse {link}: {e}")
  51. return None
  52. with concurrent.futures.ThreadPoolExecutor() as executor:
  53. future_to_link = {executor.submit(load_web_page, link): link for link in links}
  54. for future in tqdm(concurrent.futures.as_completed(future_to_link), total=len(links), desc="Loading pages"):
  55. link = future_to_link[future]
  56. try:
  57. data = future.result()
  58. if data:
  59. output.extend(data)
  60. except Exception as e:
  61. logging.error(f"Error loading page {link}: {e}")
  62. return {"doc_id": doc_id, "data": output}