diff --git a/rag/utils/web_utils.py b/rag/utils/web_utils.py index 48b1e2d..734b977 100644 --- a/rag/utils/web_utils.py +++ b/rag/utils/web_utils.py @@ -2,15 +2,7 @@ """ Утилиты для загрузки веб-страниц и PDF, извлечения основного текста, рекурсивного обхода ссылок. - -Использует aiohttp для асинхронных запросов, BeautifulSoup для парсинга, -readability-lxml для извлечения контента (опционально). - -Параметры (max_pages, max_depth) берутся из конфига при вызове из xmpp_client. - -Импортируется в: -- indexing_worker.py (fetch_any_url, crawl_url) -- возможно, в других функциях при обработке URL. +Добавлено ограничение размера загружаемого контента. """ import asyncio @@ -35,24 +27,20 @@ except ImportError: logger = logging.getLogger(__name__) -async def fetch_any_url(url: str, timeout: int = 30) -> Tuple[Optional[str], Optional[str]]: +async def fetch_any_url(url: str, timeout: int = 30, max_size_mb: int = 5) -> Tuple[Optional[str], Optional[str]]: """ - Универсальная загрузка контента по URL. + Универсальная загрузка контента по URL с ограничением размера. Аргументы: - url (str): целевой адрес. - timeout (int): таймаут запроса в секундах. + url: целевой адрес. + timeout: таймаут запроса в секундах. + max_size_mb: максимальный размер загружаемого контента в МБ. Возвращает: Tuple[Optional[str], Optional[str]]: (заголовок, текст). Для HTML: заголовок страницы и очищенный основной текст. Для PDF: имя файла (базовая часть URL) и извлечённый текст. В случае ошибки возвращает (None, None). - - Примечания: - - Для PDF используется временный файл, который удаляется после чтения. - - Поддерживаются только HTTP/HTTPS. - - Таймаут применяется к запросу в целом (чтение + соединение). """ try: async with aiohttp.ClientSession() as session: @@ -60,14 +48,22 @@ async def fetch_any_url(url: str, timeout: int = 30) -> Tuple[Optional[str], Opt if resp.status != 200: logger.error(f"HTTP {resp.status} для {url}") return None, None + content_type = resp.headers.get('Content-Type', '').lower() + limit = max_size_mb * 1024 * 1024 + + try: + data = await resp.content.read(limit=limit) + except aiohttp.ContentTooLong: + logger.error(f"Содержимое {url} превышает лимит {max_size_mb} МБ") + return None, None + # Обработка PDF if 'application/pdf' in content_type or url.lower().endswith('.pdf'): - # Сохраняем PDF во временный файл with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp: tmp_path = tmp.name async with aiofiles.open(tmp_path, 'wb') as f: - await f.write(await resp.read()) + await f.write(data) try: from pypdf import PdfReader reader = PdfReader(tmp_path) @@ -81,7 +77,7 @@ async def fetch_any_url(url: str, timeout: int = 30) -> Tuple[Optional[str], Opt os.unlink(tmp_path) else: # Обычная HTML-страница - html = await resp.text() + html = data.decode('utf-8', errors='ignore') title, text = extract_text_from_html(html, url) return title, text except asyncio.TimeoutError: @@ -93,28 +89,8 @@ async def fetch_any_url(url: str, timeout: int = 30) -> Tuple[Optional[str], Opt def extract_text_from_html(html: str, url: str = "") -> Tuple[Optional[str], str]: - """ - Извлекает заголовок и основной текст из HTML. - - Аргументы: - html (str): исходный HTML-код. - url (str): URL страницы (используется для формирования заголовка, если теги отсутствуют). - - Возвращает: - Tuple[Optional[str], str]: (заголовок, очищенный текст). - Заголовок может быть None, если не удалось извлечь. - - Алгоритм извлечения заголовка (по приоритету): - 1. - 2. <meta property="og:title"> - 3. <h1> - 4. Элементы с классами 'title', 'article-title', 'page-title', 'document-title' - 5. Последний сегмент пути URL - 6. Домен (как запасной вариант) - """ + """Извлекает заголовок и основной текст из HTML (без изменений).""" soup = BeautifulSoup(html, 'lxml') - - # ---- 1. Извлечение заголовка ---- title = None if soup.title and soup.title.string: title = soup.title.string.strip() @@ -127,14 +103,12 @@ def extract_text_from_html(html: str, url: str = "") -> Tuple[Optional[str], str if h1: title = h1.get_text(strip=True) if not title: - # Поиск по распространённым классам for cls in ['title', 'article-title', 'page-title', 'document-title']: elem = soup.find(class_=cls) if elem: title = elem.get_text(strip=True) break if not title: - # Последняя надежда: использовать последний сегмент URL parsed = urlparse(url) if parsed.path and parsed.path != '/': last_part = parsed.path.rstrip('/').split('/')[-1] @@ -142,19 +116,14 @@ def extract_text_from_html(html: str, url: str = "") -> Tuple[Optional[str], str else: title = parsed.netloc - # ---- 2. Извлечение основного текста ---- if HAS_READABILITY: - # Используем readability для выделения «читаемой» части doc = Document(html) text = doc.summary() - # Если readability дал заголовок, но мы его не получили ранее – используем if not title and doc.title(): title = doc.title() - # Очищаем полученный HTML от лишних тегов text_soup = BeautifulSoup(text, 'lxml') text = text_soup.get_text(separator='\n', strip=True) else: - # Удаляем скрипты, стили, навигацию, футеры, хедеры for script in soup(["script", "style", "nav", "footer", "header"]): script.decompose() text = soup.get_text(separator='\n', strip=True) @@ -162,7 +131,6 @@ def extract_text_from_html(html: str, url: str = "") -> Tuple[Optional[str], str logger.debug(f"Извлечён заголовок: {repr(title)} для {url}") return title, text - async def crawl_url( start_url: str, max_depth: int,