Add new file
This commit is contained in:
Markov Andrey
180
core/xmpp/client.py
Normal file
180
core/xmpp/client.py
Normal file
@@ -0,0 +1,180 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
XMPP клиент – главный класс Bot.
|
||||
Теперь это тонкий клиент, который только общается с RAG-сервером.
|
||||
Не содержит БД, Qdrant, GigaChat – всё на сервере.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from slixmpp import ClientXMPP
|
||||
|
||||
from core.utils.config_loader import BotConfig
|
||||
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: BotConfig):
|
||||
# Полный 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 = {}
|
||||
prompt_files = {
|
||||
'system': config.system_prompt_file,
|
||||
'synthesis': config.synthesis_prompt_file,
|
||||
'quality': config.quality_prompt_file,
|
||||
'rerank': config.rerank_prompt_file,
|
||||
'expand': config.expand_prompt_file,
|
||||
'intent': config.intent_prompt_file,
|
||||
'metrics': config.metrics_prompt_file,
|
||||
'summary': config.summary_prompt_file,
|
||||
'consistency': config.consistency_prompt_file,
|
||||
'critique': config.critique_prompt_file,
|
||||
'spellcheck': config.spellcheck_prompt_file,
|
||||
'generate_document': config.generate_document_prompt_file,
|
||||
}
|
||||
for name, path in prompt_files.items():
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
self.ai_prompts[name] = f.read()
|
||||
logger.debug(f"Загружен промпт {name} из {path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Не удалось загрузить промпт {name} из {path}: {e}")
|
||||
self.ai_prompts[name] = ""
|
||||
# Сохраняем промпты как атрибут для передачи в запросах
|
||||
self.prompts = self.ai_prompts
|
||||
|
||||
# ----- 3. HTTP-клиент для RAG-сервера -----
|
||||
# Адрес сервера берём из конфига (bot.conf) или используем значение по умолчанию
|
||||
rag_server_url = getattr(config, 'rag_server_url', 'http://localhost:8080')
|
||||
self.rag_client = RAGClient(server_url=rag_server_url)
|
||||
|
||||
# ----- 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())
|
||||
Reference in New Issue
Block a user