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
@@ -1,121 +0,0 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Экспертные команды: !summary, !metrics.
|
||||
Не зависят от комнаты (используют последний загруженный пользователем файл или общую БЗ).
|
||||
"""
|
||||
|
||||
import time
|
||||
import logging
|
||||
from core.commands.base import Command
|
||||
from core.functions.summarize_document import summarize_document
|
||||
from core.functions.extract_metrics import extract_metrics
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SummaryCommand(Command):
|
||||
name = "!summary"
|
||||
aliases = ["/summary", "!суммаризация", "/суммаризация"]
|
||||
|
||||
async def execute(self, msg, args: str, user_jid: str, room_jid: str = None):
|
||||
# Используем текст последнего загруженного файла (личный или комнатный – он хранится в last_file_texts)
|
||||
last_text = self.bot.last_file_texts.get(user_jid, "")
|
||||
if not last_text:
|
||||
reply = msg.reply("Нет документа для суммаризации.")
|
||||
if reply:
|
||||
reply.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить сообщение об отсутствии документа для суммаризации")
|
||||
return
|
||||
reply = msg.reply("📝 Создаю суммаризацию...")
|
||||
if reply:
|
||||
reply.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить сообщение о начале суммаризации")
|
||||
# Исправленный вызов: используем функцию из core.functions
|
||||
summary = await summarize_document(
|
||||
giga=self.bot.giga,
|
||||
text=last_text,
|
||||
title="Ваш документ",
|
||||
prompt_text=self.bot.ai_prompts.get('summary', ''),
|
||||
bot_config=self.bot.config
|
||||
)
|
||||
reply = msg.reply(f"📄 **Суммаризация:**\n\n{summary}")
|
||||
if reply:
|
||||
reply.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить суммаризацию")
|
||||
|
||||
def get_help(self) -> str:
|
||||
return "Создаёт суммаризацию последнего загруженного документа."
|
||||
|
||||
|
||||
class MetricsCommand(Command):
|
||||
name = "!metrics"
|
||||
aliases = ["/metrics", "!метрики", "/метрики"]
|
||||
cache_ttl = 300 # 5 минут
|
||||
|
||||
async def execute(self, msg, args: str, user_jid: str, room_jid: str = None):
|
||||
cache_key = self.get_cache_key(user_jid, args)
|
||||
now = time.time()
|
||||
if cache_key in self._cache:
|
||||
cached_time, cached_result = self._cache[cache_key]
|
||||
if now - cached_time < self.cache_ttl:
|
||||
reply = msg.reply(cached_result)
|
||||
if reply:
|
||||
reply.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить кэшированный результат метрик")
|
||||
return
|
||||
|
||||
reply = msg.reply("🔍 Ищу числовые показатели в ваших документах...")
|
||||
if reply:
|
||||
reply.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить сообщение о поиске метрик")
|
||||
query = "числовые показатели KPI метрики статистика проценты суммы количество"
|
||||
# Ищем по всей доступной БЗ (личной + глобальной или комнатной + глобальной)
|
||||
context = await self.bot.kb.find_relevant_info(query, user_jid, room_jid, top_k=30)
|
||||
if not context:
|
||||
reply = msg.reply("📭 В вашей базе знаний нет документов с числовыми показателями.")
|
||||
if reply:
|
||||
reply.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить сообщение об отсутствии метрик")
|
||||
return
|
||||
# Исправленный вызов: используем функцию из core.functions
|
||||
metrics = await extract_metrics(
|
||||
giga=self.bot.giga,
|
||||
context=context,
|
||||
prompt_text=self.bot.ai_prompts.get('metrics', ''),
|
||||
bot_config=self.bot.config
|
||||
)
|
||||
if metrics:
|
||||
lines = []
|
||||
for m in metrics[:20]:
|
||||
name = m.get('metric_name', 'Показатель')
|
||||
value = m.get('value', '?')
|
||||
unit = m.get('unit', '')
|
||||
period = m.get('period', '')
|
||||
line = f"- **{name}**: {value}"
|
||||
if unit:
|
||||
line += f" {unit}"
|
||||
if period:
|
||||
line += f" ({period})"
|
||||
lines.append(line)
|
||||
reply = "📊 **Извлечённые метрики из ваших документов:**\n" + "\n".join(lines)
|
||||
self._cache[cache_key] = (time.time(), reply)
|
||||
reply_obj = msg.reply(reply)
|
||||
if reply_obj:
|
||||
reply_obj.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить список метрик")
|
||||
else:
|
||||
reply_obj = msg.reply("⚠️ Не удалось извлечь структурированные метрики.")
|
||||
if reply_obj:
|
||||
reply_obj.send() # исправлено
|
||||
else:
|
||||
logger.error("Не удалось отправить сообщение об ошибке извлечения метрик")
|
||||
|
||||
def get_help(self) -> str:
|
||||
return "Извлекает KPI из базы знаний (с кэшем на 5 минут)."
|
||||
15
rag/functions/__init__.py
Normal file
15
rag/functions/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Пакет functions – независимые чистые функции без состояния.
|
||||
Каждая функция решает одну конкретную задачу:
|
||||
- intent_classify – классификация намерений,
|
||||
- expand_query – расширение запроса,
|
||||
- extract_metrics – извлечение KPI,
|
||||
- summarize_document – суммаризация,
|
||||
- check_consistency – проверка противоречий,
|
||||
- check_spelling – проверка орфографии,
|
||||
- critique_answer – самокритика,
|
||||
- rerank_context – переранжирование контекста,
|
||||
- file_processor – отправка файлов и хирургическая замена,
|
||||
- generate_document – генерация по шаблону.
|
||||
"""
|
||||
66
rag/functions/check_consistency.py
Normal file
66
rag/functions/check_consistency.py
Normal file
@@ -0,0 +1,66 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль проверки противоречий между несколькими фрагментами текста.
|
||||
Анализирует логические, числовые и терминологические расхождения.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт как аргумент.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def check_consistency(
|
||||
giga: GigaClient,
|
||||
chunks: List[str],
|
||||
query: str,
|
||||
prompt_text: str,
|
||||
bot_config,
|
||||
) -> str:
|
||||
"""
|
||||
Проверяет, есть ли противоречия между переданными фрагментами.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
chunks (List[str]): список текстовых фрагментов (минимум 2)
|
||||
query (str): исходный вопрос пользователя
|
||||
prompt_text (str): содержимое промпта проверки согласованности
|
||||
bot_config (BotConfig): объект конфигурации
|
||||
|
||||
Возвращает:
|
||||
str: "[OK]" или "[CONFLICT] ..." с описанием
|
||||
"""
|
||||
if len(chunks) < 2:
|
||||
return "Недостаточно фрагментов для проверки (нужно минимум 2)."
|
||||
if not prompt_text:
|
||||
logger.warning("Промпт проверки согласованности не загружен")
|
||||
return "Промпт проверки согласованности не загружен."
|
||||
|
||||
cons_cfg = getattr(bot_config, 'consistency', {})
|
||||
temperature = cons_cfg.get('temperature', 0.1)
|
||||
max_fragments = cons_cfg.get('max_fragments', 5)
|
||||
|
||||
if len(chunks) > max_fragments:
|
||||
logger.debug(f"Ограничиваем количество фрагментов с {len(chunks)} до {max_fragments}")
|
||||
chunks = chunks[:max_fragments]
|
||||
|
||||
combined = "\n\n---\n\n".join(
|
||||
[f"Фрагмент {i+1}:\n{chunk}" for i, chunk in enumerate(chunks)]
|
||||
)
|
||||
full_prompt = f"{prompt_text}\n\nЗапрос: {query}\n\nФрагменты:\n{combined}"
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=full_prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
return response.strip()
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка проверки согласованности: {e}", exc_info=True)
|
||||
return "Ошибка при проверке противоречий."
|
||||
135
rag/functions/check_spelling.py
Normal file
135
rag/functions/check_spelling.py
Normal file
@@ -0,0 +1,135 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль проверки орфографии и пунктуации в тексте документа.
|
||||
Разбивает документ на чанки, обрабатывает каждый чанк через GigaChat,
|
||||
собирает словарь замен (старое -> новое) и список изменений.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт и конфиг как аргументы.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from core.utils.text_utils import split_into_chunks
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def check_spelling(
|
||||
giga: GigaClient,
|
||||
original_text: str,
|
||||
prompt_text: str,
|
||||
config,
|
||||
) -> tuple[dict, list[str]]:
|
||||
"""
|
||||
Проверяет орфографию документа, возвращает замены и список изменений.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
original_text (str): исходный текст документа
|
||||
prompt_text (str): содержимое промпта проверки орфографии
|
||||
config (BotConfig): объект конфигурации
|
||||
|
||||
Возвращает:
|
||||
tuple: (словарь_замен, список_изменений_в_виде_строк)
|
||||
В случае, если парсинг изменений не удался, возвращается:
|
||||
({"_raw": сырой_ответ_LLM}, ["Сообщение для пользователя"])
|
||||
"""
|
||||
if not prompt_text:
|
||||
logger.error("Промпт проверки орфографии не загружен")
|
||||
return {}, ["❌ Промпт проверки орфографии не загружен."]
|
||||
|
||||
# Параметры чанкинга из конфига
|
||||
chunk_size_tokens = config.chunk_size_tokens
|
||||
overlap_tokens = config.overlap_tokens
|
||||
approx_chunk_chars = config.chunking_approx_chunk_chars
|
||||
approx_overlap_chars = config.chunking_approx_overlap_chars
|
||||
|
||||
# Разбиваем текст на чанки
|
||||
chunks = split_into_chunks(
|
||||
original_text,
|
||||
chunk_size_tokens=chunk_size_tokens,
|
||||
overlap_tokens=overlap_tokens,
|
||||
approx_chunk_chars=approx_chunk_chars,
|
||||
approx_overlap_chars=approx_overlap_chars,
|
||||
strategy="recursive_split_by_sentences"
|
||||
)
|
||||
logger.info(f"Разбито на {len(chunks)} чанков для проверки орфографии")
|
||||
|
||||
all_replacements = {}
|
||||
all_changes = []
|
||||
temperature = getattr(config, 'spellcheck_temperature', 0.1)
|
||||
|
||||
# Обрабатываем каждый чанк
|
||||
for idx, chunk in enumerate(chunks):
|
||||
logger.info(f"Обработка чанка {idx+1}/{len(chunks)}")
|
||||
full_prompt = f"{prompt_text}\n\nТЕКСТ ДЛЯ ПРОВЕРКИ (часть {idx+1} из {len(chunks)}):\n{chunk}"
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=full_prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
|
||||
corrected_match = re.search(r'\[CORRECTED\](.*?)\[/CORRECTED\]', response, re.DOTALL)
|
||||
changes_match = re.search(r'\[CHANGES\](.*?)\[/CHANGES\]', response, re.DOTALL)
|
||||
|
||||
if changes_match:
|
||||
changes_text = changes_match.group(1).strip()
|
||||
for line in changes_text.split('\n'):
|
||||
line = line.strip()
|
||||
if not line or line.startswith('#'):
|
||||
continue
|
||||
|
||||
patterns = [
|
||||
r'было\s*["«](.+?)["»]\s*[→-]>\s*стало\s*["«](.+?)["»]',
|
||||
r'было\s*:\s*["«](.+?)["»]\s*[→-]>\s*стало\s*:\s*["«](.+?)["»]',
|
||||
r'было\s*["«](.+?)["»]\s*[→-]>\s*(.+)',
|
||||
r'было\s*:\s*(.+?)\s*[→-]>\s*стало\s*:\s*(.+)',
|
||||
r'было\s*(.+?)\s*[→-]>\s*стало\s*(.+)'
|
||||
]
|
||||
match = None
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, line, re.IGNORECASE)
|
||||
if match:
|
||||
old, new = match.group(1).strip(), match.group(2).strip()
|
||||
break
|
||||
|
||||
if match:
|
||||
all_replacements[old] = new
|
||||
all_changes.append(f'было "{old}" → стало "{new}"')
|
||||
else:
|
||||
all_changes.append(line)
|
||||
|
||||
# Если список изменений не распарсился, но есть исправленный текст
|
||||
if corrected_match and (not changes_match or not changes_match.group(1).strip()):
|
||||
logger.warning(f"Не удалось распарсить список изменений для чанка {idx+1}, возвращаем сырой ответ")
|
||||
return {"_raw": response}, [
|
||||
f"⚠️ Не удалось распарсить список изменений (чанк {idx+1}). Вот сырой ответ LLM:\n{response}"
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при проверке чанка {idx+1}: {e}", exc_info=True)
|
||||
all_changes.append(f"⚠️ Ошибка в части {idx+1}: {e}")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Удаляем дубликаты
|
||||
unique_replacements = {}
|
||||
unique_changes = []
|
||||
seen = set()
|
||||
for old, new in all_replacements.items():
|
||||
key = f"{old}→{new}"
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique_replacements[old] = new
|
||||
for change in all_changes:
|
||||
if change not in seen:
|
||||
seen.add(change)
|
||||
unique_changes.append(change)
|
||||
|
||||
return unique_replacements, unique_changes
|
||||
68
rag/functions/critique_answer.py
Normal file
68
rag/functions/critique_answer.py
Normal file
@@ -0,0 +1,68 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль самокритики: оценивает качество сгенерированного ответа.
|
||||
Проверяет соответствие источнику, отсутствие галлюцинаций, соблюдение стиля и логику.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт как аргумент.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def critique_answer(
|
||||
giga: GigaClient,
|
||||
query: str,
|
||||
context: str,
|
||||
answer: str,
|
||||
prompt_text: str,
|
||||
bot_config,
|
||||
) -> bool:
|
||||
"""
|
||||
Оценивает ответ, возвращает True, если ответ приемлем, иначе False.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
query (str): исходный вопрос пользователя
|
||||
context (str): контекст из базы знаний
|
||||
answer (str): сгенерированный ответ
|
||||
prompt_text (str): содержимое промпта самокритики
|
||||
bot_config (BotConfig): объект конфигурации
|
||||
|
||||
Возвращает:
|
||||
bool: True если ответ хороший, False если есть замечания
|
||||
"""
|
||||
if not prompt_text:
|
||||
logger.debug("Промпт самокритики не загружен, пропускаем проверку")
|
||||
return True
|
||||
|
||||
crit_cfg = getattr(bot_config, 'critique', {})
|
||||
temperature = crit_cfg.get('temperature', 0.1)
|
||||
if isinstance(temperature, dict):
|
||||
temperature = 0.1
|
||||
|
||||
full_prompt = f"{prompt_text}\n\nВопрос: {query}\nКонтекст:\n{context}\nОтвет:\n{answer}"
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=full_prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
verdict = response.strip()
|
||||
|
||||
if re.search(r'\[OK\]|ПРИЕМЛЕМ', verdict, re.IGNORECASE):
|
||||
return True
|
||||
elif re.search(r'\[ISSUES\]', verdict, re.IGNORECASE):
|
||||
return False
|
||||
else:
|
||||
logger.warning(f"Самокритика вернула неопределённый ответ: {verdict[:200]}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка самокритики: {e}", exc_info=True)
|
||||
return True # При ошибке пропускаем проверку
|
||||
76
rag/functions/expand_query.py
Normal file
76
rag/functions/expand_query.py
Normal file
@@ -0,0 +1,76 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль расширения запроса ключевыми словами.
|
||||
Добавляет синонимы, аббревиатуры и связанные термины для увеличения релевантности.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт как аргумент.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_expand_cache = {}
|
||||
_CACHE_TTL = 300
|
||||
|
||||
|
||||
async def expand_query(
|
||||
giga: GigaClient,
|
||||
query: str,
|
||||
prompt_text: str,
|
||||
bot_config,
|
||||
) -> str:
|
||||
"""
|
||||
Расширяет исходный запрос пользователя набором ключевых слов.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
query (str): исходный запрос
|
||||
prompt_text (str): содержимое промпта расширения
|
||||
bot_config (BotConfig): объект конфигурации
|
||||
|
||||
Возвращает:
|
||||
str: расширенный запрос или исходный, если расширение не удалось
|
||||
"""
|
||||
if not prompt_text:
|
||||
logger.debug("Промпт расширения пуст, возвращаем исходный запрос")
|
||||
return query
|
||||
|
||||
cache_key = query.strip().lower()
|
||||
now = time.time()
|
||||
if cache_key in _expand_cache:
|
||||
cached_time, cached_result = _expand_cache[cache_key]
|
||||
if now - cached_time < _CACHE_TTL:
|
||||
logger.debug(f"Кэш расширения: '{cache_key}' -> '{cached_result}'")
|
||||
return cached_result
|
||||
else:
|
||||
del _expand_cache[cache_key]
|
||||
|
||||
expand_cfg = getattr(bot_config, 'expand', {})
|
||||
temperature = expand_cfg.get('temperature', 0.1)
|
||||
if isinstance(temperature, dict):
|
||||
temperature = 0.1
|
||||
|
||||
prompt = prompt_text.format(query=query)
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
expanded = response.strip()
|
||||
if expanded and expanded != query:
|
||||
logger.debug(f"Расширенный запрос: {expanded}")
|
||||
_expand_cache[cache_key] = (now, expanded)
|
||||
return expanded
|
||||
else:
|
||||
_expand_cache[cache_key] = (now, query)
|
||||
return query
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка расширения запроса: {e}", exc_info=True)
|
||||
return query
|
||||
84
rag/functions/extract_metrics.py
Normal file
84
rag/functions/extract_metrics.py
Normal file
@@ -0,0 +1,84 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль извлечения метрик (KPI) из текста.
|
||||
Использует промпт metrics_extract.txt.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт как аргумент.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Dict
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def extract_metrics(
|
||||
giga: GigaClient,
|
||||
context: str,
|
||||
prompt_text: str,
|
||||
bot_config,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Извлекает числовые показатели, KPI, статистические величины из контекста.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
context (str): текст для анализа
|
||||
prompt_text (str): содержимое промпта для извлечения метрик
|
||||
bot_config (BotConfig): объект конфигурации
|
||||
|
||||
Возвращает:
|
||||
List[Dict]: список метрик с полями metric_name, value, unit, period, source_fragment
|
||||
"""
|
||||
if not prompt_text:
|
||||
logger.warning("Промпт извлечения метрик пуст")
|
||||
return []
|
||||
|
||||
if not context or len(context) < 20:
|
||||
logger.debug("Контекст слишком короткий для извлечения метрик")
|
||||
return []
|
||||
|
||||
metrics_cfg = getattr(bot_config, 'metrics', {})
|
||||
temperature = metrics_cfg.get('temperature', 0.1)
|
||||
if isinstance(temperature, dict):
|
||||
temperature = 0.1
|
||||
|
||||
full_prompt = f"{prompt_text}\n\nТекст для анализа:\n{context}"
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=full_prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
raw_answer = response.strip()
|
||||
|
||||
# Пытаемся распарсить JSON
|
||||
try:
|
||||
metrics = json.loads(raw_answer)
|
||||
if isinstance(metrics, list):
|
||||
return metrics[:15]
|
||||
except json.JSONDecodeError:
|
||||
# Fallback: построчный разбор
|
||||
logger.debug("Не удалось распарсить JSON, пробуем построчный разбор")
|
||||
lines = raw_answer.split('\n')
|
||||
result = []
|
||||
for line in lines:
|
||||
if ':' in line:
|
||||
parts = line.split(':', 1)
|
||||
result.append({
|
||||
"metric_name": parts[0].strip(),
|
||||
"value": parts[1].strip(),
|
||||
"unit": "",
|
||||
"period": "",
|
||||
"source_fragment": ""
|
||||
})
|
||||
return result
|
||||
return []
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка извлечения метрик: {e}", exc_info=True)
|
||||
return []
|
||||
169
rag/functions/file_processor.py
Normal file
169
rag/functions/file_processor.py
Normal file
@@ -0,0 +1,169 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Обработка файлов: отправка пользователю через HTTP Upload и хирургическая замена в DOCX.
|
||||
Функции: send_file_to_user, handle_ai_file_request.
|
||||
|
||||
АДАПТАЦИЯ: эти функции уже работают через HTTP (загружают файлы в ejabberd).
|
||||
Они не требуют изменений, так как работают с локальными файлами.
|
||||
Но я добавлю комментарии для ясности.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import time
|
||||
import random
|
||||
import string
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def send_file_to_user(bot, msg, file_path, prefix="", suffix=None):
|
||||
"""
|
||||
Отправляет файл пользователю через HTTP Upload (XEP-0363).
|
||||
|
||||
Аргументы:
|
||||
bot: экземпляр бота
|
||||
msg: объект XMPP-сообщения
|
||||
file_path: путь к файлу на диске
|
||||
prefix: текст перед ссылкой
|
||||
suffix: суффикс для имени файла
|
||||
|
||||
Возвращает:
|
||||
bool: True при успехе, False при ошибке
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
logger.error(f"Файл не найден: {file_path}")
|
||||
return False
|
||||
|
||||
original_name = os.path.basename(file_path)
|
||||
name, ext = os.path.splitext(original_name)
|
||||
|
||||
# Санитизация суффикса (только буквы, цифры, подчёркивание, дефис)
|
||||
if suffix:
|
||||
suffix = re.sub(r'[^a-zA-Zа-яА-ЯёЁ0-9_-]', '', suffix)
|
||||
suffix = suffix.strip('.')
|
||||
if not suffix or suffix == '_':
|
||||
suffix = '_sanitized'
|
||||
if not suffix.startswith('_'):
|
||||
suffix = '_' + suffix
|
||||
new_name = f"{name}{suffix}{ext}"
|
||||
new_path = os.path.join(os.path.dirname(file_path), new_name)
|
||||
os.rename(file_path, new_path)
|
||||
file_path = new_path
|
||||
|
||||
try:
|
||||
# Загружаем файл через XEP-0363 (HTTP Upload)
|
||||
result = await bot.plugin['xep_0363'].upload_file(file_path)
|
||||
if isinstance(result, dict):
|
||||
file_url = result.get('get_url', result.get('url', str(result)))
|
||||
else:
|
||||
file_url = str(result)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка загрузки файла через XEP-0363: {e}", exc_info=True)
|
||||
error_reply = msg.reply("❌ Не удалось загрузить файл. Проверьте настройки HTTP Upload.")
|
||||
if error_reply:
|
||||
await error_reply.send()
|
||||
return False
|
||||
|
||||
reply = msg.reply(f"{prefix}\n📄 {file_url}")
|
||||
if reply:
|
||||
reply.send()
|
||||
else:
|
||||
logger.error("Не удалось создать reply для ссылки на файл")
|
||||
return True
|
||||
|
||||
|
||||
async def handle_ai_file_request(bot, msg, ai_response, user_jid, room_jid=None, body=None, suffix=None):
|
||||
"""
|
||||
Обрабатывает теги [SURGICAL_REPLACE], [REWRITE], [FILE] из ответа AI.
|
||||
|
||||
Поддерживает три режима:
|
||||
- SURGICAL: замена текста в существующем DOCX
|
||||
- REWRITE: создание нового DOCX из текста
|
||||
- FILE: создание нового DOCX (общий случай)
|
||||
"""
|
||||
key = (user_jid, room_jid) if room_jid is not None else user_jid
|
||||
|
||||
surgical_match = re.search(r'\[SURGICAL_REPLACE\](.*?)\[/SURGICAL_REPLACE\]', ai_response, re.DOTALL)
|
||||
rewrite_match = re.search(r'\[REWRITE\](.*?)\[/REWRITE\]', ai_response, re.DOTALL)
|
||||
file_match = re.search(r'\[FILE\](.*?)\[/FILE\]', ai_response, re.DOTALL)
|
||||
|
||||
# ---- SURGICAL_REPLACE ----
|
||||
if surgical_match:
|
||||
content = surgical_match.group(1).strip()
|
||||
replacements = {}
|
||||
for line in content.split('\n'):
|
||||
line = line.strip()
|
||||
if '|||' in line:
|
||||
old, new = line.split('|||', 1)
|
||||
replacements[old.strip()] = new.strip()
|
||||
if not replacements:
|
||||
reply = msg.reply("⚠️ Не найдено пар 'старое ||| новое' в теге SURGICAL_REPLACE.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
original_path = bot._last_file_path.get(key)
|
||||
if not original_path or not os.path.exists(original_path):
|
||||
reply = msg.reply("❌ Нет загруженного документа для замены. Сначала отправьте файл (без !learn).")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
new_path = await asyncio.to_thread(bot.files.surgical_replace, original_path, replacements)
|
||||
if new_path:
|
||||
await send_file_to_user(bot, msg, new_path, prefix="✅ Замена выполнена:", suffix="surgical")
|
||||
try:
|
||||
os.unlink(new_path)
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
|
||||
else:
|
||||
reply = msg.reply("❌ Ошибка при хирургической замене. Проверьте формат документа.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
# ---- REWRITE ----
|
||||
if rewrite_match:
|
||||
content = rewrite_match.group(1).strip()
|
||||
title = "Сгенерированный документ"
|
||||
title_match = re.search(r'\[TITLE\](.*?)\[/TITLE\]', ai_response, re.DOTALL)
|
||||
if title_match:
|
||||
title = title_match.group(1).strip()
|
||||
new_path = await asyncio.to_thread(bot.files.create_docx, content, title, user_jid, mode="REWRITE")
|
||||
if new_path:
|
||||
await send_file_to_user(bot, msg, new_path, prefix="📄 Документ создан:", suffix="rewrite")
|
||||
try:
|
||||
os.unlink(new_path)
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
|
||||
else:
|
||||
reply = msg.reply("❌ Не удалось создать документ.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
# ---- FILE ----
|
||||
if file_match:
|
||||
content = file_match.group(1).strip()
|
||||
title = "Файл"
|
||||
title_match = re.search(r'\[TITLE\](.*?)\[/TITLE\]', ai_response, re.DOTALL)
|
||||
if title_match:
|
||||
title = title_match.group(1).strip()
|
||||
new_path = await asyncio.to_thread(bot.files.create_docx, content, title, user_jid, mode="FILE")
|
||||
if new_path:
|
||||
await send_file_to_user(bot, msg, new_path, prefix="📄 Файл создан:", suffix="file")
|
||||
try:
|
||||
os.unlink(new_path)
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
|
||||
else:
|
||||
reply = msg.reply("❌ Не удалось создать файл.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
logger.debug(f"Ответ AI не содержит тегов файлов: {ai_response[:200]}")
|
||||
180
rag/functions/generate_document.py
Normal file
180
rag/functions/generate_document.py
Normal file
@@ -0,0 +1,180 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль генерации документа по шаблону.
|
||||
Заполняет плейсхолдеры {{...}} на основе данных из базы знаний.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт и конфиг как аргументы.
|
||||
"""
|
||||
|
||||
import re
|
||||
import asyncio
|
||||
import os
|
||||
import logging
|
||||
import shutil
|
||||
from docx import Document
|
||||
from core.functions.file_processor import send_file_to_user
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def generate_document_from_template(
|
||||
bot,
|
||||
msg,
|
||||
room_jid,
|
||||
template_path,
|
||||
template_name,
|
||||
user_jid,
|
||||
giga,
|
||||
kb,
|
||||
config,
|
||||
prompt_text,
|
||||
temperature=0.1,
|
||||
allow_public=False
|
||||
):
|
||||
"""
|
||||
Фоновая задача: парсит шаблон, ищет плейсхолдеры {{...}},
|
||||
для каждого делает контекстный поиск в БЗ комнаты, генерирует значение через LLM,
|
||||
заменяет в копии шаблона и отправляет готовый файл в комнату.
|
||||
|
||||
Аргументы:
|
||||
bot: экземпляр бота
|
||||
msg: объект XMPP-сообщения
|
||||
room_jid: JID комнаты
|
||||
template_path: путь к шаблону
|
||||
template_name: имя шаблона
|
||||
user_jid: JID пользователя
|
||||
giga: клиент GigaChat
|
||||
kb: сервис базы знаний
|
||||
config: объект конфигурации
|
||||
prompt_text: содержимое промпта для генерации
|
||||
temperature: температура генерации
|
||||
allow_public: разрешено ли использовать общие знания
|
||||
"""
|
||||
logger.info(f"🚀 [ГЕНЕРАЦИЯ] Запущена для шаблона '{template_name}' в комнате {room_jid}")
|
||||
|
||||
try:
|
||||
doc = Document(template_path)
|
||||
|
||||
# Собираем все плейсхолдеры и их контекст
|
||||
placeholders_info = {}
|
||||
for paragraph in doc.paragraphs:
|
||||
for match in re.finditer(r'\{\{([^}]+)\}\}', paragraph.text):
|
||||
ph = match.group(1)
|
||||
if ph not in placeholders_info:
|
||||
placeholders_info[ph] = paragraph.text
|
||||
|
||||
for table in doc.tables:
|
||||
for row in table.rows:
|
||||
for cell in row.cells:
|
||||
for match in re.finditer(r'\{\{([^}]+)\}\}', cell.text):
|
||||
ph = match.group(1)
|
||||
if ph not in placeholders_info:
|
||||
placeholders_info[ph] = cell.text
|
||||
|
||||
placeholders = set(placeholders_info.keys())
|
||||
logger.info(f"🔍 [ГЕНЕРАЦИЯ] Найдено плейсхолдеров: {len(placeholders)}")
|
||||
|
||||
if not placeholders:
|
||||
reply = msg.reply(f"❌ В шаблоне '{template_name}' нет плейсхолдеров {{...}}.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
generated = {}
|
||||
for ph, surrounding in placeholders_info.items():
|
||||
logger.info(f"📖 [ГЕНЕРАЦИЯ] Обработка плейсхолдера '{ph}'")
|
||||
|
||||
# Формируем поисковый запрос
|
||||
search_query = f"{ph} {surrounding[:100]}"
|
||||
context = await kb.find_relevant_info(search_query, user_jid, room_jid, top_k=20)
|
||||
|
||||
if not context and allow_public:
|
||||
logger.info(f"⚠️ [ГЕНЕРАЦИЯ] Контекст для '{ph}' не найден, используем общие знания")
|
||||
prompt = f"Ты — эксперт. Ответь на вопрос: {ph}. Дай краткое значение, без пояснений."
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
generated[ph] = response.strip()
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(f"❌ [ГЕНЕРАЦИЯ] Ошибка LLM для '{ph}': {e}")
|
||||
generated[ph] = "[ошибка]"
|
||||
continue
|
||||
|
||||
if not context:
|
||||
logger.warning(f"⚠️ [ГЕНЕРАЦИЯ] Нет контекста для '{ph}', и общие знания запрещены")
|
||||
generated[ph] = "не определено"
|
||||
continue
|
||||
|
||||
# Генерация значения на основе контекста
|
||||
prompt = prompt_text.format(ph=ph, context=context)
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
cleaned = re.sub(r'^```\w*\n?|```$', '', response.strip())
|
||||
generated[ph] = cleaned
|
||||
logger.info(f"✅ [ГЕНЕРАЦИЯ] Значение для '{ph}': {generated[ph][:100]}")
|
||||
except Exception as e:
|
||||
logger.error(f"❌ [ГЕНЕРАЦИЯ] Ошибка LLM для '{ph}': {e}", exc_info=True)
|
||||
generated[ph] = "[ошибка]"
|
||||
|
||||
# Замена плейсхолдеров в копии документа
|
||||
temp_dir = config.temp_dir
|
||||
new_path = temp_dir / f"generated_{template_name}_{int(asyncio.get_event_loop().time())}.docx"
|
||||
shutil.copy2(template_path, new_path)
|
||||
|
||||
new_doc = Document(new_path)
|
||||
|
||||
def replace_in_paragraph(paragraph, replacements):
|
||||
full_text = ''.join(run.text for run in paragraph.runs)
|
||||
modified = False
|
||||
for key, value in replacements.items():
|
||||
placeholder = f"{{{{{key}}}}}"
|
||||
if placeholder in full_text:
|
||||
for run in paragraph.runs:
|
||||
if placeholder in run.text:
|
||||
run.text = run.text.replace(placeholder, value)
|
||||
modified = True
|
||||
if not modified:
|
||||
for key, value in replacements.items():
|
||||
placeholder = f"{{{{{key}}}}}"
|
||||
if placeholder in paragraph.text:
|
||||
paragraph.text = paragraph.text.replace(placeholder, value)
|
||||
|
||||
def walk_document(doc_obj, replacements):
|
||||
for paragraph in doc_obj.paragraphs:
|
||||
replace_in_paragraph(paragraph, replacements)
|
||||
for table in doc_obj.tables:
|
||||
for row in table.rows:
|
||||
for cell in row.cells:
|
||||
for paragraph in cell.paragraphs:
|
||||
replace_in_paragraph(paragraph, replacements)
|
||||
|
||||
walk_document(new_doc, generated)
|
||||
new_doc.save(new_path)
|
||||
|
||||
await send_file_to_user(bot, msg, str(new_path), prefix=f"📄 Сгенерированный документ по шаблону '{template_name}':", suffix="generated")
|
||||
|
||||
try:
|
||||
os.unlink(new_path)
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
|
||||
|
||||
logger.info(f"🏁 [ГЕНЕРАЦИЯ] Завершено для шаблона '{template_name}'")
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"💥 [ГЕНЕРАЦИЯ] Критическая ошибка: {e}")
|
||||
if msg:
|
||||
reply = msg.reply(f"❌ Ошибка при генерации документа: {e}")
|
||||
if reply:
|
||||
reply.send()
|
||||
89
rag/functions/intent_classify.py
Normal file
89
rag/functions/intent_classify.py
Normal file
@@ -0,0 +1,89 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль классификации намерений пользователя.
|
||||
Определяет тип запроса: METRICS, SUMMARY, CONTRADICTION, SURGICAL, GREETING и др.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт как аргумент, а не использует глобальный.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Список всех допустимых кодов намерений
|
||||
VALID_INTENTS = {
|
||||
"METRICS", "SUMMARY", "CONTRADICTION", "FACT", "PROCEDURE",
|
||||
"COMPARISON", "CALCULATION", "SURGICAL", "GREETING", "TEMPLATE_FILL",
|
||||
"SPELLCHECK"
|
||||
}
|
||||
|
||||
# Кэш для классификации (in-memory)
|
||||
_intent_cache = {}
|
||||
_CACHE_TTL = 300 # 5 минут
|
||||
|
||||
|
||||
async def classify_intent(
|
||||
giga: GigaClient,
|
||||
query: str,
|
||||
prompt_text: str,
|
||||
bot_config,
|
||||
) -> str:
|
||||
"""
|
||||
Классифицирует намерение пользователя на основе текста запроса.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
query (str): исходный запрос пользователя
|
||||
prompt_text (str): содержимое промпта классификации (передаётся из оркестратора)
|
||||
bot_config (BotConfig): объект конфигурации бота
|
||||
|
||||
Возвращает:
|
||||
str: один из кодов VALID_INTENTS или "GENERAL"
|
||||
"""
|
||||
if not prompt_text:
|
||||
logger.warning("Промпт классификации пуст, возвращаем GENERAL")
|
||||
return "GENERAL"
|
||||
|
||||
# Проверка кэша
|
||||
cache_key = query.strip().lower()
|
||||
now = time.time()
|
||||
if cache_key in _intent_cache:
|
||||
cached_time, cached_result = _intent_cache[cache_key]
|
||||
if now - cached_time < _CACHE_TTL:
|
||||
logger.debug(f"Кэш классификации: '{cache_key}' -> {cached_result}")
|
||||
return cached_result
|
||||
else:
|
||||
del _intent_cache[cache_key]
|
||||
|
||||
# Получаем температуру из конфига
|
||||
intent_cfg = getattr(bot_config, 'intent', {})
|
||||
temperature = intent_cfg.get('temperature', 0.1)
|
||||
if isinstance(temperature, dict):
|
||||
temperature = 0.1
|
||||
|
||||
# Формируем полный промпт
|
||||
full_prompt = f"{prompt_text}\n\nВопрос: {query}"
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=full_prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
answer = response.strip().upper()
|
||||
|
||||
if answer in VALID_INTENTS:
|
||||
logger.debug(f"Классификация успешна: {answer}")
|
||||
_intent_cache[cache_key] = (now, answer)
|
||||
return answer
|
||||
else:
|
||||
logger.debug(f"Неизвестный ответ классификатора: {answer}, возвращаем GENERAL")
|
||||
return "GENERAL"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка классификации намерения: {e}", exc_info=True)
|
||||
return "GENERAL"
|
||||
80
rag/functions/rerank_context.py
Normal file
80
rag/functions/rerank_context.py
Normal file
@@ -0,0 +1,80 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль переранжирования контекста: отбирает только те фрагменты,
|
||||
которые действительно релевантны вопросу.
|
||||
Использует локальный кросс-энкодер.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает bot (для доступа к конфигу) и использует RerankerService.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from core.services.reranker_service import RerankerService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_reranker = None
|
||||
|
||||
|
||||
def get_reranker(bot_config) -> RerankerService:
|
||||
"""Возвращает глобальный экземпляр RerankerService (ленивая инициализация)."""
|
||||
global _reranker
|
||||
if _reranker is None:
|
||||
model_name = getattr(bot_config, 'rerank_model', None)
|
||||
device = getattr(bot_config, 'rerank_device', None)
|
||||
_reranker = RerankerService(model_name=model_name, device=device)
|
||||
logger.info(f"RerankerService инициализирован (модель: {model_name or 'default'})")
|
||||
return _reranker
|
||||
|
||||
|
||||
async def rerank_context(
|
||||
bot,
|
||||
query: str,
|
||||
context: str,
|
||||
prompt_text: str = None, # сохранён для обратной совместимости, но не используется
|
||||
bot_config=None,
|
||||
) -> str:
|
||||
"""
|
||||
Переранжирует контекст, оставляя только релевантные фрагменты.
|
||||
Использует кросс-энкодер вместо LLM.
|
||||
|
||||
Аргументы:
|
||||
bot: экземпляр Bot (для доступа к конфигу) — может быть None, если передан bot_config
|
||||
query: исходный запрос пользователя
|
||||
context: исходный контекст (абзацы, разделённые двойными переносами строки "\n\n")
|
||||
prompt_text: не используется
|
||||
bot_config: объект конфигурации (если bot не передан)
|
||||
|
||||
Возвращает:
|
||||
str: отфильтрованный контекст
|
||||
"""
|
||||
if not context:
|
||||
return context
|
||||
|
||||
# Получаем конфиг либо из bot, либо из bot_config
|
||||
if bot is not None and hasattr(bot, 'config'):
|
||||
config = bot.config
|
||||
elif bot_config is not None:
|
||||
config = bot_config
|
||||
else:
|
||||
logger.error("Ни bot, ни bot_config не переданы в rerank_context")
|
||||
return context
|
||||
|
||||
rerank_cfg = getattr(config, 'rerank', {})
|
||||
min_length = rerank_cfg.get('min_length', 5000)
|
||||
|
||||
# Если контекст короткий, переранжирование не требуется
|
||||
if len(context) < min_length:
|
||||
logger.debug(f"Контекст слишком короткий ({len(context)} < {min_length}), пропускаем переранжирование")
|
||||
return context
|
||||
|
||||
# Получаем сервис переранжирования
|
||||
reranker = get_reranker(config)
|
||||
|
||||
try:
|
||||
top_k = rerank_cfg.get('top_k', None)
|
||||
ranked_context = reranker.rerank_context(query, context, top_k=top_k)
|
||||
logger.debug(f"Переранжирование выполнено: исходный контекст {len(context)} символов")
|
||||
return ranked_context
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при переранжировании: {e}. Возвращаем исходный контекст.")
|
||||
return context
|
||||
77
rag/functions/summarize_document.py
Normal file
77
rag/functions/summarize_document.py
Normal file
@@ -0,0 +1,77 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Модуль суммаризации документа.
|
||||
Создаёт краткий пересказ, сохраняя ключевые факты и цифры.
|
||||
|
||||
АДАПТАЦИЯ: теперь функция принимает промпт как аргумент.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from core.services.giga_client import GigaClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def summarize_document(
|
||||
giga: GigaClient,
|
||||
text: str,
|
||||
title: str,
|
||||
prompt_text: str,
|
||||
bot_config,
|
||||
) -> str:
|
||||
"""
|
||||
Суммаризирует документ.
|
||||
|
||||
Аргументы:
|
||||
giga (GigaClient): клиент GigaChat
|
||||
text (str): полный текст документа
|
||||
title (str): название документа
|
||||
prompt_text (str): содержимое промпта суммаризации
|
||||
bot_config (BotConfig): объект конфигурации
|
||||
|
||||
Возвращает:
|
||||
str: краткое изложение документа или сообщение об ошибке
|
||||
"""
|
||||
if not prompt_text:
|
||||
logger.error("Промпт суммаризации не загружен")
|
||||
return "Не удалось загрузить промпт для суммаризации."
|
||||
|
||||
summary_cfg = getattr(bot_config, 'summary', {})
|
||||
temperature = summary_cfg.get('temperature', 0.1)
|
||||
max_chars = summary_cfg.get('max_chars', 8000)
|
||||
|
||||
# Обрезаем текст, если он слишком длинный
|
||||
if len(text) > max_chars:
|
||||
# Обрезаем по границам предложений (используем razdel)
|
||||
try:
|
||||
from razdel import sentenize
|
||||
sentences = [sent.text for sent in sentenize(text)]
|
||||
except ImportError:
|
||||
import re
|
||||
sentences = re.split(r'(?<=[.!?])\s+', text)
|
||||
|
||||
truncated = []
|
||||
current_len = 0
|
||||
for sent in sentences:
|
||||
if current_len + len(sent) <= max_chars:
|
||||
truncated.append(sent)
|
||||
current_len += len(sent)
|
||||
else:
|
||||
break
|
||||
text = ' '.join(truncated)
|
||||
logger.warning(f"Текст обрезан до {max_chars} символов по границам предложений")
|
||||
|
||||
full_prompt = f"{prompt_text}\n\nНазвание документа: {title}\nТекст:\n{text}"
|
||||
|
||||
try:
|
||||
response = await giga.chat(
|
||||
history=[],
|
||||
query=full_prompt,
|
||||
system_prompt=None,
|
||||
file_id=None,
|
||||
temperature=temperature
|
||||
)
|
||||
return response.strip()
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка суммаризации: {e}", exc_info=True)
|
||||
return "Не удалось создать суммаризацию."
|
||||
12
rag/services/__init__.py
Normal file
12
rag/services/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Пакет services – долгоживущие компоненты для работы с внешними системами.
|
||||
Содержит:
|
||||
- PostgresService – пул соединений с PostgreSQL,
|
||||
- QdrantService – векторная БД,
|
||||
- EmbeddingService – эмбеддинги GigaChat,
|
||||
- KBService – база знаний (индексация, поиск),
|
||||
- FileService – извлечение текста из файлов,
|
||||
- GigaClient – клиент GigaChat,
|
||||
- RerankerService – локальный кросс-энкодер для переранжирования.
|
||||
"""
|
||||
195
rag/services/embedding_service.py
Normal file
195
rag/services/embedding_service.py
Normal file
@@ -0,0 +1,195 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис генерации эмбеддингов через GigaChat API.
|
||||
Поддерживает одиночные запросы с кэшированием (lru_cache) и пакетную обработку.
|
||||
Все параметры (модель, таймаут, размер кэша, verify_ssl) вынесены в конфиг.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import List, Optional, Any
|
||||
|
||||
from gigachat import GigaChat
|
||||
from gigachat.exceptions import AuthenticationError, RateLimitError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from fastembed import SparseTextEmbedding
|
||||
HAS_SPARSE = True
|
||||
except ImportError:
|
||||
HAS_SPARSE = False
|
||||
logging.warning("fastembed не установлен, sparse-эмбеддинги отключены")
|
||||
|
||||
class EmbeddingService:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "GigaChat-Embeddings",
|
||||
timeout: int = 30,
|
||||
cache_size: int = 4096,
|
||||
verify_ssl: bool = False
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.cache_size = cache_size
|
||||
self.verify_ssl = verify_ssl
|
||||
self._giga = None
|
||||
# Кэширование одиночных эмбеддингов (синхронная обёртка с lru_cache)
|
||||
self.embed_sync = lru_cache(maxsize=cache_size)(self._embed_sync_impl)
|
||||
self._sparse_model = None
|
||||
|
||||
def _get_client(self) -> GigaChat:
|
||||
"""Ленивое создание клиента GigaChat (с нужными параметрами)."""
|
||||
if self._giga is None:
|
||||
self._giga = GigaChat(
|
||||
credentials=self.api_key,
|
||||
model=self.model,
|
||||
verify_ssl_certs=self.verify_ssl,
|
||||
timeout=self.timeout
|
||||
)
|
||||
logger.debug(f"Клиент GigaChat создан (модель {self.model}, timeout {self.timeout}s)")
|
||||
return self._giga
|
||||
|
||||
def _get_sparse_model(self):
|
||||
if not HAS_SPARSE:
|
||||
return None
|
||||
if self._sparse_model is None:
|
||||
self._sparse_model = SparseTextEmbedding(model_name="Qdrant/bm25")
|
||||
return self._sparse_model
|
||||
|
||||
# ---------- Синхронное получение эмбеддинга (с обрезкой) ----------
|
||||
def _embed_sync_impl(self, text: str) -> Optional[List[float]]:
|
||||
"""Синхронный вызов GigaChat API для одного текста. Результат кэшируется."""
|
||||
if not text or not text.strip():
|
||||
logger.warning("Попытка получить эмбеддинг пустого текста")
|
||||
return None
|
||||
|
||||
# Обрезка до 800 символов (безопасно для лимита 514 токенов GigaChat)
|
||||
MAX_CHARS = 800
|
||||
if len(text) > MAX_CHARS:
|
||||
text = text[:MAX_CHARS]
|
||||
logger.debug(f"Текст обрезан до {MAX_CHARS} символов для эмбеддинга")
|
||||
|
||||
client = self._get_client()
|
||||
try:
|
||||
# Совместимость с разными версиями SDK: пробуем input=[text] или просто [text]
|
||||
try:
|
||||
response = client.embeddings(input=[text])
|
||||
except TypeError:
|
||||
response = client.embeddings([text])
|
||||
if response.data and len(response.data) > 0:
|
||||
logger.debug(f"Эмбеддинг получен, размерность {len(response.data[0].embedding)}")
|
||||
return response.data[0].embedding
|
||||
else:
|
||||
logger.error(f"Пустой ответ эмбеддингов для текста: {text[:100]}")
|
||||
return None
|
||||
except (AuthenticationError, RateLimitError) as e:
|
||||
logger.error(f"Ошибка авторизации/лимитов GigaChat: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка GigaChat при эмбеддинге: {e}")
|
||||
raise
|
||||
|
||||
async def embed(self, text: str, retries: int = 3, delay: float = 1.0) -> Optional[List[float]]:
|
||||
"""Асинхронная обёртка с повторными попытками."""
|
||||
if not isinstance(text, str):
|
||||
logger.error(f"embed получил не строку, а {type(text)}: {text!r}")
|
||||
return None
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
result = await loop.run_in_executor(None, self.embed_sync, text)
|
||||
if result is not None:
|
||||
return result
|
||||
else:
|
||||
raise ValueError("Эмбеддинг вернул None")
|
||||
except Exception as e:
|
||||
wait = delay * (2 ** attempt)
|
||||
logger.warning(
|
||||
f"Попытка {attempt+1}/{retries} получить эмбеддинг не удалась: {e}. "
|
||||
f"Повтор через {wait:.1f} сек."
|
||||
)
|
||||
await asyncio.sleep(wait)
|
||||
logger.error(f"Не удалось получить эмбеддинг после {retries} попыток")
|
||||
return None
|
||||
|
||||
# ---------- Пакетная обработка (для индексации) ----------
|
||||
def _embed_batch_sync(self, texts: List[str]) -> Optional[List[List[float]]]:
|
||||
if not texts:
|
||||
return []
|
||||
MAX_CHARS = 800
|
||||
truncated_texts = [t[:MAX_CHARS] if t and len(t) > MAX_CHARS else (t or "") for t in texts]
|
||||
client = self._get_client()
|
||||
for attempt in range(5):
|
||||
try:
|
||||
try:
|
||||
response = client.embeddings(input=truncated_texts)
|
||||
except TypeError:
|
||||
response = client.embeddings(truncated_texts)
|
||||
if response.data and len(response.data) == len(truncated_texts):
|
||||
sorted_embeddings = [None] * len(truncated_texts)
|
||||
for item in response.data:
|
||||
idx = item.index
|
||||
if idx < len(sorted_embeddings):
|
||||
sorted_embeddings[idx] = item.embedding
|
||||
return sorted_embeddings
|
||||
else:
|
||||
logger.error(f"Неполный ответ эмбеддингов: ожидалось {len(truncated_texts)}, получено {len(response.data) if response.data else 0}")
|
||||
return None
|
||||
except RateLimitError as e:
|
||||
if attempt < 4:
|
||||
wait = 2 ** attempt # 1,2,4,8 секунд
|
||||
logger.warning(f"Rate limit, повтор через {wait} секунд")
|
||||
time.sleep(wait) # синхронный sleep, так как мы в синхронной функции
|
||||
else:
|
||||
logger.error("Превышено количество попыток при 429")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка пакетного эмбеддинга: {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
async def embed_batch(self, texts: List[str], retries: int = 3, delay: float = 1.0) -> Optional[List[List[float]]]:
|
||||
if not texts:
|
||||
return []
|
||||
BATCH_SIZE = 50
|
||||
results = []
|
||||
for i in range(0, len(texts), BATCH_SIZE):
|
||||
batch = texts[i:i+BATCH_SIZE]
|
||||
loop = asyncio.get_running_loop()
|
||||
vectors = await loop.run_in_executor(None, self._embed_batch_sync, batch)
|
||||
if vectors is None:
|
||||
for attempt in range(retries):
|
||||
await asyncio.sleep(delay * (2 ** attempt))
|
||||
vectors = await loop.run_in_executor(None, self._embed_batch_sync, batch)
|
||||
if vectors is not None:
|
||||
break
|
||||
else:
|
||||
return None
|
||||
results.extend(vectors)
|
||||
# ДОБАВИТЬ: пауза между батчами
|
||||
await asyncio.sleep(0.5) # 500 мс
|
||||
return results
|
||||
|
||||
async def embed_sparse_batch(self, texts: List[str]) -> List[Optional[Any]]:
|
||||
"""Возвращает список объектов SparseEmbedding для каждого текста"""
|
||||
model = self._get_sparse_model()
|
||||
if model is None:
|
||||
# Возвращаем список из None такой же длины, чтобы не нарушать индексацию
|
||||
return [None] * len(texts)
|
||||
loop = asyncio.get_running_loop()
|
||||
embeddings = await loop.run_in_executor(None, list, model.embed(texts))
|
||||
return embeddings
|
||||
|
||||
async def embed_sparse(self, text: str) -> Optional[Any]:
|
||||
model = self._get_sparse_model()
|
||||
if model is None:
|
||||
return None
|
||||
loop = asyncio.get_running_loop()
|
||||
embeddings = await loop.run_in_executor(None, lambda: list(model.embed([text])))
|
||||
if embeddings and len(embeddings) > 0:
|
||||
return embeddings[0]
|
||||
return None
|
||||
658
rag/services/file_service.py
Normal file
658
rag/services/file_service.py
Normal file
@@ -0,0 +1,658 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис обработки файлов: извлечение текста из различных форматов,
|
||||
хирургическая замена в DOCX, создание DOCX из текста,
|
||||
транскрипция аудио через SaluteSpeech.
|
||||
Все методы синхронные (используются в asyncio.to_thread).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import zipfile
|
||||
from datetime import datetime
|
||||
from typing import List, Tuple, Optional, Dict, Any
|
||||
|
||||
import requests
|
||||
from docx import Document as DocxReader
|
||||
from docx import Document as DocxWriter
|
||||
from docx.shared import Pt
|
||||
|
||||
# Опциональные библиотеки
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
HAS_PDF = True
|
||||
except ImportError:
|
||||
HAS_PDF = False
|
||||
|
||||
try:
|
||||
import openpyxl
|
||||
HAS_EXCEL = True
|
||||
except ImportError:
|
||||
HAS_EXCEL = False
|
||||
|
||||
try:
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
HAS_OCR = True
|
||||
except ImportError:
|
||||
HAS_OCR = False
|
||||
|
||||
try:
|
||||
import py7zr
|
||||
HAS_7Z = True
|
||||
except ImportError:
|
||||
HAS_7Z = False
|
||||
|
||||
try:
|
||||
from pptx import Presentation
|
||||
HAS_PPTX = True
|
||||
except ImportError:
|
||||
HAS_PPTX = False
|
||||
|
||||
# Для чанкинга используем tiktoken, но здесь оставим только извлечение текста,
|
||||
# а чанкинг вынесен в core/utils/text_utils.py
|
||||
from core.utils.config_loader import BotConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileService:
|
||||
def __init__(self, config: BotConfig):
|
||||
self.config = config
|
||||
self.temp_dir = config.temp_dir
|
||||
self.temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def cleanup_old_files(self, days: int = 3):
|
||||
"""Удаляет старые файлы из temp_dir."""
|
||||
now = datetime.now().timestamp()
|
||||
for item in self.temp_dir.iterdir():
|
||||
if item.is_file() and (now - item.stat().st_mtime) > days * 86400:
|
||||
try:
|
||||
item.unlink()
|
||||
logger.info(f"Удалён старый файл: {item.name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка удаления {item.name}: {e}")
|
||||
|
||||
# ---------- Извлечение текста из разных форматов ----------
|
||||
def extract_text_from_docx(self, file_path: str) -> str:
|
||||
"""
|
||||
Извлекает текст из DOCX, включая колонтитулы, таблицы и сноски.
|
||||
Возвращает текст, разделённый переносами строк.
|
||||
"""
|
||||
try:
|
||||
doc = DocxReader(file_path)
|
||||
full_text = []
|
||||
|
||||
# Вспомогательная функция для извлечения текста из XML-элемента
|
||||
def get_xml_text(element):
|
||||
text_nodes = element.xpath('.//*[local-name()="t"]')
|
||||
return "".join([node.text for node in text_nodes if node.text])
|
||||
|
||||
# Колонтитулы
|
||||
for section in doc.sections:
|
||||
for header in [section.header, section.first_page_header, section.even_page_header]:
|
||||
if header:
|
||||
h_text = get_xml_text(header._element)
|
||||
if h_text.strip():
|
||||
full_text.append(f"[КОЛОНТИТУЛ ВЕРХНИЙ]: {h_text}")
|
||||
for footer in [section.footer, section.first_page_footer, section.even_page_footer]:
|
||||
if footer:
|
||||
f_text = get_xml_text(footer._element)
|
||||
if f_text.strip():
|
||||
full_text.append(f"[КОЛОНТИТУЛ НИЖНИЙ]: {f_text}")
|
||||
|
||||
# Тело документа
|
||||
for block in doc.element.body:
|
||||
if block.tag.endswith('p'):
|
||||
text = get_xml_text(block)
|
||||
if text.strip():
|
||||
full_text.append(text)
|
||||
elif block.tag.endswith('tbl'):
|
||||
from docx.table import Table
|
||||
t = Table(block, doc)
|
||||
full_text.append("\n[ТАБЛИЦА]")
|
||||
for row in t.rows:
|
||||
row_data = [get_xml_text(cell._element).strip().replace('\n', ' ') for cell in row.cells]
|
||||
# Удаляем дублирование, если ячейки пустые
|
||||
clean_row = []
|
||||
for item in row_data:
|
||||
if not clean_row or item != clean_row[-1]:
|
||||
clean_row.append(item)
|
||||
full_text.append(" | ".join(clean_row))
|
||||
full_text.append("[КОНЕЦ ТАБЛИЦЫ]\n")
|
||||
|
||||
# Сноски
|
||||
try:
|
||||
for footnote in doc.footnotes:
|
||||
fn_text = get_xml_text(footnote._element)
|
||||
if fn_text.strip():
|
||||
full_text.append(f"[СНОСКА]: {fn_text}")
|
||||
except:
|
||||
pass
|
||||
|
||||
return "\n".join(full_text)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения DOCX {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_pdf(self, file_path: str) -> str:
|
||||
"""Извлекает текст из PDF через pypdf."""
|
||||
if not HAS_PDF:
|
||||
return "[Ошибка: библиотека pypdf не установлена]"
|
||||
try:
|
||||
reader = PdfReader(file_path)
|
||||
pages = []
|
||||
for i, page in enumerate(reader.pages):
|
||||
text = page.extract_text()
|
||||
if text.strip():
|
||||
pages.append(f"--- Страница {i+1} ---\n{text}")
|
||||
return "\n".join(pages)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения PDF {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_xlsx(self, file_path: str) -> str:
|
||||
"""Извлекает текст из Excel (XLSX) через openpyxl."""
|
||||
if not HAS_EXCEL:
|
||||
return "[Ошибка: библиотека openpyxl не установлена]"
|
||||
try:
|
||||
wb = openpyxl.load_workbook(file_path, data_only=True)
|
||||
output = []
|
||||
for sheet in wb.sheetnames:
|
||||
ws = wb[sheet]
|
||||
output.append(f"\n--- ЛИСТ: {sheet} ---")
|
||||
for row in ws.iter_rows(values_only=True):
|
||||
if any(c is not None for c in row):
|
||||
output.append(" | ".join(str(c) if c is not None else "" for c in row))
|
||||
return "\n".join(output)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения XLSX {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_csv(self, file_path: str) -> str:
|
||||
"""Извлекает текст из CSV."""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
reader = csv.reader(f)
|
||||
return "\n".join("\t".join(row) for row in reader)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения CSV {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_txt(self, file_path: str) -> str:
|
||||
"""Извлекает текст из TXT."""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
return f.read()
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения TXT {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_json(self, file_path: str) -> str:
|
||||
"""Извлекает текст из JSON (дамп в виде строки)."""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения JSON {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_image(self, file_path: str) -> str:
|
||||
"""Извлекает текст из изображения через OCR (Tesseract)."""
|
||||
if not HAS_OCR:
|
||||
return ""
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
text = pytesseract.image_to_string(img, lang='rus+eng')
|
||||
return f"--- ТЕКСТ С ИЗОБРАЖЕНИЯ ---\n{text}" if text.strip() else ""
|
||||
except Exception as e:
|
||||
logger.debug(f"OCR не удался: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_pptx(self, file_path: str) -> str:
|
||||
"""
|
||||
Извлекает текст из презентации .pptx:
|
||||
- текст всех слайдов (из текстовых полей, таблиц)
|
||||
- заметки к слайдам
|
||||
Возвращает строку с разделением по слайдам.
|
||||
"""
|
||||
if not HAS_PPTX:
|
||||
return "[Ошибка: библиотека python-pptx не установлена]"
|
||||
try:
|
||||
prs = Presentation(file_path)
|
||||
full_text = []
|
||||
for slide_num, slide in enumerate(prs.slides, 1):
|
||||
slide_content = []
|
||||
# Текст из фигур (надписи, автофигуры)
|
||||
for shape in slide.shapes:
|
||||
if hasattr(shape, "text") and shape.text:
|
||||
slide_content.append(shape.text.strip())
|
||||
# Таблицы
|
||||
if shape.has_table:
|
||||
for row in shape.table.rows:
|
||||
row_text = [cell.text.strip() for cell in row.cells]
|
||||
slide_content.append(" | ".join(row_text))
|
||||
# Заметки к слайду
|
||||
if slide.has_notes_slide and slide.notes_slide.notes_text_frame:
|
||||
notes = slide.notes_slide.notes_text_frame.text.strip()
|
||||
if notes:
|
||||
slide_content.append(f"[ЗАМЕТКИ]: {notes}")
|
||||
if slide_content:
|
||||
full_text.append(f"--- Слайд {slide_num} ---\n" + "\n".join(slide_content))
|
||||
return "\n\n".join(full_text)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения PPTX {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
async def transcribe_audio(self, file_path: str) -> str:
|
||||
"""
|
||||
Асинхронно транскрибирует аудиофайл через SaluteSpeech API.
|
||||
Все блокирующие вызовы (subprocess.run, requests.post) обёрнуты в asyncio.to_thread.
|
||||
Требует наличия ffmpeg и переменной окружения SALUTE_SPEECH_AUTH.
|
||||
"""
|
||||
import uuid
|
||||
|
||||
# 1. Получаем токен асинхронно
|
||||
token = await self._get_salute_token()
|
||||
if not token:
|
||||
return "[Ошибка авторизации в SaluteSpeech]"
|
||||
|
||||
# 2. Проверяем наличие ffmpeg
|
||||
ffmpeg_bin = shutil.which('ffmpeg') or '/usr/local/bin/ffmpeg'
|
||||
curl_bin = shutil.which('curl') or '/usr/local/bin/curl'
|
||||
if not shutil.which('ffmpeg') and not os.path.exists(ffmpeg_bin):
|
||||
return "[Ошибка: на сервере не установлен ffmpeg]"
|
||||
|
||||
pcm_output = file_path + ".raw"
|
||||
|
||||
try:
|
||||
# ---- ОБЁРТКА: конвертация ffmpeg (subprocess.run) ----
|
||||
def _run_ffmpeg():
|
||||
cmd_ffmpeg = [
|
||||
ffmpeg_bin, '-y', '-i', file_path,
|
||||
'-ac', '1', '-ar', '16000', '-f', 's16le', '-filter:a', 'volume=2.0', pcm_output
|
||||
]
|
||||
return subprocess.run(cmd_ffmpeg, capture_output=True)
|
||||
|
||||
result = await asyncio.to_thread(_run_ffmpeg)
|
||||
if result.returncode != 0:
|
||||
err_msg = result.stderr.decode()
|
||||
logger.error(f"FFmpeg ошибка: {err_msg}")
|
||||
return f"[Ошибка конвертации аудио: {err_msg[:100]}]"
|
||||
|
||||
if not os.path.exists(pcm_output) or os.path.getsize(pcm_output) == 0:
|
||||
return "[Ошибка: пустой аудиофайл после конвертации]"
|
||||
|
||||
# 3. Отправка запроса к SaluteSpeech через curl (subprocess.run)
|
||||
url = "https://smartspeech.sber.ru/rest/v1/speech:recognize"
|
||||
|
||||
def _run_curl():
|
||||
curl_cmd = [
|
||||
curl_bin, '-s', '-k', '--http1.1', '-X', 'POST', url,
|
||||
'-H', f'Authorization: Bearer {token}',
|
||||
'-H', 'Content-Type: audio/x-pcm;bit=16;rate=16000',
|
||||
'-H', f'RqUID: {str(uuid.uuid4())}',
|
||||
'--data-binary', f'@{pcm_output}'
|
||||
]
|
||||
return subprocess.run(curl_cmd, capture_output=True, text=True)
|
||||
|
||||
result_curl = await asyncio.to_thread(_run_curl)
|
||||
|
||||
# Удаляем временный PCM-файл
|
||||
if os.path.exists(pcm_output):
|
||||
os.remove(pcm_output)
|
||||
|
||||
if result_curl.returncode != 0:
|
||||
return f"[Ошибка сетевого соединения (curl): {result_curl.returncode}]"
|
||||
|
||||
# 4. Парсинг ответа (синхронная операция – быстрая, можно оставить без обёртки)
|
||||
try:
|
||||
response_data = json.loads(result_curl.stdout)
|
||||
if 'result' in response_data:
|
||||
if response_data['result']:
|
||||
text = " ".join(response_data['result'])
|
||||
return f"--- РАСПОЗНАННЫЙ ГОЛОС ---\n{text}"
|
||||
else:
|
||||
return "😶 Тишина на записи."
|
||||
error_msg = response_data.get('message', 'Неизвестная ошибка API')
|
||||
return f"--- ОШИБКА ГОЛОСА ---\n[Голос не распознан: {error_msg}]"
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка парсинга JSON: {result_curl.stdout}")
|
||||
return f"[Ошибка парсинга ответа: {e}]"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка в transcribe_audio: {e}")
|
||||
return f"[Ошибка распознавания: {e}]"
|
||||
|
||||
async def _get_salute_token(self) -> Optional[str]:
|
||||
"""
|
||||
Асинхронно получает токен доступа к SaluteSpeech по Basic авторизации.
|
||||
Блокирующий вызов requests.post обёрнут в asyncio.to_thread.
|
||||
"""
|
||||
import uuid
|
||||
salute_auth = self.config.salute_speech_auth
|
||||
if not salute_auth:
|
||||
logger.error("Переменная окружения SALUTE_SPEECH_AUTH не задана")
|
||||
return None
|
||||
url = "https://ngw.devices.sberbank.ru.:9443/api/v2/oauth"
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'RqUID': str(uuid.uuid4()),
|
||||
'Authorization': f'Basic {salute_auth}'
|
||||
}
|
||||
payload = {'scope': 'SALUTE_SPEECH_PERS'}
|
||||
try:
|
||||
# ---- ОБЁРТКА: синхронный requests.post выполняется в отдельном потоке ----
|
||||
def _sync_request():
|
||||
return requests.post(url, headers=headers, data=payload, verify=False)
|
||||
response = await asyncio.to_thread(_sync_request)
|
||||
token = response.json().get('access_token')
|
||||
logger.info("SaluteSpeech token получен")
|
||||
return token
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка получения токена SaluteSpeech: {e}")
|
||||
return None
|
||||
|
||||
def process_archive(self, file_path: str) -> Tuple[List[Tuple[str, str]], int]:
|
||||
"""
|
||||
Распаковывает ZIP/7z и рекурсивно обрабатывает каждый файл.
|
||||
Возвращает список (имя_файла, извлечённый_текст) и общее количество обработанных файлов.
|
||||
"""
|
||||
results = []
|
||||
total_count = 0
|
||||
extract_path = self.temp_dir / f"ext_{int(datetime.now().timestamp()*1000000)}"
|
||||
extract_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
max_total_bytes = self.config.max_file_size_mb * 2 * 1024 * 1024 # для архивов лимит удвоен
|
||||
MAX_FILES = self.config.max_archive_files
|
||||
total_files = 0
|
||||
total_size = 0
|
||||
|
||||
try:
|
||||
# ---------- ZIP ----------
|
||||
if file_path.endswith('.zip'):
|
||||
with zipfile.ZipFile(file_path, 'r') as z:
|
||||
# 1. Предварительный подсчёт суммарного несжатого размера всех файлов в архиве
|
||||
total_uncompressed = 0
|
||||
for info in z.infolist():
|
||||
if not info.is_dir():
|
||||
total_uncompressed += info.file_size
|
||||
if total_uncompressed > max_total_bytes:
|
||||
logger.warning(
|
||||
f"Архив {file_path} превышает лимит {max_total_bytes // (1024*1024)} МБ "
|
||||
f"(несжатый размер {total_uncompressed // (1024*1024)} МБ)"
|
||||
)
|
||||
return [], 0
|
||||
|
||||
# 2. Распаковка и обработка
|
||||
for info in z.infolist():
|
||||
if info.is_dir():
|
||||
continue
|
||||
total_files += 1
|
||||
if total_files > MAX_FILES:
|
||||
logger.warning(f"Архив {file_path} содержит более {MAX_FILES} файлов, обработка прервана")
|
||||
return [], 0
|
||||
|
||||
# Имя файла в правильной кодировке
|
||||
try:
|
||||
fname = info.filename.encode('cp437').decode('utf-8')
|
||||
except:
|
||||
try:
|
||||
fname = info.filename.encode('cp437').decode('cp866')
|
||||
except:
|
||||
fname = info.filename
|
||||
|
||||
target = extract_path / os.path.basename(fname)
|
||||
with open(target, 'wb') as f:
|
||||
f.write(z.read(info))
|
||||
|
||||
text, count = self.process_any_file(str(target))
|
||||
if text:
|
||||
size = len(text.encode('utf-8'))
|
||||
if total_size + size > max_total_bytes:
|
||||
logger.warning(f"Архив {file_path} превышает лимит {self.config.max_file_size_mb*2} МБ")
|
||||
return [], 0
|
||||
results.append((fname, text))
|
||||
total_count += count
|
||||
total_size += size
|
||||
|
||||
# ---------- 7Z ----------
|
||||
elif file_path.endswith('.7z') and HAS_7Z:
|
||||
with py7zr.SevenZipFile(file_path, 'r') as z:
|
||||
# Для 7z нет простого способа получить суммарный размер до распаковки,
|
||||
# поэтому распаковываем с контролем в процессе
|
||||
z.extractall(extract_path)
|
||||
for root, dirs, files in os.walk(extract_path):
|
||||
for f in files:
|
||||
total_files += 1
|
||||
if total_files > MAX_FILES:
|
||||
logger.warning(f"Архив {file_path} содержит более {MAX_FILES} файлов, обработка прервана")
|
||||
return [], 0
|
||||
fpath = os.path.join(root, f)
|
||||
text, count = self.process_any_file(fpath)
|
||||
if text:
|
||||
size = len(text.encode('utf-8'))
|
||||
if total_size + size > max_total_bytes:
|
||||
logger.warning(f"Архив {file_path} превышает лимит {self.config.max_file_size_mb*2} МБ")
|
||||
return [], 0
|
||||
results.append((f, text))
|
||||
total_count += count
|
||||
total_size += size
|
||||
else:
|
||||
# Неподдерживаемый формат архива (но по логике сюда не должны попасть)
|
||||
logger.warning(f"Неподдерживаемый тип архива: {file_path}")
|
||||
return [], 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка распаковки архива {file_path}: {e}")
|
||||
return [], 0
|
||||
finally:
|
||||
shutil.rmtree(extract_path, ignore_errors=True)
|
||||
|
||||
return results, total_count
|
||||
|
||||
def process_any_file(self, file_path: str) -> Tuple[str, int]:
|
||||
"""
|
||||
Главный метод: определяет тип файла по расширению и вызывает соответствующий
|
||||
метод извлечения текста. Возвращает (текст, количество_файлов_внутри).
|
||||
Для архивов возвращает список файлов и общее количество.
|
||||
"""
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
text = ""
|
||||
count = 0
|
||||
try:
|
||||
size = os.path.getsize(file_path)
|
||||
max_size = self.config.max_file_size_mb * 1024 * 1024
|
||||
if size > max_size:
|
||||
logger.warning(f"Файл {file_path} превышает лимит {self.config.max_file_size_mb} МБ")
|
||||
return f"[Ошибка: файл превышает максимальный размер {self.config.max_file_size_mb} МБ]", 0
|
||||
if ext == '.docx':
|
||||
text = self.extract_text_from_docx(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.pdf':
|
||||
text = self.extract_text_from_pdf(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.xlsx':
|
||||
text = self.extract_text_from_xlsx(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.csv':
|
||||
text = self.extract_text_from_csv(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.txt':
|
||||
text = self.extract_text_from_txt(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.json':
|
||||
text = self.extract_text_from_json(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff'):
|
||||
text = self.extract_text_from_image(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext in ('.ogg', '.wav', '.mp3', '.amr', '.m4a'):
|
||||
# Асинхронный метод transcribe_audio запускаем синхронно через get_event_loop
|
||||
# (так как process_any_file выполняется в отдельном потоке)
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
text = loop.run_until_complete(self.transcribe_audio(file_path))
|
||||
count = 1 if text else 0
|
||||
elif ext in ('.zip', '.7z'):
|
||||
return self.process_archive(file_path)
|
||||
elif ext == '.pptx':
|
||||
text = self.extract_text_from_pptx(file_path)
|
||||
count = 1 if text else 0
|
||||
else:
|
||||
logger.warning(f"Неподдерживаемый тип файла: {ext}")
|
||||
return "", 0
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка обработки файла {file_path}: {e}")
|
||||
return "", 0
|
||||
return text, count
|
||||
|
||||
def surgical_replace(self, file_path: str, replacements: Dict[str, str]) -> Optional[str]:
|
||||
"""
|
||||
Умная замена слов в DOCX с учётом грамматики (падеж, число, род).
|
||||
Сохраняет форматирование. При ошибке выполняет простую замену.
|
||||
"""
|
||||
import re
|
||||
from datetime import datetime
|
||||
from docx import Document
|
||||
|
||||
# ---- Попытка умной замены через mawo-pymorphy3 ----
|
||||
try:
|
||||
from mawo_pymorphy3 import create_analyzer
|
||||
morph = create_analyzer()
|
||||
logger.info("Морфологический анализатор загружен")
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось загрузить mawo-pymorphy3: {e}. Выполняю простую замену.")
|
||||
# Простая замена
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
new_path = self.temp_dir / f"surgical_{int(datetime.now().timestamp())}.docx"
|
||||
for paragraph in doc.paragraphs:
|
||||
for run in paragraph.runs:
|
||||
for old, new in replacements.items():
|
||||
if old in run.text:
|
||||
run.text = run.text.replace(old, new)
|
||||
doc.save(str(new_path))
|
||||
logger.info(f"Простая замена выполнена: {new_path}")
|
||||
return str(new_path)
|
||||
except Exception as e2:
|
||||
logger.error(f"Ошибка простой замены: {e2}")
|
||||
return None
|
||||
|
||||
# ---- Умная замена ----
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
logger.info(f"Документ загружен, параграфов: {len(doc.paragraphs)}")
|
||||
|
||||
# Кэш для разобранных слов
|
||||
parse_cache = {}
|
||||
|
||||
# Предвычисление форм старых слов
|
||||
old_forms = {}
|
||||
for old_word, new_word in replacements.items():
|
||||
old_low = old_word.lower()
|
||||
parsed = morph.parse(old_low)[0]
|
||||
forms = {parsed.normal_form}
|
||||
cases = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
|
||||
numbers = ['sing', 'plur']
|
||||
for number in numbers:
|
||||
for case in cases:
|
||||
inflected = parsed.inflect({case, number})
|
||||
if inflected:
|
||||
forms.add(inflected.word)
|
||||
old_forms[old_low] = (new_word, forms)
|
||||
|
||||
def get_new_form(new_word, old_form_text):
|
||||
if old_form_text in parse_cache:
|
||||
parsed_old = parse_cache[old_form_text]
|
||||
else:
|
||||
parsed_old = morph.parse(old_form_text)[0]
|
||||
parse_cache[old_form_text] = parsed_old
|
||||
tags = set()
|
||||
if parsed_old.tag.case:
|
||||
tags.add(parsed_old.tag.case)
|
||||
if parsed_old.tag.number:
|
||||
tags.add(parsed_old.tag.number)
|
||||
if parsed_old.tag.gender:
|
||||
tags.add(parsed_old.tag.gender)
|
||||
parsed_new = morph.parse(new_word)[0]
|
||||
inflected = parsed_new.inflect(tags)
|
||||
return inflected.word if inflected else new_word
|
||||
|
||||
total_runs = 0
|
||||
modified_runs = 0
|
||||
for paragraph in doc.paragraphs:
|
||||
for run in paragraph.runs:
|
||||
total_runs += 1
|
||||
original = run.text
|
||||
if not original:
|
||||
continue
|
||||
tokens = re.split(r'(\W+)', original)
|
||||
modified = False
|
||||
for i, token in enumerate(tokens):
|
||||
if token.isalpha():
|
||||
token_low = token.lower()
|
||||
for old_low, (new_word, forms) in old_forms.items():
|
||||
if token_low in forms:
|
||||
new_form = get_new_form(new_word, token)
|
||||
tokens[i] = new_form
|
||||
modified = True
|
||||
break
|
||||
if modified:
|
||||
run.text = ''.join(tokens)
|
||||
modified_runs += 1
|
||||
|
||||
logger.info(f"Умная замена: обработано {total_runs} runs, изменено {modified_runs}")
|
||||
new_path = self.temp_dir / f"surgical_{int(datetime.now().timestamp())}.docx"
|
||||
doc.save(str(new_path))
|
||||
logger.info(f"Умная замена выполнена: {new_path}")
|
||||
return str(new_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при умной замене: {e}. Выполняю простую замену.")
|
||||
# Простая замена (fallback)
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
new_path = self.temp_dir / f"surgical_{int(datetime.now().timestamp())}.docx"
|
||||
for paragraph in doc.paragraphs:
|
||||
for run in paragraph.runs:
|
||||
for old, new in replacements.items():
|
||||
if old in run.text:
|
||||
run.text = run.text.replace(old, new)
|
||||
doc.save(str(new_path))
|
||||
logger.info(f"Простая замена выполнена: {new_path}")
|
||||
return str(new_path)
|
||||
except Exception as e2:
|
||||
logger.error(f"Ошибка простой замены: {e2}")
|
||||
return None
|
||||
|
||||
def create_docx(self, content: str, title: str, user_name: str, mode: str = "CREATE") -> Optional[str]:
|
||||
"""
|
||||
Создаёт новый DOCX из текста и сохраняет во временную папку.
|
||||
Возвращает путь к созданному файлу или None.
|
||||
"""
|
||||
try:
|
||||
fname = f"output_{user_name}_{int(datetime.now().timestamp())}.docx"
|
||||
target = self.temp_dir / fname
|
||||
doc = DocxWriter()
|
||||
doc.add_heading(title, 0)
|
||||
for paragraph in content.split('\n'):
|
||||
if paragraph.strip():
|
||||
p = doc.add_paragraph(paragraph.strip())
|
||||
p.runs[0].font.size = Pt(11)
|
||||
doc.save(str(target))
|
||||
return str(target)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка создания DOCX: {e}")
|
||||
return None
|
||||
109
rag/services/giga_client.py
Normal file
109
rag/services/giga_client.py
Normal file
@@ -0,0 +1,109 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Простой клиент для GigaChat API.
|
||||
Только вызовы chat и upload_file, без бизнес-логики.
|
||||
Используется AIService и другими модулями.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
from gigachat import GigaChat
|
||||
from gigachat.models import Chat, Messages, MessagesRole
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GigaClient:
|
||||
def __init__(self, api_key: str, model: str, temperature: float, timeout: int, verify_ssl: bool = False):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.temperature = temperature
|
||||
self.timeout = timeout
|
||||
self.verify_ssl = verify_ssl
|
||||
self._client = None
|
||||
|
||||
def _get_client(self) -> GigaChat:
|
||||
"""Ленивое создание клиента (создаётся при первом запросе)."""
|
||||
if self._client is None:
|
||||
self._client = GigaChat(
|
||||
credentials=self.api_key,
|
||||
model=self.model,
|
||||
verify_ssl_certs=self.verify_ssl,
|
||||
timeout=self.timeout
|
||||
)
|
||||
logger.debug("GigaChat client initialized")
|
||||
return self._client
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
history: List[Dict[str, str]],
|
||||
query: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
file_id: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Отправляет запрос к GigaChat и возвращает ответ.
|
||||
|
||||
Аргументы:
|
||||
history: список предыдущих сообщений (каждое с ключами role и content).
|
||||
query: текущий запрос пользователя.
|
||||
system_prompt: системный промпт (если задан, добавляется первым сообщением).
|
||||
file_id: идентификатор загруженного файла (для Vision).
|
||||
|
||||
Возвращает:
|
||||
str: ответ модели или сообщение об ошибке.
|
||||
"""
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append(Messages(role=MessagesRole.SYSTEM, content=system_prompt))
|
||||
for msg in history[-10:]: # ограничим историю 10 сообщениями для экономии токенов
|
||||
role = MessagesRole.USER if msg['role'] == 'user' else MessagesRole.ASSISTANT
|
||||
messages.append(Messages(role=role, content=msg['content']))
|
||||
messages.append(Messages(role=MessagesRole.USER, content=query))
|
||||
|
||||
attachments = [file_id] if file_id else None
|
||||
|
||||
actual_temp = temperature if temperature is not None else self.temperature
|
||||
|
||||
try:
|
||||
# Запускаем синхронный вызов giga.chat в отдельном потоке
|
||||
response = await asyncio.to_thread(
|
||||
self._get_client().chat,
|
||||
Chat(
|
||||
messages=messages,
|
||||
model=self.model,
|
||||
temperature=actual_temp,
|
||||
n=1,
|
||||
attachments=attachments
|
||||
)
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при вызове GigaChat: {e}", exc_info=True)
|
||||
return "Извините, произошла ошибка при обращении к сервису ИИ."
|
||||
|
||||
async def upload_file(self, file_path: str) -> Optional[str]:
|
||||
"""
|
||||
Загружает файл в GigaChat для Vision-анализа (изображения).
|
||||
Возвращает file_id для использования в chat().
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
# Загружаем файл синхронно в потоке
|
||||
uploaded = await asyncio.to_thread(self._get_client().upload_file, f, "general")
|
||||
# В разных версиях SDK поле может называться id_, id или file_id
|
||||
if hasattr(uploaded, "id_"):
|
||||
return uploaded.id_
|
||||
if hasattr(uploaded, "id"):
|
||||
return uploaded.id
|
||||
if hasattr(uploaded, "file_id"):
|
||||
return uploaded.file_id
|
||||
# fallback: пробуем преобразовать в словарь
|
||||
data = uploaded.dict() if hasattr(uploaded, "dict") else {}
|
||||
return data.get("id_") or data.get("id") or data.get("file_id")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка загрузки файла в GigaChat: {e}", exc_info=True)
|
||||
return None
|
||||
426
rag/services/kb_service.py
Normal file
426
rag/services/kb_service.py
Normal file
@@ -0,0 +1,426 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис базы знаний (Knowledge Base Service).
|
||||
Реализует RAG-пайплайн:
|
||||
- Индексация документов (разбиение на чанки, получение эмбеддингов, сохранение в Qdrant и PostgreSQL).
|
||||
- Поиск релевантной информации по запросу (семантический поиск в Qdrant с фильтром по room_jid).
|
||||
- Управление доступом к документам (личные / глобальные / комнатные).
|
||||
- Поддержка обновления версий документов в комнатах (замена по хешу).
|
||||
|
||||
ВНИМАНИЕ: Поиск работает в двух режимах:
|
||||
1. Семантический (dense) — для общих вопросов, использует эмбеддинги GigaChat
|
||||
2. Полнотекстовый (keyword) — для точных запросов (пункт 1.1, фразы в кавычках)
|
||||
|
||||
Использует PostgreSQL (asyncpg) через PostgresService и Qdrant через QdrantService.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
|
||||
# Импорт сервисов
|
||||
from core.services.postgres_service import PostgresService
|
||||
from core.services.qdrant_service import QdrantService
|
||||
from core.services.embedding_service import EmbeddingService
|
||||
from core.utils.text_utils import split_into_chunks
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class KBService:
|
||||
"""
|
||||
База знаний: индексация документов, поиск, управление.
|
||||
Абстрагирует работу с PostgreSQL и Qdrant для RAG.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: PostgresService,
|
||||
qdrant: QdrantService,
|
||||
embedding: EmbeddingService,
|
||||
collection_name: str,
|
||||
chunk_size_tokens: int,
|
||||
overlap_tokens: int,
|
||||
approx_chunk_chars: int,
|
||||
approx_overlap_chars: int,
|
||||
):
|
||||
"""
|
||||
Инициализация KBService.
|
||||
|
||||
Аргументы:
|
||||
db: сервис PostgreSQL (asyncpg)
|
||||
qdrant: сервис Qdrant
|
||||
embedding: сервис эмбеддингов GigaChat
|
||||
collection_name: имя коллекции в Qdrant
|
||||
chunk_size_tokens: размер чанка в токенах (для tiktoken)
|
||||
overlap_tokens: перекрытие между чанками в токенах
|
||||
approx_chunk_chars: запасной размер чанка в символах (если tiktoken недоступен)
|
||||
approx_overlap_chars: запасное перекрытие в символах
|
||||
"""
|
||||
self.db = db
|
||||
self.qdrant = qdrant
|
||||
self.embedding = embedding
|
||||
self.collection_name = collection_name
|
||||
self.chunk_size_tokens = chunk_size_tokens
|
||||
self.overlap_tokens = overlap_tokens
|
||||
self.approx_chunk_chars = approx_chunk_chars
|
||||
self.approx_overlap_chars = approx_overlap_chars
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Поиск релевантной информации
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def find_relevant_info(self, query: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> str:
|
||||
"""
|
||||
Находит релевантные фрагменты в БЗ.
|
||||
|
||||
Режимы поиска:
|
||||
- Если запрос содержит номер пункта (например, "1.1") или кавычки —
|
||||
используется полнотекстовый поиск (search_keywords)
|
||||
- Иначе — семантический поиск (dense через GigaChat эмбеддинги)
|
||||
|
||||
При ошибке эмбеддинга или пустом результате использует fallback-поиск по source_name.
|
||||
|
||||
Аргументы:
|
||||
query: поисковый запрос пользователя
|
||||
user_jid: JID пользователя (для фильтрации личных документов)
|
||||
room_jid: JID комнаты (None для личного чата) – для фильтрации комнатных документов
|
||||
top_k: количество возвращаемых чанков
|
||||
|
||||
Возвращает:
|
||||
str: строка с объединёнными фрагментами вида "[источник: имя]\nтекст"
|
||||
"""
|
||||
try:
|
||||
# Определяем, является ли запрос "точным"
|
||||
# Точный запрос: содержит номер пункта (1.1, 2.3.4) или фразу в кавычках
|
||||
is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query
|
||||
|
||||
if is_exact_query:
|
||||
# РЕЖИМ 1: Полнотекстовый поиск для точных запросов
|
||||
logger.debug(f"Используем полнотекстовый поиск для запроса: {query}")
|
||||
search_results = self.qdrant.search_keywords(query, user_jid, room_jid, top_k=top_k)
|
||||
else:
|
||||
# Получаем dense-вектор запроса
|
||||
query_vector = await self.embedding.embed(query)
|
||||
if query_vector is None:
|
||||
logger.error("Не удалось получить dense-эмбеддинг запроса, переключаемся на fallback")
|
||||
return await self._fallback_search(query, user_jid, room_jid)
|
||||
|
||||
# Получаем sparse-вектор запроса (опционально, может быть None)
|
||||
sparse_vector = await self.embedding.embed_sparse(query)
|
||||
if sparse_vector is None:
|
||||
logger.warning("Не удалось получить sparse-вектор запроса, используем только dense-поиск")
|
||||
|
||||
# Гибридный поиск
|
||||
search_results = self.qdrant.search_hybrid(
|
||||
dense_vector=query_vector,
|
||||
sparse_vector=sparse_vector,
|
||||
user_jid=user_jid,
|
||||
room_jid=room_jid,
|
||||
top_k=top_k
|
||||
)
|
||||
|
||||
if not search_results:
|
||||
logger.info(f"По запросу '{query[:50]}' ничего не найдено в Qdrant")
|
||||
return await self._fallback_search(query, user_jid, room_jid)
|
||||
|
||||
# Формируем контекст из найденных чанков
|
||||
context_parts = []
|
||||
for res in search_results:
|
||||
payload = res.get('payload', {})
|
||||
source_name = payload.get('source_name', 'Неизвестный источник')
|
||||
chunk_text = payload.get('text', '')
|
||||
if chunk_text:
|
||||
context_parts.append(f"[источник: {source_name}]\n{chunk_text}")
|
||||
|
||||
context = "\n\n".join(context_parts)
|
||||
logger.debug(f"Сформирован контекст длиной {len(context)} символов")
|
||||
return context
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Ошибка при поиске в RAG: {e}")
|
||||
return await self._fallback_search(query, user_jid, room_jid)
|
||||
|
||||
async def _fallback_search(self, query: str, user_jid: str, room_jid: str = None) -> str:
|
||||
"""
|
||||
Полнотекстовый поиск по source_name в PostgreSQL.
|
||||
Используется, если Qdrant недоступен или не вернул результатов.
|
||||
|
||||
Аргументы:
|
||||
query: исходный запрос
|
||||
user_jid: JID пользователя
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
|
||||
Возвращает:
|
||||
str: строка с информацией о найденных документах
|
||||
"""
|
||||
if not query or len(query.strip()) < 3:
|
||||
return ""
|
||||
|
||||
# Очищаем запрос от спецсимволов для полнотекстового поиска
|
||||
clean_query = re.sub(r'[^\w\s]', ' ', query)
|
||||
|
||||
try:
|
||||
rows = await self.db.fulltext_search(clean_query, user_jid, room_jid)
|
||||
if not rows:
|
||||
return ""
|
||||
|
||||
parts = [
|
||||
f"[источник: {row['source_name']}]\n(Документ доступен, но текст не проиндексирован в Qdrant. Повторите запрос позже.)"
|
||||
for row in rows
|
||||
]
|
||||
return "\n\n".join(parts)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при fallback поиске: {e}")
|
||||
return ""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Индексация документов
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_document(
|
||||
self,
|
||||
file_name: str,
|
||||
file_text: str,
|
||||
user_jid: str,
|
||||
is_global: bool,
|
||||
title: Optional[str] = None,
|
||||
metadata: Optional[Dict] = None,
|
||||
room_jid: str = None,
|
||||
file_hash: str = None,
|
||||
update_if_exists: bool = True
|
||||
) -> Tuple[int, int]:
|
||||
"""
|
||||
Индексирует документ в базу знаний.
|
||||
|
||||
Этапы индексации:
|
||||
1. Разбивает текст на чанки
|
||||
2. Получает эмбеддинги для каждого чанка через GigaChat
|
||||
3. Сохраняет метаданные в PostgreSQL
|
||||
4. Сохраняет векторы и payload в Qdrant
|
||||
5. Если update_if_exists=True и room_jid задан, проверяет наличие документа
|
||||
с таким же source_name в комнате и при изменении хеша заменяет старую версию.
|
||||
|
||||
Аргументы:
|
||||
file_name: исходное имя файла (или URL)
|
||||
file_text: извлечённый текст документа
|
||||
user_jid: JID владельца документа
|
||||
is_global: глобальный ли документ
|
||||
title: отображаемое название (если не указано, используется file_name)
|
||||
metadata: дополнительные метаданные (сохраняются в JSONB)
|
||||
room_jid: JID комнаты (если документ привязан к комнате)
|
||||
file_hash: SHA-256 хеш содержимого (для проверки изменений)
|
||||
update_if_exists: заменять ли существующий документ в комнате при изменении хеша
|
||||
|
||||
Возвращает:
|
||||
Tuple[int, int]: (doc_id, количество_проиндексированных_чанков)
|
||||
"""
|
||||
source_name = title if title else file_name
|
||||
doc_meta = metadata or {}
|
||||
if file_hash:
|
||||
doc_meta['file_hash'] = file_hash
|
||||
|
||||
# --- ОБНОВЛЕНИЕ ДОКУМЕНТА В КОМНАТЕ (замена версии) ---
|
||||
if update_if_exists and room_jid and not is_global:
|
||||
existing = await self.db.find_document_by_name_and_room(source_name, room_jid)
|
||||
if existing:
|
||||
old_doc_id = existing['id']
|
||||
old_hash = existing.get('file_hash')
|
||||
if file_hash and old_hash == file_hash:
|
||||
logger.info(f"Документ {source_name} не изменился (хеш совпадает), пропускаем индексацию")
|
||||
return old_doc_id, 0
|
||||
# Удаляем старый документ из Qdrant и PostgreSQL
|
||||
logger.info(f"Обновление документа {source_name} в комнате {room_jid}: удаляем старую версию doc_id={old_doc_id}")
|
||||
self.qdrant.delete_by_doc_id(old_doc_id)
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM document_access WHERE doc_id = $1", old_doc_id)
|
||||
await conn.execute("DELETE FROM documents WHERE id = $1", old_doc_id)
|
||||
|
||||
# --- СОЗДАНИЕ ЗАПИСИ В POSTGRESQL ---
|
||||
doc_id = await self.db.add_document(
|
||||
source_name=source_name,
|
||||
owner_jid=user_jid,
|
||||
is_global=is_global,
|
||||
collection_name=self.collection_name,
|
||||
metadata=doc_meta,
|
||||
room_jid=room_jid,
|
||||
file_hash=file_hash
|
||||
)
|
||||
|
||||
# --- РАЗБИЕНИЕ НА ЧАНКИ ---
|
||||
chunks = split_into_chunks(
|
||||
file_text,
|
||||
chunk_size_tokens=self.chunk_size_tokens,
|
||||
overlap_tokens=self.overlap_tokens,
|
||||
approx_chunk_chars=self.approx_chunk_chars,
|
||||
approx_overlap_chars=self.approx_overlap_chars,
|
||||
strategy="recursive_split_by_sentences"
|
||||
)
|
||||
if not chunks:
|
||||
return doc_id, 0
|
||||
|
||||
# --- ПАКЕТНОЕ ПОЛУЧЕНИЕ ЭМБЕДДИНГОВ (dense + sparse) ---
|
||||
dense_vectors = await self.embedding.embed_batch(chunks)
|
||||
if dense_vectors is None or len(dense_vectors) != len(chunks):
|
||||
logger.error(f"Не удалось получить dense-эмбеддинги для документа {doc_id}")
|
||||
return doc_id, 0
|
||||
|
||||
sparse_vectors = await self.embedding.embed_sparse_batch(chunks)
|
||||
# sparse_vectors может быть None или короче, но мы продолжим (sparse опционально)
|
||||
|
||||
# --- СОХРАНЕНИЕ ЧАНКОВ В QDRANT ---
|
||||
chunk_count = 0
|
||||
for idx, chunk in enumerate(chunks):
|
||||
dense_vec = dense_vectors[idx] if idx < len(dense_vectors) else None
|
||||
if dense_vec is None:
|
||||
logger.error(f"Нет dense-вектора для чанка {idx}, пропускаем")
|
||||
continue
|
||||
|
||||
sparse_vec = None
|
||||
if sparse_vectors and idx < len(sparse_vectors):
|
||||
sparse_vec = sparse_vectors[idx]
|
||||
|
||||
payload = {
|
||||
"doc_id": doc_id,
|
||||
"user_jid": user_jid,
|
||||
"is_global": is_global,
|
||||
"chunk_index": idx,
|
||||
"text": chunk,
|
||||
"source_name": source_name,
|
||||
"room_jid": room_jid if room_jid else "",
|
||||
}
|
||||
await asyncio.to_thread(
|
||||
self.qdrant.add_chunk,
|
||||
dense=dense_vec,
|
||||
sparse=sparse_vec,
|
||||
payload=payload
|
||||
)
|
||||
chunk_count += 1
|
||||
|
||||
# --- ДЛЯ ЛИЧНЫХ ДОКУМЕНТОВ ПРЕДОСТАВЛЯЕМ ДОСТУП ВЛАДЕЛЬЦУ ---
|
||||
if not is_global and not room_jid:
|
||||
await self.db.add_access(doc_id, user_jid)
|
||||
|
||||
logger.info(f"Документ {source_name} (doc_id={doc_id}) проиндексирован, чанков: {chunk_count}")
|
||||
return doc_id, chunk_count
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Получение списков документов
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
|
||||
"""
|
||||
Возвращает список названий документов (source_name), доступных пользователю в данном контексте.
|
||||
|
||||
Аргументы:
|
||||
user_jid: JID пользователя
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
|
||||
Возвращает:
|
||||
List[str]: список названий
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name FROM documents WHERE room_jid = $1 ORDER BY source_name",
|
||||
room_jid
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL ORDER BY source_name",
|
||||
user_jid
|
||||
)
|
||||
return [row['source_name'] for row in rows]
|
||||
|
||||
async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Возвращает список документов с указанием, является ли документ глобальным.
|
||||
Используется командой !kb для раздельного отображения глобальных и локальных документов.
|
||||
|
||||
Аргументы:
|
||||
user_jid: JID пользователя
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
|
||||
Возвращает:
|
||||
List[Dict[str, Any]]: список словарей [{"source_name": ..., "is_global": bool}]
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name, is_global FROM documents WHERE room_jid = $1 ORDER BY source_name",
|
||||
room_jid
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name, is_global FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL ORDER BY source_name",
|
||||
user_jid
|
||||
)
|
||||
return [{"source_name": row['source_name'], "is_global": row['is_global']} for row in rows]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Очистка баз знаний
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def clear_user_kb(self, user_jid: str) -> None:
|
||||
"""
|
||||
Удаляет все личные документы пользователя (не глобальные и не комнатные).
|
||||
|
||||
Аргументы:
|
||||
user_jid: JID пользователя
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT id FROM documents WHERE owner_jid = $1 AND is_global = FALSE AND room_jid IS NULL",
|
||||
user_jid
|
||||
)
|
||||
doc_ids = [row['id'] for row in rows]
|
||||
|
||||
# Удаляем из Qdrant
|
||||
for doc_id in doc_ids:
|
||||
self.qdrant.delete_by_doc_id(doc_id)
|
||||
|
||||
# Удаляем из PostgreSQL (используем ANY для массового удаления)
|
||||
if doc_ids:
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
|
||||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||||
|
||||
logger.info(f"Личная БЗ пользователя {user_jid} очищена (документов: {len(doc_ids)})")
|
||||
|
||||
async def clear_room_kb(self, room_jid: str) -> None:
|
||||
"""
|
||||
Удаляет все документы указанной комнаты.
|
||||
|
||||
Аргументы:
|
||||
room_jid: JID комнаты
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT id FROM documents WHERE room_jid = $1", room_jid)
|
||||
doc_ids = [row['id'] for row in rows]
|
||||
|
||||
for doc_id in doc_ids:
|
||||
self.qdrant.delete_by_doc_id(doc_id)
|
||||
|
||||
if doc_ids:
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||||
|
||||
logger.info(f"Комнатная БЗ {room_jid} очищена (документов: {len(doc_ids)})")
|
||||
|
||||
async def clear_global_kb(self) -> None:
|
||||
"""
|
||||
Удаляет все глобальные документы.
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE")
|
||||
doc_ids = [row['id'] for row in rows]
|
||||
|
||||
for doc_id in doc_ids:
|
||||
self.qdrant.delete_by_doc_id(doc_id)
|
||||
|
||||
if doc_ids:
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||||
|
||||
logger.info("Глобальная БЗ очищена")
|
||||
356
rag/services/postgres_service.py
Normal file
356
rag/services/postgres_service.py
Normal file
@@ -0,0 +1,356 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис для работы с PostgreSQL (через asyncpg).
|
||||
Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
|
||||
управление доступом к личным документам.
|
||||
|
||||
Полностью заменяет MySQLService. Реализованы все методы, необходимые для KBService и других компонентов.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
import asyncpg
|
||||
from asyncpg import Pool, Connection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PostgresService:
|
||||
"""
|
||||
Асинхронный сервис для работы с PostgreSQL.
|
||||
Использует пул соединений asyncpg.
|
||||
"""
|
||||
|
||||
def __init__(self, host: str, port: int, user: str, password: str, db_name: str):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.db_name = db_name
|
||||
self.pool: Optional[Pool] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Подключение и инициализация таблиц
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""
|
||||
Создаёт пул соединений и инициализирует таблицы, если их нет.
|
||||
"""
|
||||
self.pool = await asyncpg.create_pool(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
user=self.user,
|
||||
password=self.password,
|
||||
database=self.db_name,
|
||||
min_size=1,
|
||||
max_size=10,
|
||||
command_timeout=60
|
||||
)
|
||||
await self._init_tables()
|
||||
logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")
|
||||
|
||||
async def _init_tables(self) -> None:
|
||||
"""
|
||||
Создаёт таблицы (если не существуют) с учётом всех необходимых колонок и индексов.
|
||||
Использует схему из schema_postgresql.sql (адаптированную для PostgreSQL).
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
# 1. Таблица documents
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id SERIAL PRIMARY KEY,
|
||||
source_name VARCHAR(255) NOT NULL,
|
||||
upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
owner_jid VARCHAR(255) NOT NULL,
|
||||
is_global BOOLEAN DEFAULT FALSE,
|
||||
collection_name VARCHAR(64) NOT NULL,
|
||||
metadata JSONB,
|
||||
room_jid VARCHAR(255) DEFAULT NULL,
|
||||
file_hash VARCHAR(64) DEFAULT NULL
|
||||
)
|
||||
""")
|
||||
# Индексы
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner ON documents(owner_jid)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_global ON documents(is_global)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room ON documents(room_jid)")
|
||||
# Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS ft_source_name
|
||||
ON documents USING GIN (to_tsvector('russian', source_name))
|
||||
""")
|
||||
|
||||
# 2. Таблица document_access
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS document_access (
|
||||
doc_id INT,
|
||||
user_jid VARCHAR(255),
|
||||
PRIMARY KEY (doc_id, user_jid)
|
||||
)
|
||||
""")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_access_user ON document_access(user_jid)")
|
||||
|
||||
# 3. Тип ENUM для history.role (если не существует)
|
||||
await conn.execute("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE history_role AS ENUM ('user', 'assistant');
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
""")
|
||||
# Таблица history
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS history (
|
||||
id SERIAL PRIMARY KEY,
|
||||
user_jid VARCHAR(255) NOT NULL,
|
||||
role history_role NOT NULL,
|
||||
content TEXT,
|
||||
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
room_jid VARCHAR(255) DEFAULT NULL
|
||||
)
|
||||
""")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_user ON history(user_jid)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_room ON history(room_jid)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_time ON history(timestamp)")
|
||||
|
||||
# 4. Таблица rooms
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS rooms (
|
||||
room_jid VARCHAR(255) PRIMARY KEY
|
||||
)
|
||||
""")
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Закрывает пул соединений."""
|
||||
if self.pool:
|
||||
await self.pool.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Документы
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_document(
|
||||
self,
|
||||
source_name: str,
|
||||
owner_jid: str,
|
||||
is_global: bool,
|
||||
collection_name: str,
|
||||
metadata: dict = None,
|
||||
room_jid: str = None,
|
||||
file_hash: str = None
|
||||
) -> int:
|
||||
"""
|
||||
Добавляет запись о документе. Возвращает doc_id.
|
||||
|
||||
:param source_name: название документа (source_name)
|
||||
:param owner_jid: JID владельца
|
||||
:param is_global: глобальный ли документ
|
||||
:param collection_name: имя коллекции Qdrant
|
||||
:param metadata: произвольные метаданные (словарь)
|
||||
:param room_jid: JID комнаты (если документ комнатный)
|
||||
:param file_hash: SHA-256 хеш содержимого (для обновления версий)
|
||||
:return: id новой записи
|
||||
"""
|
||||
metadata_json = json.dumps(metadata, ensure_ascii=False) if metadata else None
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
INSERT INTO documents (source_name, owner_jid, is_global, collection_name, metadata, room_jid, file_hash)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id
|
||||
""",
|
||||
source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash
|
||||
)
|
||||
return row['id']
|
||||
|
||||
async def get_document(self, doc_id: int) -> Optional[Dict]:
|
||||
"""Возвращает запись документа по id."""
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow("SELECT * FROM documents WHERE id = $1", doc_id)
|
||||
return dict(row) if row else None
|
||||
|
||||
async def add_access(self, doc_id: int, user_jid: str) -> None:
|
||||
"""
|
||||
Даёт пользователю доступ к личному документу (только для не-глобальных, не-комнатных).
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO document_access (doc_id, user_jid)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT (doc_id, user_jid) DO NOTHING
|
||||
""",
|
||||
doc_id, user_jid
|
||||
)
|
||||
|
||||
async def can_access(self, doc_id: int, user_jid: str, is_global: bool) -> bool:
|
||||
"""Проверяет, имеет ли пользователь доступ к документу."""
|
||||
if is_global:
|
||||
return True
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT 1 FROM document_access WHERE doc_id = $1 AND user_jid = $2",
|
||||
doc_id, user_jid
|
||||
)
|
||||
return row is not None
|
||||
|
||||
async def find_document_by_name_and_room(self, source_name: str, room_jid: str) -> Optional[Dict]:
|
||||
"""
|
||||
Поиск документа в комнате по имени (source_name). Используется для обновления версий.
|
||||
Возвращает словарь с ключами id, file_hash или None.
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT id, file_hash FROM documents WHERE source_name = $1 AND room_jid = $2 LIMIT 1",
|
||||
source_name, room_jid
|
||||
)
|
||||
return dict(row) if row else None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# История диалогов
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_history(self, user_jid: str, role: str, content: str, room_jid: str = None) -> None:
|
||||
"""
|
||||
Сохраняет сообщение в историю.
|
||||
:param role: 'user' или 'assistant'
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"INSERT INTO history (user_jid, role, content, room_jid) VALUES ($1, $2, $3, $4)",
|
||||
user_jid, role, content, room_jid
|
||||
)
|
||||
|
||||
async def get_history(self, user_jid: str, room_jid: str = None, limit: int = 100) -> List[Dict]:
|
||||
"""
|
||||
Возвращает последние limit сообщений, отсортированных от старых к новым.
|
||||
Если room_jid указан – история комнаты, иначе – личная история пользователя.
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT role, content, timestamp
|
||||
FROM history
|
||||
WHERE room_jid = $1
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT $2
|
||||
""",
|
||||
room_jid, limit
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT role, content, timestamp
|
||||
FROM history
|
||||
WHERE user_jid = $1 AND room_jid IS NULL
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT $2
|
||||
""",
|
||||
user_jid, limit
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
async def clear_user_history(self, user_jid: str) -> None:
|
||||
"""Удаляет личную историю пользователя (не трогает комнаты)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM history WHERE user_jid = $1 AND room_jid IS NULL", user_jid)
|
||||
|
||||
async def clear_room_history(self, room_jid: str) -> None:
|
||||
"""Удаляет историю указанной комнаты."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM history WHERE room_jid = $1", room_jid)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Полнотекстовый поиск (fallback, если Qdrant недоступен)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
|
||||
"""
|
||||
Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
|
||||
Возвращает список документов (id, source_name, metadata).
|
||||
"""
|
||||
if not query or len(query.strip()) < 3:
|
||||
return []
|
||||
async with self.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, source_name, metadata
|
||||
FROM documents
|
||||
WHERE room_jid = $1
|
||||
AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
|
||||
LIMIT 10
|
||||
""",
|
||||
room_jid, query
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, source_name, metadata
|
||||
FROM documents
|
||||
WHERE (owner_jid = $1 OR is_global = TRUE)
|
||||
AND room_jid IS NULL
|
||||
AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
|
||||
LIMIT 10
|
||||
""",
|
||||
user_jid, query
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Комнаты MUC
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_room(self, room_jid: str) -> None:
|
||||
"""Добавляет комнату в список (если ещё не добавлена)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"INSERT INTO rooms (room_jid) VALUES ($1) ON CONFLICT (room_jid) DO NOTHING",
|
||||
room_jid
|
||||
)
|
||||
|
||||
async def get_rooms(self) -> List[str]:
|
||||
"""Возвращает список JID комнат, в которых бот должен состоять."""
|
||||
async with self.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT room_jid FROM rooms")
|
||||
return [row['room_jid'] for row in rows]
|
||||
|
||||
async def remove_room(self, room_jid: str) -> None:
|
||||
"""Удаляет комнату из списка."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM rooms WHERE room_jid = $1", room_jid)
|
||||
|
||||
async def save_template(self, room_jid: str, name: str, file_path: str, file_hash: str) -> bool:
|
||||
"""Сохраняет шаблон в БД (если с таким именем уже есть – перезаписывает)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("""
|
||||
INSERT INTO room_templates (room_jid, name, file_path, file_hash)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (room_jid, name) DO UPDATE SET
|
||||
file_path = EXCLUDED.file_path,
|
||||
file_hash = EXCLUDED.file_hash,
|
||||
created_at = CURRENT_TIMESTAMP
|
||||
""", room_jid, name, file_path, file_hash)
|
||||
return True
|
||||
|
||||
async def get_template(self, room_jid: str, name: str) -> Optional[Dict]:
|
||||
"""Возвращает словарь {file_path, file_hash} или None."""
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow("SELECT file_path, file_hash FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
|
||||
return dict(row) if row else None
|
||||
|
||||
async def list_templates(self, room_jid: str) -> List[str]:
|
||||
"""Возвращает список названий шаблонов в комнате."""
|
||||
async with self.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT name FROM room_templates WHERE room_jid = $1 ORDER BY name", room_jid)
|
||||
return [row['name'] for row in rows]
|
||||
|
||||
async def delete_template(self, room_jid: str, name: str) -> bool:
|
||||
"""Удаляет шаблон. Возвращает True, если удалён."""
|
||||
async with self.pool.acquire() as conn:
|
||||
result = await conn.execute("DELETE FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
|
||||
return result == "DELETE 1"
|
||||
452
rag/services/qdrant_service.py
Normal file
452
rag/services/qdrant_service.py
Normal file
@@ -0,0 +1,452 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис для взаимодействия с векторной базой данных Qdrant.
|
||||
Обеспечивает создание гибридной коллекции (dense + sparse), добавление точек,
|
||||
гибридный поиск с RRF, полнотекстовый поиск и удаление по doc_id.
|
||||
|
||||
ВНИМАНИЕ: Этот сервис НЕ использует генерацию эмбеддингов на лету через models.Document.
|
||||
Вместо этого он получает готовые dense-векторы из EmbeddingService (GigaChat)
|
||||
и sparse-векторы из FastEmbed (Qdrant/bm25).
|
||||
|
||||
ВАЖНО: В версиях qdrant-client 1.10+ метод search() заменён на query_points().
|
||||
В этом файле используется правильный API для вашей версии.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from qdrant_client import QdrantClient, models
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QdrantService:
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
grpc_port: int,
|
||||
collection_name: str,
|
||||
vector_size: int,
|
||||
distance: str,
|
||||
prefer_grpc: bool = False
|
||||
):
|
||||
"""
|
||||
Инициализация сервиса Qdrant.
|
||||
|
||||
Аргументы:
|
||||
host: хост Qdrant (обычно localhost)
|
||||
port: HTTP порт Qdrant (обычно 6333)
|
||||
grpc_port: gRPC порт Qdrant (обычно 6334)
|
||||
collection_name: имя коллекции в Qdrant
|
||||
vector_size: размерность плотных векторов (1024 для GigaChat Embeddings)
|
||||
distance: метрика расстояния (Cosine или Euclid)
|
||||
prefer_grpc: использовать ли gRPC вместо HTTP (по умолчанию False)
|
||||
"""
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.collection_name = collection_name
|
||||
self.vector_size = vector_size
|
||||
self.distance = models.Distance.COSINE if distance.lower() == "cosine" else models.Distance.EUCLID
|
||||
|
||||
# Имена векторных полей в коллекции
|
||||
self.dense_vector_name = "dense"
|
||||
self.sparse_vector_name = "sparse"
|
||||
|
||||
# Создаём клиент Qdrant
|
||||
self.client = QdrantClient(host=host, port=port, grpc_port=grpc_port, prefer_grpc=prefer_grpc)
|
||||
self._ensure_collection()
|
||||
logger.info(f"QdrantService инициализирован: коллекция {collection_name}, размер dense {vector_size}")
|
||||
|
||||
def _ensure_collection(self) -> None:
|
||||
"""
|
||||
Проверяет существование коллекции. Если нет — создаёт её с поддержкой dense и sparse векторов.
|
||||
"""
|
||||
if not self.client.collection_exists(self.collection_name):
|
||||
logger.info(f"Коллекция {self.collection_name} не найдена. Создаю новую (dense + sparse)...")
|
||||
try:
|
||||
# Конфигурация dense-векторов (семантический поиск)
|
||||
vectors_config = {
|
||||
self.dense_vector_name: models.VectorParams(
|
||||
size=self.vector_size,
|
||||
distance=self.distance,
|
||||
)
|
||||
}
|
||||
# Конфигурация sparse-векторов (ключевой поиск BM25 с IDF)
|
||||
sparse_vectors_config = {
|
||||
self.sparse_vector_name: models.SparseVectorParams(
|
||||
modifier=models.Modifier.IDF
|
||||
)
|
||||
}
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=vectors_config,
|
||||
sparse_vectors_config=sparse_vectors_config
|
||||
)
|
||||
logger.info(f"Коллекция {self.collection_name} создана (dense size={self.vector_size}, sparse с IDF)")
|
||||
except Exception as e:
|
||||
# Если библиотека не поддерживает sparse (старая версия), создаём только dense
|
||||
logger.warning(f"Не удалось создать sparse-векторы: {e}. Создаю коллекцию только с dense.")
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config={
|
||||
self.dense_vector_name: models.VectorParams(
|
||||
size=self.vector_size,
|
||||
distance=self.distance,
|
||||
)
|
||||
}
|
||||
)
|
||||
logger.info(f"Коллекция {self.collection_name} создана (только dense size={self.vector_size})")
|
||||
# После создания коллекции добавляем payload-индексы
|
||||
self._create_payload_indexes()
|
||||
else:
|
||||
# Коллекция уже существует – проверяем, есть ли sparse-поле (логируем предупреждение, если нет)
|
||||
try:
|
||||
collection_info = self.client.get_collection(self.collection_name)
|
||||
if not hasattr(collection_info, 'config') or not hasattr(collection_info.config, 'params'):
|
||||
logger.debug("Не удалось проверить наличие sparse-поля – пропускаем")
|
||||
else:
|
||||
sparse_exists = (hasattr(collection_info.config.params, 'sparse_vectors') and
|
||||
self.sparse_vector_name in collection_info.config.params.sparse_vectors)
|
||||
if not sparse_exists:
|
||||
logger.warning(
|
||||
f"Коллекция {self.collection_name} не содержит sparse-векторного поля. "
|
||||
"Гибридный поиск будет работать только с dense. "
|
||||
"Рекомендуется пересоздать коллекцию с помощью --force или вручную."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Не удалось проверить конфигурацию коллекции: {e}")
|
||||
logger.debug(f"Коллекция {self.collection_name} уже существует")
|
||||
|
||||
def _create_payload_indexes(self) -> None:
|
||||
"""
|
||||
Создаёт индексы для полей payload для ускорения фильтрации и поиска.
|
||||
Индексы нужны для полей:
|
||||
- user_jid: для доступа к личным документам (keyword)
|
||||
- room_jid: для доступа к документам комнаты (keyword)
|
||||
- is_global: для доступа к глобальным документам (keyword)
|
||||
- doc_id: для быстрого удаления всех чанков документа (integer)
|
||||
- text: для полнотекстового поиска (text с настройками tokenizer=whitespace, lowercase, phrase_matching)
|
||||
"""
|
||||
logger.info("Создаю payload-индексы для коллекции...")
|
||||
|
||||
# Индексы для фильтрации (keyword)
|
||||
for field in ["user_jid", "room_jid", "is_global"]:
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name=field,
|
||||
field_schema="keyword"
|
||||
)
|
||||
logger.debug(f"Индекс keyword для поля '{field}' создан")
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось создать индекс для '{field}': {e}")
|
||||
|
||||
# Индекс для doc_id (integer)
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name="doc_id",
|
||||
field_schema="integer"
|
||||
)
|
||||
logger.debug("Индекс integer для поля 'doc_id' создан")
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось создать индекс для 'doc_id': {e}")
|
||||
|
||||
# Полнотекстовый индекс для текста с детальными параметрами
|
||||
try:
|
||||
# Используем расширенную конфигурацию (tokenizer, lowercase, phrase_matching)
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name="text",
|
||||
field_schema=models.TextIndexParams(
|
||||
tokenizer="whitespace",
|
||||
lowercase=True,
|
||||
min_token_len=1,
|
||||
max_token_len=0, # 0 = без ограничения
|
||||
phrase_matching=True
|
||||
)
|
||||
)
|
||||
logger.info("Полнотекстовый индекс для поля 'text' создан (tokenizer=whitespace, lowercase, phrase_matching)")
|
||||
except Exception as e:
|
||||
# Fallback для старых версий qdrant-client (без TextIndexParams)
|
||||
logger.warning(f"Не удалось создать расширенный индекс text: {e}. Пробую упрощённый.")
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name="text",
|
||||
field_schema="text"
|
||||
)
|
||||
logger.info("Полнотекстовый индекс для поля 'text' создан (упрощённый)")
|
||||
except Exception as e2:
|
||||
logger.error(f"Не удалось создать индекс для 'text': {e2}")
|
||||
|
||||
logger.info(f"Payload-индексы для коллекции {self.collection_name} созданы")
|
||||
|
||||
def add_chunk(self, dense: List[float], sparse: Any, payload: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Добавляет точку (чанк) в коллекцию с двумя типами векторов: dense и sparse.
|
||||
|
||||
Аргументы:
|
||||
dense: плотный вектор (1024) от GigaChat
|
||||
sparse: разрежённый вектор от FastEmbed (SparseEmbedding) или None
|
||||
payload: метаданные чанка
|
||||
|
||||
Возвращает:
|
||||
str: уникальный идентификатор точки
|
||||
"""
|
||||
point_id = str(uuid4())
|
||||
|
||||
# Преобразуем sparse-вектор из формата FastEmbed в формат Qdrant
|
||||
sparse_vector = None
|
||||
if sparse is not None:
|
||||
if hasattr(sparse, 'indices') and hasattr(sparse, 'values'):
|
||||
# FastEmbed SparseEmbedding -> SparseVector
|
||||
sparse_vector = models.SparseVector(
|
||||
indices=sparse.indices.tolist(),
|
||||
values=sparse.values.tolist()
|
||||
)
|
||||
elif isinstance(sparse, dict) and 'indices' in sparse:
|
||||
# Уже словарь
|
||||
sparse_vector = models.SparseVector(**sparse)
|
||||
else:
|
||||
logger.warning(f"Неизвестный формат sparse: {type(sparse)}")
|
||||
|
||||
# Если sparse-вектор не задан, передаём пустой (чтобы не нарушить схему коллекции)
|
||||
if sparse_vector is None:
|
||||
sparse_vector = models.SparseVector(indices=[], values=[])
|
||||
|
||||
point = models.PointStruct(
|
||||
id=point_id,
|
||||
vector={
|
||||
self.dense_vector_name: dense,
|
||||
self.sparse_vector_name: sparse_vector
|
||||
},
|
||||
payload=payload
|
||||
)
|
||||
self.client.upsert(collection_name=self.collection_name, points=[point])
|
||||
logger.debug(f"Чанк добавлен: id={point_id}, doc_id={payload.get('doc_id')}")
|
||||
return point_id
|
||||
|
||||
def search_hybrid(
|
||||
self,
|
||||
dense_vector: List[float],
|
||||
sparse_vector: Any,
|
||||
user_jid: str,
|
||||
room_jid: str = None,
|
||||
top_k: int = 15,
|
||||
rrf_k: int = 60
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Гибридный поиск: объединяет результаты dense и sparse поиска через Reciprocal Rank Fusion (RRF).
|
||||
|
||||
Аргументы:
|
||||
dense_vector: плотный вектор запроса (1024) от GigaChat
|
||||
sparse_vector: разрежённый вектор запроса от FastEmbed (SparseEmbedding) или None
|
||||
user_jid: JID пользователя для фильтрации доступа
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
top_k: количество возвращаемых результатов
|
||||
rrf_k: константа сглаживания в RRF (обычно 60)
|
||||
|
||||
Возвращает:
|
||||
List[Dict]: список результатов с ключами score, payload, id
|
||||
"""
|
||||
# --- 1. Фильтр доступа ---
|
||||
if room_jid:
|
||||
# В комнате: документы комнаты ИЛИ глобальные
|
||||
filter_cond = models.Filter(
|
||||
should=[
|
||||
models.FieldCondition(key="room_jid", match=models.MatchValue(value=room_jid)),
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
else:
|
||||
# В личном чате: личные документы пользователя ИЛИ глобальные
|
||||
filter_cond = models.Filter(
|
||||
should=[
|
||||
models.FieldCondition(key="user_jid", match=models.MatchValue(value=user_jid)),
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
|
||||
# --- 2. Dense-поиск (семантический) ---
|
||||
dense_resp = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=dense_vector,
|
||||
using=self.dense_vector_name,
|
||||
limit=top_k,
|
||||
query_filter=filter_cond,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# --- 3. Sparse-поиск (ключевой) с преобразованием формата ---
|
||||
sparse_resp = None
|
||||
if sparse_vector is not None:
|
||||
# Преобразуем sparse-вектор аналогично методу add_chunk
|
||||
sparse_query = None
|
||||
if hasattr(sparse_vector, 'indices') and hasattr(sparse_vector, 'values'):
|
||||
# FastEmbed SparseEmbedding -> SparseVector
|
||||
sparse_query = models.SparseVector(
|
||||
indices=sparse_vector.indices.tolist(),
|
||||
values=sparse_vector.values.tolist()
|
||||
)
|
||||
elif isinstance(sparse_vector, dict) and 'indices' in sparse_vector:
|
||||
sparse_query = models.SparseVector(**sparse_vector)
|
||||
elif isinstance(sparse_vector, models.SparseVector):
|
||||
sparse_query = sparse_vector
|
||||
else:
|
||||
logger.warning(f"Неизвестный формат sparse_vector: {type(sparse_vector)}")
|
||||
|
||||
if sparse_query is not None:
|
||||
sparse_resp = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=sparse_query,
|
||||
using=self.sparse_vector_name,
|
||||
limit=top_k,
|
||||
query_filter=filter_cond,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# --- 4. Reciprocal Rank Fusion (RRF) ---
|
||||
scores = {} # point_id -> {score, payload, id}
|
||||
|
||||
for rank, hit in enumerate(dense_resp.points):
|
||||
rrf_score = 1.0 / (rrf_k + rank + 1)
|
||||
scores[hit.id] = {
|
||||
"score": rrf_score,
|
||||
"payload": hit.payload,
|
||||
"id": hit.id
|
||||
}
|
||||
|
||||
if sparse_resp:
|
||||
for rank, hit in enumerate(sparse_resp.points):
|
||||
rrf_score = 1.0 / (rrf_k + rank + 1)
|
||||
if hit.id in scores:
|
||||
scores[hit.id]["score"] += rrf_score
|
||||
else:
|
||||
scores[hit.id] = {
|
||||
"score": rrf_score,
|
||||
"payload": hit.payload,
|
||||
"id": hit.id
|
||||
}
|
||||
|
||||
# --- 5. Сортировка и возврат top_k ---
|
||||
sorted_items = sorted(scores.values(), key=lambda x: x["score"], reverse=True)
|
||||
return sorted_items[:top_k]
|
||||
|
||||
def search_keywords(self, query_text: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> List[Dict]:
|
||||
"""
|
||||
Выполняет полнотекстовый поиск по полю text (для точных запросов).
|
||||
|
||||
Аргументы:
|
||||
query_text: текст запроса (ищется как подстрока в text)
|
||||
user_jid: JID пользователя для фильтрации
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
top_k: количество возвращаемых результатов
|
||||
|
||||
Возвращает:
|
||||
List[Dict]: список результатов с ключами score, payload, id
|
||||
|
||||
Используется для точных запросов типа "выведи пункт 1.1",
|
||||
где семантический поиск может не справиться.
|
||||
"""
|
||||
# Формируем фильтр доступа
|
||||
if room_jid:
|
||||
filter_cond = models.Filter(
|
||||
must=[
|
||||
models.FieldCondition(key="room_jid", match=models.MatchValue(value=room_jid))
|
||||
],
|
||||
should=[
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
else:
|
||||
filter_cond = models.Filter(
|
||||
must=[
|
||||
models.FieldCondition(key="user_jid", match=models.MatchValue(value=user_jid))
|
||||
],
|
||||
should=[
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
|
||||
# Добавляем условие полнотекстового поиска (MatchText)
|
||||
filter_cond.must.append(
|
||||
models.FieldCondition(
|
||||
key="text",
|
||||
match=models.MatchText(text=query_text)
|
||||
)
|
||||
)
|
||||
|
||||
# Получаем точки через scroll
|
||||
points = []
|
||||
offset = None
|
||||
while len(points) < top_k:
|
||||
scroll_result = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=filter_cond,
|
||||
limit=top_k,
|
||||
offset=offset,
|
||||
with_payload=True
|
||||
)
|
||||
points.extend(scroll_result[0])
|
||||
offset = scroll_result[1]
|
||||
if offset is None:
|
||||
break
|
||||
|
||||
# Простая сортировка по количеству совпавших ключевых слов (score не от Qdrant, а наша эвристика)
|
||||
query_words = set(query_text.lower().split())
|
||||
for point in points:
|
||||
text = point.payload.get('text', '').lower()
|
||||
matches = sum(1 for word in query_words if word in text)
|
||||
point._match_score = matches
|
||||
|
||||
points.sort(key=lambda p: getattr(p, '_match_score', 0), reverse=True)
|
||||
|
||||
return [
|
||||
{
|
||||
"score": getattr(p, '_match_score', 0),
|
||||
"payload": p.payload,
|
||||
"id": p.id
|
||||
}
|
||||
for p in points[:top_k]
|
||||
]
|
||||
|
||||
def delete_by_doc_id(self, doc_id: int) -> None:
|
||||
"""
|
||||
Удаляет все точки (чанки), принадлежащие указанному документу.
|
||||
|
||||
Аргументы:
|
||||
doc_id: идентификатор документа в PostgreSQL
|
||||
"""
|
||||
filter_cond = models.Filter(
|
||||
must=[
|
||||
models.FieldCondition(key="doc_id", match=models.MatchValue(value=doc_id))
|
||||
]
|
||||
)
|
||||
|
||||
points = []
|
||||
offset = None
|
||||
while True:
|
||||
scroll_result = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=filter_cond,
|
||||
limit=100,
|
||||
offset=offset
|
||||
)
|
||||
points.extend(scroll_result[0])
|
||||
offset = scroll_result[1]
|
||||
if offset is None:
|
||||
break
|
||||
|
||||
ids = [p.id for p in points]
|
||||
if ids:
|
||||
self.client.delete(collection_name=self.collection_name, points_selector=ids)
|
||||
logger.info(f"Удалено {len(ids)} чанков для doc_id={doc_id}")
|
||||
|
||||
def delete_collection(self) -> None:
|
||||
"""Удаляет всю коллекцию (опасная операция, используется редко)."""
|
||||
if self.client.collection_exists(self.collection_name):
|
||||
self.client.delete_collection(self.collection_name)
|
||||
logger.warning(f"Коллекция {self.collection_name} удалена")
|
||||
137
rag/services/reranker_service.py
Normal file
137
rag/services/reranker_service.py
Normal file
@@ -0,0 +1,137 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис переранжирования с использованием кросс-энкодеров.
|
||||
Загружает модель один раз при первом вызове и использует её для ранжирования фрагментов.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Tuple, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Импортируем sentence-transformers с обработкой ошибок
|
||||
try:
|
||||
from sentence_transformers import CrossEncoder
|
||||
HAS_CROSS_ENCODER = True
|
||||
except ImportError:
|
||||
HAS_CROSS_ENCODER = False
|
||||
logger.warning("sentence-transformers не установлен. Используйте LLM-переранжирование.")
|
||||
|
||||
|
||||
class RerankerService:
|
||||
"""
|
||||
Сервис переранжирования на основе кросс-энкодеров.
|
||||
Загружает модель при инициализации или при первом вызове.
|
||||
"""
|
||||
|
||||
# Рекомендуемая модель для русского языка
|
||||
DEFAULT_MODEL = "intfloat/multilingual-e5-reranker"
|
||||
|
||||
def __init__(self, model_name: Optional[str] = None, device: Optional[str] = None):
|
||||
"""
|
||||
Инициализация сервиса.
|
||||
|
||||
Аргументы:
|
||||
model_name: имя модели (если None, используется DEFAULT_MODEL).
|
||||
device: устройство для вычислений ('cpu' или 'cuda'). Если None, автоопределение.
|
||||
"""
|
||||
self.model_name = model_name or self.DEFAULT_MODEL
|
||||
self.device = device or ("cuda" if self._cuda_available() else "cpu")
|
||||
self._model = None
|
||||
logger.info(f"RerankerService инициализирован с моделью {self.model_name} (device={self.device})")
|
||||
|
||||
@staticmethod
|
||||
def _cuda_available() -> bool:
|
||||
"""Проверяет доступность CUDA (для GPU)."""
|
||||
try:
|
||||
import torch
|
||||
return torch.cuda.is_available()
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def _load_model(self) -> None:
|
||||
"""Загружает модель кросс-энкодера (ленивая загрузка)."""
|
||||
if not HAS_CROSS_ENCODER:
|
||||
raise RuntimeError("sentence-transformers не установлен. Установите: pip install sentence-transformers")
|
||||
|
||||
if self._model is None:
|
||||
logger.info(f"Загрузка модели {self.model_name} на устройство {self.device}...")
|
||||
self._model = CrossEncoder(self.model_name, device=self.device)
|
||||
logger.info(f"Модель {self.model_name} загружена")
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
fragments: List[str],
|
||||
top_k: Optional[int] = None
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""
|
||||
Ранжирует фрагменты по релевантности запросу.
|
||||
|
||||
Аргументы:
|
||||
query: текст запроса.
|
||||
fragments: список текстовых фрагментов.
|
||||
top_k: максимальное количество возвращаемых фрагментов (если None – все).
|
||||
|
||||
Возвращает:
|
||||
List[Tuple[str, float]]: список кортежей (фрагмент, оценка_релевантности),
|
||||
отсортированный по убыванию оценки.
|
||||
"""
|
||||
if not fragments:
|
||||
return []
|
||||
|
||||
self._load_model()
|
||||
|
||||
# Формируем пары (запрос, фрагмент) для оценки
|
||||
pairs = [(query, frag) for frag in fragments]
|
||||
|
||||
# Получаем оценки релевантности (числа от 0 до 1)
|
||||
try:
|
||||
scores = self._model.predict(pairs)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при оценке кросс-энкодером: {e}")
|
||||
# В случае ошибки возвращаем фрагменты в исходном порядке с оценкой 0.0
|
||||
return [(frag, 0.0) for frag in fragments]
|
||||
|
||||
# Собираем результаты и сортируем по убыванию оценки
|
||||
results = list(zip(fragments, scores))
|
||||
results.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Ограничиваем количество, если указано
|
||||
if top_k is not None:
|
||||
results = results[:top_k]
|
||||
|
||||
logger.debug(f"Переранжирование: из {len(fragments)} фрагментов оставлено {len(results)}")
|
||||
return results
|
||||
|
||||
def rerank_context(
|
||||
self,
|
||||
query: str,
|
||||
context: str,
|
||||
top_k: Optional[int] = None
|
||||
) -> str:
|
||||
"""
|
||||
Удобная обёртка: принимает строку контекста с фрагментами, разделёнными '\n\n',
|
||||
возвращает отранжированную строку.
|
||||
|
||||
Аргументы:
|
||||
query: текст запроса.
|
||||
context: исходный контекст (фрагменты разделены '\n\n').
|
||||
top_k: максимальное количество возвращаемых фрагментов.
|
||||
|
||||
Возвращает:
|
||||
str: отранжированный контекст (фрагменты разделены '\n\n').
|
||||
"""
|
||||
if not context:
|
||||
return ""
|
||||
|
||||
# Разбиваем контекст на фрагменты
|
||||
fragments = [f.strip() for f in context.split("\n\n") if f.strip()]
|
||||
if not fragments:
|
||||
return ""
|
||||
|
||||
# Ранжируем
|
||||
ranked = self.rerank(query, fragments, top_k)
|
||||
|
||||
# Собираем обратно в строку (только текст, без оценок)
|
||||
return "\n\n".join([frag for frag, _ in ranked])
|
||||
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