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

rag/services/postgres_service.py

@@ -5,8 +5,7 @@
 управление доступом к личным документам.
 
 Добавлена поддержка флага indexed для транзакционности индексации.
-Добавлены составные индексы для ускорения полнотекстового поиска (fallback).
+Добавлена колонка file_text для хранения исходного текста документа (для переиндексации).
-ВНИМАНИЕ: Создание таблиц теперь выполняется через Alembic, а не в коде.
 """
 
 import asyncio
@@ -58,19 +57,14 @@ class PostgresService:
             command_timeout=60
         )
         if create_if_missing:
-            # Только для тестов — создаём таблицы
             await self._init_tables()
             logger.warning("Таблицы созданы автоматически (режим тестирования). В продакшене используйте Alembic.")
         else:
-            # Проверяем существование таблиц
             await self._check_tables()
         logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")
 
     async def _check_tables(self) -> None:
-        """
+        """Проверяет наличие основных таблиц."""
-        Проверяет наличие основных таблиц. Если какая-то таблица отсутствует,
-        выводит предупреждение и советует запустить миграции Alembic.
-        """
         required_tables = ['documents', 'document_access', 'history', 'rooms', 'room_templates']
         missing = []
         async with self.pool.acquire() as conn:
@@ -90,12 +84,9 @@ class PostgresService:
             logger.debug("Все необходимые таблицы существуют")
 
     async def _init_tables(self) -> None:
-        """
+        """Создаёт таблицы (только для тестов)."""
-        Создаёт таблицы (если не существуют) — используется только для тестов.
-        В продакшене используйте Alembic.
-        """
         async with self.pool.acquire() as conn:
-            # 1. Таблица documents
+            # 1. Таблица documents с колонкой file_text
             await conn.execute("""
                 CREATE TABLE IF NOT EXISTS documents (
                     id SERIAL PRIMARY KEY,
@@ -107,18 +98,16 @@ class PostgresService:
                     metadata JSONB,
                     room_jid VARCHAR(255) DEFAULT NULL,
                     file_hash VARCHAR(64) DEFAULT NULL,
-                    indexed BOOLEAN DEFAULT FALSE
+                    indexed BOOLEAN DEFAULT FALSE,
+                    file_text TEXT DEFAULT NULL
                 )
             """)
-            # Индексы
             await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner ON documents(owner_jid)")
             await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_global ON documents(is_global)")
             await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room ON documents(room_jid)")
             await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_indexed ON documents(indexed)")
-            # Составные индексы для ускорения полнотекстового поиска (fallback)
             await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room_source ON documents(room_jid, source_name)")
             await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner_source ON documents(owner_jid, source_name)")
-            # Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
             await conn.execute("""
                 CREATE INDEX IF NOT EXISTS ft_source_name 
                 ON documents USING GIN (to_tsvector('russian', source_name))
@@ -194,18 +183,26 @@ class PostgresService:
         metadata: dict = None,
         room_jid: str = None,
         file_hash: str = None,
-        indexed: bool = False
+        indexed: bool = False,
+        file_text: Optional[str] = None
     ) -> int:
-        """Добавляет запись о документе. Возвращает doc_id."""
+        """
+        Добавляет запись о документе с возможностью сохранить исходный текст.
+        Возвращает doc_id.
+        """
         metadata_json = json.dumps(metadata, ensure_ascii=False) if metadata else None
         async with self.pool.acquire() as conn:
             row = await conn.fetchrow(
                 """
-                INSERT INTO documents (source_name, owner_jid, is_global, collection_name, metadata, room_jid, file_hash, indexed)
+                INSERT INTO documents (
-                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+                    source_name, owner_jid, is_global, collection_name,
+                    metadata, room_jid, file_hash, indexed, file_text
+                )
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                 RETURNING id
                 """,
-                source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash, indexed
+                source_name, owner_jid, is_global, collection_name,
+                metadata_json, room_jid, file_hash, indexed, file_text
             )
             return row['id']
 
@@ -326,7 +323,6 @@ class PostgresService:
         """
         Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
         Возвращает список документов (id, source_name, metadata) только с indexed = TRUE.
-        Для ускорения используются составные индексы (room_jid, source_name) и (owner_jid, source_name).
         """
         if not query or len(query.strip()) < 3:
             return []
@@ -386,7 +382,7 @@ class PostgresService:
     # ------------------------------------------------------------------
 
     async def save_template(self, room_jid: str, name: str, file_path: str, file_hash: str) -> bool:
-        """Сохраняет шаблон в БД (если с таким именем уже есть – перезаписывает)."""
+        """Сохраняет шаблон в БД."""
         async with self.pool.acquire() as conn:
             await conn.execute("""
                 INSERT INTO room_templates (room_jid, name, file_path, file_hash)
@@ -414,7 +410,7 @@ class PostgresService:
             return [row['name'] for row in rows]
 
     async def delete_template(self, room_jid: str, name: str) -> bool:
-        """Удаляет шаблон. Возвращает True, если удалён."""
+        """Удаляет шаблон."""
         async with self.pool.acquire() as conn:
             result = await conn.execute(
                 "DELETE FROM room_templates WHERE room_jid = $1 AND name = $2",
@@ -427,10 +423,7 @@ class PostgresService:
     # ------------------------------------------------------------------
 
     async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
-        """
+        """Возвращает список названий документов (source_name), доступных пользователю."""
-        Возвращает список названий документов (source_name), доступных пользователю.
-        Только те, у которых indexed = TRUE.
-        """
         async with self.pool.acquire() as conn:
             if room_jid:
                 rows = await conn.fetch(
@@ -445,10 +438,7 @@ class PostgresService:
             return [row['source_name'] for row in rows]
 
     async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]:
-        """
+        """Возвращает список документов с указанием, является ли документ глобальным."""
-        Возвращает список документов с указанием, является ли документ глобальным.
-        Только indexed = TRUE.
-        """
         async with self.pool.acquire() as conn:
             if room_jid:
                 rows = await conn.fetch(
@@ -467,7 +457,7 @@ class PostgresService:
     # ------------------------------------------------------------------
 
     async def clear_user_kb(self, user_jid: str) -> List[int]:
-        """Удаляет все личные документы пользователя. Возвращает список удалённых doc_id."""
+        """Удаляет все личные документы пользователя."""
         async with self.pool.acquire() as conn:
             rows = await conn.fetch(
                 "SELECT id FROM documents WHERE owner_jid = $1 AND is_global = FALSE AND room_jid IS NULL",
@@ -480,7 +470,7 @@ class PostgresService:
             return doc_ids
 
     async def clear_room_kb(self, room_jid: str) -> List[int]:
-        """Удаляет все документы указанной комнаты. Возвращает список doc_id."""
+        """Удаляет все документы указанной комнаты."""
         async with self.pool.acquire() as conn:
             rows = await conn.fetch("SELECT id FROM documents WHERE room_jid = $1", room_jid)
             doc_ids = [row['id'] for row in rows]
@@ -489,7 +479,7 @@ class PostgresService:
             return doc_ids
 
     async def clear_global_kb(self) -> List[int]:
-        """Удаляет все глобальные документы. Возвращает список doc_id."""
+        """Удаляет все глобальные документы."""
         async with self.pool.acquire() as conn:
             rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE")
             doc_ids = [row['id'] for row in rows]