Редактировать client.py
This commit is contained in:
Markov Andrey
@@ -7,6 +7,7 @@ XMPP клиент – главный класс Bot.
|
||||
ИСПОЛЬЗУЕТ:
|
||||
- AppConfig из rag.config_models для единой конфигурации.
|
||||
- RAGClient с передачей API-ключа для аутентификации.
|
||||
- Комнаты загружаются с сервера при старте.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -17,8 +18,7 @@ from slixmpp import ClientXMPP
|
||||
# Импорты из нового пакета rag для конфигурации
|
||||
from rag.config_models import AppConfig
|
||||
|
||||
# Остальные импорты остаются из core (старый код) – позже их нужно будет перенести в rag,
|
||||
# но пока оставляем как есть для совместимости.
|
||||
# Остальные импорты из core
|
||||
from core.services.file_service import FileService
|
||||
from core.handlers.metrics_handler import MetricsHandler
|
||||
from core.rag_client import RAGClient
|
||||
@@ -40,40 +40,19 @@ class Bot(ClientXMPP):
|
||||
Тонкий клиент: вся RAG-логика на сервере.
|
||||
"""
|
||||
|
||||
def __init__(self, config: AppConfig): # <-- изменён тип конфига
|
||||
# Полный JID: user@domain/resource
|
||||
def __init__(self, config: AppConfig):
|
||||
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():
|
||||
@@ -91,35 +70,29 @@ class Bot(ClientXMPP):
|
||||
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 с передачей ключа
|
||||
rag_server_url = config.rag_server_url
|
||||
rag_api_key = config.rag_api_key
|
||||
self.rag_client = RAGClient(
|
||||
server_url=rag_server_url,
|
||||
api_key=rag_api_key, # <-- передаём ключ
|
||||
timeout=config.http_download_timeout, # можно использовать таймаут из конфига
|
||||
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 = [] # комнаты, в которых бот состоит (загружается позже)
|
||||
self.learning_mode = {}
|
||||
self.last_file_texts = {}
|
||||
self._last_file_path = {}
|
||||
self.indexing_queue = None
|
||||
self.rooms = [] # больше не используется, загружаем с сервера
|
||||
|
||||
# ----- 5. Хендлеры -----
|
||||
self.message_handler = MessageHandler(self)
|
||||
@@ -130,13 +103,13 @@ class Bot(ClientXMPP):
|
||||
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
|
||||
self.register_plugin('xep_0030')
|
||||
self.register_plugin('xep_0199')
|
||||
self.register_plugin('xep_0045')
|
||||
self.register_plugin('xep_0054')
|
||||
self.register_plugin('xep_0172')
|
||||
self.register_plugin('xep_0066')
|
||||
self.register_plugin('xep_0363')
|
||||
|
||||
# ----- 7. Обработчики событий slixmpp -----
|
||||
self.roster_ready = asyncio.Event()
|
||||
@@ -152,24 +125,18 @@ class Bot(ClientXMPP):
|
||||
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}")
|
||||
@@ -182,31 +149,42 @@ class Bot(ClientXMPP):
|
||||
|
||||
# Отправляем запросы на подписку администраторам
|
||||
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}")
|
||||
# Загружаем комнаты с сервера и входим в них
|
||||
try:
|
||||
rooms = await self.rag_client.get_rooms()
|
||||
for room in 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}")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка загрузки комнат с сервера: {e}")
|
||||
|
||||
# Запускаем фоновые задачи:
|
||||
# - Переподключение к комнатам
|
||||
# - Очистка временных файлов
|
||||
# - Очистка словарей _last_file_path и last_file_texts
|
||||
asyncio.create_task(self.room_handler.rejoin_loop())
|
||||
asyncio.create_task(self.health_handler.cleanup_loop())
|
||||
asyncio.create_task(self._cleanup_loop())
|
||||
|
||||
async def _cleanup_loop(self):
|
||||
"""Периодическая очистка словарей _last_file_path и last_file_texts (раз в час)."""
|
||||
while True:
|
||||
await asyncio.sleep(3600)
|
||||
if len(self._last_file_path) > 100:
|
||||
self._last_file_path = dict(list(self._last_file_path.items())[-100:])
|
||||
if len(self.last_file_texts) > 100:
|
||||
self.last_file_texts = dict(list(self.last_file_texts.items())[-100:])
|
||||
|
||||
async def disconnect(self):
|
||||
"""Переопределение для закрытия HTTP-сессии."""
|
||||
await self.rag_client.close()
|
||||
await super().disconnect()
|
||||
Reference in New Issue
Block a user