diff --git a/rag/services/postgres_service.py b/rag/services/postgres_service.py index 8a4048c..f356d34 100644 --- 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" \ No newline at end of file + 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) \ No newline at end of file