zhenxun_bot/zhenxun/services/ai/context/knowledge/vector.py

from pathlib import Path
from typing import Any, Literal

import anyio
from nonebot.adapters import Bot, Event
from pydantic import BaseModel, Field

from zhenxun.services.ai.context.knowledge.base import BaseKnowledge
from zhenxun.services.ai.context.knowledge.readers import (
    BaseReader,
    CSVReader,
    TextReader,
)
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
from zhenxun.services.ai.context.rag.models import BaseRecord
from zhenxun.services.ai.core.messages import LLMMessage
from zhenxun.services.ai.llm.api import generate_structured
from zhenxun.services.ai.run import RunContext
from zhenxun.services.ai.tools.core.decorators import tool
from zhenxun.services.ai.tools.models import ToolkitConfig, ToolResult
from zhenxun.services.log import logger


class QueryAnalysis(BaseModel):
    """大模型结构化提取查询意图"""

    keywords: list[str] = Field(
        description="提取出1~3个极其简短的搜索短语或名词，严格去除所有客套话、修饰词和标点。如果用户意图跨度较大，可以拆分为多个短语。"
    )


class VectorKnowledge(BaseKnowledge):
    """
    原生语义向量知识库。
    将长文档切分、向量化并存入关系型/向量数据库，
    向大模型提供语义检索 (Semantic Search) 工具。
    """

    default_instructions = (
        "## 语义知识库\n"
        "你拥有访问外部语义向量知识库的权限。请遵循以下规则：\n"
        "1. **优先检索**：在回答专业或背景问题时，务必使用 `search_knowledge` 工具。\n"
        "2. **语义搜索**：你可以直接输入完整的问题或描述作为检索词，"
        "系统会自动进行语义匹配。\n"
        "3. **精确过滤**：如果你需要查阅特定范围，可以在 filters 参数中"
        "传入 JSON 字典进行精确匹配（如 {'source': 'local_file'}）。\n"
        "4. **基于事实**：必须仅根据检索到的内容回答，严禁编造信息。"
    )

    default_auto_inject_template = (
        "### 📚 [本地知识库自动检索结果]\n"
        "基于用户的最新提问，系统后台已自动为你检索了以下参考资料。"
        "请你务必优先结合以下资料回答用户的问题，严禁编造：\n\n"
        "{knowledge_text}"
    )

    _global_storage: Any = None

    def __init__(
        self,
        rag_client: ScopedRAGClient | None = None,
        injection_mode: Literal["tool", "auto", "smart"] = "tool",
        query_rewrite_model: str | None = None,
        auto_inject_template: str | None = None,
        query_rewrite_prompt: str | None = None,
        query_rewrite_instruction: str | None = None,
        search_limit: int = 8,
        inject_limit: int = 12,
        **kwargs: Any,
    ):
        """
        初始化向量语义知识库工具箱。

        参数:
            rag_client: RAG 基础设施客户端实例，默认 None。
            injection_mode: 知识库的介入模式。
                - "tool": 纯工具模式 (默认)。大模型需自主思考并显式调用 `search_knowledge` 工具获取信息。
                - "auto": 自动注入模式。向大模型隐藏检索工具，直接使用用户的原始输入去数据库粗筛并静默注入。
                - "smart": 智能查询模式。向大模型隐藏检索工具，先利用 LLM 对用户的提问进行改写，再查库注入，准确率最高。
            query_rewrite_model: 在 "smart" 模式下，指定用于重写查询词的大模型名称（为空则跟随当前主模型）。
            auto_inject_template: 自动/智能注入模式下向大模型提示词注入的模板字符串，默认 None。
            query_rewrite_prompt: 智能模式下对查询词进行改写时的提示词，默认 None。
            query_rewrite_instruction: 智能模式下进行查询词改写的大模型 System 提示词说明，默认 None。
            search_limit: 单次库检索的返回记录数限制，默认 8。
            inject_limit: 最终合并去重后注入给大模型的上下文片段数上限，默认 12。
            **kwargs: 传递给父类的其他关键字参数。
        """  # noqa: E501
        self.injection_mode = injection_mode
        self.query_rewrite_model = query_rewrite_model
        self.auto_inject_template = (
            auto_inject_template or self.default_auto_inject_template
        )
        self.query_rewrite_prompt = (
            query_rewrite_prompt
            or "用户原始提问：{query}\n\n请提取核心搜索词用于专业知识库向量检索。"
        )
        self.query_rewrite_instruction = (
            query_rewrite_instruction or "你是一个资深的数据检索架构师。"
        )
        self.search_limit = search_limit
        self.inject_limit = inject_limit

        if injection_mode in ("auto", "smart"):
            config = kwargs.get("config")
            if not config:
                config = ToolkitConfig()
                kwargs["config"] = config
            if config.exclude is None:
                config.exclude = []
            config.exclude.append("search_knowledge")

        super().__init__(**kwargs)

        if rag_client is None:
            from zhenxun.services.ai.context.rag.backends import DictStorageBackend
            from zhenxun.services.ai.context.rag.builder import RAGBuilder

            rag_client = RAGBuilder(DictStorageBackend()).build()

        self.rag_client = rag_client
        self.readers: dict[str, BaseReader] = {}

        txt_reader = TextReader()
        for ext in [".txt", ".md", ".json", ".log", ".yaml", ".yml", ".ini"]:
            self.readers[ext] = txt_reader
        self.readers[".csv"] = CSVReader()

    @classmethod
    def from_event(
        cls,
        event: Event | None = None,
        bot: Bot | None = None,
        isolation: Literal["group", "user"] = "group",
        **kwargs,
    ) -> "VectorKnowledge":
        """
        根据 NoneBot 的 Event 自动推导并创建一个物理隔离的向量知识库实例。
        """
        from zhenxun.services.ai.context.memory.types import Isolation
        from zhenxun.services.ai.context.rag.backends import DictStorageBackend
        from zhenxun.services.ai.context.rag.builder import RAGBuilder
        from zhenxun.services.ai.run.context import NoneBotDeps
        from zhenxun.services.ai.utils import ContextUtils

        if not bot or not event:
            deps = NoneBotDeps.get_current()
            bot = bot or (deps.bot if deps else None)
            event = event or (deps.event if deps else None)

        if not bot or not event:
            raise ValueError(
                "无法隐式获取当前对话的 Bot 或 Event 上下文,"
                "如果您在定时任务或后台线程中使用，请显式传入 bot 和 event 参数。"
            )

        scope_builder = (
            Isolation.GROUP_SHARED()
            if isolation == "group"
            else Isolation.USER_GLOBAL()
        )
        session_meta = ContextUtils.generate_session_meta(
            bot=bot, event=event, scope_builder=scope_builder, namespace="auto_kb"
        )

        if cls._global_storage is None:
            cls._global_storage = DictStorageBackend()

        client = (
            RAGBuilder(cls._global_storage)
            .with_scope(session_meta.accessible_scopes)
            .build()
        )
        return cls(rag_client=client, **kwargs)

    def register_reader(
        self, ext: str | list[str], reader: BaseReader
    ) -> "VectorKnowledge":
        """挂载自定义后缀文件解析器 (如 PDF, Docx)，支持链式调用"""
        exts = [ext] if isinstance(ext, str) else ext
        for e in exts:
            e = e.lower()
            if not e.startswith("."):
                e = f".{e}"
            self.readers[e] = reader
        return self

    def get_instructions(self) -> str | None:
        """
        如果是自动/智能注入模式，对大模型完全隐藏检索提示词，防止其误调用。
        """
        if self.injection_mode != "tool":
            return None
        return super().get_instructions()

    async def before_llm_request(
        self, context: RunContext, messages: list[Any]
    ) -> None:
        """
        生命周期钩子：在向底层 LLM 发起请求前触发。
        负责执行 "auto" 或 "smart" 模式下的前置主动检索与上下文注入。
        """
        if self.injection_mode == "tool":
            return

        user_input = context.run.user_input
        if not user_input:
            return

        queries_to_search = [user_input]

        if self.injection_mode == "smart":
            try:
                model_to_use = self.query_rewrite_model or context.run.current_model
                prompt = self.query_rewrite_prompt.format(query=user_input)
                res = await generate_structured(
                    message=prompt,
                    response_model=QueryAnalysis,
                    model=model_to_use,
                    instruction=self.query_rewrite_instruction,
                )
                if res.keywords:
                    logger.info(
                        f"✨ [Smart Knowledge] 搜索词改写成功: '{user_input}' -> "
                        f"{res.keywords}"
                    )
                    queries_to_search = res.keywords
            except Exception as e:
                logger.warning(f"[Smart Knowledge] Query 改写失败，降级使用原词: {e}")

        all_results = []
        seen_ids = set()
        for q in queries_to_search:
            results = await self.rag_client.search(query=q, limit=self.search_limit)
            for res in results:
                if res.record.id not in seen_ids:
                    seen_ids.add(res.record.id)
                    all_results.append(res)

        if not all_results:
            return

        all_results.sort(key=lambda x: x.score, reverse=True)
        all_results = all_results[: self.inject_limit]

        formatted_results = []
        for result in all_results:
            doc_name = result.record.metadata.get("name", "未命名文档")
            formatted_results.append(
                f"📄 来源: {doc_name}\n内容片段:\n{result.record.content}"
            )

        knowledge_text = "\n\n======\n\n".join(formatted_results)
        system_prompt = self.auto_inject_template.format(knowledge_text=knowledge_text)

        messages.insert(0, LLMMessage.system(system_prompt))

    async def add_document(self, document: BaseRecord) -> int:
        """
        通过注入的 Ingestion Pipeline 处理并入库文档
        返回成功入库的 Chunk 数量。
        """
        return await self.rag_client.ingest([document])

    async def add_file(self, file_path: str | Path) -> int:
        """
        读取并注入单个文件。
        """
        aio_path = anyio.Path(file_path)
        std_path = Path(file_path)
        if not await aio_path.is_file():
            logger.error(f"[VectorKnowledge] 文件不存在: {std_path}")
            return 0

        ext = std_path.suffix.lower()
        reader = self.readers.get(ext)
        if not reader:
            logger.warning(f"当前知识库未配置支持解析文件后缀: {ext}")
            return 0

        doc = await reader.read(std_path)
        if not doc:
            return 0

        return await self.rag_client.ingest([doc])

    async def add_directory(self, dir_path: str | Path) -> int:
        """扫描目录并注入所有支持的文件"""
        total_chunks = 0
        aio_path = anyio.Path(dir_path)
        async for p in aio_path.rglob("*"):
            if await p.is_file():
                total_chunks += await self.add_file(Path(p))
        return total_chunks

    @tool(
        name="search_knowledge",
        description=(
            "在语义知识库中搜索最相关的内容片段。可以通过 filters 字典进行额外过滤。"
        ),
    )
    async def search_knowledge(
        self, query: str, filters: dict[str, Any] | None = None, limit: int = 5
    ) -> ToolResult:
        results = await self.rag_client.search(
            query=query, limit=limit, metadata_filters=filters
        )

        if not results:
            return ToolResult(output=f"知识库中未找到与 '{query}' 紧密相关的内容。")

        formatted_results = []
        for result in results:
            doc_name = result.record.metadata.get("name", "未命名文档")
            formatted_results.append(
                f"📄 来源: {doc_name}\n" f"片段内容:\n{result.record.content}"
            )

        final_text = "\n\n======\n\n".join(formatted_results)
        return ToolResult(output=final_text)