Agentic RAG：当 AI Agent 接管检索-推理-行动全链路

from langgraph.graph import StateGraph, END
from typing import TypedDict, Literal

class AgentState(TypedDict):
    query: str
    route: str
    documents: list
    answer: str

def route_query(state: AgentState) -> AgentState:
    """Agent 根据查询意图决定路由"""
    query = state["query"]
    
    route_prompt = f"""分析以下用户查询，判断应该路由到哪个数据源：
    - "vector_store": 产品功能、使用方法等非结构化文档
    - "sql_database": 价格、库存、订单等结构化数据查询
    - "web_search": 需要最新信息、竞品对比等实时数据
    - "llm_direct": 通用常识问题，无需外部检索
    
    查询: {query}
    """
    
    result = llm.invoke(route_prompt)
    state["route"] = result.content.strip()
    return state

def retrieve_from_vector_store(state: AgentState) -> AgentState:
    docs = vector_store.similarity_search(state["query"], k=5)
    state["documents"] = docs
    return state

def retrieve_from_sql(state: AgentState) -> AgentState:
    sql_query = llm.invoke(f"将以下问题转换为SQL: {state['query']}")
    results = db.execute(sql_query.content)
    state["documents"] = [str(results)]
    return state

# 构建路由图
graph = StateGraph(AgentState)
graph.add_node("router", route_query)
graph.add_node("vector_retrieval", retrieve_from_vector_store)
graph.add_node("sql_retrieval", retrieve_from_sql)
graph.add_node("web_search", search_web)
graph.add_node("generate", generate_answer)

graph.set_entry_point("router")
graph.add_conditional_edges("router", lambda s: s["route"], {
    "vector_store": "vector_retrieval",
    "sql_database": "sql_retrieval",
    "web_search": "web_search",
    "llm_direct": "generate",
})

def decompose_query(state: AgentState) -> AgentState:
    """将复杂查询拆解为子查询"""
    decompose_prompt = f"""你是一个查询分析专家。请将以下复杂问题拆解为 2-4 个独立的子查询，
    每个子查询应该可以通过单次文档检索来回答。
    
    原始问题: {state["query"]}
    
    以 JSON 数组格式输出子查询列表。"""
    
    result = llm.invoke(decompose_prompt)
    state["sub_queries"] = json.loads(result.content)
    return state

def retrieve_for_sub_queries(state: AgentState) -> AgentState:
    """对每个子查询分别执行检索"""
    all_docs = []
    for sub_q in state["sub_queries"]:
        docs = vector_store.similarity_search(sub_q, k=3)
        all_docs.extend(docs)
    
    # 去重并保留最相关的文档
    state["documents"] = deduplicate_and_rank(all_docs)
    return state

def synthesize_answer(state: AgentState) -> AgentState:
    """基于多步检索的结果进行综合推理"""
    synthesis_prompt = f"""基于以下检索到的文档片段，综合回答用户的原始问题。
    注意：你需要跨多个信息源进行推理和关联分析。
    
    原始问题: {state["query"]}
    子查询及检索结果:
    {format_docs_with_sub_queries(state)}
    """
    
    state["answer"] = llm.invoke(synthesis_prompt).content
    return state

def grade_documents(state: AgentState) -> AgentState:
    """评估检索到的文档与查询的相关性"""
    graded_docs = []
    
    for doc in state["documents"]:
        grade_prompt = f"""判断以下文档是否包含回答问题所需的关键信息。
        
        问题: {state["query"]}
        文档: {doc.page_content}
        
        输出 "relevant" 或 "irrelevant"，并简要说明理由。"""
        
        result = llm.invoke(grade_prompt)
        if "relevant" in result.content.lower():
            graded_docs.append(doc)
    
    state["documents"] = graded_docs
    state["doc_quality"] = "sufficient" if len(graded_docs) >= 2 else "insufficient"
    return state

