多阶段工作流加持，高级RAG代理如何实现智能对话与精准回答 原创

uv venv rag-env
source rag-env/bin/activate

uv pip install \
    langchain \
    langgraph \
    langchain-google-genai \
    langchain-community \
    python-dotenv \
    jupyterlab \
    ipykernel

python -m ipykernel install --user --name=rag-env --display-name "RAG Agent (uv)"

GOOGLE_API_KEY=your_google_gemini_api_key_here

from dotenv import load_dotenv
load_dotenv()

# Core LangChain components
from langchain.schema import Document
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate

# Graph and state management
from typing import TypedDict, List
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

# 初始化嵌入模型
embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

# 创建全面的技术支援知识库
knowledge_documents = [
    Document(
        page_cnotallow="TechFlow Solutions offers three main service tiers: Basic Support ($29/month) includes email support and basic troubleshooting, Professional Support ($79/month) includes priority phone support and advanced diagnostics, Enterprise Support ($199/month) includes 24/7 dedicated support and custom integrations.",
        metadata={"source": "pricing_guide.pdf", "category": "pricing"},
    ),
    Document(
        page_cnotallow="Our cloud infrastructure services include: Virtual Private Servers starting at $15/month, Managed Databases from $45/month, Content Delivery Network at $0.08/GB, and Load Balancing services at $25/month. All services include 99.9% uptime guarantee.",
        metadata={"source": "infrastructure_pricing.pdf", "category": "services"},
    ),
    Document(
        page_cnotallow="TechFlow Solutions was founded in 2018 by Maria Rodriguez, a former Google engineer with 15 years of experience in cloud architecture. The company has grown from 3 employees to over 150 team members across 12 countries, specializing in enterprise cloud solutions.",
        metadata={"source": "company_history.pdf", "category": "company"},
    ),
    Document(
        page_cnotallow="Our technical support team operates 24/7 for Enterprise customers, business hours (9 AM - 6 PM EST) for Professional customers, and email-only support for Basic customers. Average response times: Enterprise (15 minutes), Professional (2 hours), Basic (24 hours).",
        metadata={"source": "support_procedures.pdf", "category": "support"},
    )
]

# 构建向量数据库
vector_store = Chroma.from_documents(knowledge_documents, embedding_model)
document_retriever = vector_store.as_retriever(search_kwargs={"k": 2})

def enhance_user_query(state: ConversationState):
    """
    根据对话历史改写用户问题，创建适合向量搜索的独立查询。
    """
    print(f"Enhancing query: {state['current_query'].content}")
    
    # 初始化新查询处理的状态
    state["retrieved_documents"] = []
    state["topic_relevance"] = ""
    state["enhanced_query"] = ""
    state["should_generate"] = False
    state["optimization_attempts"] = 0
    
    # 确保存在对话历史
    if "conversation_history" not in state or state["conversation_history"] is None:
        state["conversation_history"] = []
    
    # 如果当前问题不在对话历史中，则将其添加进去
    if state["current_query"] not in state["conversation_history"]:
        state["conversation_history"].append(state["current_query"])
    
    # 检查是否存在对话上下文
    if len(state["conversation_history"]) > 1:
        # 提取上下文和当前问题
        previous_messages = state["conversation_history"][:-1]
        current_question = state["current_query"].content
        
        # 构建上下文感知的提示
        context_messages = [
            SystemMessage(
                cnotallow="""You are an expert query reformulator. Transform the user's question into a standalone, 
                search-optimized query that incorporates relevant context from the conversation history.
                
                Guidelines:
                - Make the question self-contained and clear
                - Preserve the user's intent while adding necessary context
                - Optimize for vector database retrieval
                - Keep the reformulated query concise but comprehensive"""
            )
        ]
        context_messages.extend(previous_messages)
        context_messages.append(HumanMessage(cnotallow=f"Current question: {current_question}"))
        
        # 生成增强后的查询
        enhancement_prompt = ChatPromptTemplate.from_messages(context_messages)
        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", temperature=0.1)
        
        formatted_prompt = enhancement_prompt.format()
        response = llm.invoke(formatted_prompt)
        enhanced_question = response.content.strip()
        
        print(f"Enhanced query: {enhanced_question}")
        state["enhanced_query"] = enhanced_question
    else:
        # 对话中的第一个问题 - 直接使用原问题
        state["enhanced_query"] = state["current_query"].content
        print(f"First query - using original: {state['enhanced_query']}")
    
    return state

class TopicRelevance(BaseModel):
    """主题分类的结构化输出"""
    classification: str = Field(
        descriptinotallow="问题是否与TechFlow Solutions的服务、定价、公司信息等有关？回答'RELEVANT'或'IRRELEVANT'"
    )
    confidence: str = Field(
        descriptinotallow="置信度：'HIGH'、'MEDIUM'或'LOW'"
    )

