59feffc190
- /rag/config_models.py - /rag/utils/config_loader.py - /rag/auth.py - /rag/rag_server.py - /rag/rag_api.py - /rag/rag_orchestrator.py - /rag/__init__.py - /bots/xmpp/client.py - /bots/rag_client.py
379 lines
18 KiB
Python
379 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
HTTP-клиент для взаимодействия с RAG-сервером.
|
||
Используется ботами вместо локального RAGAPI, когда RAG-ядро вынесено
|
||
в отдельный HTTP-сервис.
|
||
|
||
Клиент отправляет асинхронные HTTP-запросы на эндпоинты /rag/query и /rag/index.
|
||
Все методы повторяют сигнатуру RAGAPI для бесшовной замены.
|
||
|
||
ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
|
||
|
||
Улучшена обработка сетевых ошибок с повторными попытками.
|
||
|
||
ДОБАВЛЕНО:
|
||
- Метод vision() для распознавания текста на изображениях
|
||
- Метод transcribe() для транскрибации аудио
|
||
- Аутентификация через API-ключ (заголовок X-API-Key)
|
||
"""
|
||
|
||
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, # <-- добавлен API-ключ
|
||
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:
|
||
# Обрабатываем 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,
|
||
message=error_msg,
|
||
headers=resp.headers
|
||
)
|
||
else:
|
||
# 4xx – не повторяем, сразу выбрасываем
|
||
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:
|
||
# Для сетевых ошибок задержка больше: 2, 4, 8 секунд
|
||
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
|
||
|
||
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.
|
||
|
||
Сигнатура полностью совпадает с RAGAPI.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,
|
||
}
|
||
# Убираем 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",
|
||
"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.
|
||
|
||
Сигнатура полностью совпадает с RAGAPI.index_document.
|
||
"""
|
||
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)}
|
||
|
||
# ==========================================================
|
||
# НОВЫЕ МЕТОДЫ: 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')}
|
||
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).
|
||
|
||
Аргументы:
|
||
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')}
|
||
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)} |