def rewrite_query(state: AgentState) -> AgentState:
    """当检索质量不足时，改写查询"""
    rewrite_prompt = f"""原始查询未能检索到足够相关的文档。请改写查询以提高召回质量。
    
    原始查询: {state["query"]}
    已检索但被判为不相关的文档主题: {state.get("irrelevant_topics", [])}
    
    输出改写后的查询。"""
    
    state["query"] = llm.invoke(rewrite_prompt).content
    state["retry_count"] = state.get("retry_count", 0) + 1
    return state

# 构建自纠错流程图
graph = StateGraph(AgentState)
graph.add_node("retrieve", retrieve_documents)
graph.add_node("grade", grade_documents)
graph.add_node("rewrite", rewrite_query)
graph.add_node("web_fallback", web_search_fallback)
graph.add_node("generate", generate_answer)

graph.set_entry_point("retrieve")
graph.add_edge("retrieve", "grade")

# 关键的条件路由逻辑
graph.add_conditional_edges("grade", decide_next_step, {
    "sufficient": "generate",      # 质量达标，直接生成
    "rewrite": "rewrite",          # 质量不足，改写查询
    "fallback": "web_fallback",    # 多次重试后仍不足，回退到 Web 搜索
})

graph.add_edge("rewrite", "retrieve")  # 形成重试循环
graph.add_edge("web_fallback", "generate")

def decide_next_step(state: AgentState) -> str:
    if state["doc_quality"] == "sufficient":
        return "sufficient"
    if state.get("retry_count", 0) >= 2:
        return "fallback"
    return "rewrite"

interface AdaptiveRAGState {
  query: string;
  queryType: 'simple' | 'complex' | 'realtime' | 'no_retrieval';
  documents: Document[];
  subQueries?: string[];
  gradeResult?: 'pass' | 'fail';
  retryCount: number;
  finalAnswer: string;
}

// 第一层决策：是否需要检索？
function classifyQuery(state: AdaptiveRAGState): AdaptiveRAGState {
  // Agent 分析查询，判断是否需要外部知识
  // "今天天气如何" -> realtime (需要搜索引擎)
  // "什么是面向对象编程" -> no_retrieval (LLM 直接回答)
  // "你们产品的退款政策" -> simple (单次检索)
  // "对比 Q1-Q3 的运营数据并分析趋势" -> complex (多步骤检索)
  return { ...state, queryType: classifiedType };
}

// 第二层决策：选择检索策略
function selectStrategy(state: AdaptiveRAGState): string {
  switch (state.queryType) {
    case 'no_retrieval': return 'direct_generate';
    case 'realtime': return 'web_search';
    case 'simple': return 'single_retrieval';
    case 'complex': return 'multi_step_retrieval';
  }
}

// 第三层决策：检索后的质量控制
function postRetrievalGate(state: AdaptiveRAGState): string {
  if (state.gradeResult === 'pass') return 'generate';
  if (state.retryCount >= MAX_RETRIES) return 'fallback_generate';
  return 'rewrite_and_retry';
}

用户查询 -> 意图分类 -> [直接回答 | Web搜索 | 单次检索 | 多步检索]
                                                  |
                                          检索结果评估 -> 通过 -> 生成答案
                                                  |
                                              不通过 -> 查询改写 -> 重新检索 (循环)
                                                  |
                                         超过重试上限 -> 降级回答 + 标记低置信度

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from typing import TypedDict, List, Optional
import json

# ========== 状态定义 ==========
class RAGState(TypedDict):
    query: str
    route: Optional[str]
    documents: List[Document]
    doc_grade: Optional[str]
    retry_count: int
    answer: Optional[str]

llm = ChatOpenAI(model="gpt-4o", temperature=0)
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(embedding_function=embeddings, collection_name="knowledge_base")

