xParse + LangGraph 构建 Agentic RAG

用户问题："如何配置数据库连接池？"
    ↓
[LangGraph 工作流]  ←  [xParse]
    ↓
[判断节点] 是否需要检索？
    ├─ 是 → [检索节点] 向量检索
    │         ↓ 
    │   [评估节点] 检索结果是否相关？
    │         ├─ 相关 → [生成节点] 生成答案
    │         └─ 不相关 → [重写节点] 重写问题
    │                           ↓
    │                    [检索节点] 再次检索
    │                           ↓
    │                    [评估节点] 再次评估
    │                           ↓
    │                    [生成节点] 生成答案
    └─ 否 → [生成节点] 直接生成答案

python -m venv .venv && source .venv/bin/activate
pip install "xparse-client>=0.2.10" langchain langchain-community langchain-core \
            langgraph langchain-milvus python-dotenv

# .env
XTI_APP_ID=your-app-id
XTI_SECRET_CODE=your-secret-code
MILVUS_DB_PATH=./agentic_rag_vectors.db
DASHSCOPE_API_KEY=your-dashscope-key

import os
from typing import TypedDict, Annotated
from dotenv import load_dotenv
from xparse_client import create_pipeline_from_config
from langchain_milvus import Milvus
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.chat_models import ChatTongyi
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, BaseMessage

# 加载环境变量
load_dotenv()

PIPELINE_CONFIG = {
    "source": {
        "type": "local",
        "directory": "./knowledge_base",  # 文档存放目录
        "pattern": ["*.pdf", "*.docx", "*.txt"]
    },
    "destination": {
        "type": "milvus",
        "db_path": os.getenv("MILVUS_DB_PATH"),
        "collection_name": "agentic_rag_docs",
        "dimension": 1024
    },
    "api_base_url": "https://api.textin.com/api/xparse",
    "api_headers": {
        "x-ti-app-id": os.getenv("XTI_APP_ID"),
        "x-ti-secret-code": os.getenv("XTI_SECRET_CODE")
    },
    "stages": [
        {
            "type": "parse",
            "config": {"provider": "textin"}
        },
        {
            "type": "chunk",
            "config": {
                "strategy": "by_title",
                "new_after_n_chars": 480,
                "max_characters": 1500,
                "overlap": 80
            }
        },
        {
            "type": "embed",
            "config": {
                "provider": "qwen",
                "model_name": "text-embedding-v4"
            }
        }
    ]
}

# 初始化 Pipeline（全局复用）
pipeline = create_pipeline_from_config(PIPELINE_CONFIG)

def build_knowledge_base():
    """构建知识库"""
    print("开始构建知识库...")
    pipeline.run()
    print("知识库构建完成！")

# 使用与 Pipeline 相同的 embedding 配置
embedding = DashScopeEmbeddings(model="text-embedding-v4")
vector_store = Milvus(
    embedding_function=embedding,
    collection_name="agentic_rag_docs",
    connection_args={"uri": os.getenv("MILVUS_DB_PATH")},
    vector_field="embeddings",
    primary_field="element_id",
    text_field="text",
    enable_dynamic_field=True
)

# 初始化大模型
llm = ChatTongyi(
    model="qwen-max",
    top_p=0.8,
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
)

class GraphState(TypedDict):
    """工作流状态"""
    messages: Annotated[list[BaseMessage], add_messages]
    question: str  # 原始问题
    rewritten_question: str  # 重写后的问题
    documents: list  # 检索到的文档
    generation: str  # 生成的答案
    next: str  # 下一步操作
    retrieval_count: int  # 检索次数

