markov/fckbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Редактировать rag_client.py

Browse Source
This commit is contained in:
Markov Andrey
2026-06-30 21:35:42 +00:00
parent 1fd55f3e11
commit 25ad77f400
1 changed files with 85 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
bots/rag_client.py
View File

@@ -15,6 +15,7 @@ HTTP-клиент для взаимодействия с RAG-сервером.
- Метод vision() для распознавания текста на изображениях
- Метод transcribe() для транскрибации аудио
- Аутентификация через API-ключ (заголовок X-API-Key)
- Методы для новых эндпоинтов (clear, list_documents, delete_global, template_*, reset_history, add_room, get_rooms, get_stats, generate_document)
"""
import os
@@ -34,7 +35,7 @@ class RAGClient:
def __init__(
self,
server_url: str,
api_key: Optional[str] = None, # <-- добавлен API-ключ
api_key: Optional[str] = None,
timeout: int = 60,
max_retries: int = 3,
retry_delay: float = 1.0
@@ -50,7 +51,7 @@ class RAGClient:
retry_delay (float): базовая задержка между попытками (секунды)
"""
self.server_url = server_url.rstrip('/')
self.api_key = api_key # <-- сохраняем ключ
self.api_key = api_key
self.timeout = timeout
self.max_retries = max_retries
self.retry_delay = retry_delay
@@ -67,7 +68,6 @@ class RAGClient:
Сессия переиспользуется для всех запросов для эффективности.
"""
if self._session is None:
# Создаём сессию с таймаутом, который применяется ко всем запросам
timeout_obj = aiohttp.ClientTimeout(total=self.timeout)
self._session = aiohttp.ClientSession(timeout=timeout_obj)
logger.debug("HTTP-сессия создана")
@@ -99,7 +99,6 @@ class RAGClient:
url = f"{self.server_url}{endpoint}"
session = await self._get_session()
# Подготавливаем заголовки
headers = {}
if self.api_key:
headers['X-API-Key'] = self.api_key
@@ -113,11 +112,9 @@ class RAGClient:
if resp.status == 200:
return await resp.json()
else:
# Обрабатываем HTTP ошибки (4xx, 5xx)
error_text = await resp.text()
error_msg = f"HTTP {resp.status}: {error_text[:200]}"
logger.error(error_msg)
# Для 5xx ошибок делаем повтор, для 4xx – нет (это ошибка клиента)
if 500 <= resp.status < 600:
raise aiohttp.ClientResponseError(
status=resp.status,
@@ -125,7 +122,6 @@ class RAGClient:
headers=resp.headers
)
else:
# 4xx – не повторяем, сразу выбрасываем
raise aiohttp.ClientResponseError(
status=resp.status,
message=error_msg,
@@ -133,13 +129,11 @@ class RAGClient:
)
except (aiohttp.ClientConnectorError, aiohttp.ServerDisconnectedError) as e:
# Ошибки подключения или разрыва соединения – повторяем с увеличенной задержкой
logger.warning(
f"Сетевая ошибка при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
)
last_exception = e
if attempt < self.max_retries - 1:
# Для сетевых ошибок задержка больше: 2, 4, 8 секунд
wait = (self.retry_delay * 2) * (2 ** attempt)
logger.info(f"Повтор через {wait:.1f} секунд")
await asyncio.sleep(wait)
@@ -147,7 +141,6 @@ class RAGClient:
logger.error(f"Все попытки к {endpoint} не удались из-за сетевых проблем")
except (aiohttp.ClientPayloadError, aiohttp.ClientResponseError) as e:
# Ошибки ответа (например, повреждённые данные)
logger.warning(
f"Ошибка ответа при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
)
@@ -160,7 +153,6 @@ class RAGClient:
logger.error(f"Все попытки к {endpoint} не удались из-за ошибок ответа")
except asyncio.TimeoutError as e:
# Таймаут
logger.warning(
f"Таймаут при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
)
@@ -173,7 +165,6 @@ class RAGClient:
logger.error(f"Все попытки к {endpoint} не удались из-за таймаута")
except aiohttp.ClientError as e:
# Все остальные ошибки клиента (общий класс)
logger.warning(
f"Ошибка клиента при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
)
@@ -186,11 +177,9 @@ class RAGClient:
logger.error(f"Все попытки к {endpoint} не удались")
except Exception as e:
# Неожиданная ошибка – не повторяем, логируем и выбрасываем
logger.error(f"Неожиданная ошибка при запросе к {endpoint}: {e}", exc_info=True)
raise
# Если мы вышли из цикла, значит, все попытки провалились
error_msg = f"Не удалось выполнить запрос к {endpoint} после {self.max_retries} попыток"
if last_exception:
error_msg += f": {last_exception}"
@@ -223,6 +212,10 @@ class RAGClient:
logger.error(f"Ошибка при health check: {e}")
return False
# ============================================================
# ОСНОВНЫЕ МЕТОДЫ (query, index_document, vision, transcribe)
# ============================================================
async def query(
self,
query: str,
@@ -235,9 +228,6 @@ class RAGClient:
) -> Dict[str, Any]:
"""
Выполняет RAG-запрос через HTTP-вызов /rag/query.
Сигнатура полностью совпадает с RAGAPI.query для бесшовной замены.
ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
"""
payload = {
"query": query,
@@ -248,16 +238,13 @@ class RAGClient:
"last_file_path": last_file_path,
"last_file_text": last_file_text,
}
# Убираем None-значения, чтобы не засорять запрос
payload = {k: v for k, v in payload.items() if v is not None}
try:
result = await self._post("/rag/query", payload)
logger.debug(f"RAGClient.query успешно выполнен для {user_jid}")
return result
except Exception as e:
logger.error(f"Ошибка в RAGClient.query: {e}", exc_info=True)
# Возвращаем структурированный ответ с ошибкой
return {
"answer": f"⚠️ Ошибка связи с RAG-сервером: {str(e)}",
"intent": "ERROR",
@@ -281,8 +268,6 @@ class RAGClient:
) -> Dict[str, Any]:
"""
Индексирует документ через HTTP-вызов /rag/index.
Сигнатура полностью совпадает с RAGAPI.index_document.
"""
payload = {
"file_name": file_name,
@@ -296,7 +281,6 @@ class RAGClient:
"update_if_exists": update_if_exists,
}
payload = {k: v for k, v in payload.items() if v is not None}
try:
result = await self._post("/rag/index", payload)
logger.info(f"Документ {file_name} проиндексирован через RAG-сервер")
@@ -305,31 +289,17 @@ class RAGClient:
logger.error(f"Ошибка в RAGClient.index_document: {e}", exc_info=True)
return {"doc_id": None, "chunk_count": 0, "error": str(e)}
# ==========================================================
# НОВЫЕ МЕТОДЫ: OCR и транскрибация
# ==========================================================
async def vision(self, file_path: str) -> Dict[str, Any]:
"""
Отправляет изображение на RAG-сервер для распознавания текста (OCR).
Аргументы:
file_path: путь к файлу изображения
Возвращает:
Dict с ключами: text (str), status (str), error (опционально)
"""
if not os.path.exists(file_path):
return {"text": "", "status": "file_not_found", "error": "Файл не найден"}
session = await self._get_session()
url = f"{self.server_url}/rag/vision"
# Добавляем заголовок с ключом
headers = {}
if self.api_key:
headers['X-API-Key'] = self.api_key
try:
with open(file_path, 'rb') as f:
files = {'file': (os.path.basename(file_path), f, 'image/jpeg')}
@@ -347,23 +317,14 @@ class RAGClient:
async def transcribe(self, file_path: str) -> Dict[str, Any]:
"""
Отправляет аудиофайл на RAG-сервер для транскрибации (SaluteSpeech).
Аргументы:
file_path: путь к аудиофайлу
Возвращает:
Dict с ключами: text (str), status (str), error (опционально)
"""
if not os.path.exists(file_path):
return {"text": "", "status": "file_not_found", "error": "Файл не найден"}
session = await self._get_session()
url = f"{self.server_url}/rag/transcribe"
headers = {}
if self.api_key:
headers['X-API-Key'] = self.api_key
try:
with open(file_path, 'rb') as f:
files = {'file': (os.path.basename(file_path), f, 'audio/ogg')}
@@ -377,3 +338,80 @@ class RAGClient:
except Exception as e:
logger.error(f"Ошибка вызова транскрибации: {e}", exc_info=True)
return {"text": "", "status": "error", "error": str(e)}
# ============================================================
# НОВЫЕ МЕТОДЫ ДЛЯ КОМАНД БОТА
# ============================================================
async def clear(self, user_jid: str, room_jid: Optional[str] = None, is_global: bool = False) -> Dict:
"""Очистка базы знаний."""
return await self._post("/rag/clear", {
"user_jid": user_jid,
"room_jid": room_jid,
"is_global": is_global
})
async def list_documents(self, user_jid: str, room_jid: Optional[str] = None) -> Dict:
"""Получение списка документов."""
return await self._post("/rag/list", {
"user_jid": user_jid,
"room_jid": room_jid
})
async def delete_global(self, source_name: str) -> Dict:
"""Удаление глобального документа."""
return await self._post("/rag/delete_global", {
"source_name": source_name
})
async def template_save(self, room_jid: str, name: str, file_path: str, file_hash: str) -> Dict:
"""Сохранение шаблона."""
return await self._post("/rag/template/save", {
"room_jid": room_jid,
"name": name,
"file_path": file_path,
"file_hash": file_hash
})
async def template_list(self, room_jid: str) -> Dict:
"""Список шаблонов комнаты."""
return await self._post("/rag/template/list", {
"room_jid": room_jid
})
async def template_delete(self, room_jid: str, name: str) -> Dict:
"""Удаление шаблона."""
return await self._post("/rag/template/delete", {
"room_jid": room_jid,
"name": name
})
async def reset_history(self, user_jid: str, room_jid: Optional[str] = None) -> Dict:
"""Сброс истории диалога."""
return await self._post("/rag/reset_history", {
"user_jid": user_jid,
"room_jid": room_jid
})
async def add_room(self, room_jid: str) -> Dict:
"""Добавление комнаты."""
return await self._post("/rag/room/add", {
"room_jid": room_jid
})
async def get_rooms(self) -> List[str]:
"""Получение списка комнат."""
result = await self._post("/rag/room/list", {})
return result.get("rooms", [])
async def get_stats(self) -> Dict:
"""Получение статистики."""
return await self._post("/rag/stats", {})
async def generate_document(self, template_name: str, user_jid: str, room_jid: str) -> Dict:
"""Генерация документа по шаблону."""
return await self._post("/rag/generate_document", {
"template_name": template_name,
"user_jid": user_jid,
"room_jid": room_jid
})