Agentic RAG 实战：让检索拥有「规划能力」的 2026 进阶指南

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel
from typing import TypedDict, List, Optional
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.documents import Document

# 定义 Agent 状态
class AgentState(TypedDict):
    question: str
    plan: List[str]           # 分解后的子任务
    current_task: str         # 当前处理的任务
    documents: List[Document] # 已检索到的文档
    answer: str               # 最终答案
    iterations: int           # 迭代轮次

# 检索工具
@tool
def retrieve(query: str) -> List[Document]:
    """从向量数据库检索相关文档"""
    return vector_store.similarity_search(query, k=5)

@tool
def web_search(query: str) -> str:
    """搜索外部网络（用于动态信息）"""
    return search_api.run(query)

# 规划节点：分解问题
def planner_node(state: AgentState) -> AgentState:
    llm = ChatOpenAI(model="gpt-4o")
    question = state["question"]
    
    # LLM 自主分解任务
    plan = llm.bind_tools([retrieve, web_search]).invoke(
        f"将以下问题分解为检索步骤：{question}"
    )
    
    return {**state, "plan": plan}

# 检索质量评估节点
def grade_node(state: AgentState) -> AgentState:
    """判断检索结果是否足够回答问题"""
    llm = ChatOpenAI(model="gpt-4o")
    
    context = "\n".join([doc.page_content for doc in state["documents"]])
    question = state["question"]
    
    grade_prompt = f"""
    问题：{question}
    检索到的上下文：{context}
    
    这个上下文是否足够完整地回答问题？回答 YES 或 NO，并说明理由。
    """
    
    response = llm.invoke(grade_prompt)
    
    if response.content.startswith("YES"):
        return {**state, "quality": "adequate"}
    else:
        # 提取改进建议，重新检索
        return {**state, "quality": "inadequate", "iterations": state.get("iterations", 0) + 1}

# 生成答案节点
def generate_node(state: AgentState) -> AgentState:
    llm = ChatOpenAI(model="gpt-4o")
    
    context = "\n".join([doc.page_content for doc in state["documents"]])
    
    answer = llm.invoke(
        f"基于以下上下文回答问题。\n\n上下文：{context}\n\n问题：{state['question']}"
    )
    
    return {**state, "answer": answer.content}

# 构建图
graph = StateGraph(AgentState)

graph.add_node("planner", planner_node)
graph.add_node("retrieve", ToolNode([retrieve, web_search]))
graph.add_node("grade", grade_node)
graph.add_node("generate", generate_node)

# 边
graph.set_entry_point("planner")
graph.add_edge("planner", "retrieve")
graph.add_edge("retrieve", "grade")

# 条件边：质量不足则重新检索（最多3次）
def should_retry(state: AgentState) -> str:
    if state.get("quality") == "adequate":
        return "generate"
    elif state.get("iterations", 0) >= 3:
        return "generate"  # 达到上限，强制生成
    else:
        return "retrieve" # 重新检索

graph.add_conditional_edges("grade", should_retry)
graph.add_edge("generate", END)

app = graph.compile()

def context_pruner(state: AgentState, max_tokens: int = 120000) -> AgentState:
    """动态修剪上下文，保留最相关的内容"""
    current_tokens = estimate_tokens(state["documents"])
    
    if current_tokens <= max_tokens:
        return state
    
    # 按相关度重新排序，裁剪低相关度文档
    scored = []
    for doc in state["documents"]:
        score = embedding_similarity(state["question"], doc.page_content)
        scored.append((score, doc))
    
    scored.sort(reverse=True)
    
    kept = []
    token_count = 0
    for score, doc in scored:
        doc_tokens = estimate_tokens(doc.page_content)
        if token_count + doc_tokens <= max_tokens:
            kept.append(doc)
            token_count += doc_tokens
    
    return {**state, "documents": kept}

def filtered_retrieve(query: str, user_roles: List[str]) -> List[Document]:
    """按用户角色过滤后才返回文档"""
    candidates = vector_store.similarity_search(query, k=20)
    
    filtered = []
    for doc in candidates:
        acl = doc.metadata.get("acl", {})  # 文档元数据携带权限信息
        if any(role in acl.get("allowed_roles", []) for role in user_roles):
            filtered.append(doc)
    
    return filtered[:5]