Spaces:
Running
Running
| import asyncio | |
| import time | |
| from typing import Optional | |
| from urllib.parse import urljoin, urlparse | |
| import polars as pl | |
| from crawl4ai import AsyncWebCrawler | |
| from src.config import Config | |
| from src.utils import content_hash, normalize_url | |
| class Crawler: | |
| def __init__(self, config: Config): | |
| self.config = config | |
| self.visited: set[str] = set() | |
| self.pages: list[dict] = [] | |
| self.domain: str = "" | |
| def _get_domain(self, url: str) -> str: | |
| parsed = urlparse(url) | |
| return parsed.netloc | |
| def _is_same_domain(self, url: str) -> bool: | |
| return self._get_domain(url) == self.domain | |
| async def crawl(self, url: str, depth: int, max_pages: int, progress=None) -> pl.DataFrame: | |
| url = normalize_url(url) | |
| self.domain = self._get_domain(url) | |
| self.visited.clear() | |
| self.pages.clear() | |
| queue: list[tuple[str, int]] = [(url, 0)] | |
| async with AsyncWebCrawler() as crawler: | |
| while queue and len(self.pages) < max_pages: | |
| current_url, current_depth = queue.pop(0) | |
| if current_url in self.visited: | |
| continue | |
| if current_depth > depth: | |
| continue | |
| if self._get_domain(current_url) and not self._is_same_domain(current_url): | |
| continue | |
| self.visited.add(current_url) | |
| try: | |
| result = await crawler.arun( | |
| url=current_url, | |
| word_count_threshold=10, | |
| bypass_cache=True, | |
| ) | |
| if result.success: | |
| page = { | |
| "url": current_url, | |
| "title": result.metadata.get("title", "") if result.metadata else "", | |
| "markdown": result.markdown or "", | |
| "html": result.html or "", | |
| "depth": current_depth, | |
| "timestamp": time.time(), | |
| "content_hash": content_hash(result.markdown or ""), | |
| } | |
| self.pages.append(page) | |
| if progress is not None: | |
| progress( | |
| min(len(self.pages) / max_pages, 1.0), | |
| desc=f"Crawled {len(self.pages)} pages (depth {current_depth})...", | |
| ) | |
| if current_depth < depth: | |
| links = self._extract_links(result.html or "", current_url) | |
| for link in links: | |
| if link not in self.visited: | |
| queue.append((link, current_depth + 1)) | |
| except Exception: | |
| continue | |
| if not self.pages: | |
| return pl.DataFrame( | |
| schema={ | |
| "url": pl.Utf8, | |
| "title": pl.Utf8, | |
| "markdown": pl.Utf8, | |
| "html": pl.Utf8, | |
| "depth": pl.Int32, | |
| "timestamp": pl.Float64, | |
| "content_hash": pl.Utf8, | |
| } | |
| ) | |
| df = pl.DataFrame(self.pages) | |
| df = df.unique(subset=["content_hash"], keep="first") | |
| return df | |
| def _extract_links(self, html: str, base_url: str) -> list[str]: | |
| from bs4 import BeautifulSoup | |
| soup = BeautifulSoup(html, "lxml") | |
| links = [] | |
| for a_tag in soup.find_all("a", href=True): | |
| href = a_tag["href"].strip() | |
| if not href or href.startswith("#") or href.startswith("javascript:"): | |
| continue | |
| absolute_url = urljoin(base_url, href) | |
| absolute_url = normalize_url(absolute_url) | |
| if self._is_same_domain(absolute_url): | |
| links.append(absolute_url) | |
| return links | |