обновлённые методы

2026-06-30 12:00:00 +00:00
parent e5bb3c2900
commit 58ac79feac
1 changed files with 109 additions and 17 deletions
--- a/rag/services/postgres_service.py
+++ b/rag/services/postgres_service.py
@@ -4,7 +4,7 @@
 Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
 управление доступом к личным документам.

-Полностью заменяет MySQLService. Реализованы все методы, необходимые для KBService и других компонентов.
+Добавлена поддержка флага indexed для транзакционности индексации.
 """

 import asyncio
@@ -70,13 +70,15 @@ class PostgresService:
                    collection_name VARCHAR(64) NOT NULL,
                    metadata JSONB,
                    room_jid VARCHAR(255) DEFAULT NULL,
-                    file_hash VARCHAR(64) DEFAULT NULL
+                    file_hash VARCHAR(64) DEFAULT NULL,
+                    indexed BOOLEAN DEFAULT FALSE
                )
            """)
            # Индексы
            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)")
            # Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS ft_source_name 
@@ -140,32 +142,33 @@ class PostgresService:
        collection_name: str,
        metadata: dict = None,
        room_jid: str = None,
-        file_hash: str = None
+        file_hash: str = None,
+        indexed: bool = False   # новый параметр
    ) -> int:
        """
        Добавляет запись о документе. Возвращает doc_id.
-
-        :param source_name: название документа (source_name)
-        :param owner_jid: JID владельца
-        :param is_global: глобальный ли документ
-        :param collection_name: имя коллекции Qdrant
-        :param metadata: произвольные метаданные (словарь)
-        :param room_jid: JID комнаты (если документ комнатный)
-        :param file_hash: SHA-256 хеш содержимого (для обновления версий)
-        :return: id новой записи
+        :param indexed: начальное состояние индексации (по умолчанию False)
        """
        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)
-                VALUES ($1, $2, $3, $4, $5, $6, $7)
+                INSERT INTO documents (source_name, owner_jid, is_global, collection_name, metadata, room_jid, file_hash, indexed)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                RETURNING id
                """,
-                source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash
+                source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash, indexed
            )
            return row['id']

+    async def update_document_indexed(self, doc_id: int, indexed: bool) -> None:
+        """Обновляет флаг indexed для документа."""
+        async with self.pool.acquire() as conn:
+            await conn.execute(
+                "UPDATE documents SET indexed = $1 WHERE id = $2",
+                indexed, doc_id
+            )
+
    async def get_document(self, doc_id: int) -> Optional[Dict]:
        """Возвращает запись документа по id."""
        async with self.pool.acquire() as conn:
@@ -271,7 +274,7 @@ class PostgresService:
    async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
        """
        Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
-        Возвращает список документов (id, source_name, metadata).
+        Возвращает список документов (id, source_name, metadata) только с indexed = TRUE.
        """
        if not query or len(query.strip()) < 3:
            return []
@@ -282,6 +285,7 @@ class PostgresService:
                    SELECT id, source_name, metadata
                    FROM documents
                    WHERE room_jid = $1
+                      AND indexed = TRUE
                      AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
                    LIMIT 10
                    """,
@@ -294,6 +298,7 @@ class PostgresService:
                    FROM documents
                    WHERE (owner_jid = $1 OR is_global = TRUE)
                      AND room_jid IS NULL
+                      AND indexed = TRUE
                      AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
                    LIMIT 10
                    """,
@@ -324,6 +329,10 @@ class PostgresService:
        async with self.pool.acquire() as conn:
            await conn.execute("DELETE FROM rooms WHERE room_jid = $1", room_jid)

+    # ------------------------------------------------------------------
+    # Шаблоны
+    # ------------------------------------------------------------------
+
    async def save_template(self, room_jid: str, name: str, file_path: str, file_hash: str) -> bool:
        """Сохраняет шаблон в БД (если с таким именем уже есть – перезаписывает)."""
        async with self.pool.acquire() as conn:
@@ -353,4 +362,87 @@ class PostgresService:
        """Удаляет шаблон. Возвращает True, если удалён."""
        async with self.pool.acquire() as conn:
            result = await conn.execute("DELETE FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
-            return result == "DELETE 1"
+            return result == "DELETE 1"
+
+    # ------------------------------------------------------------------
+    # Получение списков документов для пользователей (только indexed = TRUE)
+    # ------------------------------------------------------------------
+
+    async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
+        """
+        Возвращает список названий документов (source_name), доступных пользователю.
+        Только те, у которых indexed = TRUE.
+        """
+        async with self.pool.acquire() as conn:
+            if room_jid:
+                rows = await conn.fetch(
+                    "SELECT source_name FROM documents WHERE room_jid = $1 AND indexed = TRUE ORDER BY source_name",
+                    room_jid
+                )
+            else:
+                rows = await conn.fetch(
+                    "SELECT source_name FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL AND indexed = TRUE ORDER BY source_name",
+                    user_jid
+                )
+            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(
+                    "SELECT source_name, is_global FROM documents WHERE room_jid = $1 AND indexed = TRUE ORDER BY source_name",
+                    room_jid
+                )
+            else:
+                rows = await conn.fetch(
+                    "SELECT source_name, is_global FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL AND indexed = TRUE ORDER BY source_name",
+                    user_jid
+                )
+            return [{"source_name": row['source_name'], "is_global": row['is_global']} for row in rows]
+
+    # ------------------------------------------------------------------
+    # Очистка баз знаний (удаление всех документов, независимо от indexed)
+    # ------------------------------------------------------------------
+
+    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",
+                user_jid
+            )
+            doc_ids = [row['id'] for row in rows]
+            if doc_ids:
+                await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
+                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
+            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]
+            if doc_ids:
+                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
+            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]
+            if doc_ids:
+                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
+            return doc_ids
+
+    async def delete_documents_by_ids(self, doc_ids: List[int]) -> None:
+        """Удаляет документы по списку id (используется в global_remove)."""
+        if not doc_ids:
+            return
+        async with self.pool.acquire() as conn:
+            await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
+            await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)