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

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. <title>
-        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,