def should_retrieve(state: GraphState) -> GraphState:
    """判断是否需要检索"""
    question = state["question"]
    
    # 使用 LLM 判断问题是否需要检索知识库
    prompt = f"""判断以下问题是否需要从知识库中检索信息才能回答。

            问题：{question}

            如果问题需要特定的文档、数据或知识库信息才能回答，返回 "retrieve"。
            如果问题是一般性对话、问候或不需要特定信息的简单问题，返回 "generate"。

            只返回 "retrieve" 或 "generate"，不要返回其他内容。"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    decision = response.content.strip().lower()
    
    next_step = "retrieve" if "retrieve" in decision else "generate"
    
    return {
        **state,
        "next": next_step
    }

def retrieve(state: GraphState) -> GraphState:
    """检索相关文档"""
    question = state.get("rewritten_question") or state["question"]
    retrieval_count = state.get("retrieval_count", 0)
    
    # 向量检索
    docs = vector_store.similarity_search(question, k=5)
    
    # 格式化检索结果
    documents = []
    for doc in docs:
        documents.append({
            "content": doc.page_content,
            "metadata": doc.metadata
        })
    
    return {
        **state,
        "documents": documents,
        "retrieval_count": retrieval_count + 1  # 增加检索计数
    }

def grade_documents(state: GraphState) -> GraphState:
    """评估检索结果的相关性"""
    question = state.get("rewritten_question") or state["question"]
    documents = state["documents"]
    retrieval_count = state.get("retrieval_count", 0)
    
    # 限制最多重写 2 次（即最多检索 3 次）
    if retrieval_count >= 2:
        return {
            **state,
            "next": "generate"
        }
    
    if not documents:
        return {
            **state,
            "next": "rewrite"
        }
    
    # 构建评估提示
    docs_text = "\n\n".join([
        f"文档 {i+1}:\n{doc['content'][:300]}..."
        for i, doc in enumerate(documents[:3])
    ])
    
    prompt = f"""评估以下检索到的文档是否与问题相关。

            问题：{question}

            检索到的文档：
            {docs_text}

            如果文档与问题高度相关，能够回答问题，返回 "generate"。
            如果文档与问题不相关或相关性很低，返回 "rewrite"。

            只返回 "generate" 或 "rewrite"，不要返回其他内容。"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    decision = response.content.strip().lower()
    
    next_step = "generate" if "generate" in decision else "rewrite"
    
    return {
        **state,
        "next": next_step
    }

def rewrite_question(state: GraphState) -> GraphState:
    """重写问题"""
    question = state["question"]
    documents = state.get("documents", [])
    
    if documents:
        # 基于检索结果重写问题
        docs_summary = "\n".join([
            f"- {doc['content'][:200]}..."
            for doc in documents[:2]
        ])
        prompt = f"""原始问题：{question}

                当前检索到的文档摘要：
                {docs_summary}

                这些文档与问题不够相关。请重写问题，使其能够更好地匹配知识库中的内容。
                重写时应该：
                1. 保持问题的核心意图
                2. 使用更具体的关键词
                3. 考虑知识库可能使用的术语

                只返回重写后的问题，不要返回其他内容。"""
    else:
        # 首次重写
        prompt = f"""原始问题：{question}

                请重写这个问题，使其更具体、更清晰，便于在知识库中检索相关信息。
                重写时应该：
                1. 保持问题的核心意图
                2. 使用更具体的关键词
                3. 考虑知识库可能使用的术语

                只返回重写后的问题，不要返回其他内容。"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    rewritten = response.content.strip()
    
    return {
        **state,
        "rewritten_question": rewritten
    }

def generate(state: GraphState) -> GraphState:
    """生成答案"""
    question = state["question"]
    documents = state.get("documents", [])
    
    if documents:
        # 基于检索结果生成答案
        context = "\n\n".join([
            f"文档来源：{doc['metadata'].get('filename', '未知')}\n内容：{doc['content']}"
            for i, doc in enumerate(documents)
        ])
        
        prompt = f"""基于以下文档内容回答用户问题。

                文档内容：
                {context}

                用户问题：{question}

                请基于文档内容回答问题。如果文档中没有相关信息，请说明。
                在回答中引用具体的文档来源。"""
    else:
        # 直接生成答案（不需要检索的情况）
        prompt = f"""回答以下问题：{question}"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    
    return {
        **state,
        "generation": response.content
    }

