codekingpro/portable-devtools

1"""Scrapfly Web Reader."""2 3import logging4from typing import Iterator, List, Literal, Optional5 6from langchain_core.document_loaders import BaseLoader7from langchain_core.documents import Document8from langchain_core.utils import get_from_env9 10logger = logging.getLogger(__file__)11 12 13class ScrapflyLoader(BaseLoader):14    """Turn a url to llm accessible markdown with `Scrapfly.io`.15 16    For further details, visit: https://scrapfly.io/docs/sdk/python17    """18 19    def __init__(20        self,21        urls: List[str],22        *,23        api_key: Optional[str] = None,24        scrape_format: Literal["markdown", "text"] = "markdown",25        scrape_config: Optional[dict] = None,26        continue_on_failure: bool = True,27    ) -> None:28        """Initialize client.29 30        Args:31            urls: List of urls to scrape.32            api_key: The Scrapfly API key. If not specified must have env var33                SCRAPFLY_API_KEY set.34            scrape_format: Scrape result format, one or "markdown" or "text".35            scrape_config: Dictionary of ScrapFly scrape config object.36            continue_on_failure: Whether to continue if scraping a url fails.37        """38        try:39            from scrapfly import ScrapflyClient40        except ImportError:41            raise ImportError(42                "`scrapfly` package not found, please run `pip install scrapfly-sdk`"43            )44        if not urls:45            raise ValueError("URLs must be provided.")46        api_key = api_key or get_from_env("api_key", "SCRAPFLY_API_KEY")47        self.scrapfly = ScrapflyClient(key=api_key)48        self.urls = urls49        self.scrape_format = scrape_format50        self.scrape_config = scrape_config51        self.continue_on_failure = continue_on_failure52 53    def lazy_load(self) -> Iterator[Document]:54        from scrapfly import ScrapeConfig55 56        scrape_config = self.scrape_config if self.scrape_config is not None else {}57        for url in self.urls:58            try:59                response = self.scrapfly.scrape(60                    ScrapeConfig(url, format=self.scrape_format, **scrape_config)61                )62                yield Document(63                    page_content=response.scrape_result["content"],64                    metadata={"url": url},65                )66            except Exception as e:67                if self.continue_on_failure:68                    logger.error(f"Error fetching data from {url}, exception: {e}")69                else:70                    raise e71