codekingpro/portable-devtools

1from typing import Iterator, List, Union2 3import requests4from langchain_core.documents import Document5 6from langchain_community.document_loaders.base import BaseLoader7 8 9class BrowserlessLoader(BaseLoader):10    """Load webpages with `Browserless` /content endpoint."""11 12    def __init__(13        self, api_token: str, urls: Union[str, List[str]], text_content: bool = True14    ):15        """Initialize with API token and the URLs to scrape"""16        self.api_token = api_token17        """Browserless API token."""18        self.urls = urls19        """List of URLs to scrape."""20        self.text_content = text_content21 22    def lazy_load(self) -> Iterator[Document]:23        """Lazy load Documents from URLs."""24 25        for url in self.urls:26            if self.text_content:27                response = requests.post(28                    "https://chrome.browserless.io/scrape",29                    params={30                        "token": self.api_token,31                    },32                    json={33                        "url": url,34                        "elements": [35                            {36                                "selector": "body",37                            }38                        ],39                    },40                )41                yield Document(42                    page_content=response.json()["data"][0]["results"][0]["text"],43                    metadata={44                        "source": url,45                    },46                )47            else:48                response = requests.post(49                    "https://chrome.browserless.io/content",50                    params={51                        "token": self.api_token,52                    },53                    json={54                        "url": url,55                    },56                )57 58                yield Document(59                    page_content=response.text,60                    metadata={61                        "source": url,62                    },63                )64