Редактировать web_utils.py

rag/utils/web_utils.py

@@ -2,7 +2,11 @@
 """
 Утилиты для загрузки веб-страниц и PDF, извлечения основного текста,
 рекурсивного обхода ссылок.
-Добавлено ограничение размера загружаемого контента.
+
+Использует aiohttp для асинхронных запросов, BeautifulSoup для парсинга,
+readability-lxml для извлечения контента (опционально).
+
+Добавлено ограничение размера загружаемого контента (max_size_mb).
 """
 
 import asyncio
@@ -29,18 +33,24 @@ logger = logging.getLogger(__name__)
 
 async def fetch_any_url(url: str, timeout: int = 30, max_size_mb: int = 5) -> Tuple[Optional[str], Optional[str]]:
     """
-    Универсальная загрузка контента по URL с ограничением размера.
+    Универсальная загрузка контента по URL.
 
     Аргументы:
-        url: целевой адрес.
+        url (str): целевой адрес.
-        timeout: таймаут запроса в секундах.
+        timeout (int): таймаут запроса в секундах.
-        max_size_mb: максимальный размер загружаемого контента в МБ.
+        max_size_mb (int): максимальный размер загружаемого контента в МБ.
 
     Возвращает:
         Tuple[Optional[str], Optional[str]]: (заголовок, текст).
         Для HTML: заголовок страницы и очищенный основной текст.
         Для PDF: имя файла (базовая часть URL) и извлечённый текст.
         В случае ошибки возвращает (None, None).
+
+    Примечания:
+        - Для PDF используется временный файл, который удаляется после чтения.
+        - Поддерживаются только HTTP/HTTPS.
+        - Таймаут применяется к запросу в целом (чтение + соединение).
+        - Размер контента ограничен параметром max_size_mb.
     """
     try:
         async with aiohttp.ClientSession() as session:
@@ -48,10 +58,10 @@ async def fetch_any_url(url: str, timeout: int = 30, max_size_mb: int = 5) -> Tu
                 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
 
+                # Ограничение размера: читаем только limit байт
                 try:
                     data = await resp.content.read(limit=limit)
                 except aiohttp.ContentTooLong:
@@ -60,6 +70,7 @@ async def fetch_any_url(url: str, timeout: int = 30, max_size_mb: int = 5) -> Tu
 
                 # Обработка 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:
@@ -89,8 +100,28 @@ async def fetch_any_url(url: str, timeout: int = 30, max_size_mb: int = 5) -> Tu
 
 
 def extract_text_from_html(html: str, url: str = "") -> Tuple[Optional[str], str]:
-    """Извлекает заголовок и основной текст из HTML (без изменений)."""
+    """
+    Извлекает заголовок и основной текст из HTML.
+
+    Аргументы:
+        html (str): исходный HTML-код.
+        url (str): URL страницы (используется для формирования заголовка, если теги отсутствуют).
+
+    Возвращает:
+        Tuple[Optional[str], str]: (заголовок, очищенный текст).
+        Заголовок может быть None, если не удалось извлечь.
+
+    Алгоритм извлечения заголовка (по приоритету):
+        1. <title>
+        2. <meta property="og:title">
+        3. <h1>
+        4. Элементы с классами 'title', 'article-title', 'page-title', 'document-title'
+        5. Последний сегмент пути URL
+        6. Домен (как запасной вариант)
+    """
     soup = BeautifulSoup(html, 'lxml')
+
+    # ---- 1. Извлечение заголовка ----
     title = None
     if soup.title and soup.title.string:
         title = soup.title.string.strip()
@@ -103,12 +134,14 @@ 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]
@@ -116,14 +149,19 @@ 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)
@@ -131,11 +169,13 @@ 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,
     max_pages: int,
     domain_restrict: bool = True,
+    max_size_mb: int = 5,
 ) -> List[Tuple[str, str, str]]:
     """
     Рекурсивно обходит ссылки на том же домене (или не ограничивая, если domain_restrict=False).
@@ -145,6 +185,7 @@ async def crawl_url(
         max_depth (int): максимальная глубина перехода по ссылкам.
         max_pages (int): максимальное количество страниц для обработки.
         domain_restrict (bool): ограничиваться ли тем же доменом (по умолчанию True).
+        max_size_mb (int): максимальный размер загружаемой страницы (МБ).
 
     Возвращает:
         List[Tuple[str, str, str]]: список кортежей (url, title, text).
@@ -159,7 +200,7 @@ async def crawl_url(
         if depth > max_depth:
             continue
         logger.info(f"Обработка {url}, глубина {depth}")
-        title, text = await fetch_any_url(url)
+        title, text = await fetch_any_url(url, max_size_mb=max_size_mb)
         if text:
             results.append((url, title or url, text))
         else: