Source code for langchain_community.document_loaders.url_selenium

"""使用Selenium加载页面的加载器，然后使用unstructured加载HTML。
"""
import logging
from typing import TYPE_CHECKING, List, Literal, Optional, Union

if TYPE_CHECKING:
    from selenium.webdriver import Chrome, Firefox

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader

logger = logging.getLogger(__name__)


class SeleniumURLLoader(BaseLoader):
    """使用`Selenium`加载`HTML`页面，并使用`Unstructured`进行解析。

    这对于加载需要JavaScript渲染的页面非常有用。

    属性:
        urls (List[str]): 要加载的URL列表。
        continue_on_failure (bool): 如果为True，则在失败时继续加载其他URL。
        browser (str): 要使用的浏览器，可以是'chrome'或'firefox'。
        binary_location (Optional[str]): 浏览器二进制文件的位置。
        executable_path (Optional[str]): 浏览器可执行文件的路径。
        headless (bool): 如果为True，则浏览器将以无头模式运行。
        arguments [List[str]]: 要传递给浏览器的参数列表。"""

    def __init__(
        self,
        urls: List[str],
        continue_on_failure: bool = True,
        browser: Literal["chrome", "firefox"] = "chrome",
        binary_location: Optional[str] = None,
        executable_path: Optional[str] = None,
        headless: bool = True,
        arguments: List[str] = [],
    ):
        """使用Selenium和非结构化方式加载URL列表。"""
        try:
            import selenium  # noqa:F401
        except ImportError:
            raise ImportError(
                "selenium package not found, please install it with "
                "`pip install selenium`"
            )

        try:
            import unstructured  # noqa:F401
        except ImportError:
            raise ImportError(
                "unstructured package not found, please install it with "
                "`pip install unstructured`"
            )

        self.urls = urls
        self.continue_on_failure = continue_on_failure
        self.browser = browser
        self.binary_location = binary_location
        self.executable_path = executable_path
        self.headless = headless
        self.arguments = arguments

    def _get_driver(self) -> Union["Chrome", "Firefox"]:
        """根据指定的浏览器创建并返回一个WebDriver实例。

抛出:
    ValueError: 如果指定了无效的浏览器。

返回:
    Union[Chrome, Firefox]: 指定浏览器的WebDriver实例。
"""
        if self.browser.lower() == "chrome":
            from selenium.webdriver import Chrome
            from selenium.webdriver.chrome.options import Options as ChromeOptions
            from selenium.webdriver.chrome.service import Service

            chrome_options = ChromeOptions()

            for arg in self.arguments:
                chrome_options.add_argument(arg)

            if self.headless:
                chrome_options.add_argument("--headless")
                chrome_options.add_argument("--no-sandbox")
            if self.binary_location is not None:
                chrome_options.binary_location = self.binary_location
            if self.executable_path is None:
                return Chrome(options=chrome_options)
            return Chrome(
                options=chrome_options,
                service=Service(executable_path=self.executable_path),
            )
        elif self.browser.lower() == "firefox":
            from selenium.webdriver import Firefox
            from selenium.webdriver.firefox.options import Options as FirefoxOptions
            from selenium.webdriver.firefox.service import Service

            firefox_options = FirefoxOptions()

            for arg in self.arguments:
                firefox_options.add_argument(arg)

            if self.headless:
                firefox_options.add_argument("--headless")
            if self.binary_location is not None:
                firefox_options.binary_location = self.binary_location
            if self.executable_path is None:
                return Firefox(options=firefox_options)
            return Firefox(
                options=firefox_options,
                service=Service(executable_path=self.executable_path),
            )
        else:
            raise ValueError("Invalid browser specified. Use 'chrome' or 'firefox'.")

    def _build_metadata(self, url: str, driver: Union["Chrome", "Firefox"]) -> dict:
        from selenium.common.exceptions import NoSuchElementException
        from selenium.webdriver.common.by import By

        """Build metadata based on the contents of the webpage"""
        metadata = {
            "source": url,
            "title": "No title found.",
            "description": "No description found.",
            "language": "No language found.",
        }
        if title := driver.title:
            metadata["title"] = title
        try:
            if description := driver.find_element(
                By.XPATH, '//meta[@name="description"]'
            ):
                metadata["description"] = (
                    description.get_attribute("content") or "No description found."
                )
        except NoSuchElementException:
            pass
        try:
            if html_tag := driver.find_element(By.TAG_NAME, "html"):
                metadata["language"] = (
                    html_tag.get_attribute("lang") or "No language found."
                )
        except NoSuchElementException:
            pass
        return metadata

    def load(self) -> List[Document]:
        """使用Selenium加载指定的URL，并创建Document实例。

返回：
    List[Document]: 加载内容的Document实例列表。
"""
        from unstructured.partition.html import partition_html

        docs: List[Document] = list()
        driver = self._get_driver()

        for url in self.urls:
            try:
                driver.get(url)
                page_content = driver.page_source
                elements = partition_html(text=page_content)
                text = "\n\n".join([str(el) for el in elements])
                metadata = self._build_metadata(url, driver)
                docs.append(Document(page_content=text, metadata=metadata))
            except Exception as e:
                if self.continue_on_failure:
                    logger.error(f"Error fetching or processing {url}, exception: {e}")
                else:
                    raise e

        driver.quit()
        return docs