# 创建状态图
workflow = StateGraph(GraphState)

# 添加节点
workflow.add_node("should_retrieve", should_retrieve)
workflow.add_node("retrieve", retrieve)
workflow.add_node("grade_documents", grade_documents)
workflow.add_node("rewrite_question", rewrite_question)
workflow.add_node("generate", generate)

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

# 添加条件边
workflow.add_conditional_edges(
    "should_retrieve",
    lambda state: state.get("next", "generate"),
    {
        "retrieve": "retrieve",
        "generate": "generate"
    }
)

workflow.add_edge("retrieve", "grade_documents")

workflow.add_conditional_edges(
    "grade_documents",
    lambda state: state.get("next", "generate"),
    {
        "generate": "generate",
        "rewrite": "rewrite_question"
    }
)

workflow.add_edge("rewrite_question", "retrieve")
workflow.add_edge("generate", END)

# 编译工作流
app = workflow.compile()

def ask_question(question: str) -> str:
    """提问并获取答案"""
    initial_state = {
        "messages": [HumanMessage(content=question)],
        "question": question,
        "rewritten_question": "",
        "documents": [],
        "generation": "",
        "next": "",
        "retrieval_count": 0  # 初始化检索计数
    }
    
    # 运行工作流
    result = app.invoke(initial_state)
    
    return result["generation"]

if __name__ == "__main__":
    # 首次运行需要构建知识库
    build_knowledge_base()
    
    # 提问示例
    questions = [
        "如何配置数据库连接池？",
        "产品的定价策略是什么？",
        "你好"  # 简单问候，不需要检索
    ]
    
    for question in questions:
        print("=" * 60)
        print(f"问题：{question}")
        print("=" * 60)
        answer = ask_question(question)
        print(f"回答：{answer}\n")

用户问题 → 向量检索 → 生成答案

用户问题 → 判断是否需要检索
         ├─ 需要 → 检索 → 评估相关性
         │         ├─ 相关 → 生成答案
         │         └─ 不相关 → 重写问题 → 重新检索 → ...
         └─ 不需要 → 直接生成答案

class GraphState(TypedDict):
    # ... 其他字段
    retrieval_count: int  # 检索次数

def grade_documents(state: GraphState) -> GraphState:
    """评估检索结果的相关性"""
    # 限制最多重写 2 次
    if state.get("retrieval_count", 0) >= 2:
        return {
            **state,
            "next": "generate"
        }
    
    # ... 原有的评估逻辑

class GraphState(TypedDict):
    # ... 其他字段
    chat_history: list[BaseMessage]  # 对话历史

def generate(state: GraphState) -> GraphState:
    """生成答案"""
    # 在生成时考虑对话历史
    messages = state.get("chat_history", []) + [
        HumanMessage(content=state["question"])
    ]
    # ... 生成逻辑

def retrieve(state: GraphState) -> GraphState:
    """检索相关文档"""
    question = state.get("rewritten_question") or state["question"]
    
    # 向量检索
    vector_docs = vector_store.similarity_search(question, k=3)
    
    # 关键词检索（如果 Milvus 支持）
    # keyword_docs = keyword_search(question, k=2)
    
    # 合并结果
    documents = []
    for doc in vector_docs:
        documents.append({
            "content": doc.page_content,
            "metadata": doc.metadata,
            "source": "vector"
        })
    
    return {
        **state,
        "documents": documents
    }

def generate(state: GraphState) -> GraphState:
    """生成答案"""
    # ... 生成逻辑
    
    # 添加引用来源
    if documents:
        sources = [
            doc['metadata'].get('filename', '未知')
            for doc in documents
        ]
        generation_with_sources = f"{generation}\n\n📚 参考来源：{', '.join(set(sources))}"
        return {
            **state,
            "generation": generation_with_sources
        }