codekingpro/portable-devtools

1"""Loader that uses Selenium to load a page, then uses unstructured to load the html."""2 3import logging4from typing import TYPE_CHECKING, List, Literal, Optional, Union5 6if TYPE_CHECKING:7    from selenium.webdriver import Chrome, Firefox8 9from langchain_core.documents import Document10 11from langchain_community.document_loaders.base import BaseLoader12 13logger = logging.getLogger(__name__)14 15 16class SeleniumURLLoader(BaseLoader):17    """Load `HTML` pages with `Selenium` and parse with `Unstructured`.18 19    This is useful for loading pages that require javascript to render.20 21    Attributes:22        urls (List[str]): List of URLs to load.23        continue_on_failure (bool): If True, continue loading other URLs on failure.24        browser (str): The browser to use, either 'chrome' or 'firefox'.25        binary_location (Optional[str]): The location of the browser binary.26        executable_path (Optional[str]): The path to the browser executable.27        headless (bool): If True, the browser will run in headless mode.28        arguments [List[str]]: List of arguments to pass to the browser.29    """30 31    def __init__(32        self,33        urls: List[str],34        continue_on_failure: bool = True,35        browser: Literal["chrome", "firefox"] = "chrome",36        binary_location: Optional[str] = None,37        executable_path: Optional[str] = None,38        headless: bool = True,39        arguments: List[str] = [],40    ):41        """Load a list of URLs using Selenium and unstructured."""42        try:43            import selenium  # noqa:F40144        except ImportError:45            raise ImportError(46                "selenium package not found, please install it with "47                "`pip install selenium`"48            )49 50        try:51            import unstructured  # noqa:F40152        except ImportError:53            raise ImportError(54                "unstructured package not found, please install it with "55                "`pip install unstructured`"56            )57 58        self.urls = urls59        self.continue_on_failure = continue_on_failure60        self.browser = browser61        self.binary_location = binary_location62        self.executable_path = executable_path63        self.headless = headless64        self.arguments = arguments65 66    def _get_driver(self) -> Union["Chrome", "Firefox"]:67        """Create and return a WebDriver instance based on the specified browser.68 69        Raises:70            ValueError: If an invalid browser is specified.71 72        Returns:73            Union[Chrome, Firefox]: A WebDriver instance for the specified browser.74        """75        if self.browser.lower() == "chrome":76            from selenium.webdriver import Chrome77            from selenium.webdriver.chrome.options import Options as ChromeOptions78            from selenium.webdriver.chrome.service import Service79 80            chrome_options = ChromeOptions()81 82            for arg in self.arguments:83                chrome_options.add_argument(arg)84 85            if self.headless:86                chrome_options.add_argument("--headless")87                chrome_options.add_argument("--no-sandbox")88            if self.binary_location is not None:89                chrome_options.binary_location = self.binary_location90            if self.executable_path is None:91                return Chrome(options=chrome_options)92            return Chrome(93                options=chrome_options,94                service=Service(executable_path=self.executable_path),95            )96        elif self.browser.lower() == "firefox":97            from selenium.webdriver import Firefox98            from selenium.webdriver.firefox.options import Options as FirefoxOptions99            from selenium.webdriver.firefox.service import Service100 101            firefox_options = FirefoxOptions()102 103            for arg in self.arguments:104                firefox_options.add_argument(arg)105 106            if self.headless:107                firefox_options.add_argument("--headless")108            if self.binary_location is not None:109                firefox_options.binary_location = self.binary_location110            if self.executable_path is None:111                return Firefox(options=firefox_options)112            return Firefox(113                options=firefox_options,114                service=Service(executable_path=self.executable_path),115            )116        else:117            raise ValueError("Invalid browser specified. Use 'chrome' or 'firefox'.")118 119    def _build_metadata(self, url: str, driver: Union["Chrome", "Firefox"]) -> dict:120        from selenium.common.exceptions import NoSuchElementException121        from selenium.webdriver.common.by import By122 123        """Build metadata based on the contents of the webpage"""124        metadata = {125            "source": url,126            "title": "No title found.",127            "description": "No description found.",128            "language": "No language found.",129        }130        if title := driver.title:131            metadata["title"] = title132        try:133            if description := driver.find_element(134                By.XPATH, '//meta[@name="description"]'135            ):136                metadata["description"] = (137                    description.get_attribute("content") or "No description found."138                )139        except NoSuchElementException:140            pass141        try:142            if html_tag := driver.find_element(By.TAG_NAME, "html"):143                metadata["language"] = (144                    html_tag.get_attribute("lang") or "No language found."145                )146        except NoSuchElementException:147            pass148        return metadata149 150    def load(self) -> List[Document]:151        """Load the specified URLs using Selenium and create Document instances.152 153        Returns:154            List[Document]: A list of Document instances with loaded content.155        """156        from unstructured.partition.html import partition_html157 158        docs: List[Document] = list()159        driver = self._get_driver()160 161        for url in self.urls:162            try:163                driver.get(url)164                page_content = driver.page_source165                elements = partition_html(text=page_content)166                text = "\n\n".join([str(el) for el in elements])167                metadata = self._build_metadata(url, driver)168                docs.append(Document(page_content=text, metadata=metadata))169            except Exception as e:170                if self.continue_on_failure:171                    logger.error(f"Error fetching or processing {url}, exception: {e}")172                else:173                    raise e174 175        driver.quit()176        return docs177