AI项目实战（5）企业知识库问答 Agent

"""项目二：领域模型与数据类型。

定义 RAG 系统中使用的数据结构、状态类型和配置常量。
"""
from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum, auto
from pathlib import Path
from typing import Any

from core.config import DATA_DIR, get_settings

# =========================================================================
# 目录配置
# =========================================================================
DOCS_DIR = DATA_DIR / "knowledge_base"
DOCS_DIR.mkdir(parents=True, exist_ok=True)


def get_vector_store_dir() -> Path:
    """获取向量存储目录路径。"""
    settings = get_settings()
    return Path(settings.vector_store_dir) / "knowledge_qa"


# =========================================================================
# 数据传输对象 (DTO)
# =========================================================================
@dataclass
class DocumentChunk:
    """文档片段数据模型。"""
    text: str
    source: str = "未知文档"
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.text.strip():
            raise ValueError("文档内容不能为空")


@dataclass
class RetrievalResult:
    """检索结果数据模型。"""
    text: str
    score: float
    source: str

    @classmethod
    def from_node(cls, node: Any) -> "RetrievalResult":
        """从 LlamaIndex 的 Node 对象构建检索结果。"""
        source = node.metadata.get("file_name", "未知文档")
        text = node.get_text()
        score = round(node.score or 0.0, 4)
        return cls(text=text, score=score, source=source)

    def to_dict(self) -> dict[str, Any]:
        """转换为字典格式。"""
        return {
            "text": self.text,
            "score": self.score,
            "source": self.source,
        }


# =========================================================================
# 索引状态
# =========================================================================
class IndexState(Enum):
    """索引状态枚举。"""
    NOT_LOADED = "not_loaded"
    LOADING = "loading"
    LOADED = "loaded"
    ERROR = "error"

"""项目二：Prompt 定义。

集中管理所有系统提示词和生成模板，便于统一优化和版本管理。
"""
from __future__ import annotations

from typing import Final

# =========================================================================
# RAG 系统提示词
# =========================================================================
RAG_SYSTEM_PROMPT: Final[str] = """你是企业知识库问答助手，基于公司内部文档回答员工问题。

工作规则：
1. 回答任何问题前，必须先调用 retrieve_knowledge 检索相关文档。
2. 严格基于检索到的文档内容回答，不得添加文档中没有的信息。
3. 如果检索结果为空，回答"未在企业知识库中找到相关文档，建议联系 HR 部门确认"。
4. 回答时，在关键信息后标注出处，例如：【来源：员工手册.pdf】
5. 回答用中文，简洁清晰。
6. 如果用户要求添加文档，调用 add_document_to_kb 工具。
"""

# =========================================================================
# 答案生成 Prompt（用于直接生成模式）
# =========================================================================
ANSWER_GENERATION_PROMPT: Final[str] = """基于以下检索到的文档内容，回答用户的问题。

【检索到的文档】
{context}

【用户问题】
{query}

【回答要求】
1. 只使用上面文档中出现的信息回答
2. 不添加任何文档中没有的内容
3. 如果文档中没有相关信息，回答"未在企业知识库中找到相关文档，建议联系 HR 部门确认"
4. 回答结尾标注信息来源，例如：【来源：员工手册.pdf】
5. 用中文回答，简洁明了
"""

# =========================================================================
# 工具描述
# =========================================================================
RETRIEVE_KNOWLEDGE_DESCRIPTION: Final[str] = """从企业知识库中检索与查询相关的文档内容。

适用于：政策查询、流程说明、制度问答、产品文档等。
返回最相关的文档片段及出处。"""

ADD_DOCUMENT_DESCRIPTION: Final[str] = """将文档文件添加到知识库中（支持 PDF / Word / Markdown / TXT）。

file_path 是文档的绝对路径或相对于知识库目录的路径。"""

"""项目二：仓储层实现。

封装 LlamaIndex 向量索引的底层操作，提供清晰的仓储接口。
处理文档加载、索引构建、检索、持久化等数据访问逻辑。
"""
from __future__ import annotations

from pathlib import Path
from typing import Any

from core.config import get_settings
from core.logging_conf import get_logger
from .models import (
    DocumentChunk,
    IndexState,
    RetrievalResult,
    get_vector_store_dir,
)

logger = get_logger("p02.knowledge_qa.repository")


# =========================================================================
# RAG 仓储接口
# =========================================================================
class RAGRepository:
    """RAG 仓储层，封装向量索引的所有操作。

    职责：
    - 管理索引的生命周期（创建/加载/持久化）
    - 文档索引与增量更新
    - 语义检索与结果格式化
    """

    def __init__(self) -> None:
        self._index: Any | None = None
        self._state = IndexState.NOT_LOADED

    def _ensure_loaded(self) -> None:
        """惰性加载索引（避免启动时加载大索引）。

        Raises:
            RuntimeError: 索引加载失败时抛出
        """
        if self._state == IndexState.LOADED:
            return

        if self._state == IndexState.LOADING:
            logger.warning("索引正在加载中，等待完成...")
            return

        self._state = IndexState.LOADING
        try:
            self._build_index()
            self._state = IndexState.LOADED
        except Exception as e:
            self._state = IndexState.ERROR
            logger.error("索引加载失败: %s", e)
            raise RuntimeError(f"RAG 索引加载失败: {e}") from e

    def _build_index(self) -> None:
        """构建或加载向量索引。

        如果已有持久化索引则加载，否则创建空索引。

        注意：此方法会临时修改 LlamaIndex 全局 Settings（embed_model 和 node_parser），
        操作完成后会恢复原始值，以避免污染其他模块的配置。
        """
        from llama_index.core import (
            Settings,
            StorageContext,
            VectorStoreIndex,
            load_index_from_storage,
        )
        from llama_index.core.node_parser import SentenceSplitter
        from llama_index.embeddings.huggingface import HuggingFaceEmbedding

        settings = get_settings()

        # 保存原始 Settings，构建索引后恢复（防止全局污染）
        original_embed_model = Settings.embed_model
        original_node_parser = Settings.node_parser

        try:
            # 临时配置 Embedding 模型和文档切分器
            logger.info("初始化 Embedding 模型: %s", settings.embedding_model)
            Settings.embed_model = HuggingFaceEmbedding(
                model_name=settings.embedding_model,
                embed_batch_size=32,
            )

            Settings.node_parser = SentenceSplitter(
                chunk_size=800,
                chunk_overlap=100,
            )

            persist_dir = get_vector_store_dir()
            storage_dir = Path(persist_dir)

            if (storage_dir / "docstore.json").exists():
                # 已存在索引，加载
                logger.info("加载已有索引: %s", persist_dir)
                storage_context = StorageContext.from_defaults(
                    persist_dir=str(persist_dir),
                )
                self._index = load_index_from_storage(storage_context)
                logger.info("RAG 索引加载完成")
            else:
                # 新建空索引
                logger.info("创建新索引: %s", persist_dir)
                self._index = VectorStoreIndex([])
                storage_dir.mkdir(parents=True, exist_ok=True)
                self._index.storage_context.persist(persist_dir=str(persist_dir))
                logger.info("RAG 空索引已创建")
        finally:
            # 恢复原始 Settings，防止污染其他模块
            Settings.embed_model = original_embed_model
            Settings.node_parser = original_node_parser
            logger.debug("已恢复 LlamaIndex Settings 原始值")

    def add_document(self, file_path: str | Path) -> int:
        """追加文档到索引（增量，不重建）。

        Args:
            file_path: 文档路径，支持单个文件或目录

        Returns:
            新增的文档片段数量

        Raises:
            FileNotFoundError: 文件不存在时抛出
            RuntimeError: 索引加载失败时抛出
        """
        self._ensure_loaded()

        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"文档不存在: {file_path}")

        from llama_index.core import SimpleDirectoryReader

        if path.is_dir():
            documents = SimpleDirectoryReader(str(path)).load_data()
        else:
            documents = SimpleDirectoryReader(
                input_files=[str(path)]
            ).load_data()

        if not documents:
            logger.warning("文档为空或无内容: %s", file_path)
            return 0

        logger.info("正在索引文档: %s (%d 页)", file_path, len(documents))

        for doc in documents:
            self._index.insert(doc)

        # 持久化更新
        persist_dir = get_vector_store_dir()
        self._index.storage_context.persist(persist_dir=str(persist_dir))

        logger.info("文档索引完成: %s, %d 个片段", file_path, len(documents))
        return len(documents)

    def retrieve(
        self,
        query: str,
        top_k: int | None = None,
    ) -> list[RetrievalResult]:
        """检索最相关的文档片段。

        Args:
            query: 查询文本
            top_k: 返回结果数量，None 则使用配置值

        Returns:
            检索结果列表，按相关度降序排列

        Raises:
            RuntimeError: 索引加载失败时抛出
        """
        if not query.strip():
            logger.warning("空查询，返回空结果")
            return []

        self._ensure_loaded()

        settings = get_settings()
        k = top_k or settings.top_k_retrieval
        threshold = settings.retrieval_similarity_threshold

        logger.debug("执行检索: query=%s, top_k=%d, threshold=%.2f",
                    query[:50], k, threshold)

        retriever = self._index.as_retriever(
            similarity_top_k=k,
            similarity_cutoff=threshold,
        )
        nodes = retriever.retrieve(query)

        results = []
        for node in nodes:
            result = RetrievalResult.from_node(node)
            if result.text.strip():
                results.append(result)

        logger.debug("检索完成: %d 条结果", len(results))
        return results

    def retrieve_as_dict(
        self,
        query: str,
        top_k: int | None = None,
    ) -> list[dict[str, Any]]:
        """检索并返回字典格式的结果（工具函数兼容）。"""
        results = self.retrieve(query, top_k)
        return [r.to_dict() for r in results]

    def clear_index(self) -> None:
        """清空当前索引（用于测试）。"""
        from llama_index.core import VectorStoreIndex
        self._index = VectorStoreIndex([])
        self._state = IndexState.LOADED


