⛓️ LangChain RAG企业级开发

# ============ 步骤1：安装依赖 ============
pip install langchain langchain-openai langchain-community
pip install chromadb faiss-cpu rank-bm25

# ============ 步骤2：导入模块 ============
import os
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.document_loaders import DirectoryLoader, TextLoader
from langchain.retrievers import BM25Retriever, EnsembleRetriever

os.environ["OPENAI_API_KEY"] = "sk-your-api-key"

# ============ 步骤3：加载文档 ============
loader = DirectoryLoader('./docs', glob="**/*.txt", loader_cls=TextLoader)
documents = loader.load()

# 文档分块（中文优化）
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", "。", "！", "？", ".", "!", "?"]
)
chunks = text_splitter.split_documents(documents)
print(f"📄 加载了 {len(documents)} 个文档，分块为 {len(chunks)} 个")

# ============ 步骤4：配置Embedding ============
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# ============ 步骤5：创建向量存储 ============
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./chroma_langchain"
)
print("✅ 向量存储创建完成")

# ============ 步骤6：混合检索（核心特色） ============
# 向量检索器
vector_retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5}
)

# BM25关键词检索器
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = 5

# 组合检索器
ensemble_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.5, 0.5]
)
print("✅ 混合检索器配置完成")

# ============ 步骤7：配置LLM ============
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)

# ============ 步骤8：自定义Prompt ============
custom_prompt = PromptTemplate(
    template="""你是一个专业的问答助手。请基于以下参考文档回答问题。

参考文档：
{context}

问题：{question}

要求：
1. 仅基于参考文档回答
2. 回答要准确、简洁
3. 说明信息来源

回答：""",
    input_variables=["context", "question"]
)

# ============ 步骤9：创建QA Chain ============
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=ensemble_retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": custom_prompt}
)

print("✅ QA Chain创建完成")

# ============ 步骤10：查询 ============
def query_rag(question):
    result = qa_chain({"query": question})
    
    print(f"\n❓ 问题：{question}")
    print(f"💡 回答：{result['result']}")
    print(f"\n📚 参考文档（{len(result['source_documents'])}个）：")
    for i, doc in enumerate(result['source_documents'], 1):
        print(f"  {i}. {doc.page_content[:80]}...")
    
    return result

# ============ 使用示例 ============
if __name__ == "__main__":
    query_rag("什么是RAG系统？")
    query_rag("RAG有哪些优势？")

print("\n🎉 LangChain RAG系统开发完成！")