fckbot/rag/utils/web_utils.py

# -*- coding: utf-8 -*-
"""
Утилиты для загрузки веб-страниц и PDF, извлечения основного текста,
рекурсивного обхода ссылок.

Использует aiohttp для асинхронных запросов, BeautifulSoup для парсинга,
readability-lxml для извлечения контента (опционально).

Параметры (max_pages, max_depth) берутся из конфига при вызове из xmpp_client.

Импортируется в:
- indexing_worker.py (fetch_any_url, crawl_url)
- возможно, в других функциях при обработке URL.
"""

import asyncio
import aiohttp
import logging
from typing import Optional, List, Tuple
from urllib.parse import urljoin, urlparse

from bs4 import BeautifulSoup
import aiofiles
import tempfile
import os

# Опциональная библиотека readability-lxml улучшает извлечение основного текста
try:
    from readability import Document
    HAS_READABILITY = True
except ImportError:
    HAS_READABILITY = False
    logging.warning("readability-lxml не установлен, используем BeautifulSoup без очистки")

logger = logging.getLogger(__name__)


async def fetch_any_url(url: str, timeout: int = 30) -> Tuple[Optional[str], Optional[str]]:
    """
    Универсальная загрузка контента по URL.

    Аргументы:
        url (str): целевой адрес.
        timeout (int): таймаут запроса в секундах.

    Возвращает:
        Tuple[Optional[str], Optional[str]]: (заголовок, текст).
        Для HTML: заголовок страницы и очищенный основной текст.
        Для PDF: имя файла (базовая часть URL) и извлечённый текст.
        В случае ошибки возвращает (None, None).

    Примечания:
        - Для PDF используется временный файл, который удаляется после чтения.
        - Поддерживаются только HTTP/HTTPS.
        - Таймаут применяется к запросу в целом (чтение + соединение).
    """
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, timeout=timeout) as resp:
                if resp.status != 200:
                    logger.error(f"HTTP {resp.status} для {url}")
                    return None, None
                content_type = resp.headers.get('Content-Type', '').lower()
                # Обработка 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())
                    try:
                        from pypdf import PdfReader
                        reader = PdfReader(tmp_path)
                        text = "\n".join([page.extract_text() or "" for page in reader.pages])
                        title = os.path.basename(url)
                        return title, text
                    except ImportError:
                        logger.error("pypdf не установлен")
                        return None, None
                    finally:
                        os.unlink(tmp_path)
                else:
                    # Обычная HTML-страница
                    html = await resp.text()
                    title, text = extract_text_from_html(html, url)
                    return title, text
    except asyncio.TimeoutError:
        logger.error(f"Таймаут при загрузке {url}")
        return None, None
    except Exception as e:
        logger.error(f"Ошибка загрузки {url}: {e}")
        return None, None


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. Домен (как запасной вариант)
    """
    soup = BeautifulSoup(html, 'lxml')

    # ---- 1. Извлечение заголовка ----
    title = None
    if soup.title and soup.title.string:
        title = soup.title.string.strip()
    if not title:
        og_title = soup.find('meta', property='og:title')
        if og_title and og_title.get('content'):
            title = og_title['content'].strip()
    if not title:
        h1 = soup.find('h1')
        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]
            title = last_part.replace('-', ' ').replace('_', ' ').title()
        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)

    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,
) -> List[Tuple[str, str, str]]:
    """
    Рекурсивно обходит ссылки на том же домене (или не ограничивая, если domain_restrict=False).

    Аргументы:
        start_url (str): начальная страница.
        max_depth (int): максимальная глубина перехода по ссылкам.
        max_pages (int): максимальное количество страниц для обработки.
        domain_restrict (bool): ограничиваться ли тем же доменом (по умолчанию True).

    Возвращает:
        List[Tuple[str, str, str]]: список кортежей (url, title, text).
    """
    visited = set()
    to_visit = [(start_url, 0)]
    results = []
    visited.add(start_url)

    while to_visit and len(results) < max_pages:
        url, depth = to_visit.pop(0)
        if depth > max_depth:
            continue
        logger.info(f"Обработка {url}, глубина {depth}")
        title, text = await fetch_any_url(url)
        if text:
            results.append((url, title or url, text))
        else:
            continue
        if depth < max_depth:
            # Извлекаем ссылки из текущей страницы
            links = await extract_links_from_url(url, domain_restrict)
            for link in links:
                if link not in visited:
                    visited.add(link)
                    to_visit.append((link, depth + 1))
    return results


async def extract_links_from_url(url: str, domain_restrict: bool = True) -> List[str]:
    """
    Извлекает все ссылки (href) из HTML-страницы по заданному URL.

    Аргументы:
        url (str): адрес страницы.
        domain_restrict (bool): если True, оставлять только ссылки на том же домене.

    Возвращает:
        List[str]: список абсолютных URL.
    """
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, timeout=10) as resp:
                if resp.status != 200:
                    return []
                html = await resp.text()
                soup = BeautifulSoup(html, 'lxml')
                base_domain = urlparse(url).netloc
                links = set()
                for a in soup.find_all('a', href=True):
                    href = a['href']
                    full_url = urljoin(url, href)
                    parsed = urlparse(full_url)
                    if domain_restrict and parsed.netloc != base_domain:
                        continue
                    if parsed.scheme in ('http', 'https'):
                        # Удаляем фрагмент (#)
                        full_url = parsed._replace(fragment='').geturl()
                        links.add(full_url)
                return list(links)
    except Exception as e:
        logger.error(f"Ошибка извлечения ссылок с {url}: {e}")
        return []