def validate_topic_relevance(state: ConversationState):
    """
    判断用户问题是否在我们的知识领域内。
    使用增强后的查询以提高分类准确性。
    """
    print("Validating topic relevance...")
    
    classification_prompt = SystemMessage(
        cnotallow="""You are a topic classifier for TechFlow Solutions support system.
        
        RELEVANT topics include:
        - TechFlow Solutions services (cloud infrastructure, migration, DevOps)
        - Pricing for any TechFlow Solutions products or services
        - Company information (history, team, locations)
        - Support procedures and response times
        - Security and compliance features
        - Technical specifications and capabilities
        
        IRRELEVANT topics include:
        - General technology questions not specific to TechFlow
        - Other companies' products or services
        - Personal questions unrelated to business
        - Weather, news, or general knowledge queries
        
        根据包含对话上下文的增强查询进行分类。"""
    )
    
    user_question = HumanMessage(
        cnotallow=f"Enhanced query to classify: {state['enhanced_query']}"
    )
    
    # 创建分类链
    classification_chain = ChatPromptTemplate.from_messages([classification_prompt, user_question])
    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", temperature=0)
    
    structured_llm = llm.with_structured_output(TopicRelevance)
    classifier = classification_chain | structured_llm
    
    result = classifier.invoke({})
    state["topic_relevance"] = result.classification.strip()
    
    print(f"Topic classification: {state['topic_relevance']} (Confidence: {result.confidence})")
    return state

def fetch_relevant_content(state: ConversationState):
    """
    使用增强后的查询从知识库中检索文档。
    """
    print("Fetching relevant documents...")
    
    # 使用增强后的查询进行检索
    retrieved_docs = document_retriever.invoke(state["enhanced_query"])
    
    print(f"Retrieved {len(retrieved_docs)} documents")
    for i, doc in enumerate(retrieved_docs):
        print(f"   Document {i+1}: {doc.page_content[:50]}...")
    
    state["retrieved_documents"] = retrieved_docs
    return state

class DocumentRelevance(BaseModel):
    """文档相关性评估的结构化输出"""
    relevance: str = Field(
        descriptinotallow="该文档是否与回答问题相关？回答'RELEVANT'或'IRRELEVANT'"
    )
    reasoning: str = Field(
        descriptinotallow="简要说明文档相关或不相关的原因"
    )

def assess_document_relevance(state: ConversationState):
    """
    评估每个检索到的文档，确定它是否与回答用户问题相关。
    """
    print("Assessing document relevance...")
    
    assessment_prompt = SystemMessage(
        cnotallow="""You are a document relevance assessor. Evaluate whether each document 
        contains information that can help answer the user's question.
        
        A document is RELEVANT if it contains:
        - Direct answers to the question
        - Supporting information that contributes to a complete answer
        - Context that helps understand the topic
        
        A document is IRRELEVANT if it:
        - Contains no information related to the question
        - Discusses completely different topics
        - Provides no value for answering the question
        
        在评估时要严格但公正。"""
    )
    
    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", temperature=0)
    structured_llm = llm.with_structured_output(DocumentRelevance)
    
    relevant_documents = []
    
    for i, doc in enumerate(state["retrieved_documents"]):
        assessment_query = HumanMessage(
            cnotallow=f"""Question: {state['enhanced_query']}
            
            Document to assess:
            {doc.page_content}
            
            Is this document relevant for answering the question?"""
        )
        
        assessment_chain = ChatPromptTemplate.from_messages([assessment_prompt, assessment_query])
        assessor = assessment_chain | structured_llm
        
        result = assessor.invoke({})
        
        print(f"Document {i+1}: {result.relevance} - {result.reasoning}")
        
        if result.relevance.strip().upper() == "RELEVANT":
            relevant_documents.append(doc)
    
    # 更新状态，保留相关文档
    state["retrieved_documents"] = relevant_documents
    state["should_generate"] = len(relevant_documents) > 0
    
    print(f"Final relevant documents: {len(relevant_documents)}")
    return state

def generate_contextual_response(state: ConversationState):
    """
    根据对话历史和相关文档生成最终的回答。
    """
    print("Generating contextual response...")
    
    if "conversation_history" not in state or state["conversation_history"] is None:
        raise ValueError("对话历史是回答生成所必需的")
    
    # 提取回答生成所需的组件
    conversation_context = state["conversation_history"]
    relevant_docs = state["retrieved_documents"]
    enhanced_question = state["enhanced_query"]
    
    # 创建全面的回答模板
    response_template = """You are a knowledgeable TechFlow Solutions support agent. Generate a helpful, 
    accurate response based on the conversation history and retrieved documents.
    
    Guidelines:
    - Use information from the provided documents to answer the question
    - Maintain conversation context and refer to previous exchanges when relevant
    - Be conversational and helpful in tone
    - If the documents don't fully answer the question, acknowledge limitations
    - Provide specific details when available (prices, timeframes, etc.)
    
    Conversation History:
    {conversation_history}
    
    Retrieved Knowledge:
    {document_context}
    
    Current Question: {current_question}
    
    Generate a helpful response:"""
    
    response_prompt = ChatPromptTemplate.from_template(response_template)
    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", temperature=0.3)
    
    # 创建回答生成链
    response_chain = response_prompt | llm
    
    # 生成回答
    response = response_chain.invoke({
        "conversation_history": conversation_context,
        "document_context": relevant_docs,
        "current_question": enhanced_question
    })
    
    generated_response = response.content.strip()
    
    # 将回答添加到对话历史中
    state["conversation_history"].append(AIMessage(cnotallow=generated_response))
    
    print(f"Generated response: {generated_response[:100]}...")
    return state

def optimize_search_query(state: ConversationState):
    """
    当初始检索未能成功时，优化搜索查询。
    包含循环预防机制，避免无限优化循环。
    """
    print("Optimizing search query...")
    
    current_attempts = state.get("optimization_attempts", 0)
    
    # 防止无限优化循环
    if current_attempts >= 2:
        print("⚠Maximum optimization attempts reached")
        return state
    
    current_query = state["enhanced_query"]
    
    optimization_prompt = SystemMessage(
        cnotallow="""You are a search query optimizer. The current query didn't retrieve relevant documents.
        
        Create an improved version that:
        - Uses different keywords or synonyms
        - Adjusts the query structure for better matching
        - Maintains the original intent while improving searchability
        - Considers alternative ways to express the same concept
        
        只提供优化后的查询，无需解释。"""
    )
    
    optimization_request = HumanMessage(
        cnotallow=f"Current query that needs optimization: {current_query}"
    )
    
    optimization_chain = ChatPromptTemplate.from_messages([optimization_prompt, optimization_request])
    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", temperature=0.2)
    
    formatted_prompt = optimization_chain.format()
    response = llm.invoke(formatted_prompt)
    optimized_query = response.content.strip()
    
    # 更新状态
    state["enhanced_query"] = optimized_query
    state["optimization_attempts"] = current_attempts + 1
    
    print(f"Optimized query (attempt {current_attempts + 1}): {optimized_query}")
    return state

def route_by_topic(state: ConversationState):
    """根据主题相关性分类进行路由"""
    print("Routing based on topic relevance...")
    
    relevance = state.get("topic_relevance", "").strip().upper()
    
    if relevance == "RELEVANT":
        print("   → Proceeding to content retrieval")
        return "fetch_content"
    else:
        print("   → Routing to off-topic handler")
        return "handle_off_topic"

def route_by_document_quality(state: ConversationState):
    """根据文档相关性评估进行路由"""
    print("Routing based on document quality...")
    
    optimization_attempts = state.get("optimization_attempts", 0)
    
    if state.get("should_generate", False):
        print("   → Generating response with relevant documents")
        return "generate_response"
    elif optimization_attempts >= 2:
        print("   → Maximum optimization attempts reached")
        return "handle_no_results"
    else:
        print("   → Optimizing query for better results")
        return "optimize_query"

# 边缘情况的辅助函数
def handle_off_topic_queries(state: ConversationState):
    """处理超出知识领域的问题"""
    print("Handling off-topic query...")
    
    if "conversation_history" not in state or state["conversation_history"] is None:
        state["conversation_history"] = []
    
    off_topic_response = """我专注于帮助解答有关TechFlow Solutions服务、定价和公司信息的问题。 
    您的问题似乎超出了我的专业领域。 
    
    我可以帮助您了解：
    - 我们的云基础设施服务和定价
    - 支持流程和响应时间  
    - 公司信息和团队详情
    - 安全和合规特性
    
    您是否有关于TechFlow Solutions的具体问题需要帮助？"""
    
    state["conversation_history"].append(AIMessage(cnotallow=off_topic_response))
    return state

def handle_no_relevant_results(state: ConversationState):
    """处理在优化后仍未找到相关文档的情况"""
    print("No relevant results found after optimization...")
    
    if "conversation_history" not in state or state["conversation_history"] is None:
        state["conversation_history"] = []
    
    no_results_response = """抱歉，我在当前知识库中未能找到能够回答您问题的具体信息。 
    
    这可能是因为：
    - 该信息在我们的文档中不可用
    - 您的问题可能需要进一步明确
    - 您可能需要直接联系我们的支持团队
    
    如需立即帮助，您可以联系我们的支持团队，邮箱为support@techflow.com，或拨打1-800-TECHFLOW。"""
    
    state["conversation_history"].append(AIMessage(cnotallow=no_results_response))
    return state

# 初始化对话记忆
conversation_memory = MemorySaver()

# 创建工作流图
workflow = StateGraph(ConversationState)

# 添加所有处理节点
workflow.add_node("enhance_query", enhance_user_query)
workflow.add_node("validate_topic", validate_topic_relevance)
workflow.add_node("handle_off_topic", handle_off_topic_queries)
workflow.add_node("fetch_content", fetch_relevant_content)
workflow.add_node("assess_relevance", assess_document_relevance)
workflow.add_node("generate_response", generate_contextual_response)
workflow.add_node("optimize_query", optimize_search_query)
workflow.add_node("handle_no_results", handle_no_relevant_results)

# 定义工作流连接
workflow.add_edge("enhance_query", "validate_topic")

# 根据主题相关性进行条件路由
workflow.add_conditional_edges(
    "validate_topic",
    route_by_topic,
    {
        "fetch_content": "fetch_content",
        "handle_off_topic": "handle_off_topic",
    },
)

# 内容处理流程
workflow.add_edge("fetch_content", "assess_relevance")

# 根据文档质量进行条件路由
workflow.add_conditional_edges(
    "assess_relevance",
    route_by_document_quality,
    {
        "generate_response": "generate_response",
        "optimize_query": "optimize_query", 
        "handle_no_results": "handle_no_results",
    },
)

# 优化循环
workflow.add_edge("optimize_query", "fetch_content")

# 终止节点
workflow.add_edge("generate_response", END)
workflow.add_edge("handle_no_results", END)
workflow.add_edge("handle_off_topic", END)

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

# 编译工作流
advanced_rag_agent = workflow.compile(checkpointer=conversation_memory)

print("🧪 Testing Advanced RAG Agent\n")

# 测试1：超出主题范围的问题
print("=== Test 1: Off-Topic Query ===")
test_input = {"current_query": HumanMessage(cnotallow="What's the weather like today?")}
result = advanced_rag_agent.invoke(
    input=test_input, 
    cnotallow={"configurable": {"thread_id": "test_session_1"}}
)
print(f"Response: {result['conversation_history'][-1].content}\n")

Testing Advanced RAG Agent

=== Test 1: Off-Topic Query ===
 Enhancing query: What's the weather like today?
 First query - using original: What's the weather like today?
 Validating topic relevance...
 Topic classification: IRRELEVANT (Confidence: HIGH)
 Routing based on topic relevance...
   → Routing to off-topic handler
 Handling off-topic query...
Response: 我专注于帮助解答有关TechFlow Solutions服务、定价和公司信息的问题。 
    您的问题似乎超出了我的专业领域。 
    
    我可以帮助您了解：
    - 我们的云基础设施服务和定价
    - 支持流程和响应时间  
    - 公司信息和团队详情
    - 安全和合规特性
    
    您是否有关于TechFlow Solutions的具体问题需要帮助？

# 测试2：关于定价的在主题范围内的问题
print("=== Test 2: Service Pricing Query ===")
test_input = {"current_query": HumanMessage(cnotallow="What are your support service pricing options?")}
result = advanced_rag_agent.invoke(
    input=test_input,
    cnotallow={"configurable": {"thread_id": "test_session_2"}}
)
print(f"Response: {result['conversation_history'][-1].content}\n")

=== Test 2: Service Pricing Query ===
 Enhancing query: What are your support service pricing options?
📝 First query - using original: What are your support service pricing options?
🎯 Validating topic relevance...
🏷️  Topic classification: RELEVANT (Confidence: HIGH)
🚦 Routing based on topic relevance...
   → Proceeding to content retrieval
📚 Fetching relevant documents...
📄 Retrieved 2 documents
   Document 1: TechFlow Solutions offers three main service tiers...
   Document 2: Our cloud infrastructure services include: Virtual...
🔍 Assessing document relevance...
📋 Document 1: RELEVANT - The document directly answers the question by listing the names, features, and prices of the support service tiers offered by TechFlow Solutions.
📋 Document 2: IRRELEVANT - The document describes pricing options for cloud infrastructure services, not support services. Therefore, it's not relevant to the question about support service pricing.
✅ Final relevant documents: 1
🚦 Routing based on document quality...
   → Generating response with relevant documents
💬 Generating contextual response...
📝 Generated response: We have three support service tiers available. Basic Support is $29 per month and includes email sup...
Response: We have three support service tiers available. Basic Support is $29 per month and includes email support and basic troubleshooting. Professional Support is $79 per month, providing priority phone support and advanced diagnostics. Finally, Enterprise Support, at $199 per month, includes 24/7 dedicated support and custom integrations.