feat: update chromium loader

This commit is contained in:
Marco Vinciguerra 2024-10-09 10:39:29 +02:00
parent 7797631eac
commit 4f816f3b04

View File

@ -1,10 +1,12 @@
""" """"
Chromium module chromium module
""" """
import asyncio import asyncio
from typing import Any, AsyncIterator, Iterator, List, Optional from typing import Any, AsyncIterator, Iterator, List, Optional
from langchain_community.document_loaders.base import BaseLoader from langchain_community.document_loaders.base import BaseLoader
from langchain_core.documents import Document from langchain_core.documents import Document
import aiohttp
import async_timeout
from ..utils import Proxy, dynamic_import, get_logger, parse_or_search_proxy from ..utils import Proxy, dynamic_import, get_logger, parse_or_search_proxy
logger = get_logger("web-loader") logger = get_logger("web-loader")
@ -21,6 +23,9 @@ class ChromiumLoader(BaseLoader):
urls: A list of URLs to scrape content from. urls: A list of URLs to scrape content from.
""" """
RETRY_LIMIT = 3
TIMEOUT = 10
def __init__( def __init__(
self, self,
urls: List[str], urls: List[str],
@ -66,17 +71,29 @@ class ChromiumLoader(BaseLoader):
Returns: Returns:
str: The scraped HTML content or an error message if an exception occurs. str: The scraped HTML content or an error message if an exception occurs.
""" """
import undetected_chromedriver as uc import undetected_chromedriver as uc
logger.info(f"Starting scraping with {self.backend}...") logger.info(f"Starting scraping with {self.backend}...")
results = "" results = ""
try: attempt = 0
driver = uc.Chrome(headless=self.headless)
results = driver.get(url).page_content while attempt < self.RETRY_LIMIT:
except Exception as e: try:
results = f"Error: {e}" async with async_timeout.timeout(self.TIMEOUT):
driver = uc.Chrome(headless=self.headless)
driver.get(url)
results = driver.page_content
logger.info(f"Successfully scraped {url}")
break
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
attempt += 1
logger.error(f"Attempt {attempt} failed: {e}")
if attempt == self.RETRY_LIMIT:
results = f"Error: Network error after {self.RETRY_LIMIT} attempts - {e}"
finally:
driver.quit()
return results return results
async def ascrape_playwright(self, url: str) -> str: async def ascrape_playwright(self, url: str) -> str:
@ -88,28 +105,36 @@ class ChromiumLoader(BaseLoader):
Returns: Returns:
str: The scraped HTML content or an error message if an exception occurs. str: The scraped HTML content or an error message if an exception occurs.
""" """
from playwright.async_api import async_playwright from playwright.async_api import async_playwright
from undetected_playwright import Malenia from undetected_playwright import Malenia
logger.info(f"Starting scraping with {self.backend}...") logger.info(f"Starting scraping with {self.backend}...")
results = "" results = ""
async with async_playwright() as p: attempt = 0
browser = await p.chromium.launch(
headless=self.headless, proxy=self.proxy, **self.browser_config while attempt < self.RETRY_LIMIT:
)
try: try:
context = await browser.new_context() async with async_playwright() as p, async_timeout.timeout(self.TIMEOUT):
await Malenia.apply_stealth(context) browser = await p.chromium.launch(
page = await context.new_page() headless=self.headless, proxy=self.proxy, **self.browser_config
await page.goto(url, wait_until="domcontentloaded") )
await page.wait_for_load_state(self.load_state) context = await browser.new_context()
results = await page.content() # Simply get the HTML content await Malenia.apply_stealth(context)
logger.info("Content scraped") page = await context.new_page()
except Exception as e: await page.goto(url, wait_until="domcontentloaded")
results = f"Error: {e}" await page.wait_for_load_state(self.load_state)
await browser.close() results = await page.content()
logger.info("Content scraped")
break
except (aiohttp.ClientError, asyncio.TimeoutError, Exception) as e:
attempt += 1
logger.error(f"Attempt {attempt} failed: {e}")
if attempt == self.RETRY_LIMIT:
results = f"Error: Network error after {self.RETRY_LIMIT} attempts - {e}"
finally:
await browser.close()
return results return results
def lazy_load(self) -> Iterator[Document]: def lazy_load(self) -> Iterator[Document]:
@ -121,7 +146,6 @@ class ChromiumLoader(BaseLoader):
Yields: Yields:
Document: The scraped content encapsulated within a Document object. Document: The scraped content encapsulated within a Document object.
""" """
scraping_fn = getattr(self, f"ascrape_{self.backend}") scraping_fn = getattr(self, f"ascrape_{self.backend}")