# 全局单例仓储实例
_repository: RAGRepository | None = None


def get_rag_repository() -> RAGRepository:
    """获取 RAG 仓储单例。"""
    global _repository
    if _repository is None:
        _repository = RAGRepository()
    return _repository


def set_rag_repository(repo: RAGRepository | None) -> None:
    """设置 RAG 仓储实例（用于测试 Mock）。"""
    global _repository
    _repository = repo

"""项目二：服务层与流程编排。

使用 LangGraph StateGraph 实现完整的 RAG 工作流：
问题 -> 检索 -> 生成答案 -> 返回结果

状态包含 question、contexts、answer，节点包括检索和生成。
"""
from __future__ import annotations

from typing import Any, TypedDict

from langchain.tools import tool
from langgraph.graph import END, StateGraph

from core.logging_conf import get_logger
from projects.p02_knowledge_qa.models import RetrievalResult
from projects.p02_knowledge_qa.prompts import (
    ADD_DOCUMENT_DESCRIPTION,
    ANSWER_GENERATION_PROMPT,
    RETRIEVE_KNOWLEDGE_DESCRIPTION,
    RAG_SYSTEM_PROMPT,
)
from projects.p02_knowledge_qa.repositories import get_rag_repository

logger = get_logger("p02.knowledge_qa.service")


# =========================================================================
# LangGraph RAG State 定义
# =========================================================================
class RAGState(TypedDict):
    """RAG 工作流状态。"""
    question: str  # 用户问题
    contexts: list[str]  # 检索到的文档片段列表
    answer: str  # 生成的答案


# =========================================================================
# 工具函数（保持向后兼容）
# =========================================================================
@tool
def retrieve_knowledge(query: str) -> str:
    """从企业知识库中检索与查询相关的文档内容。

    适用于：政策查询、流程说明、制度问答、产品文档等。
    返回最相关的文档片段及出处。
    """
    try:
        repo = get_rag_repository()
        results = repo.retrieve(query)
        return _format_retrieval_results(results)
    except Exception as e:
        logger.error("知识库检索失败: query=%s error=%s", query, e)
        return "知识库检索服务暂时不可用，请稍后重试。"


