mai-bot/src/chat/knowledge/lpmm_ops.py

import asyncio
import os
from functools import partial
from typing import List, Callable, Any
from src.chat.knowledge.embedding_store import EmbeddingManager
from src.chat.knowledge.kg_manager import KGManager
from src.chat.knowledge.qa_manager import QAManager
from src.common.logger import get_logger
from src.config.config import global_config
from src.chat.knowledge import get_qa_manager, lpmm_start_up

logger = get_logger("LPMM-Plugin-API")


class LPMMOperations:
    """
    LPMM 内部操作接口。
    封装了 LPMM 的核心操作，供插件系统 API 或其他内部组件调用。
    """

    def __init__(self):
        self._initialized = False

    async def _run_cancellable_executor(self, func: Callable, *args, **kwargs) -> Any:
        """
        在线程池中执行可取消的同步操作。
        当任务被取消时（如 Ctrl+C），会立即响应并抛出 CancelledError。
        注意：线程池中的操作可能仍在运行，但协程会立即返回，不会阻塞主进程。

        Args:
            func: 要执行的同步函数
            *args: 函数的位置参数
            **kwargs: 函数的关键字参数

        Returns:
            函数的返回值

        Raises:
            asyncio.CancelledError: 当任务被取消时
        """
        loop = asyncio.get_event_loop()
        # 在线程池中执行，当协程被取消时会立即响应
        # 虽然线程池中的操作可能仍在运行，但协程不会阻塞
        return await loop.run_in_executor(None, func, *args, **kwargs)

    async def _get_managers(self) -> tuple[EmbeddingManager, KGManager, QAManager]:
        """获取并确保 LPMM 管理器已初始化"""
        qa_mgr = get_qa_manager()
        if qa_mgr is None:
            # 如果全局没初始化，尝试初始化
            if not global_config.lpmm_knowledge.enable:
                logger.warning("LPMM 知识库在全局配置中未启用，操作可能受限。")

            lpmm_start_up()
            qa_mgr = get_qa_manager()

        if qa_mgr is None:
            raise RuntimeError("无法获取 LPMM QAManager，请检查 LPMM 是否已正确安装和配置。")

        return qa_mgr.embed_manager, qa_mgr.kg_manager, qa_mgr

    async def add_content(self, text: str, auto_split: bool = True) -> dict:
        """
        向知识库添加新内容。

        Args:
            text: 原始文本。
            auto_split: 是否自动按双换行符分割段落。
                - True: 自动分割（默认），支持多段文本（用双换行分隔）
                - False: 不分割，将整个文本作为完整一段处理

        Returns:
            dict: {"status": "success/error", "count": 导入段落数, "message": "描述"}
        """
        try:
            embed_mgr, kg_mgr, _ = await self._get_managers()

            # 1. 分段处理
            if auto_split:
                # 自动按双换行符分割
                paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
            else:
                # 不分割，作为完整一段
                text_stripped = text.strip()
                if not text_stripped:
                    return {"status": "error", "message": "文本内容为空"}
                paragraphs = [text_stripped]

            if not paragraphs:
                return {"status": "error", "message": "文本内容为空"}

            # 2. 实体与三元组抽取 (内部调用大模型)
            from src.chat.knowledge.ie_process import IEProcess
            from src.services.llm_service import LLMServiceClient

            llm_ner = LLMServiceClient(
                task_name="lpmm_entity_extract", request_type="lpmm.entity_extract"
            )
            llm_rdf = LLMServiceClient(
                task_name="lpmm_rdf_build", request_type="lpmm.rdf_build"
            )
            ie_process = IEProcess(llm_ner, llm_rdf)

            logger.info(f"[Plugin API] 正在对 {len(paragraphs)} 段文本执行信息抽取...")
            extracted_docs = await ie_process.process_paragraphs(paragraphs)

            # 3. 构造并导入数据
            # 这里我们手动实现导入逻辑，不依赖外部脚本
            # a. 准备段落
            raw_paragraphs = {doc["idx"]: doc["passage"] for doc in extracted_docs}
            # b. 准备三元组
            triple_list_data = {doc["idx"]: doc["extracted_triples"] for doc in extracted_docs}

            # 向量化并入库
            # 注意：此处模仿 import_openie.py 的核心逻辑
            # 1. 先进行去重检查，只处理新段落
            # store_new_data_set 期望的格式：raw_paragraphs 的键是段落hash（不带前缀），值是段落文本
            new_raw_paragraphs = {}
            new_triple_list_data = {}

            for pg_hash, passage in raw_paragraphs.items():
                key = f"paragraph-{pg_hash}"
                if key not in embed_mgr.stored_pg_hashes:
                    new_raw_paragraphs[pg_hash] = passage
                    new_triple_list_data[pg_hash] = triple_list_data[pg_hash]

            if not new_raw_paragraphs:
                return {"status": "success", "count": 0, "message": "内容已存在，无需重复导入"}

            # 2. 使用 EmbeddingManager 的标准方法存储段落、实体和关系的嵌入
            # store_new_data_set 会自动处理嵌入生成和存储
            # 将同步阻塞操作放到线程池中执行，避免阻塞事件循环
            await self._run_cancellable_executor(embed_mgr.store_new_data_set, new_raw_paragraphs, new_triple_list_data)

            # 3. 构建知识图谱（只需要三元组数据和embedding_manager）
            await self._run_cancellable_executor(kg_mgr.build_kg, new_triple_list_data, embed_mgr)

            # 4. 持久化
            await self._run_cancellable_executor(embed_mgr.rebuild_faiss_index)
            await self._run_cancellable_executor(embed_mgr.save_to_file)
            await self._run_cancellable_executor(kg_mgr.save_to_file)

            return {
                "status": "success",
                "count": len(new_raw_paragraphs),
                "message": f"成功导入 {len(new_raw_paragraphs)} 条知识",
            }

        except asyncio.CancelledError:
            logger.warning("[Plugin API] 导入操作被用户中断")
            return {"status": "cancelled", "message": "导入操作已被用户中断"}
        except Exception as e:
            logger.error(f"[Plugin API] 导入知识失败: {e}", exc_info=True)
            return {"status": "error", "message": str(e)}

    async def search(self, query: str, top_k: int = 3) -> List[str]:
        """
        检索知识库。

        Args:
            query: 查询问题。
            top_k: 返回最相关的条目数。

        Returns:
            List[str]: 相关文段列表。
        """
        try:
            _, _, qa_mgr = await self._get_managers()
            # 直接调用 QAManager 的检索接口
            knowledge = qa_mgr.get_knowledge(query, top_k=top_k)
            # 返回通常是拼接好的字符串，这里我们可以尝试按其内部规则切分回列表，或者直接返回
            return [knowledge] if knowledge else []
        except Exception as e:
            logger.error(f"[Plugin API] 检索知识失败: {e}")
            return []

    async def delete(self, keyword: str, exact_match: bool = False) -> dict:
        """
        根据关键词或完整文段删除知识库内容。

        Args:
            keyword: 匹配关键词或完整文段。
            exact_match: 是否使用完整文段匹配（True=完全匹配，False=关键词模糊匹配）。

        Returns:
            dict: {"status": "success/info", "deleted_count": 删除条数, "message": "描述"}
        """
        try:
            embed_mgr, kg_mgr, _ = await self._get_managers()

            # 1. 查找匹配的段落
            to_delete_keys = []
            to_delete_hashes = []

            for key, item in embed_mgr.paragraphs_embedding_store.store.items():
                if exact_match:
                    # 完整文段匹配
                    if item.str.strip() == keyword.strip():
                        to_delete_keys.append(key)
                        to_delete_hashes.append(key.replace("paragraph-", "", 1))
                else:
                    # 关键词模糊匹配
                    if keyword in item.str:
                        to_delete_keys.append(key)
                        to_delete_hashes.append(key.replace("paragraph-", "", 1))

            if not to_delete_keys:
                match_type = "完整文段" if exact_match else "关键词"
                return {"status": "info", "deleted_count": 0, "message": f"未找到匹配的内容（{match_type}匹配）"}

            # 2. 执行删除
            # 将同步阻塞操作放到线程池中执行，避免阻塞事件循环

            # a. 从向量库删除
            deleted_count, _ = await self._run_cancellable_executor(
                embed_mgr.paragraphs_embedding_store.delete_items, to_delete_keys
            )
            embed_mgr.stored_pg_hashes = set(embed_mgr.paragraphs_embedding_store.store.keys())

            # b. 从知识图谱删除
            # 注意：必须使用关键字参数，避免 True 被误当作 ent_hashes 参数
            # 使用 partial 来传递关键字参数，因为 run_in_executor 不支持 **kwargs
            delete_func = partial(
                kg_mgr.delete_paragraphs, to_delete_hashes, ent_hashes=None, remove_orphan_entities=True
            )
            await self._run_cancellable_executor(delete_func)

            # 3. 持久化
            await self._run_cancellable_executor(embed_mgr.rebuild_faiss_index)
            await self._run_cancellable_executor(embed_mgr.save_to_file)
            await self._run_cancellable_executor(kg_mgr.save_to_file)

            match_type = "完整文段" if exact_match else "关键词"
            return {
                "status": "success",
                "deleted_count": deleted_count,
                "message": f"已成功删除 {deleted_count} 条相关知识（{match_type}匹配）",
            }

        except asyncio.CancelledError:
            logger.warning("[Plugin API] 删除操作被用户中断")
            return {"status": "cancelled", "message": "删除操作已被用户中断"}
        except Exception as e:
            logger.error(f"[Plugin API] 删除知识失败: {e}", exc_info=True)
            return {"status": "error", "message": str(e)}

    async def clear_all(self) -> dict:
        """
        清空整个LPMM知识库（删除所有段落、实体、关系和知识图谱数据）。

        Returns:
            dict: {"status": "success/error", "message": "描述", "stats": {...}}
        """
        try:
            embed_mgr, kg_mgr, _ = await self._get_managers()

            # 记录清空前的统计信息
            before_stats = {
                "paragraphs": len(embed_mgr.paragraphs_embedding_store.store),
                "entities": len(embed_mgr.entities_embedding_store.store),
                "relations": len(embed_mgr.relation_embedding_store.store),
                "kg_nodes": len(kg_mgr.graph.get_node_list()),
                "kg_edges": len(kg_mgr.graph.get_edge_list()),
            }

            # 将同步阻塞操作放到线程池中执行，避免阻塞事件循环

            # 1. 清空所有向量库
            # 获取所有keys
            para_keys = list(embed_mgr.paragraphs_embedding_store.store.keys())
            ent_keys = list(embed_mgr.entities_embedding_store.store.keys())
            rel_keys = list(embed_mgr.relation_embedding_store.store.keys())

            # 删除所有段落向量
            para_deleted, _ = await self._run_cancellable_executor(
                embed_mgr.paragraphs_embedding_store.delete_items, para_keys
            )
            embed_mgr.stored_pg_hashes.clear()

            # 删除所有实体向量
            if ent_keys:
                ent_deleted, _ = await self._run_cancellable_executor(
                    embed_mgr.entities_embedding_store.delete_items, ent_keys
                )
            else:
                ent_deleted = 0

            # 删除所有关系向量
            if rel_keys:
                rel_deleted, _ = await self._run_cancellable_executor(
                    embed_mgr.relation_embedding_store.delete_items, rel_keys
                )
            else:
                rel_deleted = 0

            # 2. 清空所有 embedding store 的索引和映射
            # 确保 faiss_index 和 idx2hash 也被重置，并删除旧的索引文件
            def _clear_embedding_indices():
                # 清空段落索引
                embed_mgr.paragraphs_embedding_store.faiss_index = None
                embed_mgr.paragraphs_embedding_store.idx2hash = None
                embed_mgr.paragraphs_embedding_store.dirty = False
                # 删除旧的索引文件
                if os.path.exists(embed_mgr.paragraphs_embedding_store.index_file_path):
                    os.remove(embed_mgr.paragraphs_embedding_store.index_file_path)
                if os.path.exists(embed_mgr.paragraphs_embedding_store.idx2hash_file_path):
                    os.remove(embed_mgr.paragraphs_embedding_store.idx2hash_file_path)

                # 清空实体索引
                embed_mgr.entities_embedding_store.faiss_index = None
                embed_mgr.entities_embedding_store.idx2hash = None
                embed_mgr.entities_embedding_store.dirty = False
                # 删除旧的索引文件
                if os.path.exists(embed_mgr.entities_embedding_store.index_file_path):
                    os.remove(embed_mgr.entities_embedding_store.index_file_path)
                if os.path.exists(embed_mgr.entities_embedding_store.idx2hash_file_path):
                    os.remove(embed_mgr.entities_embedding_store.idx2hash_file_path)

                # 清空关系索引
                embed_mgr.relation_embedding_store.faiss_index = None
                embed_mgr.relation_embedding_store.idx2hash = None
                embed_mgr.relation_embedding_store.dirty = False
                # 删除旧的索引文件
                if os.path.exists(embed_mgr.relation_embedding_store.index_file_path):
                    os.remove(embed_mgr.relation_embedding_store.index_file_path)
                if os.path.exists(embed_mgr.relation_embedding_store.idx2hash_file_path):
                    os.remove(embed_mgr.relation_embedding_store.idx2hash_file_path)

            await self._run_cancellable_executor(_clear_embedding_indices)

            # 3. 清空知识图谱
            # 获取所有段落hash
            all_pg_hashes = list(kg_mgr.stored_paragraph_hashes)
            if all_pg_hashes:
                # 删除所有段落节点（这会自动清理相关的边和孤立实体）
                # 注意：必须使用关键字参数，避免 True 被误当作 ent_hashes 参数
                # 使用 partial 来传递关键字参数，因为 run_in_executor 不支持 **kwargs
                delete_func = partial(
                    kg_mgr.delete_paragraphs, all_pg_hashes, ent_hashes=None, remove_orphan_entities=True
                )
                await self._run_cancellable_executor(delete_func)

            # 完全清空KG：创建新的空图（无论是否有段落hash都要执行）
            from quick_algo import di_graph

            kg_mgr.graph = di_graph.DiGraph()
            kg_mgr.stored_paragraph_hashes.clear()
            kg_mgr.ent_appear_cnt.clear()

            # 4. 保存所有数据（此时所有store都是空的，索引也是None）
            # 注意：即使store为空，save_to_file也会保存空的DataFrame，这是正确的
            await self._run_cancellable_executor(embed_mgr.save_to_file)
            await self._run_cancellable_executor(kg_mgr.save_to_file)

            after_stats = {
                "paragraphs": len(embed_mgr.paragraphs_embedding_store.store),
                "entities": len(embed_mgr.entities_embedding_store.store),
                "relations": len(embed_mgr.relation_embedding_store.store),
                "kg_nodes": len(kg_mgr.graph.get_node_list()),
                "kg_edges": len(kg_mgr.graph.get_edge_list()),
            }

            return {
                "status": "success",
                "message": f"已成功清空LPMM知识库（删除 {para_deleted} 个段落、{ent_deleted} 个实体、{rel_deleted} 个关系）",
                "stats": {
                    "before": before_stats,
                    "after": after_stats,
                },
            }

        except asyncio.CancelledError:
            logger.warning("[Plugin API] 清空操作被用户中断")
            return {"status": "cancelled", "message": "清空操作已被用户中断"}
        except Exception as e:
            logger.error(f"[Plugin API] 清空知识库失败: {e}", exc_info=True)
            return {"status": "error", "message": str(e)}


# 内部使用的单例
lpmm_ops = LPMMOperations()