59feffc190
- /rag/config_models.py - /rag/utils/config_loader.py - /rag/auth.py - /rag/rag_server.py - /rag/rag_api.py - /rag/rag_orchestrator.py - /rag/__init__.py - /bots/xmpp/client.py - /bots/rag_client.py
212 lines
12 KiB
Python
212 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
XMPP клиент – главный класс Bot.
|
||
Теперь это тонкий клиент, который только общается с RAG-сервером.
|
||
Не содержит БД, Qdrant, GigaChat – всё на сервере.
|
||
|
||
ИСПОЛЬЗУЕТ:
|
||
- AppConfig из rag.config_models для единой конфигурации.
|
||
- RAGClient с передачей API-ключа для аутентификации.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from datetime import datetime
|
||
from slixmpp import ClientXMPP
|
||
|
||
# Импорты из нового пакета rag для конфигурации
|
||
from rag.config_models import AppConfig
|
||
|
||
# Остальные импорты остаются из core (старый код) – позже их нужно будет перенести в rag,
|
||
# но пока оставляем как есть для совместимости.
|
||
from core.services.file_service import FileService
|
||
from core.handlers.metrics_handler import MetricsHandler
|
||
from core.rag_client import RAGClient
|
||
from core.commands.registry import register_commands
|
||
from core.handlers.message_handler import MessageHandler
|
||
from core.handlers.file_handler import FileHandler
|
||
from core.handlers.room_handler import RoomHandler
|
||
from core.handlers.subscription_handler import SubscriptionHandler
|
||
from core.handlers.health_handler import HealthHandler
|
||
from core.workers.indexing_worker import IndexingWorker
|
||
from core.functions.file_processor import handle_ai_file_request as _handle_ai_file_request
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class Bot(ClientXMPP):
|
||
"""
|
||
Основной класс XMPP бота.
|
||
Тонкий клиент: вся RAG-логика на сервере.
|
||
"""
|
||
|
||
def __init__(self, config: AppConfig): # <-- изменён тип конфига
|
||
# Полный JID: user@domain/resource
|
||
full_jid = f"{config.jid}/{config.xmpp_resource}"
|
||
super().__init__(full_jid, config.xmpp_password)
|
||
self.config = config
|
||
|
||
# ----- 1. Локальные сервисы (только для файлов и метрик) -----
|
||
# FileService нужен для извлечения текста из файлов перед отправкой на сервер
|
||
self.files = FileService(config)
|
||
# Метрики (счётчики сообщений, ошибок и т.д.)
|
||
self.metrics_handler = MetricsHandler()
|
||
|
||
# ----- 2. Загрузка промптов (остаются локально) -----
|
||
# Промпты загружаются из файлов и передаются в каждом запросе к серверу.
|
||
self.ai_prompts = {}
|
||
# В AppConfig теперь есть свойства для путей к файлам промптов,
|
||
# но они называются так же, как в BotConfig, поэтому доступ через config.<имя> работает.
|
||
# Однако в BotConfig они назывались system_prompt_file и т.д. – они есть в AppConfig?
|
||
# В AppConfig мы добавили прокси-свойства только для основных параметров,
|
||
# но не для путей к промптам. Поэтому здесь нужно либо использовать config.prompts_dir_path
|
||
# и формировать пути вручную, либо оставить как было, если в AppConfig есть соответствующие свойства.
|
||
# Для обратной совместимости мы добавили свойства в AppConfig, но не для prompt_file.
|
||
# Поэтому временно оставим как было: будем использовать config.prompts_dir_path и имена файлов.
|
||
# Лучше прочитать имена файлов из config.ai.prompts и загрузить их.
|
||
# Но для простоты пока оставим старый способ, если свойства есть.
|
||
# Однако в AppConfig нет system_prompt_file и т.д. – они были в BotConfig.
|
||
# В BotConfig они были путями к файлам. В AppConfig мы не стали их добавлять,
|
||
# чтобы не дублировать. Вместо этого мы загружаем все промпты в config.prompts_content.
|
||
# Поэтому мы можем использовать config.prompts_content, который уже содержит все промпты.
|
||
# Используем его:
|
||
self.ai_prompts = config.prompts_content.copy() if config.prompts_content else {}
|
||
if not self.ai_prompts:
|
||
logger.warning("Промпты не загружены из конфига, пробуем загрузить из файлов вручную...")
|
||
# Запасной вариант – загружаем из файлов, используя config.prompts_dir_path
|
||
prompts_dir = config.prompts_dir_path
|
||
if prompts_dir and prompts_dir.exists():
|
||
for prompt_name, filename in config.ai.prompts.items():
|
||
file_path = prompts_dir / filename
|
||
if file_path.exists():
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
self.ai_prompts[prompt_name] = f.read()
|
||
logger.debug(f"Загружен промпт {prompt_name} из {file_path}")
|
||
except Exception as e:
|
||
logger.error(f"Не удалось загрузить промпт {prompt_name} из {file_path}: {e}")
|
||
self.ai_prompts[prompt_name] = ""
|
||
else:
|
||
logger.warning(f"Файл промпта {filename} не найден в {prompts_dir}")
|
||
self.ai_prompts[prompt_name] = ""
|
||
else:
|
||
logger.error("Каталог промптов не найден, промпты не загружены")
|
||
|
||
# Сохраняем промпты как атрибут для передачи в запросах
|
||
self.prompts = self.ai_prompts
|
||
|
||
# ----- 3. HTTP-клиент для RAG-сервера -----
|
||
# Адрес сервера и API-ключ берём из конфига
|
||
rag_server_url = config.rag_server_url # теперь это свойство
|
||
rag_api_key = config.rag_api_key # получаем API-ключ
|
||
|
||
# Создаём RAGClient с передачей ключа
|
||
self.rag_client = RAGClient(
|
||
server_url=rag_server_url,
|
||
api_key=rag_api_key, # <-- передаём ключ
|
||
timeout=config.http_download_timeout, # можно использовать таймаут из конфига
|
||
max_retries=3,
|
||
retry_delay=1.0
|
||
)
|
||
|
||
# ----- 4. Команды и состояния -----
|
||
# Регистрируем все команды бота (из core/commands/)
|
||
self.commands = register_commands(self)
|
||
self._command_names = [cmd.name for cmd in self.commands.values() if cmd.name]
|
||
|
||
# Временные состояния (хранятся в памяти)
|
||
self.learning_mode = {} # user_jid -> (mode, room_jid)
|
||
self.last_file_texts = {} # user_jid -> текст последнего файла
|
||
self._last_file_path = {} # key -> путь к последнему файлу
|
||
self.indexing_queue = None # очередь для фоновой индексации
|
||
self.rooms = [] # комнаты, в которых бот состоит (загружается позже)
|
||
|
||
# ----- 5. Хендлеры -----
|
||
self.message_handler = MessageHandler(self)
|
||
self.file_handler = FileHandler(self)
|
||
self.room_handler = RoomHandler(self)
|
||
self.subscription_handler = SubscriptionHandler(self)
|
||
self.health_handler = HealthHandler(self)
|
||
self.indexing_worker = IndexingWorker(self)
|
||
|
||
# ----- 6. XMPP плагины -----
|
||
self.register_plugin('xep_0030') # Service Discovery
|
||
self.register_plugin('xep_0199') # Ping
|
||
self.register_plugin('xep_0045') # Multi-User Chat
|
||
self.register_plugin('xep_0054') # VCard
|
||
self.register_plugin('xep_0172') # User Nickname
|
||
self.register_plugin('xep_0066') # Out of Band Data
|
||
self.register_plugin('xep_0363') # HTTP File Upload
|
||
|
||
# ----- 7. Обработчики событий slixmpp -----
|
||
self.roster_ready = asyncio.Event()
|
||
self.add_event_handler("session_start", self.session_start)
|
||
self.add_event_handler("message", self.message_handler.handle)
|
||
self.add_event_handler("groupchat_invite", self.room_handler.accept_invite)
|
||
self.add_event_handler("roster_subscription_request", self.subscription_handler.handle_subscription)
|
||
self.add_event_handler('roster_update', self._on_roster_ready)
|
||
self.add_event_handler("muc::kicked", self.room_handler.on_kicked)
|
||
self.add_event_handler("muc::banned", self.room_handler.on_banned)
|
||
|
||
# Время старта бота (для команды !stats)
|
||
self.start_time = datetime.now()
|
||
|
||
def _on_roster_ready(self, *args):
|
||
"""Вызывается, когда ростер загружен."""
|
||
self.roster_ready.set()
|
||
|
||
def is_message_encrypted(self, msg) -> bool:
|
||
"""Проверяет, зашифровано ли сообщение (OMEMO). Заглушка."""
|
||
return False
|
||
|
||
async def handle_ai_file_request(self, msg, ai_response: str, user_jid: str, room_jid: str = None, body: str = None):
|
||
"""
|
||
Обрабатывает теги [SURGICAL_REPLACE], [REWRITE], [FILE] из ответа AI.
|
||
"""
|
||
await _handle_ai_file_request(self, msg, ai_response, user_jid, room_jid, body=body)
|
||
|
||
async def session_start(self, event):
|
||
"""
|
||
Обработчик успешного подключения и авторизации.
|
||
"""
|
||
# Устанавливаем статус присутствия
|
||
self.send_presence(pshow='chat', pstatus=self.config.status_message)
|
||
self.get_roster()
|
||
logger.info(f"Бот {self.config.name} запущен, JID: {self.boundjid}")
|
||
|
||
# Ждём загрузки ростера (максимум 10 секунд)
|
||
try:
|
||
await asyncio.wait_for(self.roster_ready.wait(), timeout=10.0)
|
||
except asyncio.TimeoutError:
|
||
logger.warning("Таймаут ожидания ростера")
|
||
|
||
# Отправляем запросы на подписку администраторам
|
||
for admin_jid in self.config.admin_jids:
|
||
if admin_jid in self.roster:
|
||
item = self.roster[admin_jid]
|
||
if item and item.get('subscription') in ('both', 'from'):
|
||
continue
|
||
self.send_presence(pto=admin_jid, ptype='subscribe')
|
||
logger.info(f"Отправлен запрос подписки для администратора {admin_jid}")
|
||
|
||
# Создаём очередь индексации и запускаем фоновый воркер
|
||
self.indexing_queue = asyncio.Queue()
|
||
asyncio.create_task(self.indexing_worker.run())
|
||
|
||
# Загружаем комнаты из конфига или БД. Пока оставляем пустым.
|
||
# Бот будет входить в комнаты только по приглашениям.
|
||
self.rooms = []
|
||
|
||
# Входим в комнаты (если есть)
|
||
for room in self.rooms:
|
||
try:
|
||
await self.plugin['xep_0045'].join_muc(room, self.boundjid.user)
|
||
logger.info(f"Вошли в комнату {room}")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка входа в комнату {room}: {e}")
|
||
|
||
# Запускаем фоновые задачи:
|
||
# - Переподключение к комнатам
|
||
# - Очистка временных файлов
|
||
asyncio.create_task(self.room_handler.rejoin_loop())
|
||
asyncio.create_task(self.health_handler.cleanup_loop()) |