@tool
def add_document_to_kb(file_path: str) -> str:
    """将文档文件添加到知识库中（支持 PDF / Word / Markdown / TXT）。

    file_path 是文档的绝对路径或相对于知识库目录的路径。
    """
    try:
        repo = get_rag_repository()
        cnt = repo.add_document(file_path)
        return f"已成功将《{file_path}》添加到知识库，共索引 {cnt} 个知识片段。"
    except FileNotFoundError as e:
        return f"文档《{file_path}》不存在，请检查路径是否正确。"
    except Exception as e:
        logger.error("文档索引失败: file=%s error=%s", file_path, e)
        return f"文档《{file_path}》添加失败: {e}"


def _format_retrieval_results(results: list[RetrievalResult]) -> str:
    """将检索结果格式化为可读的字符串。"""
    if not results:
        return "未在企业知识库中找到与您问题相关的文档。"

    lines = []
    for i, result in enumerate(results, 1):
        lines.append(
            f"[{i}] 来源《{result.source}》（相关度: {result.score}）\n"
            f"{result.text}"
        )
    return "\n\n---\n\n".join(lines)


# =========================================================================
# LangGraph RAG 工作流节点
# =========================================================================
def retrieve_node(state: RAGState) -> dict[str, Any]:
    """检索节点：根据用户问题检索相关文档。

    Args:
        state: 当前工作流状态，包含 question

    Returns:
        更新后的状态，包含检索到的 contexts
    """
    question = state["question"]
    logger.debug("RAG 工作流 - 执行检索: %s", question[:50])

    repo = get_rag_repository()
    results = repo.retrieve(question)
    contexts = [f"来源《{r.source}》：{r.text}" for r in results]

    logger.debug("RAG 工作流 - 检索完成: %d 条结果", len(contexts))
    return {"contexts": contexts}


def generate_node(state: RAGState) -> dict[str, Any]:
    """生成节点：基于检索结果生成答案。

    Args:
        state: 当前工作流状态，包含 question 和 contexts

    Returns:
        更新后的状态，包含生成的 answer
    """
    from core import build_chat_model

    question = state["question"]
    contexts = state["contexts"]

    logger.debug("RAG 工作流 - 执行生成: %s", question[:50])

    if not contexts:
        answer = "未在企业知识库中找到相关文档，建议联系 HR 部门确认。"
    else:
        context_text = "\n\n".join(contexts)
        prompt = ANSWER_GENERATION_PROMPT.format(
            context=context_text,
            query=question,
        )
        model = build_chat_model()
        response = model.invoke(prompt)
        answer = response.content

    logger.debug("RAG 工作流 - 生成完成")
    return {"answer": answer}


# =========================================================================
# LangGraph RAG 工作流构建
# =========================================================================
def build_rag_workflow_graph() -> StateGraph:
    """构建 RAG 工作流图。

    工作流：
    question -> retrieve_node -> contexts -> generate_node -> answer -> END
    """
    workflow = StateGraph(RAGState)

    # 添加节点
    workflow.add_node("retrieve", retrieve_node)
    workflow.add_node("generate", generate_node)

    # 设置入口点
    workflow.set_entry_point("retrieve")

    # 添加边：检索 -> 生成 -> 结束
    workflow.add_edge("retrieve", "generate")
    workflow.add_edge("generate", END)

    return workflow


# =========================================================================
# RAGWorkflow 类（保持向后兼容）
# =========================================================================
class RAGWorkflow:
    """RAG 工作流编排，基于 LangGraph StateGraph 实现。

    完整的检索-生成流水线：
    1. retrieve_node: 根据问题检索相关文档
    2. generate_node: 基于检索结果生成答案

    同时保持与 create_agent 的向后兼容性。
    """

    @staticmethod
    def get_tools() -> list[Any]:
        """获取 RAG 工具列表（用于 create_agent 兼容）。"""
        return [retrieve_knowledge, add_document_to_kb]

    @staticmethod
    def get_system_prompt() -> str:
        """获取系统提示词（用于 create_agent 兼容）。"""
        return RAG_SYSTEM_PROMPT

    @staticmethod
    def compile_graph() -> Any:
        """编译并返回 LangGraph 工作流图。

        Returns:
            编译后的 LangGraph 可执行图
        """
        workflow = build_rag_workflow_graph()
        return workflow.compile()

    @staticmethod
    def run(question: str) -> str:
        """运行 RAG 工作流回答问题。

        Args:
            question: 用户问题

        Returns:
            生成的答案
        """
        graph = RAGWorkflow.compile_graph()
        result = graph.invoke({"question": question})
        return result["answer"]


