417 lines
18 KiB
Python
417 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
HTTP-клиент для взаимодействия с RAG-сервером.
|
||
Используется ботами вместо локального RAGAPI, когда RAG-ядро вынесено
|
||
в отдельный HTTP-сервис.
|
||
|
||
Клиент отправляет асинхронные HTTP-запросы на эндпоинты /rag/query и /rag/index.
|
||
Все методы повторяют сигнатуру RAGAPI для бесшовной замены.
|
||
|
||
ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
|
||
|
||
Улучшена обработка сетевых ошибок с повторными попытками.
|
||
|
||
ДОБАВЛЕНО:
|
||
- Метод vision() для распознавания текста на изображениях
|
||
- Метод transcribe() для транскрибации аудио
|
||
- Аутентификация через API-ключ (заголовок X-API-Key)
|
||
- Методы для новых эндпоинтов (clear, list_documents, delete_global, template_*, reset_history, add_room, get_rooms, get_stats, generate_document)
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
from typing import Optional, Dict, List, Any
|
||
import aiohttp
|
||
import asyncio
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RAGClient:
|
||
"""
|
||
HTTP-клиент для RAG-сервера с поддержкой повторных попыток при сетевых ошибках.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
server_url: str,
|
||
api_key: Optional[str] = None,
|
||
timeout: int = 60,
|
||
max_retries: int = 3,
|
||
retry_delay: float = 1.0
|
||
):
|
||
"""
|
||
Инициализация клиента.
|
||
|
||
Аргументы:
|
||
server_url (str): базовый URL RAG-сервера (например, http://localhost:8080)
|
||
api_key (Optional[str]): API-ключ для аутентификации (передаётся в заголовке X-API-Key)
|
||
timeout (int): таймаут на запрос в секундах
|
||
max_retries (int): максимальное количество повторных попыток при ошибке
|
||
retry_delay (float): базовая задержка между попытками (секунды)
|
||
"""
|
||
self.server_url = server_url.rstrip('/')
|
||
self.api_key = api_key
|
||
self.timeout = timeout
|
||
self.max_retries = max_retries
|
||
self.retry_delay = retry_delay
|
||
self._session: Optional[aiohttp.ClientSession] = None
|
||
logger.info(
|
||
f"RAGClient инициализирован с сервером {server_url}, "
|
||
f"таймаут={timeout}с, макс. попыток={max_retries}, "
|
||
f"API-ключ {'задан' if api_key else 'не задан'}"
|
||
)
|
||
|
||
async def _get_session(self) -> aiohttp.ClientSession:
|
||
"""
|
||
Возвращает сессию aiohttp (создаёт при первом вызове).
|
||
Сессия переиспользуется для всех запросов для эффективности.
|
||
"""
|
||
if self._session is None:
|
||
timeout_obj = aiohttp.ClientTimeout(total=self.timeout)
|
||
self._session = aiohttp.ClientSession(timeout=timeout_obj)
|
||
logger.debug("HTTP-сессия создана")
|
||
return self._session
|
||
|
||
async def close(self):
|
||
"""Закрывает HTTP-сессию."""
|
||
if self._session:
|
||
await self._session.close()
|
||
self._session = None
|
||
logger.debug("RAGClient сессия закрыта")
|
||
|
||
async def _post(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""
|
||
Внутренний метод для отправки POST-запроса с повторными попытками.
|
||
Автоматически добавляет заголовок X-API-Key, если ключ задан.
|
||
|
||
Аргументы:
|
||
endpoint (str): путь эндпоинта (например, '/rag/query')
|
||
payload (Dict): тело запроса в формате JSON
|
||
|
||
Возвращает:
|
||
Dict: ответ сервера
|
||
|
||
Исключения:
|
||
aiohttp.ClientError: если все попытки не удались
|
||
asyncio.TimeoutError: если все попытки истекли по таймауту
|
||
"""
|
||
url = f"{self.server_url}{endpoint}"
|
||
session = await self._get_session()
|
||
|
||
headers = {}
|
||
if self.api_key:
|
||
headers['X-API-Key'] = self.api_key
|
||
|
||
last_exception = None
|
||
|
||
for attempt in range(self.max_retries):
|
||
try:
|
||
logger.debug(f"Попытка {attempt + 1}/{self.max_retries} к {endpoint}")
|
||
async with session.post(url, json=payload, headers=headers) as resp:
|
||
if resp.status == 200:
|
||
return await resp.json()
|
||
else:
|
||
error_text = await resp.text()
|
||
error_msg = f"HTTP {resp.status}: {error_text[:200]}"
|
||
logger.error(error_msg)
|
||
if 500 <= resp.status < 600:
|
||
raise aiohttp.ClientResponseError(
|
||
status=resp.status,
|
||
message=error_msg,
|
||
headers=resp.headers
|
||
)
|
||
else:
|
||
raise aiohttp.ClientResponseError(
|
||
status=resp.status,
|
||
message=error_msg,
|
||
headers=resp.headers
|
||
)
|
||
|
||
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:
|
||
wait = (self.retry_delay * 2) * (2 ** attempt)
|
||
logger.info(f"Повтор через {wait:.1f} секунд")
|
||
await asyncio.sleep(wait)
|
||
else:
|
||
logger.error(f"Все попытки к {endpoint} не удались из-за сетевых проблем")
|
||
|
||
except (aiohttp.ClientPayloadError, aiohttp.ClientResponseError) as e:
|
||
logger.warning(
|
||
f"Ошибка ответа при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
|
||
)
|
||
last_exception = e
|
||
if attempt < self.max_retries - 1:
|
||
wait = self.retry_delay * (2 ** attempt)
|
||
logger.info(f"Повтор через {wait:.1f} секунд")
|
||
await asyncio.sleep(wait)
|
||
else:
|
||
logger.error(f"Все попытки к {endpoint} не удались из-за ошибок ответа")
|
||
|
||
except asyncio.TimeoutError as e:
|
||
logger.warning(
|
||
f"Таймаут при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
|
||
)
|
||
last_exception = e
|
||
if attempt < self.max_retries - 1:
|
||
wait = self.retry_delay * (2 ** attempt)
|
||
logger.info(f"Повтор через {wait:.1f} секунд")
|
||
await asyncio.sleep(wait)
|
||
else:
|
||
logger.error(f"Все попытки к {endpoint} не удались из-за таймаута")
|
||
|
||
except aiohttp.ClientError as e:
|
||
logger.warning(
|
||
f"Ошибка клиента при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
|
||
)
|
||
last_exception = e
|
||
if attempt < self.max_retries - 1:
|
||
wait = self.retry_delay * (2 ** attempt)
|
||
logger.info(f"Повтор через {wait:.1f} секунд")
|
||
await asyncio.sleep(wait)
|
||
else:
|
||
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}"
|
||
logger.error(error_msg)
|
||
raise aiohttp.ClientError(error_msg)
|
||
|
||
async def health_check(self) -> bool:
|
||
"""
|
||
Проверяет доступность RAG-сервера через эндпоинт /health.
|
||
Возвращает True, если сервер отвечает со статусом 200.
|
||
"""
|
||
try:
|
||
session = await self._get_session()
|
||
headers = {}
|
||
if self.api_key:
|
||
headers['X-API-Key'] = self.api_key
|
||
async with session.get(f"{self.server_url}/health", timeout=10, headers=headers) as resp:
|
||
if resp.status == 200:
|
||
data = await resp.json()
|
||
if data.get('status') == 'healthy':
|
||
logger.debug("RAG-сервер доступен (health check пройден)")
|
||
return True
|
||
else:
|
||
logger.warning(f"RAG-сервер вернул нездоровый статус: {data}")
|
||
return False
|
||
else:
|
||
logger.warning(f"Health check вернул HTTP {resp.status}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при health check: {e}")
|
||
return False
|
||
|
||
# ============================================================
|
||
# ОСНОВНЫЕ МЕТОДЫ (query, index_document, vision, transcribe)
|
||
# ============================================================
|
||
|
||
async def query(
|
||
self,
|
||
query: str,
|
||
user_jid: str,
|
||
room_jid: Optional[str],
|
||
prompts: Optional[Dict[str, str]] = None,
|
||
intent_override: Optional[str] = None,
|
||
last_file_path: Optional[str] = None,
|
||
last_file_text: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Выполняет RAG-запрос через HTTP-вызов /rag/query.
|
||
"""
|
||
payload = {
|
||
"query": query,
|
||
"user_jid": user_jid,
|
||
"room_jid": room_jid,
|
||
"prompts": prompts or {},
|
||
"intent_override": intent_override,
|
||
"last_file_path": last_file_path,
|
||
"last_file_text": last_file_text,
|
||
}
|
||
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",
|
||
"context": "",
|
||
"sources": [],
|
||
"confidence": None,
|
||
"error": str(e)
|
||
}
|
||
|
||
async def index_document(
|
||
self,
|
||
file_name: str,
|
||
file_text: str,
|
||
user_jid: str,
|
||
room_jid: Optional[str],
|
||
is_global: bool = False,
|
||
title: Optional[str] = None,
|
||
metadata: Optional[Dict] = None,
|
||
file_hash: Optional[str] = None,
|
||
update_if_exists: bool = True
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Индексирует документ через HTTP-вызов /rag/index.
|
||
"""
|
||
payload = {
|
||
"file_name": file_name,
|
||
"file_text": file_text,
|
||
"user_jid": user_jid,
|
||
"room_jid": room_jid,
|
||
"is_global": is_global,
|
||
"title": title,
|
||
"metadata": metadata,
|
||
"file_hash": file_hash,
|
||
"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-сервер")
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"Ошибка в RAGClient.index_document: {e}", exc_info=True)
|
||
return {"doc_id": None, "chunk_count": 0, "error": str(e)}
|
||
|
||
async def vision(self, file_path: str) -> Dict[str, Any]:
|
||
"""
|
||
Отправляет изображение на RAG-сервер для распознавания текста (OCR).
|
||
"""
|
||
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')}
|
||
async with session.post(url, data=files, headers=headers) as resp:
|
||
if resp.status == 200:
|
||
return await resp.json()
|
||
else:
|
||
error_text = await resp.text()
|
||
logger.error(f"OCR ошибка: {resp.status} - {error_text[:200]}")
|
||
return {"text": "", "status": "error", "error": error_text}
|
||
except Exception as e:
|
||
logger.error(f"Ошибка вызова OCR: {e}", exc_info=True)
|
||
return {"text": "", "status": "error", "error": str(e)}
|
||
|
||
async def transcribe(self, file_path: str) -> Dict[str, Any]:
|
||
"""
|
||
Отправляет аудиофайл на RAG-сервер для транскрибации (SaluteSpeech).
|
||
"""
|
||
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')}
|
||
async with session.post(url, data=files, headers=headers) as resp:
|
||
if resp.status == 200:
|
||
return await resp.json()
|
||
else:
|
||
error_text = await resp.text()
|
||
logger.error(f"Транскрибация ошибка: {resp.status} - {error_text[:200]}")
|
||
return {"text": "", "status": "error", "error": error_text}
|
||
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
|
||
}) |