Update 86 files
- /rag/commands/expert.py - /core/services/__init__.py - /core/services/embedding_service.py - /core/services/file_service.py - /core/services/giga_client.py - /core/services/kb_service.py - /core/services/postgres_service.py - /core/services/qdrant_service.py - /core/services/reranker_service.py - /core/functions/__init__.py - /core/functions/check_consistency.py - /core/functions/check_spelling.py - /core/functions/critique_answer.py - /core/functions/expand_query.py - /core/functions/extract_metrics.py - /core/functions/file_processor.py - /core/functions/generate_document.py - /core/functions/intent_classify.py - /core/functions/rerank_context.py - /core/functions/summarize_document.py - /core/utils/__init__.py - /core/utils/arg_parser.py - /core/utils/config_loader.py - /core/utils/layout_converter.py - /core/utils/logger.py - /core/utils/text_utils.py - /core/utils/web_utils.py - /bots/commands/__init__.py - /bots/commands/base.py - /bots/commands/create.py - /core/commands/global_remove.py - /core/commands/help.py - /core/commands/info.py - /core/commands/kb.py - /core/commands/learn.py - /core/commands/other.py - /core/commands/registry.py - /core/commands/stats.py - /core/commands/template.py - /core/handlers/__init__.py - /core/handlers/message_handler.py - /core/workers/__init__.py - /core/workers/indexing_worker.py - /core/xmpp/__init__.py - /core/xmpp/client.py - /bots/commands/global_remove.py - /bots/commands/help.py - /bots/commands/info.py - /bots/commands/kb.py - /bots/commands/registry.py - /bots/commands/other.py - /bots/commands/template.py - /bots/commands/learn.py - /bots/commands/stats.py - /bots/handlers/message_handler.py - /bots/handlers/__init__.py - /bots/workers/__init__.py - /bots/workers/indexing_worker.py - /bots/xmpp/__init__.py - /bots/xmpp/client.py - /rag/services/qdrant_service.py - /rag/services/giga_client.py - /rag/services/embedding_service.py - /rag/services/reranker_service.py - /rag/services/__init__.py - /rag/services/file_service.py - /rag/services/postgres_service.py - /rag/services/kb_service.py - /rag/functions/check_spelling.py - /rag/functions/__init__.py - /rag/functions/extract_metrics.py - /rag/functions/generate_document.py - /rag/functions/expand_query.py - /rag/functions/critique_answer.py - /rag/functions/file_processor.py - /rag/functions/check_consistency.py - /rag/functions/summarize_document.py - /rag/functions/rerank_context.py - /rag/functions/intent_classify.py - /rag/utils/__init__.py - /rag/utils/layout_converter.py - /rag/utils/text_utils.py - /rag/utils/arg_parser.py - /rag/utils/config_loader.py - /rag/utils/logger.py - /rag/utils/web_utils.py
This commit is contained in:
11
rag/utils/__init__.py
Normal file
11
rag/utils/__init__.py
Normal file
@@ -0,0 +1,11 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Пакет utils – вспомогательные утилиты без бизнес-логики.
|
||||
Содержит:
|
||||
- config_loader – загрузка конфигурации,
|
||||
- logger – настройка логирования,
|
||||
- layout_converter – конвертер раскладки клавиатуры,
|
||||
- text_utils – разбиение текста на чанки,
|
||||
- web_utils – веб-скрапинг,
|
||||
- arg_parser – парсинг аргументов команд.
|
||||
"""
|
||||
62
rag/utils/arg_parser.py
Normal file
62
rag/utils/arg_parser.py
Normal file
@@ -0,0 +1,62 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Утилита для разбора аргументов командной строки в стиле --flag value или --flag.
|
||||
Используется командами бота для извлечения флагов (--global, --title и т.д.) и позиционных аргументов.
|
||||
"""
|
||||
|
||||
import shlex
|
||||
from typing import Tuple, Dict, List, Any
|
||||
|
||||
|
||||
def parse_args(args_str: str) -> Tuple[str, Dict[str, Any], List[str]]:
|
||||
"""
|
||||
Разбирает строку аргументов.
|
||||
|
||||
Возвращает кортеж: (команда, словарь флагов, список позиционных аргументов).
|
||||
|
||||
Пример:
|
||||
"!learn --global --title 'Мой документ' file.docx"
|
||||
-> ("!learn", {"global": True, "title": "Мой документ"}, ["file.docx"])
|
||||
|
||||
Алгоритм:
|
||||
1. Разбить строку на части, учитывая кавычки (через shlex.split).
|
||||
2. Если shlex не справился (непарные кавычки), разбить по пробелам (грубо).
|
||||
3. Первая часть – имя команды.
|
||||
4. Далее идём по частям: если часть начинается с '--', то это флаг.
|
||||
- Если следующий элемент не начинается с '--', то это значение флага.
|
||||
- Иначе флаг булевый (True).
|
||||
5. Всё, что не начинается с '--', – позиционные аргументы.
|
||||
"""
|
||||
if not args_str:
|
||||
return "", {}, []
|
||||
|
||||
# Пытаемся разбить с учётом кавычек (shlex)
|
||||
try:
|
||||
parts = shlex.split(args_str)
|
||||
except ValueError:
|
||||
# fallback: просто по пробелам (может быть неточно, но работает для простых случаев)
|
||||
parts = args_str.split()
|
||||
|
||||
if not parts:
|
||||
return "", {}, []
|
||||
|
||||
cmd = parts[0] # первая часть – имя команды
|
||||
flags = {}
|
||||
positionals = []
|
||||
i = 1
|
||||
while i < len(parts):
|
||||
p = parts[i]
|
||||
if p.startswith('--'):
|
||||
key = p[2:] # убираем '--'
|
||||
# Если следующий элемент есть и он не флаг, то это значение
|
||||
if i + 1 < len(parts) and not parts[i+1].startswith('--'):
|
||||
flags[key] = parts[i+1]
|
||||
i += 2
|
||||
else:
|
||||
flags[key] = True
|
||||
i += 1
|
||||
else:
|
||||
positionals.append(p)
|
||||
i += 1
|
||||
|
||||
return cmd, flags, positionals
|
||||
290
rag/utils/config_loader.py
Normal file
290
rag/utils/config_loader.py
Normal file
@@ -0,0 +1,290 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Загрузчик конфигурации бота из bot.conf и переменных окружения.
|
||||
Собирает все настройки в единый объект BotConfig.
|
||||
"""
|
||||
|
||||
import os
|
||||
import yaml
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Any, Dict
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
class BotConfig:
|
||||
"""
|
||||
Единый объект конфигурации для одного профиля бота.
|
||||
Все параметры читаются из bot.conf, core.conf и переменных окружения.
|
||||
При отсутствии любого параметра – ошибка.
|
||||
"""
|
||||
|
||||
def __init__(self, profile_dir: str):
|
||||
self.profile_dir = Path(profile_dir).resolve()
|
||||
self.bot_config_path = self.profile_dir / "bot.conf"
|
||||
if not self.bot_config_path.exists():
|
||||
raise FileNotFoundError(f"Файл профиля не найден: {self.bot_config_path}")
|
||||
|
||||
with open(self.bot_config_path, 'r', encoding='utf-8') as f:
|
||||
self.bot_raw = yaml.safe_load(f)
|
||||
|
||||
# Загружаем core.conf
|
||||
core_config_path = self.profile_dir.parent / "core" / "core.conf"
|
||||
if core_config_path.exists():
|
||||
with open(core_config_path, 'r', encoding='utf-8') as f:
|
||||
self.core_raw = yaml.safe_load(f) or {}
|
||||
else:
|
||||
self.core_raw = {}
|
||||
|
||||
# Объединяем: bot.conf имеет приоритет над core.conf
|
||||
self.raw = self._merge_configs(self.core_raw, self.bot_raw)
|
||||
|
||||
# ----- 1. Обязательные поля (без них бот не запустится) -----
|
||||
self.name: str = self._get_required("name")
|
||||
self.jid: str = self._get_required("jid")
|
||||
self.password_env: str = self._get_required("password_env")
|
||||
self.xmpp_password: str = self._get_env(self.password_env)
|
||||
|
||||
# ----- 2. XMPP параметры (обязательные) -----
|
||||
self.xmpp_resource = self._get_required("resource")
|
||||
self.status_message = self._get_required("status_message")
|
||||
self.xmpp_server = self._get_required("xmpp_server")
|
||||
self.xmpp_port = int(self._get_required("xmpp_port"))
|
||||
self.xmpp_use_ssl = bool(self._get_required("xmpp_use_ssl"))
|
||||
|
||||
# ----- 3. Настройки базы данных (PostgreSQL) -----
|
||||
db_cfg = self._get_required_section("database")
|
||||
self.db_name = self._get_required_from_section(db_cfg, "database")
|
||||
self.db_type = self._get_required_from_section(db_cfg, "type")
|
||||
self.db_host = self._get_required_from_section(db_cfg, "host")
|
||||
self.db_port = int(self._get_required_from_section(db_cfg, "port"))
|
||||
self.db_user = self._get_required_from_section(db_cfg, "user")
|
||||
self.db_password_env = self._get_required_from_section(db_cfg, "password_env")
|
||||
self.db_password = self._get_env(self.db_password_env)
|
||||
|
||||
# SSL параметры (опциональны, но если секция ssl не задана – будет ошибка, если нет)
|
||||
self.db_ssl = self._get_optional_from_section(db_cfg, "ssl", default=None)
|
||||
if self.db_ssl is None:
|
||||
raise ValueError("В секции database отсутствует обязательный параметр 'ssl'")
|
||||
self.db_ssl_mode = self._get_optional_from_section(db_cfg, "ssl_mode", default=None)
|
||||
if self.db_ssl_mode is None:
|
||||
raise ValueError("В секции database отсутствует обязательный параметр 'ssl_mode'")
|
||||
|
||||
# ----- 4. Настройки Qdrant -----
|
||||
qdrant_cfg = self._get_required_section("qdrant")
|
||||
self.qdrant_collection = self._get_required_from_section(qdrant_cfg, "collection")
|
||||
self.qdrant_host = self._get_required_from_section(qdrant_cfg, "host")
|
||||
self.qdrant_port = int(self._get_required_from_section(qdrant_cfg, "port"))
|
||||
self.qdrant_grpc_port = int(self._get_required_from_section(qdrant_cfg, "grpc_port"))
|
||||
self.qdrant_vector_size = int(self._get_required_from_section(qdrant_cfg, "vector_size"))
|
||||
self.qdrant_distance = self._get_required_from_section(qdrant_cfg, "distance")
|
||||
|
||||
# ----- 5. Настройки GigaChat -----
|
||||
gc_cfg = self._get_required_section("gigachat")
|
||||
self.gigachat_credentials_env = self._get_required_from_section(gc_cfg, "credentials_env")
|
||||
self.gigachat_api_key = self._get_env(self.gigachat_credentials_env)
|
||||
self.gigachat_model_embedding = self._get_required_from_section(gc_cfg, "model_embedding")
|
||||
self.gigachat_model_generation = self._get_required_from_section(gc_cfg, "model_generation")
|
||||
self.gigachat_timeout = int(self._get_required_from_section(gc_cfg, "timeout"))
|
||||
self.gigachat_temperature = float(self._get_required_from_section(gc_cfg, "temperature"))
|
||||
|
||||
# ----- 6. Кэш эмбеддингов -----
|
||||
emb_cfg = self._get_required_section("embedding")
|
||||
self.embedding_cache_size = int(self._get_required_from_section(emb_cfg, "cache_size"))
|
||||
self.embedding_model = self._get_required_from_section(emb_cfg, "model")
|
||||
self.embedding_timeout = int(self._get_required_from_section(emb_cfg, "timeout"))
|
||||
self.embedding_verify_ssl = bool(self._get_required_from_section(emb_cfg, "verify_ssl_certs"))
|
||||
|
||||
# ----- 7. Разбиение документов на чанки -----
|
||||
chunk_cfg = self._get_required_section("chunking")
|
||||
self.chunking_enabled = bool(self._get_required_from_section(chunk_cfg, "enabled"))
|
||||
self.chunk_size_tokens = int(self._get_required_from_section(chunk_cfg, "chunk_size_tokens"))
|
||||
self.overlap_tokens = int(self._get_required_from_section(chunk_cfg, "overlap_tokens"))
|
||||
self.chunking_strategy = self._get_required_from_section(chunk_cfg, "strategy")
|
||||
self.chunking_approx_chunk_chars = int(self._get_required_from_section(chunk_cfg, "approx_chunk_chars"))
|
||||
self.chunking_approx_overlap_chars = int(self._get_required_from_section(chunk_cfg, "approx_overlap_chars"))
|
||||
|
||||
# ----- 8. Пути к файлам и каталогам -----
|
||||
paths = self._get_required_section("paths")
|
||||
self.data_dir = (self.profile_dir / paths.get("data_dir")).resolve()
|
||||
self.temp_dir = (self.profile_dir / paths.get("temp_dir")).resolve()
|
||||
self.prompts_dir = (self.profile_dir / paths.get("prompts_dir")).resolve()
|
||||
self.upload_dir = Path(paths.get("upload_dir")).resolve()
|
||||
self.log_dir = Path(paths.get("log_dir")).resolve()
|
||||
log_filename = paths.get("log_filename")
|
||||
if not log_filename:
|
||||
raise ValueError("В секции paths отсутствует log_filename")
|
||||
self.log_file = self.log_dir / log_filename.format(name=self.name)
|
||||
|
||||
# ----- 9. HTTP Upload -----
|
||||
http_up = self._get_required_section("http_upload")
|
||||
self.upload_base_url = http_up.get("base_url")
|
||||
|
||||
# ----- 10. Настройки AI (промпты, модель) -----
|
||||
ai_cfg = self._get_required_section("ai")
|
||||
self.ai_model = ai_cfg.get("model")
|
||||
if not self.ai_model:
|
||||
self.ai_model = self.gigachat_model_generation
|
||||
self.ai_temperature = float(ai_cfg.get("temperature", self.gigachat_temperature))
|
||||
self.ai_timeout = int(ai_cfg.get("timeout", self.gigachat_timeout))
|
||||
|
||||
prompts_cfg = ai_cfg.get("prompts")
|
||||
if not prompts_cfg:
|
||||
raise ValueError("В секции ai отсутствует подсекция prompts")
|
||||
|
||||
# промпты – обязательные файлы
|
||||
self.system_prompt_file = self.prompts_dir / prompts_cfg.get("system")
|
||||
self.synthesis_prompt_file = self.prompts_dir / prompts_cfg.get("synthesis")
|
||||
self.quality_prompt_file = self.prompts_dir / prompts_cfg.get("quality")
|
||||
self.rerank_prompt_file = self.prompts_dir / prompts_cfg.get("rerank")
|
||||
self.expand_prompt_file = self.prompts_dir / prompts_cfg.get("expand")
|
||||
self.intent_prompt_file = self.prompts_dir / prompts_cfg.get("intent")
|
||||
self.metrics_prompt_file = self.prompts_dir / prompts_cfg.get("metrics")
|
||||
self.summary_prompt_file = self.prompts_dir / prompts_cfg.get("summary")
|
||||
self.consistency_prompt_file = self.prompts_dir / prompts_cfg.get("consistency")
|
||||
self.critique_prompt_file = self.prompts_dir / prompts_cfg.get("critique")
|
||||
self.spellcheck_prompt_file = self.prompts_dir / prompts_cfg.get("spellcheck")
|
||||
self.generate_document_prompt_file = self.prompts_dir / prompts_cfg.get("generate_document")
|
||||
|
||||
# ----- 11. Параметры RAG -----
|
||||
rag_cfg = self._get_required_section("rag")
|
||||
self.rag_default_top_k = int(rag_cfg.get("default_top_k"))
|
||||
self.rag_metrics_top_k = int(rag_cfg.get("metrics_top_k"))
|
||||
self.rag_contradiction_top_k = int(rag_cfg.get("contradiction_top_k"))
|
||||
self.max_context_tokens = int(rag_cfg.get("max_context_tokens"))
|
||||
|
||||
# ----- 12. Веб-скрапер -----
|
||||
ws_cfg = self._get_required_section("web_scraper")
|
||||
self.ws_max_pages = int(ws_cfg.get("max_pages"))
|
||||
self.ws_max_depth = int(ws_cfg.get("max_depth"))
|
||||
|
||||
# ----- 13. Включение/отключение функций -----
|
||||
feat = self._get_required_section("features")
|
||||
self.file_processing = bool(feat.get("file_processing"))
|
||||
self.surgical_replace = bool(feat.get("surgical_replace"))
|
||||
self.voice_recognition = bool(feat.get("voice_recognition"))
|
||||
self.vision = bool(feat.get("vision"))
|
||||
self.archive_support = bool(feat.get("archive_support"))
|
||||
self.ocr = bool(feat.get("ocr"))
|
||||
self.allow_public_knowledge = bool(feat.get("allow_public_knowledge"))
|
||||
self.max_file_size_mb = int(feat.get("max_file_size_mb"))
|
||||
self.max_archive_files = int(feat.get("max_archive_files"))
|
||||
self.enable_intent_classification = bool(feat.get("enable_intent_classification"))
|
||||
self.enable_self_critique = bool(feat.get("enable_self_critique"))
|
||||
self.surgical_keywords = feat.get("surgical_keywords")
|
||||
if not isinstance(self.surgical_keywords, list):
|
||||
raise ValueError("surgical_keywords должен быть списком строк")
|
||||
self.mention_keyword = feat.get("mention_keyword", self.name) # если не указано, используем имя бота
|
||||
|
||||
# ----- 14. Очистка старых файлов -----
|
||||
clean_cfg = self._get_required_section("cleanup")
|
||||
self.temp_cleanup_days = int(clean_cfg.get("temp_days"))
|
||||
self.upload_cleanup_days = int(clean_cfg.get("upload_days"))
|
||||
self.cleanup_interval = int(clean_cfg.get("interval_seconds"))
|
||||
self.user_data_cleanup_days = int(clean_cfg.get("user_data_days"))
|
||||
|
||||
# ----- 15. Таймауты HTTP -----
|
||||
http_cfg = self._get_required_section("http")
|
||||
self.http_download_timeout = int(http_cfg.get("download_timeout"))
|
||||
self.http_upload_timeout = int(http_cfg.get("upload_timeout"))
|
||||
|
||||
# ----- 16. Права на создаваемые файлы -----
|
||||
fp_cfg = self._get_required_section("file_permissions")
|
||||
self.file_mode = int(fp_cfg.get("mode"), 8) # восьмеричное
|
||||
|
||||
# ----- 17. Администраторы -----
|
||||
self.admin_jids = self._get_required("admin_jids")
|
||||
if not isinstance(self.admin_jids, list):
|
||||
raise ValueError("admin_jids должен быть списком JID")
|
||||
|
||||
# ----- 18. SaluteSpeech (из .env) -----
|
||||
self.salute_speech_auth = os.getenv("SALUTE_SPEECH_AUTH", "")
|
||||
if not self.salute_speech_auth:
|
||||
raise ValueError("Переменная окружения SALUTE_SPEECH_AUTH не задана")
|
||||
|
||||
# ----- 19. Log level -----
|
||||
log_level_str = self._get_required("log_level")
|
||||
self.log_level = getattr(logging, log_level_str.upper(), logging.INFO)
|
||||
|
||||
# ----- 20. Настройки функций (intent, expand, metrics, summary, consistency, critique, rerank) -----
|
||||
# Все они читаются из соответствующих секций, но сами секции могут отсутствовать?
|
||||
# Поскольку мы требуем всё явно, то каждая секция должна быть.
|
||||
self.intent_temperature = float(self._get_required_from_section(self._get_required_section("intent"), "temperature"))
|
||||
self.expand_temperature = float(self._get_required_from_section(self._get_required_section("expand"), "temperature"))
|
||||
self.metrics_temperature = float(self._get_required_from_section(self._get_required_section("metrics"), "temperature"))
|
||||
self.summary_temperature = float(self._get_required_from_section(self._get_required_section("summary"), "temperature"))
|
||||
self.summary_max_chars = int(self._get_required_from_section(self._get_required_section("summary"), "max_chars"))
|
||||
self.consistency_temperature = float(self._get_required_from_section(self._get_required_section("consistency"), "temperature"))
|
||||
self.consistency_max_fragments = int(self._get_required_from_section(self._get_required_section("consistency"), "max_fragments"))
|
||||
self.critique_temperature = float(self._get_required_from_section(self._get_required_section("critique"), "temperature"))
|
||||
self.rerank_temperature = float(self._get_required_from_section(self._get_required_section("rerank"), "temperature"))
|
||||
self.rerank_min_length = int(self._get_required_from_section(self._get_required_section("rerank"), "min_length"))
|
||||
self.generate_document_temperature = float(self._get_required_from_section(self._get_required_section("generate_document"), "temperature"))
|
||||
|
||||
# ----- 21. Совместимость с функциями: создаём вложенные объекты -----
|
||||
self.intent = {"temperature": self.intent_temperature}
|
||||
self.expand = {"temperature": self.expand_temperature}
|
||||
self.metrics = {"temperature": self.metrics_temperature}
|
||||
self.summary = {"temperature": self.summary_temperature, "max_chars": self.summary_max_chars}
|
||||
self.consistency = {"temperature": self.consistency_temperature, "max_fragments": self.consistency_max_fragments}
|
||||
self.critique = {"temperature": self.critique_temperature}
|
||||
self.rerank = {"temperature": self.rerank_temperature, "min_length": self.rerank_min_length}
|
||||
|
||||
# Создаём директории
|
||||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.prompts_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.upload_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Вспомогательные методы
|
||||
# ------------------------------------------------------------------
|
||||
def _merge_configs(self, core: dict, bot: dict) -> dict:
|
||||
"""Рекурсивное слияние: bot имеет приоритет над core."""
|
||||
merged = core.copy()
|
||||
for key, value in bot.items():
|
||||
if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
|
||||
merged[key] = self._merge_configs(merged[key], value)
|
||||
else:
|
||||
merged[key] = value
|
||||
return merged
|
||||
|
||||
def _get_required(self, key: str) -> Any:
|
||||
"""Возвращает значение ключа из объединённого конфига. Если нет – ошибка."""
|
||||
if key not in self.raw:
|
||||
raise KeyError(f"В конфигурации (bot.conf или core.conf) отсутствует обязательный ключ: '{key}'")
|
||||
return self.raw[key]
|
||||
|
||||
def _get_required_section(self, section: str) -> dict:
|
||||
"""Возвращает словарь-секцию. Если секции нет – ошибка."""
|
||||
if section not in self.raw:
|
||||
raise KeyError(f"В конфигурации отсутствует обязательная секция: '{section}'")
|
||||
if not isinstance(self.raw[section], dict):
|
||||
raise TypeError(f"Секция '{section}' должна быть словарём")
|
||||
return self.raw[section]
|
||||
|
||||
def _get_required_from_section(self, section_dict: dict, key: str) -> Any:
|
||||
"""Возвращает значение ключа внутри секции. Если нет – ошибка."""
|
||||
if key not in section_dict:
|
||||
raise KeyError(f"В секции отсутствует обязательный ключ: '{key}'")
|
||||
return section_dict[key]
|
||||
|
||||
def _get_optional_from_section(self, section_dict: dict, key: str, default=None):
|
||||
"""Возвращает значение или default, если ключ отсутствует (не вызывает ошибку)."""
|
||||
return section_dict.get(key, default)
|
||||
|
||||
def _get_env(self, env_var: str) -> str:
|
||||
"""Читает переменную окружения. Если нет – ошибка."""
|
||||
value = os.getenv(env_var)
|
||||
if value is None or value.strip() == "":
|
||||
raise ValueError(
|
||||
f"Переменная окружения {env_var} не задана или пуста. "
|
||||
f"Бот не может запуститься без неё."
|
||||
)
|
||||
return value.strip()
|
||||
71
rag/utils/layout_converter.py
Normal file
71
rag/utils/layout_converter.py
Normal file
@@ -0,0 +1,71 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Конвертер раскладки клавиатуры: определяет, в какой раскладке (русской или английской)
|
||||
набран текст, и при необходимости переключает её.
|
||||
Используется для исправления сообщений, набранных по ошибке не в той раскладке.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def convert_layout(text: str) -> str:
|
||||
"""
|
||||
Определяет, в какой раскладке набран текст, и переключает, если нужно.
|
||||
Команды (начинающиеся с ! или /) и ссылки (http://) не трогает.
|
||||
|
||||
Алгоритм:
|
||||
1. Если текст начинается с http://, https://, ! или / – возвращается без изменений.
|
||||
2. Составляем таблицы перевода: английские символы -> русские и наоборот.
|
||||
3. Если в тексте есть и русские, и английские символы – считаем, что всё правильно (не переключаем).
|
||||
4. Считаем количество гласных букв в английской и русской раскладках для исходного текста
|
||||
и для его переведённых вариантов.
|
||||
5. Выбираем вариант с наибольшим количеством гласных (вероятная правильная раскладка).
|
||||
6. Если гласных поровну или нет ни одной, возвращаем исходный текст.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
# Не трогаем команды и ссылки
|
||||
if text.startswith('http') or text.startswith(('!', '/')):
|
||||
return text
|
||||
|
||||
# Таблицы соответствия символов (раскладка QWERTY / ЙЦУКЕН)
|
||||
en_chars = "qwertyuiop[]asdfghjkl;'zxcvbnm,./QWERTYUIOP{}ASDFGHJKL:\"ZXCVBNM<>?~"
|
||||
ru_chars = "йцукенгшщзхъфывапролджэячсмитьбю.ЙЦУКЕНГШЩЗХЪФЫВАПРОЛДЖЭЯЧСМИТЬБЮ,Ё"
|
||||
to_ru = str.maketrans(en_chars, ru_chars) # таблица для перевода английской раскладки в русскую
|
||||
to_en = str.maketrans(ru_chars, en_chars) # таблица для перевода русской раскладки в английскую
|
||||
|
||||
# Проверяем, есть ли символы обеих раскладок в тексте
|
||||
has_russian = any('а' <= c <= 'я' or 'А' <= c <= 'Я' for c in text)
|
||||
has_english = any('a' <= c <= 'z' or 'A' <= c <= 'Z' for c in text)
|
||||
if has_russian and has_english:
|
||||
# Смесь – вероятно, уже правильно
|
||||
return text
|
||||
|
||||
# Гласные буквы для подсчёта (в английской и русской раскладках)
|
||||
eng_vowels = "aeiouyAEIOUY"
|
||||
ru_vowels = "аеёиоуыэюяАЕЁИОУЫЭЮЯ"
|
||||
|
||||
orig_en_vowels = sum(1 for c in text if c in eng_vowels)
|
||||
orig_ru_vowels = sum(1 for c in text if c in ru_vowels)
|
||||
|
||||
# Переводим текст в другую раскладку
|
||||
translated_en = text.translate(to_en) # русская -> английская
|
||||
translated_ru = text.translate(to_ru) # английская -> русская
|
||||
|
||||
trans_en_vowels = sum(1 for c in translated_en if c in eng_vowels)
|
||||
trans_ru_vowels = sum(1 for c in translated_ru if c in ru_vowels)
|
||||
|
||||
# Если в исходном тексте нет русских символов, но при переводе в русскую раскладку
|
||||
# появляется больше гласных, значит, текст был набран в английской раскладке.
|
||||
if not has_russian and trans_ru_vowels > orig_en_vowels:
|
||||
return translated_ru
|
||||
|
||||
# Если в исходном тексте нет английских символов, но при переводе в английскую раскладку
|
||||
# появляется больше гласных, значит, текст был набран в русской раскладке.
|
||||
if has_russian and trans_en_vowels > orig_ru_vowels:
|
||||
return translated_en
|
||||
|
||||
return text
|
||||
53
rag/utils/logger.py
Normal file
53
rag/utils/logger.py
Normal file
@@ -0,0 +1,53 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Настройка логирования для бота.
|
||||
Создаёт логгер с выводом в консоль и в файл.
|
||||
Уровень логирования может быть установлен в конфиге (DEBUG для разработки, INFO для production).
|
||||
|
||||
Изменения: удалён aiomysql из списка шумных библиотек (заменён на asyncpg).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def setup_logging(profile_name: str, log_file: Path, log_level: int = logging.INFO):
|
||||
"""
|
||||
Инициализирует логирование:
|
||||
- вывод в консоль (StreamHandler)
|
||||
- вывод в файл (FileHandler)
|
||||
- снижает уровень шума для сторонних библиотек (slixmpp, asyncpg, qdrant_client и др.)
|
||||
"""
|
||||
log_format = '%(asctime)s [%(levelname)s] %(name)s: %(message)s'
|
||||
# Формат даты: 2026-05-20 15:30:00
|
||||
date_format = '%Y-%m-%d %H:%M:%S'
|
||||
|
||||
# Настройка корневого логгера
|
||||
logging.basicConfig(
|
||||
level=log_level,
|
||||
format=log_format,
|
||||
datefmt=date_format,
|
||||
handlers=[
|
||||
logging.StreamHandler(sys.stdout), # вывод в консоль
|
||||
logging.FileHandler(log_file, encoding='utf-8') # вывод в файл
|
||||
]
|
||||
)
|
||||
|
||||
# Уменьшаем многословность некоторых библиотек (они пишут много DEBUG)
|
||||
noisy_loggers = [
|
||||
"slixmpp",
|
||||
"asyncpg", # вместо aiomysql (PostgreSQL)
|
||||
"qdrant_client",
|
||||
"httpx",
|
||||
"gigachat",
|
||||
"urllib3",
|
||||
"asyncio",
|
||||
"aiohttp"
|
||||
]
|
||||
for lib in noisy_loggers:
|
||||
logging.getLogger(lib).setLevel(logging.WARNING)
|
||||
|
||||
# Для особо шумных можно ещё снизить, но оставляем WARNING
|
||||
logging.getLogger("requests").setLevel(logging.WARNING)
|
||||
logging.getLogger("urllib3").setLevel(logging.WARNING)
|
||||
226
rag/utils/text_utils.py
Normal file
226
rag/utils/text_utils.py
Normal file
@@ -0,0 +1,226 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Утилиты для работы с текстом:
|
||||
- Разбиение на чанки (токенизация) с помощью tiktoken для индексации документов.
|
||||
- Подсчёт количества токенов в тексте (для обрезки истории диалога).
|
||||
- Fallback-разбиение по символам, если tiktoken не установлен.
|
||||
|
||||
Используется в:
|
||||
- KBService при индексации документов (split_into_chunks).
|
||||
- RAG-пайплайне для обрезки истории диалога по токенам (count_tokens).
|
||||
|
||||
Особенности:
|
||||
- Разбиение на предложения использует библиотеку razdel для корректной обработки
|
||||
аббревиатур (т.е., и.т.д.) и других сложных случаев русского языка.
|
||||
- Если razdel недоступен, используется fallback-регулярное выражение.
|
||||
- Если tiktoken недоступен, используется грубая оценка токенов (3 символа на токен).
|
||||
"""
|
||||
|
||||
import re
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
# ---- Инициализация токенизатора tiktoken (для точного подсчёта токенов) ----
|
||||
try:
|
||||
import tiktoken
|
||||
# Используем кодировку cl100k_base, которая используется в моделях GPT-4, GPT-3.5 и GigaChat
|
||||
TOKENIZER = tiktoken.get_encoding("cl100k_base")
|
||||
except ImportError:
|
||||
TOKENIZER = None
|
||||
logging.warning("tiktoken не установлен. Для точного подсчёта токенов установите: pip install tiktoken")
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
# ---- Импорт razdel для умного разбиения по предложениям (русский язык) ----
|
||||
try:
|
||||
from razdel import sentenize
|
||||
HAS_RAZDEL = True
|
||||
except ImportError:
|
||||
HAS_RAZDEL = False
|
||||
logging.warning("razdel не установлен. Разбиение по предложениям будет неточным. Установите: pip install razdel")
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""
|
||||
Подсчитывает количество токенов в тексте, используя tiktoken (если доступен).
|
||||
Если tiktoken недоступен, использует грубую оценку: 1 токен ≈ 3 символа.
|
||||
|
||||
Эта функция используется в RAG-пайплайне для обрезки истории диалога
|
||||
до заданного лимита токенов (max_context_tokens).
|
||||
|
||||
Аргументы:
|
||||
text (str): текст для подсчёта токенов.
|
||||
|
||||
Возвращает:
|
||||
int: количество токенов (или оценка).
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
|
||||
if TOKENIZER is not None:
|
||||
# Точный подсчёт через tiktoken
|
||||
return len(TOKENIZER.encode(text))
|
||||
else:
|
||||
# Грубая оценка: для русского языка в среднем 1 токен ≈ 2.5–3 символа.
|
||||
# Используем 3 как консервативную оценку (чтобы не превысить лимит).
|
||||
return len(text) // 3 + 1
|
||||
|
||||
|
||||
def split_into_chunks(
|
||||
text: str,
|
||||
chunk_size_tokens: int = 200,
|
||||
overlap_tokens: int = 50,
|
||||
approx_chunk_chars: int = 500,
|
||||
approx_overlap_chars: int = 100,
|
||||
strategy: str = "recursive_split_by_sentences"
|
||||
) -> List[str]:
|
||||
"""
|
||||
Разбивает текст на чанки (фрагменты) для последующей индексации в векторной БД.
|
||||
Сначала пытается использовать tiktoken для точного разбиения по токенам.
|
||||
Если tiktoken недоступен, использует приближённое разбиение по символам.
|
||||
|
||||
Аргументы:
|
||||
text (str): исходный текст.
|
||||
chunk_size_tokens (int): целевой размер чанка в токенах (по умолчанию 200).
|
||||
overlap_tokens (int): перекрытие между соседними чанками в токенах (по умолчанию 50).
|
||||
approx_chunk_chars (int): размер чанка в символах для fallback-режима (по умолчанию 500).
|
||||
approx_overlap_chars (int): перекрытие в символах для fallback-режима (по умолчанию 100).
|
||||
strategy (str): стратегия разбиения (пока только "recursive_split_by_sentences").
|
||||
|
||||
Возвращает:
|
||||
List[str]: список чанков (текстовых фрагментов).
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
# Нормализуем пробелы: заменяем множественные пробелы, табуляции и переносы на один пробел.
|
||||
# Это упрощает дальнейшую обработку и улучшает качество разбиения.
|
||||
text = re.sub(r'\s+', ' ', text).strip()
|
||||
|
||||
# Если токенизатор недоступен, используем fallback-разбиение по символам.
|
||||
if TOKENIZER is None:
|
||||
logger.debug("tiktoken недоступен, используем разбиение по символам")
|
||||
return _split_by_chars(text, approx_chunk_chars, approx_overlap_chars)
|
||||
|
||||
# Основной метод: разбиение по предложениям с учётом токенов.
|
||||
return _split_by_tokens(text, chunk_size_tokens, overlap_tokens)
|
||||
|
||||
|
||||
def _split_by_tokens(text: str, chunk_size: int, overlap: int) -> List[str]:
|
||||
"""
|
||||
Разбивает текст на чанки, используя tiktoken для точного подсчёта токенов.
|
||||
Алгоритм:
|
||||
1. Разбить текст на предложения (с помощью razdel для русского языка).
|
||||
2. Накопить предложения, пока сумма токенов не превысит chunk_size.
|
||||
3. Когда превышает – сохранить текущий чанк и создать перекрытие из последних предложений (overlap токенов).
|
||||
4. Продолжить с оставшимся хвостом.
|
||||
5. В конце объединить очень короткие чанки (< chunk_size/3) с предыдущим,
|
||||
чтобы избежать множества мелких фрагментов.
|
||||
|
||||
Аргументы:
|
||||
text (str): исходный текст.
|
||||
chunk_size (int): максимальный размер чанка в токенах.
|
||||
overlap (int): количество токенов перекрытия между соседними чанками.
|
||||
|
||||
Возвращает:
|
||||
List[str]: список чанков.
|
||||
"""
|
||||
# Вспомогательная функция для подсчёта токенов в строке.
|
||||
def _count_tokens(s: str) -> int:
|
||||
return len(TOKENIZER.encode(s))
|
||||
|
||||
# ---- 1. Разбиение на предложения ----
|
||||
if HAS_RAZDEL:
|
||||
# razdel возвращает итератор объектов Sentence, у которых есть атрибут .text
|
||||
sentences = [sent.text for sent in sentenize(text)]
|
||||
else:
|
||||
# Fallback: разбиение по точкам, вопросительным и восклицательным знакам.
|
||||
# Это может быть неточным для аббревиатур (например, "т.е.", "и.т.д.").
|
||||
sentences = re.split(r'(?<=[.!?])\s+', text)
|
||||
# -------------------------------------
|
||||
|
||||
chunks = []
|
||||
current_chunk = [] # Список предложений текущего чанка
|
||||
current_len = 0 # Суммарное количество токенов в текущем чанке
|
||||
|
||||
for sent in sentences:
|
||||
sent_tokens = _count_tokens(sent) # Токены в текущем предложении
|
||||
|
||||
# Проверяем, помещается ли предложение в текущий чанк
|
||||
if current_len + sent_tokens <= chunk_size:
|
||||
# Помещается – добавляем
|
||||
current_chunk.append(sent)
|
||||
current_len += sent_tokens
|
||||
else:
|
||||
# Не помещается – сохраняем текущий чанк
|
||||
if current_chunk:
|
||||
chunks.append(' '.join(current_chunk))
|
||||
|
||||
# --- Формируем перекрытие для следующего чанка ---
|
||||
# Берём последние предложения текущего чанка, сумма токенов которых
|
||||
# не превышает overlap (чтобы сохранить контекст между чанками).
|
||||
overlap_sentences = []
|
||||
overlap_len = 0
|
||||
# Идём с конца текущего чанка (от последнего предложения к первому)
|
||||
for prev_sent in reversed(current_chunk):
|
||||
prev_len = _count_tokens(prev_sent)
|
||||
if overlap_len + prev_len <= overlap:
|
||||
overlap_sentences.insert(0, prev_sent) # Вставляем в начало, чтобы сохранить порядок
|
||||
overlap_len += prev_len
|
||||
else:
|
||||
break
|
||||
# -------------------------------------------------
|
||||
|
||||
# Новый чанк начинается с перекрытия + текущее предложение
|
||||
current_chunk = overlap_sentences + [sent]
|
||||
current_len = overlap_len + sent_tokens
|
||||
|
||||
# Добавляем последний чанк, если он не пуст
|
||||
if current_chunk:
|
||||
chunks.append(' '.join(current_chunk))
|
||||
|
||||
# ---- Склеиваем очень короткие чанки с предыдущим ----
|
||||
# Это предотвращает создание большого количества мелких фрагментов,
|
||||
# которые могут ухудшить качество поиска.
|
||||
merged = []
|
||||
for chunk in chunks:
|
||||
# Если текущий чанк слишком короткий (< 1/3 от chunk_size) и есть предыдущий чанк,
|
||||
# присоединяем его к предыдущему.
|
||||
if merged and _count_tokens(chunk) < chunk_size // 3:
|
||||
merged[-1] = merged[-1] + " " + chunk
|
||||
else:
|
||||
merged.append(chunk)
|
||||
# -------------------------------------------------
|
||||
|
||||
return merged
|
||||
|
||||
|
||||
def _split_by_chars(text: str, chunk_size: int, overlap: int) -> List[str]:
|
||||
"""
|
||||
Fallback-метод: разбиение по символам (при отсутствии tiktoken).
|
||||
Использует простое скользящее окно с фиксированной длиной chunk_size и перекрытием overlap.
|
||||
|
||||
Аргументы:
|
||||
text (str): исходный текст.
|
||||
chunk_size (int): размер чанка в символах.
|
||||
overlap (int): перекрытие в символах.
|
||||
|
||||
Возвращает:
|
||||
List[str]: список чанков.
|
||||
"""
|
||||
if chunk_size <= 0:
|
||||
return [text]
|
||||
|
||||
chunks = []
|
||||
start = 0
|
||||
while start < len(text):
|
||||
end = min(start + chunk_size, len(text))
|
||||
chunks.append(text[start:end])
|
||||
start = end - overlap
|
||||
if start < 0:
|
||||
start = 0
|
||||
if start >= len(text):
|
||||
break
|
||||
return chunks
|
||||
242
rag/utils/web_utils.py
Normal file
242
rag/utils/web_utils.py
Normal file
@@ -0,0 +1,242 @@
|
||||
# -*- 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 []
|
||||
Reference in New Issue
Block a user