# 导出工具函数供外部使用
__all__ = [
    "retrieve_knowledge",
    "add_document_to_kb",
    "RAGWorkflow",
    "RAGState",
    "retrieve_node",
    "generate_node",
    "build_rag_workflow_graph",
]

"""项目二：项目定义。

定义 KnowledgeQAProject 类，继承 BaseProject，
实现 Agent 构建和项目注册。
"""
from __future__ import annotations

from typing import Any

from core import BaseProject, build_chat_model, registry
from projects.p02_knowledge_qa.graph import RAGWorkflow


class KnowledgeQAProject(BaseProject):
    """企业知识库问答 Agent 项目。

    提供基于 RAG 技术的企业内部文档问答能力，支持：
    - 文档索引与增量更新
    - 语义检索
    - 带出处的精准回答
    - 防幻觉机制
    """

    id = "p02_knowledge_qa"
    name = "企业知识库问答 Agent"
    description = "基于 LangGraph+LlamaIndex 的 RAG 问答，附带引用出处。"
    capabilities = ["RAG", "LangGraph", "LlamaIndex"]

    def build_agent(self) -> Any:
        """构建 RAG Agent 实例。

        使用 create_agent 统一构建方式，集成：
        - 知识库检索工具
        - 文档添加工具
        - RAG 专用系统提示词
        """
        from langchain.agents import create_agent

        return create_agent(
            model=build_chat_model(),
            tools=RAGWorkflow.get_tools(),
            system_prompt=RAGWorkflow.get_system_prompt(),
        )


# 注册项目实例
registry.register(KnowledgeQAProject())

"""项目二：企业知识库问答 Agent（生产级 RAG）

主讲能力：RAG（检索增强生成），LangGraph + LlamaIndex

业务场景：企业内部文档问答，员工可以问"年假有几天""报销流程是什么"，
Agent 基于上传的 PDF/Word/Markdown 文档精准回答，附带引用出处。

生产级特性：
- LlamaIndex 向量索引（支持混合检索：关键词 BM25 + 语义 Embedding）
- LangGraph 编排检索-重排-生成流水线
- 引用溯源（每个答案标注来源于哪个文档的哪一段）
- 增量更新（新增文档无需重建全量索引）
- 防幻觉（资料中没有就说"未找到"，绝不编造）
- 多格式文档支持（PDF / Word / Markdown / TXT）

三层架构：
- models.py: 领域模型与数据类型
- prompts.py: 系统提示词与答案生成模板
- repositories.py: RAG 仓储层（索引、检索、持久化）
- graph.py: 服务层与工具函数
- project.py: 项目定义与 Agent 构建
"""
from __future__ import annotations

# 向后兼容：重新导出旧版 API，保持导入路径不变
from projects.p02_knowledge_qa.graph import (
    RAGWorkflow,
    add_document_to_kb,
    retrieve_knowledge,
)
from projects.p02_knowledge_qa.models import (
    DOCS_DIR,
    DocumentChunk,
    IndexState,
    RetrievalResult,
)
from projects.p02_knowledge_qa.project import KnowledgeQAProject
from projects.p02_knowledge_qa.repositories import (
    RAGRepository,
    get_rag_repository,
    set_rag_repository,
)

# 为了完全兼容旧版，提供 RAGEngine 别名
RAGEngine = RAGRepository


def _get_rag() -> RAGRepository:
    """向后兼容：旧版代码使用 _get_rag() 获取 RAG 引擎实例。"""
    return get_rag_repository()


