diff --git a/bots/rag_client.py b/bots/rag_client.py index ded340c..32bc58d 100644 --- a/bots/rag_client.py +++ b/bots/rag_client.py @@ -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')} @@ -376,4 +337,81 @@ class RAGClient: 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)} \ No newline at end of file + 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 + }) \ No newline at end of file