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:
Markov Andrey
2026-06-30 10:33:28 +00:00
parent ea0d42de9c
commit 63900feece
45 changed files with 0 additions and 265 deletions

11
rag/utils/__init__.py Normal file
View 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
View 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
View 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()

View 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
View 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
View 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.53 символа.
# Используем 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
View 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 []