# 导出所有公共 API
__all__ = [
    # 数据模型
    "DocumentChunk",
    "RetrievalResult",
    "IndexState",
    "DOCS_DIR",
    # 仓储层
    "RAGRepository",
    "RAGEngine",  # 兼容别名
    "get_rag_repository",
    "set_rag_repository",
    "_get_rag",  # 兼容旧版
    # 工具与服务
    "retrieve_knowledge",
    "add_document_to_kb",
    "RAGWorkflow",
    # 项目
    "KnowledgeQAProject",
]

"""项目二：企业知识库问答 Agent - 测试套件。

覆盖三层架构各层的单元测试和集成测试：
- models: 数据模型验证
- repositories: 仓储层测试（索引、检索、持久化）
- graph: 工具函数测试
- project: Agent 集成测试

运行：
  cd backend
  pytest projects/p02_knowledge_qa/test_agent.py -v
"""
from __future__ import annotations

import os
import tempfile
from pathlib import Path

import pytest


# =========================================================================
# 模型层测试
# =========================================================================
class TestModels:
    """测试领域模型与数据类型。"""

    def test_index_state_enum(self) -> None:
        """测试 IndexState 是真正的 Enum 类型。"""
        from enum import Enum
        from projects.p02_knowledge_qa.models import IndexState

        assert issubclass(IndexState, Enum), "IndexState 应该是 Enum 类型"
        assert IndexState.NOT_LOADED.value == "not_loaded"
        assert IndexState.LOADING.value == "loading"
        assert IndexState.LOADED.value == "loaded"
        assert IndexState.ERROR.value == "error"

    def test_document_chunk_validation(self) -> None:
        """测试 DocumentChunk 不能为空内容验证。"""
        from projects.p02_knowledge_qa.models import DocumentChunk

        # 正常情况
        chunk = DocumentChunk(text="测试内容", source="test.txt")
        assert chunk.text == "测试内容"
        assert chunk.source == "test.txt"

        # 空内容应该抛出异常
        with pytest.raises(ValueError, match="不能为空"):
            DocumentChunk(text="   ", source="test.txt")

    def test_retrieval_result_conversion(self) -> None:
        """测试 RetrievalResult 的字典转换。"""
        from projects.p02_knowledge_qa.models import RetrievalResult

        result = RetrievalResult(text="测试内容", score=0.95, source="test.txt")
        result_dict = result.to_dict()

        assert result_dict["text"] == "测试内容"
        assert result_dict["score"] == 0.95
        assert result_dict["source"] == "test.txt"


# =========================================================================
# 仓储层测试
# =========================================================================
class TestRAGRepository:
    """测试 RAG 仓储层核心功能。"""

    def test_add_and_retrieve(self) -> None:
        """测试文档索引和检索流程。"""
        from projects.p02_knowledge_qa.repositories import RAGRepository

        # 创建新的仓储实例，避免使用全局单例
        repo = RAGRepository()
        repo.clear_index()  # 确保是干净的索引

        # 创建临时文档
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            f.write("年假政策：入职满 1 年享 5 天年假，满 3 年享 10 天年假。")
            tmp = f.name

        try:
            # 添加文档
            cnt = repo.add_document(tmp)
            assert cnt > 0, "索引应产出至少 1 个 chunk"

            # 检索
            results = repo.retrieve("年假有几天")
            assert len(results) > 0
            assert any("年假" in r.text for r in results)
        finally:
            Path(tmp).unlink(missing_ok=True)

    def test_retrieve_empty_query(self) -> None:
        """测试空查询或无关查询应返回空结果。"""
        from projects.p02_knowledge_qa.repositories import RAGRepository

        repo = RAGRepository()
        repo.clear_index()

        # 空查询
        results = repo.retrieve("")
        assert len(results) == 0, "空查询应返回空结果"

        # 无关查询（空索引下）
        results = repo.retrieve("xyz不存在的查询abc123")
        assert len(results) == 0, "空索引下无关查询应返回空"

    def test_add_nonexistent_file(self) -> None:
        """测试添加不存在的文件应抛出 FileNotFoundError。"""
        from projects.p02_knowledge_qa.repositories import RAGRepository

        repo = RAGRepository()
        repo.clear_index()

        with pytest.raises(FileNotFoundError):
            repo.add_document("/nonexistent/path/file.txt")

    def test_get_rag_repository_singleton(self) -> None:
        """测试仓储单例模式。"""
        from projects.p02_knowledge_qa.repositories import (
            get_rag_repository,
            set_rag_repository,
        )

        # 正常获取
        repo1 = get_rag_repository()
        repo2 = get_rag_repository()
        assert repo1 is repo2, "应该返回同一个单例实例"

        # 测试设置自定义实例
        set_rag_repository(None)
        repo3 = get_rag_repository()
        assert repo3 is not repo1, "重置后应该创建新实例"


