chatbot-rag / src /crawler.py
scythe327's picture
Upload src/crawler.py with huggingface_hub
ef05780 verified
Raw
History Blame Contribute Delete
4.01 kB
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