# ========== 节点定义 ==========
def analyze_and_route(state: RAGState) -> RAGState:
    """入口节点：分析查询并决定路由"""
    response = llm.invoke(
        f"""分析查询意图，输出JSON格式：
        {{"route": "vector_db" | "web_search" | "direct"}}
        查询：{state['query']}"""
    )
    route_data = json.loads(response.content)
    return {**state, "route": route_data["route"]}

def vector_retrieve(state: RAGState) -> RAGState:
    """从向量数据库检索"""
    docs = vectorstore.similarity_search(state["query"], k=4)
    return {**state, "documents": docs}

def grade_retrieval(state: RAGState) -> RAGState:
    """评估检索质量"""
    if not state["documents"]:
        return {**state, "doc_grade": "fail"}
    
    doc_summaries = "\n".join([d.page_content[:200] for d in state["documents"]])
    response = llm.invoke(
        f"""评估这些文档是否能充分回答问题。输出 "pass" 或 "fail"。
        问题：{state['query']}
        文档摘要：{doc_summaries}"""
    )
    grade = "pass" if "pass" in response.content.lower() else "fail"
    return {**state, "doc_grade": grade}

def rewrite_and_retry(state: RAGState) -> RAGState:
    """查询改写并重新检索"""
    response = llm.invoke(
        f"改写以下查询以提高检索效果，直接输出改写结果：{state['query']}"
    )
    new_query = response.content.strip()
    docs = vectorstore.similarity_search(new_query, k=4)
    return {
        **state,
        "query": new_query,
        "documents": docs,
        "retry_count": state["retry_count"] + 1,
    }

def generate_answer(state: RAGState) -> RAGState:
    """生成最终回答"""
    context = "\n\n".join([d.page_content for d in state["documents"]])
    response = llm.invoke(
        f"""基于以下上下文回答问题。如果上下文不足以完全回答，请如实说明。
        问题：{state['query']}
        上下文：{context}"""
    )
    return {**state, "answer": response.content}

def direct_answer(state: RAGState) -> RAGState:
    """LLM 直接回答（无需检索）"""
    response = llm.invoke(state["query"])
    return {**state, "answer": response.content}

# ========== 路由函数 ==========
def route_decision(state: RAGState) -> str:
    return state["route"]

def grade_decision(state: RAGState) -> str:
    if state["doc_grade"] == "pass":
        return "pass"
    if state["retry_count"] >= 2:
        return "max_retries"
    return "retry"

# ========== 构建图 ==========
workflow = StateGraph(RAGState)

workflow.add_node("analyze", analyze_and_route)
workflow.add_node("vector_retrieve", vector_retrieve)
workflow.add_node("grade", grade_retrieval)
workflow.add_node("rewrite", rewrite_and_retry)
workflow.add_node("generate", generate_answer)
workflow.add_node("direct", direct_answer)

workflow.set_entry_point("analyze")

workflow.add_conditional_edges("analyze", route_decision, {
    "vector_db": "vector_retrieve",
    "web_search": "vector_retrieve",  # 简化：web_search 也走向量库
    "direct": "direct",
})

workflow.add_edge("vector_retrieve", "grade")

workflow.add_conditional_edges("grade", grade_decision, {
    "pass": "generate",
    "retry": "rewrite",
    "max_retries": "generate",  # 超过重试上限，用现有文档尽力生成
})

workflow.add_edge("rewrite", "grade")
workflow.add_edge("generate", END)
workflow.add_edge("direct", END)

# 编译并运行
app = workflow.compile()

result = app.invoke({
    "query": "对比上季度和本季度的客户流失率变化趋势",
    "documents": [],
    "retry_count": 0,
})
print(result["answer"])

# 必须设置硬性上限
MAX_RETRIES = 3
MAX_TOTAL_STEPS = 10

def should_continue(state: RAGState) -> str:
    if state["retry_count"] >= MAX_RETRIES:
        return "force_generate"
    if state.get("total_steps", 0) >= MAX_TOTAL_STEPS:
        return "force_generate"
    return "continue"