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")

# 創(chuàng)建全面的技術(shù)支援知識庫
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"},
    )
]

# 構(gòu)建向量數(shù)據(jù)庫
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):
    """
    根據(jù)對話歷史改寫用戶問題，創(chuàng)建適合向量搜索的獨立查詢。
    """
    print(f"Enhancing query: {state['current_query'].content}")
    
    # 初始化新查詢處理的狀態(tài)
    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"] = []
    
    # 如果當(dāng)前問題不在對話歷史中，則將其添加進(jìn)去
    if state["current_query"] not in state["conversation_history"]:
        state["conversation_history"].append(state["current_query"])
    
    # 檢查是否存在對話上下文
    if len(state["conversation_history"]) > 1:
        # 提取上下文和當(dāng)前問題
        previous_messages = state["conversation_history"][:-1]
        current_question = state["current_query"].content
        
        # 構(gòu)建上下文感知的提示
        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}"))
        
        # 生成增強(qiáng)后的查詢
        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):
    """主題分類的結(jié)構(gòu)化輸出"""
    classification: str = Field(
        descriptinotallow="問題是否與TechFlow Solutions的服務(wù)、定價、公司信息等有關(guān)？回答'RELEVANT'或'IRRELEVANT'"
    )
    confidence: str = Field(
        descriptinotallow="置信度：'HIGH'、'MEDIUM'或'LOW'"
    )

def validate_topic_relevance(state: ConversationState):
    """
    判斷用戶問題是否在我們的知識領(lǐng)域內(nèi)。
    使用增強(qiáng)后的查詢以提高分類準(zhǔn)確性。
    """
    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
        
        根據(jù)包含對話上下文的增強(qiáng)查詢進(jìn)行分類。"""
    )
    
    user_question = HumanMessage(
        cnotallow=f"Enhanced query to classify: {state['enhanced_query']}"
    )
    
    # 創(chuàng)建分類鏈
    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):
    """
    使用增強(qiáng)后的查詢從知識庫中檢索文檔。
    """
    print("Fetching relevant documents...")
    
    # 使用增強(qiáng)后的查詢進(jìn)行檢索
    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):
    """文檔相關(guān)性評估的結(jié)構(gòu)化輸出"""
    relevance: str = Field(
        descriptinotallow="該文檔是否與回答問題相關(guān)？回答'RELEVANT'或'IRRELEVANT'"
    )
    reasoning: str = Field(
        descriptinotallow="簡要說明文檔相關(guān)或不相關(guān)的原因"
    )

def assess_document_relevance(state: ConversationState):
    """
    評估每個檢索到的文檔，確定它是否與回答用戶問題相關(guān)。
    """
    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
        
        在評估時要嚴(yán)格但公正。"""
    )
    
    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)
    
    # 更新狀態(tài)，保留相關(guān)文檔
    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):
    """
    根據(jù)對話歷史和相關(guān)文檔生成最終的回答。
    """
    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"]
    
    # 創(chuàng)建全面的回答模板
    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)
    
    # 創(chuàng)建回答生成鏈
    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):
    """
    當(dāng)初始檢索未能成功時，優(yōu)化搜索查詢。
    包含循環(huán)預(yù)防機(jī)制，避免無限優(yōu)化循環(huán)。
    """
    print("Optimizing search query...")
    
    current_attempts = state.get("optimization_attempts", 0)
    
    # 防止無限優(yōu)化循環(huán)
    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
        
        只提供優(yōu)化后的查詢，無需解釋。"""
    )
    
    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()
    
    # 更新狀態(tài)
    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):
    """根據(jù)主題相關(guān)性分類進(jìn)行路由"""
    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):
    """根據(jù)文檔相關(guān)性評估進(jìn)行路由"""
    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"

# 邊緣情況的輔助函數(shù)
def handle_off_topic_queries(state: ConversationState):
    """處理超出知識領(lǐng)域的問題"""
    print("Handling off-topic query...")
    
    if "conversation_history" not in state or state["conversation_history"] is None:
        state["conversation_history"] = []
    
    off_topic_response = """我專注于幫助解答有關(guān)TechFlow Solutions服務(wù)、定價和公司信息的問題。 
    您的問題似乎超出了我的專業(yè)領(lǐng)域。 
    
    我可以幫助您了解：
    - 我們的云基礎(chǔ)設(shè)施服務(wù)和定價
    - 支持流程和響應(yīng)時間  
    - 公司信息和團(tuán)隊詳情
    - 安全和合規(guī)特性
    
    您是否有關(guān)于TechFlow Solutions的具體問題需要幫助？"""
    
    state["conversation_history"].append(AIMessage(cnotallow=off_topic_response))
    return state

def handle_no_relevant_results(state: ConversationState):
    """處理在優(yōu)化后仍未找到相關(guān)文檔的情況"""
    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 = """抱歉，我在當(dāng)前知識庫中未能找到能夠回答您問題的具體信息。 
    
    這可能是因為：
    - 該信息在我們的文檔中不可用
    - 您的問題可能需要進(jìn)一步明確
    - 您可能需要直接聯(lián)系我們的支持團(tuán)隊
    
    如需立即幫助，您可以聯(lián)系我們的支持團(tuán)隊，郵箱為support@techflow.com，或撥打1-800-TECHFLOW。"""
    
    state["conversation_history"].append(AIMessage(cnotallow=no_results_response))
    return state

# 初始化對話記憶
conversation_memory = MemorySaver()

# 創(chuàng)建工作流圖
workflow = StateGraph(ConversationState)

# 添加所有處理節(jié)點
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")

# 根據(jù)主題相關(guān)性進(jìn)行條件路由
workflow.add_conditional_edges(
    "validate_topic",
    route_by_topic,
    {
        "fetch_content": "fetch_content",
        "handle_off_topic": "handle_off_topic",
    },
)

# 內(nèi)容處理流程
workflow.add_edge("fetch_content", "assess_relevance")

# 根據(jù)文檔質(zhì)量進(jìn)行條件路由
workflow.add_conditional_edges(
    "assess_relevance",
    route_by_document_quality,
    {
        "generate_response": "generate_response",
        "optimize_query": "optimize_query", 
        "handle_no_results": "handle_no_results",
    },
)

# 優(yōu)化循環(huán)
workflow.add_edge("optimize_query", "fetch_content")

# 終止節(jié)點
workflow.add_edge("generate_response", END)
workflow.add_edge("handle_no_results", END)
workflow.add_edge("handle_off_topic", END)

# 設(shè)置入口點
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: 我專注于幫助解答有關(guān)TechFlow Solutions服務(wù)、定價和公司信息的問題。 
    您的問題似乎超出了我的專業(yè)領(lǐng)域。 
    
    我可以幫助您了解：
    - 我們的云基礎(chǔ)設(shè)施服務(wù)和定價
    - 支持流程和響應(yīng)時間  
    - 公司信息和團(tuán)隊詳情
    - 安全和合規(guī)特性
    
    您是否有關(guān)于TechFlow Solutions的具體問題需要幫助？

# 測試2：關(guān)于定價的在主題范圍內(nèi)的問題
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.