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

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')}
@@ -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
+        })