# =========================================================================
# 工具函数测试
# =========================================================================
class TestTools:
    """测试工具函数。"""

    def test_retrieve_knowledge_empty(self) -> None:
        """测试空知识库检索。"""
        from projects.p02_knowledge_qa.graph import retrieve_knowledge

        out = retrieve_knowledge.invoke({"query": "xyz不存在xyz"})
        assert "未在" in out or "未找到" in out, "空检索应返回提示信息"

    def test_add_document_tool_success(self) -> None:
        """测试添加文档工具成功场景。"""
        from projects.p02_knowledge_qa.graph import add_document_to_kb

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            f.write("报销流程：提交申请→直属领导审批→财务审核→打款，约 5 个工作日。")
            tmp = f.name

        try:
            out = add_document_to_kb.invoke({"file_path": tmp})
            assert "已成功" in out or "索引" in out
        finally:
            Path(tmp).unlink(missing_ok=True)

    def test_add_document_tool_not_found(self) -> None:
        """测试添加不存在的文档。"""
        from projects.p02_knowledge_qa.graph import add_document_to_kb

        out = add_document_to_kb.invoke({"file_path": "/nonexistent/file.txt"})
        assert "不存在" in out, "文件不存在应该返回错误提示"

    def test_add_document_tool_error_handling(self) -> None:
        """测试添加文档时的错误处理。"""
        from projects.p02_knowledge_qa.graph import add_document_to_kb

        # 传入无效路径（不是文件也不是目录）
        out = add_document_to_kb.invoke({"file_path": ""})
        assert "失败" in out or "不存在" in out


# =========================================================================
# Prompt 测试
# =========================================================================
class TestPrompts:
    """测试系统提示词和约束条件。"""

    def test_rag_system_prompt_contains_retrieve_constraint(self) -> None:
        """验证系统提示词包含必须先检索的约束。"""
        from projects.p02_knowledge_qa.prompts import RAG_SYSTEM_PROMPT

        assert "必须先调用 retrieve_knowledge" in RAG_SYSTEM_PROMPT, \
            "系统提示词应包含必须先调用检索的约束"

    def test_rag_system_prompt_contains_source_citation(self) -> None:
        """验证系统提示词包含标注出处的要求。"""
        from projects.p02_knowledge_qa.prompts import RAG_SYSTEM_PROMPT

        assert "出处" in RAG_SYSTEM_PROMPT or "来源" in RAG_SYSTEM_PROMPT, \
            "系统提示词应包含标注出处的要求"

    def test_rag_system_prompt_contains_no_hallucination(self) -> None:
        """验证系统提示词包含不知道就说不知道的防幻觉约束。"""
        from projects.p02_knowledge_qa.prompts import RAG_SYSTEM_PROMPT

        assert "未在企业知识库中找到" in RAG_SYSTEM_PROMPT or "未找到" in RAG_SYSTEM_PROMPT, \
            "系统提示词应包含防幻觉约束（未找到就说未找到）"

    def test_rag_system_prompt_all_constraints_present(self) -> None:
        """验证系统提示词包含所有核心约束。"""
        from projects.p02_knowledge_qa.prompts import RAG_SYSTEM_PROMPT

        # 检查三大核心约束都存在
        constraints = [
            "必须先调用 retrieve_knowledge",  # 1. 必须先检索
            "出处",  # 2. 标注出处
            "未在企业知识库中找到",  # 3. 不知道就说不知道
        ]
        for constraint in constraints:
            assert constraint in RAG_SYSTEM_PROMPT, \
                f"系统提示词缺少核心约束: {constraint}"

    def test_rag_workflow_get_system_prompt(self) -> None:
        """验证 RAGWorkflow.get_system_prompt() 返回正确的提示词。"""
        from projects.p02_knowledge_qa.graph import RAGWorkflow
        from projects.p02_knowledge_qa.prompts import RAG_SYSTEM_PROMPT

        prompt = RAGWorkflow.get_system_prompt()
        assert prompt == RAG_SYSTEM_PROMPT
        assert "必须先调用 retrieve_knowledge" in prompt


