Add new file
This commit is contained in:
Markov Andrey
173
bots/handlers/file_handler.py
Normal file
173
bots/handlers/file_handler.py
Normal file
@@ -0,0 +1,173 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Обработчик файлов: скачивание файлов по URL (из OOB), сохранение во временную директорию,
|
||||
постановка задачи в очередь индексации.
|
||||
|
||||
Ключевое изменение для групповой работы:
|
||||
- Если файл пришёл в групповой чат (room_jid не None), он автоматически индексируется
|
||||
в базу знаний комнаты, без необходимости команды !learn.
|
||||
- Для личных чатов – поведение зависит от режима обучения (learning_mode) и флага --global.
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
import urllib.parse
|
||||
import requests
|
||||
from datetime import datetime
|
||||
import re
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileHandler:
|
||||
"""
|
||||
Отвечает за скачивание файлов, переданных через ссылку в сообщении (XEP-0066 OOB).
|
||||
"""
|
||||
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
async def _wait_for_queue(self):
|
||||
"""
|
||||
Ожидает создания очереди индексации (защита от гонки при старте бота).
|
||||
|
||||
Проблема: очередь `self.bot.indexing_queue` создаётся в `session_start` после
|
||||
подключения к БД. Если сообщение с файлом придёт ДО этого момента,
|
||||
`indexing_queue` будет `None`, и `put()` вызовет исключение.
|
||||
Решение: ждём, пока очередь станет доступна, проверяя каждые 0.5 секунды.
|
||||
"""
|
||||
while self.bot.indexing_queue is None:
|
||||
logger.debug("Очередь индексации ещё не создана, ожидание...")
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
async def handle_file(self, msg, url: str, user_jid: str, room_jid: str = None):
|
||||
"""
|
||||
Скачивает файл и создаёт задачу для indexing_worker.
|
||||
|
||||
Аргументы:
|
||||
msg: stanza сообщения (нужна для получения user, ответных сообщений).
|
||||
url (str): HTTP/HTTPS ссылка на файл (из атрибута url элемента oob).
|
||||
user_jid (str): JID отправителя (без ресурса).
|
||||
room_jid (str | None): JID комнаты, если сообщение из groupchat; иначе None.
|
||||
|
||||
Логика:
|
||||
1. Извлечь имя файла из URL, декодировать, сгенерировать уникальное имя во временной папке.
|
||||
2. Скачать файл (синхронно, в отдельном потоке, с таймаутом из конфига).
|
||||
3. Определить режим индексации:
|
||||
- Если room_jid не None -> mode = 'room', is_global = False (всегда).
|
||||
- Если личный чат -> проверить learning_mode пользователя.
|
||||
4. Извлечь --title из тела сообщения (если есть).
|
||||
5. Положить задачу в очередь индексации (asyncio.Queue).
|
||||
6. Ответить пользователю, что файл принят в обработку.
|
||||
"""
|
||||
self.bot.metrics_handler.increment("files_processed")
|
||||
try:
|
||||
# --- 1. Извлечение имени файла из URL ---
|
||||
raw_name = url.split('/')[-1].split('?')[0]
|
||||
original_fname = urllib.parse.unquote(raw_name)
|
||||
ext = os.path.splitext(original_fname)[1].lower()
|
||||
fname = f"{msg['from'].user}_{int(datetime.now().timestamp())}{ext}"
|
||||
target_path = os.path.join(self.bot.config.temp_dir, fname)
|
||||
|
||||
# --- 2. Скачивание файла ---
|
||||
# HEAD-запрос для получения размера
|
||||
try:
|
||||
head_resp = requests.head(url, timeout=self.bot.config.http_download_timeout)
|
||||
content_length = head_resp.headers.get('Content-Length')
|
||||
if content_length is not None and int(content_length) > self.bot.config.max_file_size_mb * 1024 * 1024:
|
||||
reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Скачивание с прерыванием по лимиту
|
||||
def download_with_limit():
|
||||
r = requests.get(url, stream=True, timeout=self.bot.config.http_download_timeout)
|
||||
if r.status_code != 200:
|
||||
return False
|
||||
downloaded = 0
|
||||
limit = self.bot.config.max_file_size_mb * 1024 * 1024
|
||||
with open(target_path, 'wb') as f:
|
||||
for chunk in r.iter_content(8192):
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
if downloaded > limit:
|
||||
os.unlink(target_path)
|
||||
raise Exception("File too large")
|
||||
return True
|
||||
|
||||
try:
|
||||
success = await asyncio.to_thread(download_with_limit)
|
||||
if not success:
|
||||
reply = msg.reply("❌ Не удалось загрузить файл. Проверьте доступность ссылки.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
except Exception as e:
|
||||
if "File too large" in str(e):
|
||||
reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.")
|
||||
else:
|
||||
reply = msg.reply("❌ Ошибка при скачивании файла.")
|
||||
if reply:
|
||||
reply.send()
|
||||
return
|
||||
|
||||
# --- 3. Определение режима индексации ---
|
||||
if room_jid:
|
||||
mode = 'room'
|
||||
is_global = False
|
||||
title = None
|
||||
else:
|
||||
mode_info = self.bot.learning_mode.get(user_jid)
|
||||
if mode_info is None:
|
||||
mode = None
|
||||
is_global = False
|
||||
title = None
|
||||
else:
|
||||
if isinstance(mode_info, tuple):
|
||||
mode, saved_room = mode_info
|
||||
is_global = (mode == 'global')
|
||||
else:
|
||||
mode = mode_info
|
||||
is_global = (mode == 'global')
|
||||
title = None
|
||||
|
||||
# --- 4. Парсинг --title из тела сообщения ---
|
||||
raw_body = msg['body'].strip() if msg['body'] else ""
|
||||
if raw_body:
|
||||
match = re.search(r'--title "([^"]+)"', raw_body)
|
||||
if match:
|
||||
title = match.group(1)
|
||||
|
||||
# --- 5. Создание задачи для очереди ---
|
||||
await self._wait_for_queue()
|
||||
task = {
|
||||
"user_jid": user_jid,
|
||||
"room_jid": room_jid,
|
||||
"file_path": target_path,
|
||||
"original_fname": original_fname,
|
||||
"title": title,
|
||||
"is_global": is_global,
|
||||
"mode": mode,
|
||||
"ext": ext
|
||||
}
|
||||
await self.bot.indexing_queue.put(task)
|
||||
|
||||
# --- 6. Уведомление пользователя ---
|
||||
if room_jid:
|
||||
reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для КОМНАТНОЙ БЗ.")
|
||||
else:
|
||||
scope = "глобальную" if is_global else "личную"
|
||||
reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для {scope} БЗ.")
|
||||
if reply:
|
||||
reply.send()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Ошибка при скачивании файла: {e}")
|
||||
self.bot.metrics_handler.increment("errors_total")
|
||||
reply = msg.reply("❌ Ошибка при обработке файла. Проверьте логи.")
|
||||
if reply:
|
||||
reply.send()
|
||||
Reference in New Issue
Block a user