# =========================================================================
# LangGraph 工作流测试
# =========================================================================
class TestRAGWorkflowGraph:
    """测试 LangGraph RAG 工作流实现。"""

    def test_state_graph_structure(self) -> None:
        """验证 StateGraph 结构正确。"""
        from langgraph.graph import StateGraph
        from projects.p02_knowledge_qa.graph import build_rag_workflow_graph

        graph = build_rag_workflow_graph()
        assert isinstance(graph, StateGraph), "应该返回 StateGraph 实例"

    def test_rag_state_definition(self) -> None:
        """验证 RAGState 包含必要字段。"""
        from projects.p02_knowledge_qa.graph import RAGState
        from typing import TypedDict

        assert issubclass(RAGState, TypedDict)
        # 验证字段存在（通过 TypedDict 注解检查）
        annotations = RAGState.__annotations__
        assert "question" in annotations, "RAGState 应包含 question 字段"
        assert "contexts" in annotations, "RAGState 应包含 contexts 字段"
        assert "answer" in annotations, "RAGState 应包含 answer 字段"

    def test_nodes_exist(self) -> None:
        """验证工作流包含 retrieve 和 generate 节点。"""
        from projects.p02_knowledge_qa.graph import build_rag_workflow_graph

        graph = build_rag_workflow_graph()
        nodes = graph.nodes
        assert "retrieve" in nodes, "工作流应包含 retrieve 节点"
        assert "generate" in nodes, "工作流应包含 generate 节点"

    def test_retrieve_node_function(self) -> None:
        """测试 retrieve_node 函数签名和行为。"""
        from projects.p02_knowledge_qa.graph import retrieve_node

        result = retrieve_node({"question": "年假有几天"})
        assert "contexts" in result, "retrieve_node 应该返回 contexts"
        assert isinstance(result["contexts"], list), "contexts 应该是列表"

    def test_compile_graph(self) -> None:
        """测试 RAGWorkflow.compile_graph() 返回可执行的图。"""
        from projects.p02_knowledge_qa.graph import RAGWorkflow

        compiled = RAGWorkflow.compile_graph()
        assert hasattr(compiled, "invoke"), "编译后的图应该有 invoke 方法"


# =========================================================================
# 集成测试
# =========================================================================
@pytest.mark.skipif(
    not os.getenv("ANTHROPIC_API_KEY"),
    reason="需要 ANTHROPIC_API_KEY",
)
class TestAgentIntegration:
    """Agent 集成测试。"""

    def test_rag_agent_uses_retrieval(self) -> None:
        """测试 Agent 能正确使用检索工具回答问题。"""
        from projects.p02_knowledge_qa import KnowledgeQAProject, _get_rag

        # 先添加一条知识
        rag = _get_rag()
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            f.write("远程办公：每周最多可申请 2 天，需主管审批。")
            tmp = f.name

        try:
            rag.add_document(tmp)
            project = KnowledgeQAProject()
            reply = project.run("远程办公怎么申请")
            assert "2 天" in reply or "主管" in reply
        finally:
            Path(tmp).unlink(missing_ok=True)

    def test_rag_agent_no_hallucination(self) -> None:
        """测试 Agent 不会编造知识库中没有的信息。"""
        from projects.p02_knowledge_qa import KnowledgeQAProject

        project = KnowledgeQAProject()
        reply = project.run("火星上的外星人长什么样？")
        # 应该说"未找到"而不是编造
        assert "未在" in reply or "未找到" in reply or "不知道" in reply.lower()